SDXSv3 — looped diffusion transformer (Krea-style, weight-tied)

Compact text-to-image model: a single-stream DiT with a looped middle, at 2.16B params but 28 forward passes. The transformer is trained on the butterfly dataset (checkpoint saved locally to transformer/); this repo has the architecture, components, dataset/training scripts, and a finished smoke test.

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

Model size

• Stored (unique) weights: 2.16B = 4.0 GB (bf16) — SD3-Medium / SDXL-UNet class.
• Synthetic (effective) size: ~4.0B — the 13 middle blocks run twice (looped middle, weight-tied), so the forward is 28 passes = a ~28-block DiT at hidden 3328; if those 28 blocks were independent, that would be 28 × ~144M ≈ 4.0B params.
• Compute: ~120 GFLOP/token ≈ SD3.5-Medium (2.5B, 24 blocks).

Architecture

Krea-2 building blocks (RMSNorm + QK-norm, SwiGLU 8/3, sigmoid-gated attention, light bias modulation, 3D axial RoPE), with two changes:

hidden 3328  (26 heads × 128, GQA 2:1)
depth:  15 unique blocks = 1 input + 13 middle + 1 output
        middle blocks applied TWICE (weight-tied) → 28 passes
text:   Qwen3.5-0.8B (hidden 1024), 4 слоя (2, 8, 14, 22) → text-fusion → 3328
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 keeps the depth of a 28-block DiT while halving the block parameters. Text: hidden states of Qwen layers 2, 8, 14, 22 are stacked → layerwise attention over the layer axis → projector 4→1 → refiner blocks → project to 3328. (No final norm — that was a CLIP-specific trick in sdxs.)

Files

transformer_sdxsv3.py   # SDXSv3Transformer + config (looped forward)
pipeline_sdxsv3.py      # custom pipeline (text → DiT → VAE)
generate.py             # inference (wiring test)
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,
                      #   resolution buckets min_size/max_size/step)
make_test_dataset.py  # one-image dataset for the smoke test (datasets/test)
train.py              # REAL training on the butterfly dataset (buckets 320-640,
                      #   Accelerator, resolution sampler, word/caption dropout,
                      #   flow matching + shift, EMA loss, save_pretrained)
train_test.py         # one-image smoke training (PSNR in PIXEL space, per-step
                      #   loss, wandb, checkpoints)
train_utils.py        # pack/unpack helpers (patch=2)

smoke_test_images/    # finished smoke test: gt.png, gen_2000.png, rec_2000.png

Smoke test (finished)

One photo (448×640 bucket), looped DiT 2.16B, patch=2, asymmetric VAE (32ch), 2000 steps, batch 2, lr 1e-4, shift 5, flow matching velocity, cfg-dropout 10%.

step	loss	rec PSNR	gen PSNR
200	—	16.6	13.7
600	—	25.0	31.1
1400	—	26.3	35.3
2000	0.012	25.1	31.77

PSNR is in pixel space (decoded images), independent of the VAE. smoke_test_images/: gt.png (original), gen_2000.png (from pure noise), rec_2000.png (from 50% noise). See wandb.ai/recoilme/sdxsv3-smoke.

Usage

# 1. build the butterfly dataset (resolution buckets)
python dataset.py

# 2. build the one-image test dataset
python make_test_dataset.py

# 3. smoke train (one image, 2000 steps, pixel PSNR, wandb)
python train_test.py --wandb --model-path transformer

# 4. real training (butterfly, 1000 samples, batch 8)
python train.py --batch-size 8 --lr 1e-4 --epochs 40 --wandb
#   --model-path: saves diffusers format (config.json + safetensors),
#   overwrites the model there on every checkpoint (like sdxs unet.save_pretrained)

Training recipe (flow matching, from recoilme/sdxs)

• velocity prediction, MSE, noisy = (1-t)·x0 + t·noise, target noise - x0
• timestep shift: t = shift·u/(1+(shift-1)·u), shift = 5.0
• cfg-dropout 10%, word-dropout 5%, AdamW8bit lr 1e-4, bf16, gradient checkpointing, clip 1.0
• EMA loss (0.99) logged to tqdm/wandb/console alongside raw loss
• samples: step 0 = generation from noise, step 10 = VAE decode of a dataset sample (dataset sanity check), then every 60 min
• checkpoint: save_pretrained every 30 min, overwrite (diffusers format)

Status

Real butterfly training in progress (batch 8, ~1000 samples, buckets 320-640). Checkpoints are saved locally to transformer/ (overwrite) and are not pushed to this repo. The transformer as uploaded is untrained (random init); generate.py runs end-to-end but produces random output until trained weights are released.