- Wrap training loop in try/finally so _unapply_lora always runs.
Without this, an exception mid-training would leave LoRALinear wrappers
in the cached DiTWrapper; a subsequent training run would then apply LoRA
on top of existing LoRA, silently doubling the effective rank.
- Fix misleading comment: diffusion.model is DiTWrapper (not DiffusionTransformer).
DiffusionTransformer is at diffusion.model.model; _apply_lora reaches it
recursively but the direct attribute is the wrapper.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Ethanfel
4f40e15db3