SDXSv3 — looped diffusion transformer (multi-layer text fusion)

Compact text-to-image model: a single-stream DiT with a looped middle, at 1.652B params but 28 forward passes. Checkpoint saved locally to transformer/ (multi-layer text, txtlayers=3).

Published (public): https://huggingface.co/recoilme/sdxsv3

Model size

• Stored (unique) weights: 1.652B = 3.3 GB (bf16).
• Depth: 28 forward passes = 2 input + 12 middle (×2 loop) + 2 output; effectively a 28-block DiT in compute, with 16 unique blocks (the middle is weight-tied, per-loop RMSNorm + loop_emb). Compute ≈ 1.652B × 28 ≈ 92 GFLOP/token.

Architecture

Single-stream DiT (Krea/FLUX-style blocks: RMSNorm + QK-norm, SwiGLU 8/3, AdaNorm bias modulation, 3D axial RoPE):

hidden 2560  (20 heads × 128, FULL attention — kvheads = 20, no GQA)
depth:  16 unique blocks = 2 input + 12 middle + 2 output
        middle blocks applied TWICE (weight-tied, loop_emb + per-loop RMSNorm)
        → 28 passes
text:   Qwen3.5-2B (hidden 2048), 3 equally-spaced hidden layers
        (hidden_states[2, 12, 22], incl. penultimate — FLUX.2-klein style)
        → text-fusion (n=3): layerwise attention + projector → 2560
patch:  2  (image tokens = latent/2² per axis)
VAE:    AsymmetricAutoencoderKL, 32 latent channels, encoder f8 / decoder f16
        (latents_mean/std applied), 2× upscale → generate at 2× train res

The looped middle (Universal-Transformer style — a per-pass loop_emb embedding tells the block which iteration it is) gives the depth of a 30-block DiT while keeping only 16 block-weights. Text conditioning uses multi-layer fusion: 3 equally-spaced hidden states ([2,12,22]) are stacked and passed through layerwise attention + a learned projector (FLUX.2-klein style), instead of a single layer.

Files

transformer_sdxsv3.py   # SDXSv3Transformer + config (looped forward, full attention)
pipeline_sdxsv3.py      # custom pipeline (text → DiT → VAE)
generate.py             # inference
model_index.json, transformer/config.json
vae/  text_encoder/  tokenizer/  scheduler/

# data & training (adapted from recoilme/sdxs)
dataset.py            # images + .txt -> HF dataset (VAE latents + text + size)
train.py              # training (Accelerator, resolution sampler, flow-matching shift 5,
                      #   EMA loss, save_pretrained)
make_model.py         # build a fresh random-init SDXSv3 with a given config

one_sample_train/     # legacy 1-photo smoke test (make_test_dataset.py, train_test.py)
                      #   moved out of the main tree

# distillation (FLUX.2-klein teacher + Qwen text adapter)
train_distill.py      # FLUX.2-klein (quantized SDNQ) → SDXSv3 via Qwen text adapter
    --txt-layers 3    # rebuild student to multi-layer text fusion ([2,12,22])
                      #   and transfer pretrain weights (only projector re-init)

# FLUX.2-klein adapter (our text encoder -> FLUX.2-klein)
qwen_adapter_2b_4b/   # Qwen3.5-2B [2,12,22] → Qwen3-4B [9,18,27] text adapter (native FLUX.2)
                      #   (make_teacher_cache.py, make_adapter.py, generate_flux2_adapter.py,
                      #    adapter.pt, samples/) — README внутри папки

Setup / dependencies

Install the Python deps:

pip install torch diffusers transformers accelerate datasets bitsandbytes einops wandb sdnq

Local models / data (paths used by the code):

What	Path
SDXSv3 (student)	repo root transformer/, vae/, text_encoder/, tokenizer/, scheduler/
Teacher (quantized SDNQ)	/workspace/.hf_home/hub/models--Disty0--FLUX.2-klein-4B-SDNQ-4bit-dynamic/snapshots/<hash>/
Qwen adapter	qwen_adapter_2b_4b/adapter.pt
Distill datasets	/workspace/ds (auto-merged: alchemist + civitai, HF arrow)

Dataset

--ds-path may point at a folder with one or more HF arrow datasets (dataset_info.json / state.json). train_distill.py finds every sub-dataset, loads it and concatenate_datasetss them. Each must have columns [vae, text, width, height].

Usage

# 1. build a dataset from a folder of images + .txt captions (saves VAE latents)
python dataset.py

# 2. build a random-init model (defaults match the trained config: txtlayers=3, n_mid=12)
python make_model.py

# 3. distill FLUX.2-klein -> SDXSv3 (see "Distillation from FLUX.2-klein" below).
#    After `accelerate config` (MULTI_GPU, num_processes=8), or pass the flags yourself.
#    Example for 8 GPU, ~10-12 h:
accelerate launch train_distill.py \
    --ds-path /workspace/ds --batch-size 80 --epochs 1400 --wandb

# 4. inference
python generate.py --prompt "a majestic deer in a forest"

Distillation from FLUX.2-klein teacher

# teacher: quantized FLUX.2-klein (SDNQ) -> student: SDXSv3
# teacher is conditioned through the Qwen text adapter (2B -> 4B)
# 8 GPU (uses your `accelerate config`). Teacher path defaults to the local SDNQ snapshot;
# override with --teacher-dir if needed.
accelerate launch train_distill.py \
    --ds-path /workspace/ds \
    --adapter qwen_adapter_2b_4b/adapter.pt \
    --model-path transformer \
    --batch-size 80 --lr 1e-4 --epochs 1400 \
    --lambda-gt 0.1 --wandb

• Teacher is frozen (Flux2Transformer2DModel, SDNQ-quantized, compact on GPU).
• Student is SDXSv3Transformer; the teacher sees the same text states as the student via the Qwen adapter (Qwen3.5-2B [2,12,22] -> 6144 -> adapter -> 7680).
• Loss = MSE(student velocity, teacher velocity) + lambda_gt · MSE(student, ground-truth).
• Shared VAE encoder (the decoder was retrained to 16×, so the latent space is compatible between the two models).
• Timestep shift = 5.0 (same t in [0,1] for teacher & student; the teacher scales by 1000 inside).
• --cache-emb pre-encodes text (student [2,12,22] + teacher via adapter) once and frees the text encoder → allows a larger batch; without it the encoder runs every batch.
• Defaults: --sample-every-minutes 30, --save-every-minutes 60 (override on the CLI).
• --limit N — debug: keep only the first N samples after merging.
• Multi-GPU specifics: the teacher is loaded on each rank's own device (not cuda:0), and cp.checkpoint(..., use_reentrant=True) is used with the weight-tied loop.
• Samples (teacher_N.png, student_N.png, gt_N.png) saved to samples_distill/ every --sample-every-minutes.

Checkpoints are saved with save_pretrained (diffusers format) into transformer/ (overwrite).

Training recipe (flow matching)

• velocity prediction, MSE, noisy = (1-t)·x0 + t·noise, target noise - x0
• timestep shift: t = shift·u/(1+(shift-1)·u), shift = 5.0
• AdamW8bit lr 1e-4, bf16, gradient checkpointing, clip 1.0
• text: Qwen3.5-2B frozen, multi-layer fusion hidden_states[[2, 12, 22]] (3 equally-spaced layers incl. penultimate; FLUX.2-klein style)
• samples: step 0 = generation from noise, step 10 = VAE decode of a dataset sample, then every 30 min (default); EMA loss in tqdm/wandb

Rebuild to multi-layer text

# rebuild student from the pretrain to 3-layer text fusion and run distillation
accelerate launch train_distill.py --ds-path /workspace/ds --epochs 1 \
    --txt-layers 3 --txt-slices "2,12,22" --wandb

• --txt-layers N — rebuild the model to N text layers and transfer the pretrain weights (only txtfusion.projector is re-initialized).
• --txt-slices — which hidden states to stack (comma-separated indices).

Status

• Model: 1.652B params, 28 passes, multi-layer text fusion (txtlayers=3, layers [2,12,22]).
• Weights pushed to HF: recoilme/sdxsv3 (transformer/, 3.3GB).
• The active training path is distillation from FLUX.2-klein (quantized SDNQ) through the Qwen text adapter (train_distill.py, --cache-emb optional). Teacher & student share the same VAE encoder, so their latent spaces are compatible.
• The old butterfly pretraining and the 1-photo smoke test are kept for reference (one_sample_train/), not part of the active pipeline.