Lora

[mi804]

PR Description: Replace peft with an in-repo LoRA training implementation

Summary

This PR removes the peft dependency. LoRA training is now handled by a 60-line implementation inside the repository (diffsynth/core/lora_train/), which is behavior-equivalent to peft while dropping the default adapter-name segment from checkpoint keys.

Changes

1. In-repo LoRA injection

New diffsynth/core/lora_train/patch.py (60 lines, depends only on torch):

• LoRALinear: wraps the target layer as base_layer, plus lora_A / lora_B as plain nn.Linear(bias=False). No ModuleDict, so state_dict keys are lora_A.weight instead of peft's lora_A.default.weight.
• inject_lora_into_model(): replaces matched modules in place, freezes everything else, and marks only the LoRA submodules trainable.
• Optional lora_bias (zero-initialized bias on lora_B, mirroring peft) is supported at the layer level; nothing wires it yet.

Details kept intentionally identical to peft, because they change results:

• Initialization order: kaiming_uniform_(A, a=√5) + zeros_(B) performed on CPU in fp32, then cast — matching peft's RNG stream and rounding.
• Accumulation order: base_out + lora_B(lora_A(x)) * scaling promotes to the LoRA dtype and only casts back at the end. Rounding the delta first differs by ~0.03 in bf16.
• Target matching: a plain string is a re.fullmatch pattern, a list matches by exact name or .-suffix.

Two deliberate deviations from peft:

• A single-element target list is no longer treated as a regex. Previously --lora_target_modules "to_q" matched nothing and raised.
• A matched non-linear module is reported and skipped instead of aborting the run; only "no linear module matched at all" is an error.

2. Checkpoint keys no longer contain default

• DiffusionTrainingModule.add_lora_to_model() delegates to the new implementation; LoRA weights are created directly in the requested dtype instead of being injected and cast afterwards.
• mapping_lora_state_dict() now strips .default., so --lora_checkpoint still accepts checkpoints produced before this change.
• Runtime-generated LoRA keys in z_image_image2lora.py / qwen_image_image2lora.py follow the new naming.
• SdxlLoRAConverter / Krea2LoRAConverter opensource-format exporters (--align_to_opensource_format) updated to the new naming. Under the old naming these two crashed with UnboundLocalError, so this path is fixed rather than merely migrated.
• Removed the SingleMMDiTConfig.get() stub in krea2_dit.py, which existed only for peft's config probing.
• peft removed from pyproject.toml dependencies.

Inference-side loading is unchanged and accepts both namings: GeneralLoRALoader.get_name_dict() already dropped the adapter-name segment, which is why previously trained LoRAs still load for inference and for --lora_checkpoint resume.

3. Incidental fix: SDXL alpha key

In SdxlLoRAConverter.align_to_diffsynth_format(), .replace("lora_A.weight", ".alpha") was missing a leading dot, so the alpha entry collided with the lora_A key. When .alpha was iterated after .lora_down.weight, the alpha scalar overwrote the lora_A weight (reproduced: shape became ()). Fixed, so alpha is now applied as weight * alpha / rank.

Verification

Unit level:

• 7 functional cases: key naming, output-preserving injection, gradient routed to LoRA only, AutoWrappedLinear (VRAM management) as base layer, mixed dtype, fp8 base weights, optional lora_bias.
• Bit-exact parity against peft in 4 dtype/device combinations (cpu fp32/fp32, cpu bf16/bf16, cpu bf16/fp32, cuda bf16/bf16): weights, forward outputs and gradients all compare equal with torch.equal.
• Checkpoint interoperability: 476 keys map one-to-one with a peft checkpoint modulo .default.; both namings resolve to the same 238 target modules.
• Base-layer probe: the patched layer is a plain nn.Linear in default training, bitsandbytes.nn.Linear4bit (uint8 weights) under nf4, and AutoWrappedLinear only under --fp8_models / --offload_models.

Full trainings (official scripts, unmodified hyperparameters, 5 epochs):

Model	LoRA modules	Checkpoint	Inference
Z-Image-Turbo	238	476 keys, bf16, no default	fused, matches training data
Qwen-Image-2512	720	1440 keys	fused
Krea-2-Turbo	264	528 keys (opensource format)	fused; also exercises the exporter round-trip
Wan2.2-TI2V-5B	300	600 keys	image-to-video, matching the training condition

Also run: Z-Image nf4 quantized training on 2 GPUs with DeepSpeed-free DDP, plus fuse / hot-load / legacy-checkpoint hot-load inference paths.

peft regression baselines (same official scripts, 5 epochs, injection monkeypatched back to inject_adapter_in_model):

• Z-Image-Turbo, Qwen-Image-2512 and Wan2.2-TI2V-5B checkpoints match key-for-key after stripping .default., with no shape or dtype differences.
• Step-level comparison on Z-Image: step 0 loss is bit-identical between the two implementations; later-step drift is the same magnitude as running one implementation twice, i.e. CUDA/bf16 nondeterminism.
• The Krea-2 peft baseline crashes in the opensource exporter, which is the pre-existing bug fixed in §2.

Training flags and special paths (7 runs): --enable_lora_hot_loading, --lora_checkpoint resume from both new and legacy checkpoints, --preset_lora_path, --fp8_models, DeepSpeed ZeRO-2, and task=trajectory_imitation (which deep-copies the whole pipeline after injection). All pass.

Notes

• Trained LoRA checkpoints produced before this change still work for inference and for --lora_checkpoint. They are not accepted by --resume_from_checkpoint, which does not normalize keys.
• examples/lingbot_video/model_training/scripts/prompt_rewriter.py still imports peft lazily for a third-party VLM adapter; using that rewriter now requires installing peft manually.