fix: guard model cleanup in try/finally and fix DiTWrapper comments

- 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>
This commit is contained in:
2026-03-28 15:49:04 +01:00
parent 08d73773c5
commit 4f40e15db3
2 changed files with 63 additions and 59 deletions
+1 -1
View File
@@ -95,7 +95,7 @@ class PrismAudioLoRALoader:
# Merge LoRA weights in-place into the DiT's base linear layers.
# ComfyUI re-executes the upstream ModelLoader on the next queue run
# when inputs change, providing a fresh base model as needed.
dit = model["model"].model # DiffusionTransformer
dit = model["model"].model # DiTWrapper
if strength == 0.0:
print("[PrismAudio] LoRA strength=0.0 — skipping merge, base model unchanged.", flush=True)