Commit Graph

2 Commits

Author SHA1 Message Date
Ethanfel 4f40e15db3 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>
2026-03-28 15:49:04 +01:00
Ethanfel 08d73773c5 feat: LoRA trainer and loader nodes for PrismAudio DiT fine-tuning
Adds PrismAudioLoRATrainer and PrismAudioLoRALoader nodes enabling
low-rank adaptation of the DiT on paired (video features + audio) datasets.

- LoRALinear wraps nn.Linear with trainable lora_A/lora_B matrices
- Rectified flow training loop with fp16 GradScaler, AdamW, cfg dropout
- Checkpoint saving every N steps + _config.json metadata alongside weights
- _unapply_lora restores base model state after training completes
- Weight-merge loader: delta_W added in-place, no deep copy overhead
- Three target presets: attn_only, attn_ffn (default), full

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-28 12:18:50 +01:00