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.

python3 train_distill.py --wandb --project micro-distill
# defaults: --ds-path datasets/ --batch-size 32 --teacher-fp8 --epochs 40,
#           save every 1000 steps -> ./transformers, previews -> ./samples

• 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, student also at the early step 10): the local teacher is the base klein (klein-4b/transformer is byte-identical to klein-4b-base/transformer), so it needs --teacher-cfg 4.0 and a negative prompt; the distilled weights in klein-4b/transformer-distiled want --teacher-cfg 1.0 --teacher-steps 4 instead;
• latent frames: our latents are normalised per channel by the VAE's latents_mean/std, klein's DiT is trained on latents normalised by its VAE BatchNorm over the packed 128-dim space (the encoders are the same encoder — 106 tensors, max relative diff 2.4e-3 = bf16 rounding). The teacher therefore gets the raw latent of the same point and its velocity is mapped with v_ours = v_teacher · bn_std / our_std (--no-teacher-frame-convert keeps it for an A/B).

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

--batch-size	peak	step
32	26.0 GiB — all resolutions pass	~5 s
36 / 40	OOM	—

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.

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

• Teacher frames and the teacher choice. The velocity mapping in distillation follows from the chain rule, but a free-fit scale on real batches lands at ≈0.74 — which is what a marginal field looks like next to a single-sample GT target, so it neither confirms nor refutes the factor (--no-teacher-frame-convert exists for the A/B). Which teacher to distil from is open too: base klein with CFG 4, or transformer-distiled at 4 steps.
• 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.