Z Image Train From Scratch

A from-scratch, dependency-light LoRA training implementation for Z-Image-Turbo - the companion to Z_Image_Inference_From_Scratch.

This project performs SFT LoRA fine-tuning of the 6B Z-Image-Turbo diffusion transformer using hand-written model definitions and a hand-written LoRA, with no external training-framework runtime dependency. The model code is shared with the inference project.

---

🎯 Trains LoRA on a single 24 GB GPU (RTX 3090)

The 6B DiT, a 4B Qwen3 text encoder, and a Flux VAE do not fit together on a 24 GB card. This trainer avoids that in the simplest way possible:

1. Precompute. The frozen text encoder and VAE are loaded to GPU one at a time, used to encode every (image, prompt) pair into (latent, prompt-embedding), then freed. The cached results are tiny.
2. Train. Only the DiT stays on GPU. A LoRA adapter is injected into the attention/FFN linears; only the LoRA weights are trainable. Gradient checkpointing over every transformer block keeps activation memory low enough for 24 GB.

No accelerate, no peft, no DeepSpeed - just PyTorch.

---

✨ Highlights

• From scratch. The DiT, text encoder, VAE (encoder and decoder), flow-match scheduler, state-dict converters, and the LoRA adapter are all hand-written in plain PyTorch. No diffusers, no peft.
• Faithful training recipe. Flow-matching SFT loss, 1000-step BSMNTW per-timestep weighting, the same to_q,to_k,to_v,to_out.0,w1,w2,w3 LoRA targets at rank 32, AdamW (lr 1e-4, wd 0.01), 5 epochs over a 50×-repeated dataset.
• Compatible LoRA format. The saved epoch-<n>.safetensors uses the standard <path>.lora_A.weight / <path>.lora_B.weight layout, so it loads straight into any compatible LoRA loader (and the from-scratch inference_lora.py here).
• Zero-cost model loading. Models are built on the meta device and weights loaded with assign=True (same trick as the inference project).

---

🚀 Quick Start

Train

bash Z-Image-Turbo-Train-lora.sh

or directly:

python train_lora.py \
  --dataset_base_path  /path/to/Z-Image-Turbo \
  --dataset_metadata_path /path/to/Z-Image-Turbo/metadata.csv \
  --max_pixels 1048576 --dataset_repeat 50 \
  --learning_rate 1e-4 --num_epochs 5 \
  --lora_target_modules "to_q,to_k,to_v,to_out.0,w1,w2,w3" \
  --lora_rank 32 --lora_alpha 32 \
  --use_gradient_checkpointing \
  --output_path output/Z-Image-Turbo_lora

Each epoch writes output/Z-Image-Turbo_lora/epoch-<n>.safetensors.

Generate with the trained LoRA

bash Z-Image-Turbo-Train-lora.sh

or directly:

python inference_lora.py \
  --lora_path output/Z-Image-Turbo_lora/epoch-4.safetensors \
  --prompt "a dog" --output lora_dog.jpg 

---

🖼️ Results

The bundled example_images/ set is five dog photos, each paired with a short prompt in example_images/metadata.csv. This is the entire fine-tuning set - repeated 50× per epoch, trained at LoRA rank 32.

Training data

Generated with the trained LoRA

After fine-tuning, generating with the simple prompt a dog from the epoch-4 checkpoint (1024×1024, 8 steps):

The LoRA steers the base Z-Image-Turbo toward the dog concept captured by the five reference photos: the generated dog reflects the appearance and style of the training set, while the underlying 6B DiT retains its general image priors.

---

🧠 How the training step works

For each sample (batch size 1), following the flow-match SFT loss:

1. Sample a random timestep t from the 1000-step Z-Image flow-match schedule (shift = 3.0).
2. Add flow-matching noise: noisy = (1 - σ) · latent + σ · noise.
3. Target velocity: target = noise − latent.
4. Predict velocity with the DiT: pred = model_fn_z_image_turbo(dit, noisy, t, prompt_embeds).
5. Loss: MSE(pred, target) · w(t), where w(t) is the BSMNTW bell-shaped weight.

Only the LoRA A/B matrices receive gradients; the 6B base DiT is frozen. AdamW updates the LoRA weights.

---

🔧 Notes & Limitations

• Precomputed latents. All (latent, prompt-embed) pairs are encoded up-front and kept on CPU RAM. This is ideal for small/medium datasets (the bundled 5-image set, a few hundred images). For very large datasets you'd want on-the-fly encoding instead.
• Single GPU. The script targets one GPU (CUDA_VISIBLE_DEVICES). Multi-GPU would need accelerate, which is deliberately out of scope for "as simple as possible".
• bf16 LoRA weights. LoRA params are kept in bf16; the optimizer states follow the parameter dtype.

---

🙏 Acknowledgements

Built directly on Z_Image_Inference_From_Scratch (which itself credits DiffSynth-Studio and Z-Image). The training recipe, flow-match loss, BSMNTW weighting, and VAE-encoder architecture follow DiffSynth-Studio's implementation. Model weights are from the official Tongyi-MAI/Z-Image-Turbo release.

📄 License

An independent, from-scratch reimplementation for research and educational purposes. The Z-Image-Turbo model weights remain under their own license - please respect the terms set by Tongyi-MAI/Z-Image-Turbo.