recoilme/sdxs-micro

Micro diffusion — 1.24B latent DiT with a patch-1 refinement stage

A text-to-image flow-matching model: a single-stream diffusion transformer (20 main blocks + 8 fine blocks) over cached latents of an asymmetric 32-channel VAE, with Qwen3-0.6B as the text encoder. Two training paths live in here: from scratch (train_micro.py — plain flow matching, optional VAE-REPA) and distillation from FLUX.2-klein-4B (train_distill.py — the teacher is driven by the same 0.6B text encoder through a text adapter, so student and teacher read one conditioning).

transformer	1.235B params = 20 main blocks @2048 + 8 fine blocks @640 (bf16 → 2.47 GB)
text encoder	Qwen3-0.6B (596M, hidden 1024), hidden states [2,9,14,18,23,27] stacked (6×1024 = 6144)
VAE	AsymmetricAutoencoderKL, 32 latent channels, encoder f8 / decoder f16
tokens per image	800 patch-2 tokens (latent 40×80 at 320×640) + 256 text, then 3200 patch-1 tokens in the fine stage
compute	3.23 TFLOP forward per sample (512²-class grid)
training	batch 16 + torch.compile → ~45 min/epoch on 17k images @320×640 (estimated), peak 13.0 GiB

Status of this snapshot

This is a snapshot at step 15 000 of a 40-epoch from-scratch run on 84 983 text–image pairs with the VAE-REPA alignment below switched on — kept as a backup, not the finished model. The current work distils klein into a fresh student (train_distill.py) and writes to transformers/ (plural), so this snapshot is never overwritten by it.

Download & folder layout

Everything needed to run the model is in this repo — clone it as one folder and the relative paths inside the scripts resolve by themselves:

git lfs install
git clone https://huggingface.co/recoilme/sdxs-micro
cd sdxs-micro

No git/git-lfs? Then download the same tree with the Hub client (parallel, resumable):

hf download recoilme/sdxs-micro --local-dir sdxs-micro

sdxs-micro/
  transformer_micro.py     the model: config dataclass, blocks, fine stage
  pipeline_micro.py        MicroPipeline + build_pipeline()
  train_micro.py           flow-matching training (VAE-REPA alignment optional)
  train_distill.py         distillation training: FLUX.2-klein-4B (fp8) -> this DiT
  adapter/                 Qwen3-0.6B -> klein text adapter (v14) + its training code
  ref/train_distill.py     the reference distillation script this port started from
  generate.py              CLI: prompt -> image
  dataset.py               images + .txt captions -> arrow dataset (VAE latents + text)
  loss_watch.py            live loss/log watcher
  run_repa.sh              the exact command of the from-scratch REPA run (backup era)
  requirements.txt
  transformer/   config.json + 8 safetensors shards + index.json   (2.35 GB, bf16)
  vae/           AsymmetricAutoencoderKL, 32 latent channels       (366 MB)
  text_encoder/  Qwen3-0.6B in plain transformers format           (1.2 GB)
  tokenizer/     its tokenizer
  scheduler/     FlowMatchEulerDiscreteScheduler config (shift 5.0)

transformer/ is stored sharded — diffusion_pytorch_model-0000N-of-00008.safetensors plus diffusion_pytorch_model.safetensors.index.json. That is the ordinary diffusers layout for checkpoints above ~1 GB: from_pretrained reads the index and picks the shards up on its own, there is nothing to concatenate by hand. (The trainer writes a single diffusion_pytorch_model.safetensors; sharding is only how this snapshot is stored.)

MicroTransformer is a custom class defined in transformer_micro.py, not a model registered inside diffusers, so DiffusionPipeline.from_pretrained("recoilme/sdxs-micro") does not work. Load it through the shipped pipeline with the repo folder on sys.path — see Generation below.

Architecture

latent → first(128→2048)
text (Qwen3-0.6B, 6 layers × 1024) → cross-layer fusion → txtmlp(→2048)
        │
        ▼  single stream: [text tokens | image tokens], 2D axial RoPE
   20 × Block(2048, 16 heads × 128)          SwiGLU, QK-norm, sigmoid-gated attention,
        │                                     3D→2D RoPE, per-block timestep modulation
        ▼
   last: RMSNorm+mod → Linear(2048 → 128)    128 = 2×2 patch × 32 channels
        │
        ├── unfold (free reshape)            (B, 800, 128) → (B, 3200, 32)
        ├── concat the unfolded input latent (B, 3200, 64)
        ▼
    8 × Block(640, 10 heads × 64)            image-only, own RoPE on the 80×40 grid
        │                                    no text, no mask → flash-attention path
        ▼
   fine_last: Linear(640 → 32), zero-init → fold back → **add** to the main prediction

Design decisions worth knowing:

• Two RoPE axes, half the head dim each (row, column). One global coordinate frame, so a given phase means the same pixel distance in any block.
• Full attention (kvheads == heads): a diffusion loop has no KV-cache, so GQA would only shrink wk/wv.
• The fine stage is a residual refinement head, not a layer stack in the middle: it sees the patch-2 prediction plus the input latent, both unfolded to native latent resolution, and adds a correction. Its output projection is zero-initialized, so at step 0 the model is bit-identical to the same model without the stage — it can be attached to, or removed from, a trained checkpoint. Its purpose is detail: a patch-2 token spends its 128 outputs on a 2×2 latent cell, so sub-cell structure and seams between neighbouring tokens can only be fixed at patch-1 resolution. Measured cost: +44M params (+3.5%), +13% step time.
• Honest status: the fine stage trains faster and fits a training image better, but on a single-image overfit test it did not consistently beat the plain stack (the config ranking flipped between checkpoints). Its benefit can only be settled on held-out real data — see "Open questions".

Layout

transformer_micro.py            the model (config dataclass + blocks + fine stage)
pipeline_micro.py              inference: MicroPipeline + build_pipeline()
train_micro.py                 from-scratch flow-matching training
train_distill.py               distillation from FLUX.2-klein-4B (fp8 teacher + text adapter)
ref/train_distill.py           reference distillation script (recoilme/sdxs) this port came from
generate.py                    CLI: prompt -> image
one_sample_train/
  train_test_micro.py          single-image overfit smoke test (loss + PSNR/HF)
  make_test_dataset.py         1-sample HF dataset from one photo
dataset.py                     images + .txt captions -> HF arrow dataset (VAE latents + text)
vae/                           AsymmetricAutoencoderKL (32ch, f8/f16) — needed for decode
scheduler/                     FlowMatchEulerDiscreteScheduler config (shift 5.0 for this resolution)
text_encoder/                  Qwen3-0.6B (transformers format)
tokenizer/                     its tokenizer (files copied out of the HF snapshot)
transformer/                   DiT checkpoint shipped in this repo (sharded)
dataset/                       cached latents — build with dataset.py (not in this repo)

Install

pip install -r requirements.txt

The text encoder is already in this repo (text_encoder/, tokenizer/) — nothing to rebuild. Should you ever want to reproduce it from upstream:

hf download Qwen/Qwen3-0.6B --local-dir /tmp/qwen3_06b_raw
python3 -c "from transformers import AutoModel, AutoTokenizer; \
  AutoModel.from_pretrained('/tmp/qwen3_06b_raw').save_pretrained('text_encoder'); \
  AutoTokenizer.from_pretrained('/tmp/qwen3_06b_raw').save_pretrained('tokenizer')"

flash-attn is not required — attention goes through torch.nn.functional.scaled_dot_product_attention. Note: a bool attention mask forces the mem_efficient backend (flash rejects arbitrary masks); that costs ~2-3% here, because attention is only ~7% of a block at these sequence lengths.

Dataset

An HF arrow dataset with columns vae (float16 latent, 32×80×40 for 320×640), text, width, height. Build it with dataset.py from a folder of images + paired .txt captions. Latents are precomputed, so the VAE is not needed during training.

--ds-path accepts one arrow dataset or a folder of several — load_any_dataset walks it and concatenates the chunks. The local pool is datasets/ (7 chunks: testg, pd12m, vklbd, ae3, allphoto, alchemist, civitai = 1 187 239 samples, 11 resolutions, 320×640 … 640×640; ~324 GB, not in the repo). One epoch over all of it is ≈49 k batches at batch 32 — pass --limit or point --ds-path at a single dataset if you want short epochs.

Training

# from scratch, ~45 min/epoch on a 16 GB card
python3 train_micro.py --ds-path dataset/testg_p1 --epochs 40 \
    --compile --compile-mode default \
    --sample-every-steps 500 --save-every-steps 1000

# resume: point --model-path at the checkpoint (config is restored from config.json)
python3 train_micro.py --ds-path dataset/testg_p1 --model-path transformer --epochs 40 --compile

Useful flags: --depth (main blocks), --fine-depth 0 (disable the refinement stage), --batch-size (16 fits in 16 GB at ~13 GiB peak; 32 OOMs), --max-length, --lr, --warmup, --word-dropout, --caption-dropout, --t-detail-bias, --compile-mode, --limit (debug).

Measured on an RTX 4080 16 GB, batch 8, 1056-token sequences, 18 main blocks (no fine stage):

setting	ms/step	ms/image	peak VRAM
eager	1657	207	8.6 GiB
--compile	1091	136	7.8 GiB
--compile --compile-mode max-autotune-no-cudagraphs	1023	128	7.8 GiB

torch.compile gives −34…−38%; autotune spends ~90 s once per sequence shape (use default if you plan to train on many resolutions). Gradient checkpointing stays on — without it batch 8 OOMs. The fine stage adds ~13% step time. At batch 16 the default config (20+8) peaks at 13.0 GiB, so batch 16 fits but with little headroom; batch 8 is the safe setting.

VAE-REPA alignment (optional, off by default)

The checkpoint in this repo was trained with an auxiliary alignment loss on top of flow matching: the image-token hidden states after block 4 of 20 are projected by a 5-layer MLP (17 M params, 2048→2048→2048→2048→128) and matched against the clean VAE latent of the same sample (smooth-ℓ1, summed over the 128 feature dims, mean over tokens). The target is already in the batch (column vae), so there is no external encoder and no extra forward pass — measured cost: 0 % step time, +0.1 GiB VRAM.

python3 train_micro.py --ds-path dataset/testg_p1 --model-path transformer \
    --epochs 40 --batch-size 12 --compile --compile-mode default \
    --repa-coeff 0.15 --repa-depth 4 --repa-warmup-steps 500 --seed 43

flag	meaning
--repa-coeff	weight λ, 0 disables the loss entirely (default 0)
--repa-depth	which main block to tap (4 = 20 % depth; the reference recipe uses layer 2/12)
--repa-layers, --repa-width	projector MLP depth (default 5) and width (default = features)
--repa-beta	smooth-ℓ1 β (default 0.05)
--repa-warmup-steps	linear ramp of λ — the projector starts random, so a resume warms it in

λ is not a copy of the paper value: on a real batch at λ=1.0 the alignment gradient on the tapped blocks is 2.36× the flow gradient there, so 0.15 puts it at ≈0.35× (guides, does not steer). flow and repa are logged separately — the printed total is dominated by repa and says nothing about quality; watch the flow= column.

The projector lives inside the checkpoint (config.repa = true, repa_proj.*, 34 MB of the 2.35 GB) but is only used in training and is ignored at inference: in eval() mode forward returns a plain tensor, so generate.py and pipeline_micro.py need no changes. To ship a checkpoint without it, load it, set model.config.repa = False, and save_pretrained again.

Distillation from FLUX.2-klein-4B (train_distill.py)

The teacher is frozen and driven by our text encoder: Qwen3-0.6B layers 2,9,14,18,23,27 → the adapter (AiArtLab/qwen3-0.6b-4b-adapter, 220M: per-layer-norm MLP 6144→8192→8192→7680 + a residual attention branch) → klein's 7680-dim joint space. The student reads the same 6144-dim feature stack, so both models see one conditioning; the leading 5 tokens (template attention sinks, ‖y‖≈6000 against ≈180 for content) are cut from the teacher condition and kept for the student.

Which klein is which (verified by sha256 against the HF repos, not by folder names):

file	sha256	what it is
klein-4b/transformer	e1096746…	black-forest-labs/FLUX.2-klein-base-4B — the base (undistilled) model
klein-4b/transformer-distiled	9f29f9ed…	black-forest-labs/FLUX.2-klein-4B — the 4-step distilled model
klein-4b-base/	symlinks onto klein-4b/*	hand-made alias with is_distilled: false

klein-4b/model_index.json claims is_distilled: true although the folder holds the base weights. That flag is not cosmetic: Flux2KleinPipeline disables CFG for distilled models, so running the base through klein-4b silently used guidance 1.0 and produced washed-out images. Use klein-4b-base (or our own calls) whenever CFG matters.

Working recipe (RTX 5090, base teacher, adapter conditioning, CFG 4.0, negative prompt bad quality, the project's own shift 5 schedule — no dynamic shifting, no mu): gens/deer_ADAPTER_shift5_30st_cfg4.png (1024×1024, 30 steps) and gens/deer_base_shift5_50st_cfg4.png (50 steps), gens/deer_base_50st_cfg4_neg_badquality.png (native Qwen3-4B encoder, klein's own dynamic-shift schedule), gens/card_example_cat_base_50st_cfg4.png (BFL's own card example). Both sides run on shift 5 — one schedule for the whole distillation.

Resolution matters more than anything else in the previews. With identical settings the base gives a soft, smeared deer on the dataset's own buckets (320×640 → latent 40×20, 800 tokens: gens/base_adapter_320x640_shift5_50st_cfg4.png) and a sharp one at 1024² (gens/base_adapter_1024_shift5_50st_cfg4.png); the dataset ground truth at the same bucket is a sharp photo (gens/dataset_gt_320x640.png). Distilling on the small buckets therefore means chasing a soft teacher — previews (and, if it shows up in the field agreement, training) should move to bigger buckets.

The teacher's field is only valid near the teacher's own noise prior. Same field, same sampler (our shift 5, 30 steps, CFG 4), same frame conversion, only the starting noise differs: gens/ideal_teacher_field_kleinnoise_shift5.png — a clean deer (mean 71, Laplacian variance 640) when the run starts from a klein-frame standard normal mapped into our frame; gens/ideal_teacher_field_ournoise_shift5.png — a flat dark blob (mean 21, variance 4) when it starts from our own std-1 noise (which is 0.93σ in klein's frame). So the trainer builds the path in our packed frame but draws the noise in klein's frame and maps it in: z_ours = (z_kl·bn_std + bn_mean − our_mean) / our_std. That makes the teacher see exactly its own interpolation (1−t)·lat_klein + t·z_klein while the student sees the same point in our units. The student preview starts from the same mapped noise.

⚠ Consequence for later fine-tuning: the distilled student's noise prior is klein's (≈1.11σ with a small offset in our units), not the std-1 prior of train_micro.py. Generating from a distilled checkpoint must start from the mapped noise (student_sample does); switching to plain GT training means changing the prior, so either fine-tune that explicitly or keep sampling distilled weights with the mapped noise.

The preview path is the pipeline. An independent base+adapter run with the teacher preview's settings (same prompt, seed 1000, same bucket, shift 5, 50 steps, CFG 4, negative bad quality) reproduces the trainer's teacher preview pixel for pixel — the conditioning, the frame conversion and the decode in train_distill.py are correct.

python3 train_distill.py --wandb --project micro-distill
# defaults: --ds-path datasets/ --batch-size 24 --teacher-fp8 --epochs 40,
#           previews at 50 steps for both sides (teacher and student),
#           save every 500 steps -> ./transformers, previews -> ./samples

One-image smoke test (the plumbing test — it does not prove distillation works, see below):

python3 train_distill.py --distill --lambda-gt 0 --limit 1 --min-batch 1 --batch-size 1 \
    --epochs 1000 --no-save --out /tmp/distill_smoke

• loss MSE(v_student, v_teacher) + 0.01 · MSE(v_student, v_gt) — the ground-truth term only anchors (the GT field is noisier than the teacher's), --lambda-gt 0 for pure distillation;
• timesteps: same sampler and shift map as train_micro.py;
• teacher in fp8 (diffusers layerwise casting: fp8 storage, bf16 compute) — 7.22 → 3.63 GiB resident, which is what buys the bigger batch;
• previews (3 samples by default, also logged to wandb as samples_gt / samples_teacher / samples_student, the student also at the early step 10): the teacher is the base klein, which needs CFG (--teacher-cfg 4.0) and a negative prompt to generate at all; the distilled weights in klein-4b/transformer-distiled are run with --teacher-cfg 1.0 --teacher-steps 4 instead;
• latent frames (measured, not assumed): the VAE encoders are identical (106 tensors, max relative diff 2.4e-3 = bf16 rounding; quant_conv/post_quant_conv match too), but the two normalisations are not the same numbers — the dataset latents come out at std 1.073 in our frame and 0.930 in klein's (the VAE BatchNorm of the packed 128-dim space), a ~13 % scale difference plus small offsets (up to 0.59σ in the mean). They are linked by an exact affine map, and the trainer applies it in both directions: the teacher gets the point in its frame, the returned velocity comes back with v_ours = v_teacher · bn_std / our_std (--no-teacher-frame-convert for the A/B). Since the map is affine, (1−t)·lat + t·z maps to (1−t)·lat_kl + t·z_kl exactly, so the teacher is always evaluated on its own training distribution. Dropping the map puts it ~1.8σ off (if the raw latent is fed) and the student learns a field that generates nothing.

Measured on one RTX 5090 32 GB (fp8 teacher resident, gradient checkpointing on, all 11 dataset resolutions swept):

--batch-size	peak	step
24	safe maximum — a full run holds ~29 GiB	~4.9 s
32	26.0 GiB on a cold one-step-per-resolution sweep, but OOMs in a real run (previews at step 0/10 keep CFG tensors alive, wandb, and allocator fragmentation over hours)	—
36 / 40	OOM even in the sweep	—

Two independent knobs make it survivable: run_distill.sh restarts the run and lowers the batch by 8 on OutOfMemoryError (resuming from the checkpoint), and watch_distill.sh <pid> appends step/loss/ checkpoint/VRAM to monitor_distill.log every 5 minutes.

For comparison, the same run with a bf16 teacher (7.22 GiB resident) capped out at batch 24 measured on the 640×640 bucket alone, i.e. the two numbers are not methodologically identical — re-measure with the same sweep if the exact delta matters. torch.compile was measured here and hurts distillation: 11.3 s/step against 5.6 s eager at the same batch (and it drops the max batch), so leave it off.

Generation

python3 generate.py --prompt "a girl" --steps 24 --cfg 4.0 --out girl.png
python3 generate.py --prompt "a cat" --prompt "a fox" --out two.png     # concatenated

Run it from the repo root (ROOT is the script's own folder, so the defaults ./transformer, ./vae, ./text_encoder, ./tokenizer, ./scheduler just work); --model-path, --te-path, --tok-path override them.

The decoder is f16 while the encoder is f8, so --height 640 --width 320 (the training resolution) produces a 640×1280 px image.

⚠ generate.py starts from a standard normal in our frame (the train_micro.py prior). A distilled checkpoint expects the mapped klein-frame noise instead, otherwise the sample collapses to a flat blob — see the distillation section. Only train_distill.py's own previews implement that today.

Smoke test (one image)

python3 one_sample_train/train_test_micro.py --steps 1000 --sample-every 250 --out one_sample_train/out

Verifies the whole loop — text encode, forward/backward, sampling, VAE decode, PSNR — and writes gt.png, gen_XXXX.png and a per-step line with pixel PSNR, high-frequency (Laplacian) PSNR and mean abs error. What was measured: loss 2.5 → 0.03 within ~400 steps, and at 1000 steps PSNR ≈ 34 dB / HF ≈ 31 dB — i.e. the model can reproduce details up to the VAE decoder's ceiling.

Open questions

• Is the distilled student actually close to the teacher? Not shown yet. The plumbing is verified (frames, noise prior, preview path = pipeline, bit-exact), but no trained student has been compared with the teacher's field. The one-image tests — 1200 steps and then 3000 steps with the mapped noise prior, --lambda-gt 0 — end at a recognizable but broken deer, which for a single sample with no GT anchor is expected and neither proves nor disproves the pipeline. The teacher loss settles at ~0.05 (velocity units) on that sample; whether that is small enough for a faithful 50-step trajectory is not measured. Decisive (planned, not run): walk one sampling trajectory and compare v_student against v_teacher at every step — conditional, unconditional and the CFG-combined one, because CFG ×4 amplifies whatever the unconditional branch gets wrong.
• Does the fine stage pay for itself? On one overfit image the ranking of (18 main), (22 main) and (22 main + 6 fine) flipped between checkpoints and stayed inside run-to-run noise, so the test cannot decide. It needs a held-out comparison on the real dataset (validation loss / PSNR at matched steps). train_micro.py currently has no validation split — that is the missing piece.
• Depth vs width vs patch size at a fixed 1.2B budget: +54.8M params and ~+11% step time per main block; a patch-4 stack would cut the sequence 4× (2.3× faster) but each token would cover a 32×32 px region, which is the opposite of what the detail goal needs.

Previous work

This repository used to hold an SDXS-v3 style distilled 2B model (single-stream DiT + Qwen3.5-2B + teacher distillation). The transformer, text encoder, adapter, teacher pipeline and its training scripts were removed; the VAE, the scheduler and the dataset tooling are unchanged. The old files are still recoverable from git history.