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

pip install torch safetensors sentence-transformers pillow numpy huggingface-hub

Inference (standalone — loads from Hugging Face)

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 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

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

---