Update README.md

README.md

@@ -27,7 +27,7 @@ Random Noise → [VAE Encoder] → Latent (16ch) → [DiT (12 blocks)] → Denoi
 
 ### VAE (Autoencoder)
 
-The VAE compresses 256×256 RGB images into a **16-channel latent** with spatial size 32×32 (downsampled by 8×). It uses residual blocks with GroupNorm and SiLU activations. During training, a small KL penalty keeps latents close to a standard normal distribution.
+The VAE compresses 256×256 RGB images into a **16-channel latent** with spatial size 32×32 (downsampled by 8×). It uses residual blocks with GroupNorm and SiLU activations. During training, a KL penalty (β=0.1) keeps latents close to a standard normal distribution.
 
 ### DiT (Diffusion Transformer)
 
@@ -35,28 +35,47 @@ The DiT operates on patched latents (patch size 4 → 8×8 = 64 patches). Each b
 - **Self-attention** for spatial relationships
 - **Cross-attention** conditioned on text embeddings
 - **Adaptive timestep conditioning** via an MLP-projected sinusoidal embedding
+- **Dropout** (0.1) in attention and MLP for regularization
 
 ### Diffusion Process
 
-Standard DDPM (Denoising Diffusion Probabilistic Model) with 1000 timesteps, cosine-like linear beta schedule (β₁=1e‑4, βᵀ=0.02). The model is trained to predict the added noise ε.
+Standard DDPM (Denoising Diffusion Probabilistic Model) with 1000 timesteps and a linear beta schedule (β₁=1e‑4, βᵀ=0.02). The model is trained to predict the added noise ε. Classifier-free guidance (CFG) is used during training with a text-conditioning dropout probability of 15%.
 
 ---
 
 ## Training
 
-- **Dataset**: Midjourney v5 202304 Clean (~20K image-text pairs, with procedural fallback)
-- **Image size**: 256×256
-- **Batch size**: 8
-- **Optimizer**: Adam (β₁=0.9, β₂=0.999, lr=1e‑4)
-- **LR schedule**: Linear warmup (500 steps) + Cosine annealing
-- **Gradient clipping**: 1.0
-- **Mixed precision**: FP16 via GradScaler
-- **VAE pretraining**: 1 epoch of reconstruction + KL loss (lr=1e‑3) before diffusion training
+### Dataset
+
+**[jackyhate/text-to-image-2M](https://huggingface.co/datasets/jackyhate/text-to-image-2M)** — \~2M high-quality text-image pairs in webdataset format. Loaded via the `datasets` library with **streaming** to avoid materializing the full 2TB+ dataset into memory. A rotating **2000-image in-memory buffer** (\~400 MB RAM) is refreshed each epoch from a fresh random stream to provide shuffle diversity without disk I/O bottlenecks.
+
+### Data Augmentation
+
+Random horizontal flip (p=0.5), color jitter (brightness/contrast/saturation ±0.2, hue ±0.05), and random affine transforms (rotation ±10°, translation ±5%, scale 0.9–1.1) via `torchvision`.
+
+### Hyperparameters
+
+| Parameter | Value |
+|-----------|-------|
+| **Image size** | 256×256 |
+| **Batch size** | 8 |
+| **Optimizer** | AdamW (β₁=0.9, β₂=0.999, lr=1e‑4, weight decay=0.01) |
+| **LR schedule** | Linear warmup (500 steps) + Cosine annealing |
+| **Gradient clipping** | 1.0 (norm) |
+| **Mixed precision** | FP16 via `torch.cuda.amp.GradScaler` |
+| **Dropout** | 0.1 (DiT attention + MLP) |
+| **EMA** | Exponential moving average (decay=0.999) applied at every step; EMA weights used for validation and final checkpoint |
+| **Early stopping** | Patience of 8 epochs on validation loss (10% held-out split) |
+| **KL weight** | 0.1 (β-VAE style) |
+
+### VAE Pretraining
+
+Before diffusion training, the VAE is pretrained for **10 epochs** on reconstruction + KL loss with a higher learning rate (lr=1e‑3) to establish a meaningful latent space. During this phase the DiT and text encoder are frozen.
 
 ### Training Phases
 
-1. **VAE Pretraining** — Train encoder + decoder on image reconstruction to establish a meaningful latent space.
-2. **DiT Diffusion Training** — Freeze VAE, train DiT to denoise latents conditioned on text embeddings.
+1. **VAE Pretraining** (10 epochs) — Train encoder + decoder on image reconstruction to establish a meaningful latent space. DiT and text encoder are frozen.
+2. **DiT Diffusion Training** (up to 30 epochs, early-stopped) — Freeze VAE, train DiT to denoise latents conditioned on text embeddings. CFG dropout randomly replaces text embeddings with a learned null embedding to enable classifier-free guidance at inference time.
 
 ---
 
@@ -68,6 +87,11 @@ Standard DDPM (Denoising Diffusion Probabilistic Model) with 1000 timesteps, cos
 pip install torch safetensors sentence-transformers pillow numpy huggingface-hub
 ```
 
+For training, also install:
+```bash
+pip install torchvision datasets
+```
+
 ### Inference (standalone — loads from Hugging Face)
 
 ```python
@@ -78,33 +102,21 @@ image = pipe.generate("a serene lake surrounded by mountains")
 image.save("output.png")
 ```
 
-`ShellDInference` will automatically download the weights from Hugging Face using `huggingface_hub` on first use, then cache them locally.
-
-See [`inference.py`](./inference.py) for the complete standalone inference script.
-
----
-
-## Model Files
+`ShellDInference` automatically downloads weights from Hugging Face via `huggingface_hub` on first use and caches them locally.
 
-| File | Description |
-|------|-------------|
-| `config.json` | Model hyperparameters (ShellDConfig) |
-| `model.safetensors` | All weights: VAE, DiT, and text encoder |
+### Streaming Generation
 
-### Loading the weights manually
+View the diffusion process unfold step-by-step:
 
 ```python
-from huggingface_hub import snapshot_download
-from safetensors.torch import load_file
-import json
-
-model_path = snapshot_download("FlameF0X/ShellD")
-
-with open(f"{model_path}/config.json") as f:
-    config = json.load(f)
-
-state_dict = load_file(f"{model_path}/model.safetensors")
-print(state_dict.keys())  # prefixed: vae.*, dit.*, text_encoder.*
+for img, step_info in pipe.generate_stream(
+    prompt="a futuristic city at night",
+    num_steps=250,
+    cfg_scale=3.0,
+    display_every=25,  # emit an image every 25 steps
+):
+    print(f"Step {step_info['step']}/{step_info['total']}")
+    img.save(f"progress_{step_info['step']:04d}.png")
 ```
 
 ---
@@ -119,6 +131,5 @@ print(state_dict.keys())  # prefixed: vae.*, dit.*, text_encoder.*
 
 - 256×256 resolution only (no upscaling built in)
 - Limited prompt understanding due to small DiT and frozen lightweight text encoder
-- Procedural training data substitutes real images — quality depends on the real/fine-tuned checkpoints
-
+- Quality depends on training data distribution — may not match large-scale models like SDXL or Flux
 ---