Skip to content
Commit 1f586669 authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files

loras: native Z-Image (DiT) LoRA training (_train_dit)



coderai's trainer only targeted SDXL/SD1.x U-Nets, producing unet.-prefixed LoRAs
that silently no-op on Z-Image's DiT (ZImageTransformer2DModel, transformer. keys,
Qwen3 text encoder) — "No LoRA keys associated to ZImageTransformer2DModel found
with prefix='transformer'". The lora_train_base_model=SDXL workaround can't make a
LoRA that loads on Z-Image.

Add _train_dit: a native flow-matching DiT LoRA trainer for Z-Image, reusing the
existing job/progress/checkpoint/queue infra (modeled on the Wan video DiT trainer,
which is video-only). All deps (ZImageTransformer2DModel, ZImagePipeline, PEFT,
Qwen3) are already in the main venv — no ai-toolkit, no separate venv.

Reverse-engineered from diffusers pipeline_z_image.py so training matches inference:
chat-template + Qwen3 hidden_states[-2] masked per-sample embed list; AutoencoderKL
latents scaled (lat-shift)*scale; list-based transformer I/O with normalized
timestep (1000-t)/1000; RAW target = x0-noise (Z-Image negates the model output);
timesteps sampled from the discrete turbo schedule (set_timesteps mu-shift) to keep
the distilled model sharp; save via ZImagePipeline.save_lora_weights (transformer.
keys) so _apply_loras' load_lora_weights applies it.

_train_lora_sync routes a Z-Image DiT base to _train_dit (detected via
model_index/transformer config _class_name); other DiTs (Flux/SD3) still raise with
guidance. v1 — needs one validation train; recipe + knobs in
docs/zimage-lora-training.md.

Co-Authored-By: default avatarClaude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RdMufYvtTbtGDWsiZVoXce
parent 88063940
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment