FluxForge — Text-to-Image Diffusion Transformer

A from-scratch implementation of a Flux.1 Pro–style text-to-image generation system built on Diffusion Transformers (DiT) with Flow Matching.

Architecture Overview

Text Prompt ──► CLIP + T5-XXL Encoders ──► Cross-Attention
                                              │
Noise z_t ──► DiT Blocks (24 layers) ────────┤──► Predicted v(z_t, t)
                                              │
Timestep t ──► Sinusoidal + Adaln ────────────┘
                     │
                     ▼
            VAE Decoder ──► RGB Image (512–1024px)

Features

Quick Start

# Install dependencies
pip install -r requirements.txt

# Generate an image (after training or loading weights)
python generate.py --prompt "a majestic dragon flying over mountains at sunset" \
                   --steps 28 --guidance 7.5 --seed 42

# Launch Gradio UI
python serve.py --ui gradio --port 7860

# Launch FastAPI
python serve.py --ui api --port 8000

# Start training
python train.py --config config/train_default.yaml

Project Structure

flux-from-scratch/
├── config/              # Model & training configurations
├── models/              # Core neural network modules
│   ├── dit.py           # Diffusion Transformer
│   ├── vae.py           # Variational Autoencoder
│   ├── text_encoder.py  # CLIP + T5-XXL encoders
│   ├── flow_matching.py # Flow Matching scheduler
│   ├── lora.py          # LoRA adapters
│   └── controlnet.py    # ControlNet conditioning
├── training/            # Training pipeline
│   ├── trainer.py       # Main training loop
│   ├── dataset.py       # Data loading (LAION-5B)
│   └── losses.py        # Loss functions
├── inference/           # Inference pipeline
│   ├── pipeline.py      # End-to-end generation
│   ├── samplers.py      # ODE/SDE samplers
│   └── safety.py        # NSFW + watermarking
├── ui/                  # Serving interfaces
│   ├── gradio_app.py    # Gradio UI
│   └── api.py           # FastAPI REST API
├── utils/               # Shared utilities
├── train.py             # Training entry point
├── generate.py          # CLI generation
└── serve.py             # Server entry point

Hardware Requirements

License

Apache 2.0