• Stefy Lanza (nextime / spora )'s avatar
    loras: native Z-Image (DiT) LoRA training (_train_dit) · 1f586669
    Stefy Lanza (nextime / spora ) authored
    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: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
    Claude-Session: https://claude.ai/code/session_01RdMufYvtTbtGDWsiZVoXce
    1f586669
Name
Last commit
Last update
codai Loading commit data...
docs Loading commit data...
packaging Loading commit data...
samples Loading commit data...
tests Loading commit data...
tools Loading commit data...
.dockerignore Loading commit data...
.gitignore Loading commit data...
AI.PROMPT Loading commit data...
CODERAI_API_DOCUMENTATION.md Loading commit data...
CoderAI.gif Loading commit data...
DISTRIBUTION.md Loading commit data...
LICENSE.md Loading commit data...
MULTIMODAL_CAPABILITIES.md Loading commit data...
MULTIMODAL_UI_EXAMPLES.md Loading commit data...
README.md Loading commit data...
build-oci.sh Loading commit data...
build.ps1 Loading commit data...
build.sh Loading commit data...
coderai Loading commit data...
coderai-broker-implementation-reference.md Loading commit data...
coderai-integration.md Loading commit data...
commands Loading commit data...
osxbuild.sh Loading commit data...
package-oci.sh Loading commit data...
package-tarball.sh Loading commit data...
requirements-nvidia.txt Loading commit data...
requirements-vulkan.txt Loading commit data...
requirements.txt Loading commit data...
run-oci.sh Loading commit data...
smoke-test-oci.sh Loading commit data...
todo.md Loading commit data...
video_editor.config.json Loading commit data...