A modular training framework for compact flow-matching diffusion models with dual expert distillation.
- Flow Matching: Rectified flow training with flux-style timestep shifting
- Dual Expert Distillation: Learn from Lune (trajectory) and Sol (attention) teachers
- Feature Caching: Extract expensive teacher features once, reuse forever
- Modular Design: Swap components, customize losses, extend easily
- HuggingFace Integration: Upload checkpoints and samples automatically
- Memory Efficient: Gradient checkpointing, mixed precision, model offloading
pip install torch torchvision
pip install transformers diffusers accelerate
pip install safetensors huggingface_hub datasetsfrom tinyflux.model.zoo import ModelZoo
from tinyflux.trainer.cache_experts import DatasetCache
# Load extraction models
zoo = ModelZoo(device="cuda")
zoo.load_all()
# Build cache (one-time)
cache = DatasetCache.build(
zoo=zoo,
images=my_images, # List of PIL Images
prompts=my_prompts, # List of strings
name="my_dataset",
)
cache.save("my_cache.pt")
zoo.unload_all()from tinyflux.model.model import TinyFluxConfig, TinyFluxDeep
from tinyflux.trainer.trainer import Trainer, TrainerConfig
from tinyflux.trainer.cache_experts import DatasetCache, MultiSourceCache
from tinyflux.trainer.data import CachedDataset, collate_fn
from torch.utils.data import DataLoader
# Load cache
cache = DatasetCache.load("my_cache.pt")
# Create model
model = TinyFluxDeep(TinyFluxConfig()).to("cuda")
# Setup data
dataset = CachedDataset(cache)
loader = DataLoader(dataset, batch_size=8, shuffle=True, collate_fn=collate_fn)
multi_cache = MultiSourceCache()
multi_cache.add(cache, dataset_id=0)
# Train
trainer = Trainer(model, TrainerConfig(total_steps=100000))
trainer.setup(loader, multi_cache)
trainer.train()from tinyflux.model.loader import load_model
from tinyflux.trainer.sampling import Sampler
model = load_model("path/to/checkpoint.safetensors")
sampler = Sampler(zoo, model)
images = sampler.generate(
prompts=["a cat", "a dog"],
num_steps=28,
guidance_scale=5.0,
)┌─────────────────────────────────────────────────────────────┐
│ TinyFlux Pipeline │
├─────────────────────────────────────────────────────────────┤
│ │
│ Phase 1: Cache Building (One-Time) │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │ Images │───▶│ ModelZoo│───▶│ Cache │ │
│ │ Prompts │ │ (VAE, │ │(latents,│ │
│ └─────────┘ │T5,CLIP, │ │ expert │ │
│ │Lune,Sol)│ │features)│ │
│ └─────────┘ └─────────┘ │
│ │
│ Phase 2: Training (Iterative) │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │DataLoader───▶│ Trainer │───▶│Checkpoint │
│ └─────────┘ │(TinyFlux│ │ + EMA │ │
│ │ │ Deep) │ └─────────┘ │
│ │ └─────────┘ │
│ │ ▲ │
│ └──── Cache ───┘ │
│ (Lune + Sol) │
│ │
│ Phase 3: Inference │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │ Prompt │───▶│TinyFlux │───▶│ Image │ │
│ └─────────┘ │(no cache│ └─────────┘ │
│ │ needed) │ │
│ └─────────┘ │
└─────────────────────────────────────────────────────────────┘
TinyFlux learns from two teacher experts:
| Expert | Source | What it Teaches |
|---|---|---|
| Lune | SD1.5 mid-block | Trajectory guidance, semantic structure |
| Sol | SD1.5 attention | Spatial importance, attention patterns |
During training, teacher features are cached and used as targets. At inference, internal predictors replace the teachers - no external models needed.
TinyFluxConfig(
hidden_size=512,
num_attention_heads=4,
attention_head_dim=128,
num_double_layers=15,
num_single_layers=25,
use_lune_expert=True,
use_sol_prior=True,
)TrainerConfig(
learning_rate=1e-4,
total_steps=100000,
optimizer="adamw",
lr_scheduler="cosine",
# Expert distillation
enable_lune=True,
lune_weight=0.1,
lune_dropout=0.1,
enable_sol=True,
sol_weight=0.05,
sol_dropout=0.1,
# Checkpointing
save_every_steps=5000,
hf_repo_id="username/my-model",
)tinyflux/
├── model/
│ ├── model.py # TinyFluxConfig, TinyFluxDeep
│ ├── zoo.py # ModelZoo (extraction models)
│ └── loader.py # Unified checkpoint loading
├── trainer/
│ ├── cache_experts.py # Feature caching
│ ├── trainer.py # Trainer, TrainerConfig
│ ├── losses.py # Loss functions
│ ├── schedules.py # Timestep/LR schedules
│ ├── ema.py # EMA tracking
│ └── sampling.py # Inference
└── util/
└── predictions.py # Flow matching math
- Architecture Blueprint - System design and data flow
- Technical Reference - Component API documentation
- Applications Guide - Use cases and examples
- Expert Cache System - Lune/Sol distillation details
- Conversion Checklist - Migration from train_v4.py
TinyFlux uses rectified flow matching:
# Interpolation: straight line from noise to data
x_t = (1 - t) * noise + t * data
# Target: velocity = direction of flow
v_target = data - noise
# Model predicts velocity
v_pred = model(x_t, t)
loss = MSE(v_pred, v_target)Predictors must work without teachers at inference. Force this during training:
# 10% of steps: drop teacher features
if random.random() < lune_dropout:
lune_features = None # Predictor must work aloneTeacher features vary with timestep. Cache at 10 points, interpolate:
t_buckets = [0.05, 0.15, 0.25, ..., 0.95] # 10 buckets
# For any t, linearly interpolate between nearest buckets| Phase | Components | VRAM (A100) |
|---|---|---|
| Cache Building | VAE + T5 + CLIP + Lune + Sol | ~6 GB |
| Training | Model + Optimizer + Activations | ~8-12 GB |
| Inference | Model + VAE + T5 + CLIP | ~3 GB |
MIT License - see LICENSE file.
- Flow matching: Lipman et al.
- Flux architecture: Black Forest Labs
- Stable Diffusion: Stability AI