Update README.md

README.md

@@ -1,3 +1,124 @@
 ---
 license: apache-2.0
+pipeline_tag: text-to-image
 ---
+# ShellD (Shell Diffusion)
+
+**Small DiT-based Text-to-Image Latent Diffusion Model**
+
+ShellD is a lightweight text-to-image model that generates **256×256** images from natural language prompts. It uses a Diffusion Transformer (DiT) backbone operating in a compact VAE latent space, making it feasible to train and run on consumer GPUs.
+
+---
+
+## Model Architecture
+
+```
+Text Prompt → [MiniLM-L6-v2 (frozen)] → Text Embedding (384-d)
+                                                      ↓
+Random Noise → [VAE Encoder] → Latent (16ch) → [DiT (12 blocks)] → Denoised Latent → [VAE Decoder] → 256×256 Image
+```
+
+| Component | Details | Params |
+|-----------|---------|--------|
+| **Text Encoder** | `sentence-transformers/all-MiniLM-L6-v2` (frozen) | 22.71M |
+| **VAE** | Encoder + Decoder with residual blocks, 3 down/up stages, latent dim=16 | 23.43M |
+| **DiT** | 12-layer Transformer with self-attention, cross-attention (text), and adaptive timestep conditioning. Patch size=4, hidden dim=256, 8 heads | 20.80M |
+| **Total** | | **66.95M** (trainable: **44.23M**) |
+
+### 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.
+
+### DiT (Diffusion Transformer)
+
+The DiT operates on patched latents (patch size 4 → 8×8 = 64 patches). Each block includes:
+- **Self-attention** for spatial relationships
+- **Cross-attention** conditioned on text embeddings
+- **Adaptive timestep conditioning** via an MLP-projected sinusoidal embedding
+
+### 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 ε.
+
+---
+
+## 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
+
+### 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.
+
+---
+
+## Usage
+
+### Requirements
+
+```bash
+pip install torch safetensors sentence-transformers pillow numpy huggingface-hub
+```
+
+### Inference (standalone — loads from Hugging Face)
+
+```python
+from inference import ShellDInference
+
+pipe = ShellDInference("FlameF0X/ShellD")
+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
+
+| File | Description |
+|------|-------------|
+| `config.json` | Model hyperparameters (ShellDConfig) |
+| `model.safetensors` | All weights: VAE, DiT, and text encoder |
+
+### Loading the weights manually
+
+```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.*
+```
+
+---
+
+## Intended Use
+
+- Educational exploration of diffusion transformers
+- Lightweight text-to-image generation on consumer hardware
+- Starting point for fine-tuning on custom datasets
+
+## Limitations
+
+- 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
+
+---