From e36cdd79472b39aefba382e8b4708f9cdd34c80d Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Wed, 8 Apr 2026 23:08:13 +0200 Subject: [PATCH] fix(ti-trainer): fix gradient flow and spectral metric shapes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace in-place text_clip assignment with torch.cat so the computation graph correctly links text_input → learned_tokens; in-place assignment into a requires_grad=False leaf severs the graph and learned_tokens receives no gradients - _spectral_metrics(wav, sr): was passing wav.unsqueeze(0) [1,1,L] instead of wav [1,L]; stft mean(dim=1) would return wrong shape [1,T] not [n_freqs] - _save_spectrogram(wav, sr, ...): was passing wav.squeeze(0) [L] (1D) instead of wav [1,L] as the function expects Co-Authored-By: Claude Sonnet 4.6 --- nodes/selva_textual_inversion_trainer.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/nodes/selva_textual_inversion_trainer.py b/nodes/selva_textual_inversion_trainer.py index 41c0bdf..59767f4 100644 --- a/nodes/selva_textual_inversion_trainer.py +++ b/nodes/selva_textual_inversion_trainer.py @@ -282,15 +282,19 @@ class SelvaTextualInversionTrainer: sync_f = torch.stack([x.squeeze(0) for x in sync_list]).to(device, dtype) text_clip = torch.stack([x.squeeze(0) for x in text_list]).to(device, dtype).clone() - # Inject learned tokens into last n_tokens positions - text_clip[:, -n_tokens:, :] = learned_tokens.unsqueeze(0).expand(batch_size, -1, -1) + # Inject learned tokens into last n_tokens positions. + # Must use torch.cat (not in-place assignment) so the computation graph + # links text_input → learned_tokens and gradients flow correctly. + text_front = text_clip[:, :-n_tokens, :].detach() # [B, 77-K, D], no grad + tokens_expanded = learned_tokens.unsqueeze(0).expand(batch_size, -1, -1) # [B, K, D] + text_input = torch.cat([text_front, tokens_expanded], dim=1) # [B, 77, D] with grad x1 = generator.normalize(x1) t = torch.rand(batch_size, device=device, dtype=dtype) x0 = torch.randn_like(x1) xt = fm.get_conditional_flow(x0, x1, t) - v_pred = generator.forward(xt, clip_f, sync_f, text_clip, t) + v_pred = generator.forward(xt, clip_f, sync_f, text_input, t) loss = fm.loss(v_pred, x0, x1).mean() loss.backward() @@ -340,9 +344,9 @@ class SelvaTextualInversionTrainer: import soundfile as sf sf.write(str(wav_path), wav.squeeze(0).numpy(), sr) - metrics = _spectral_metrics(wav.unsqueeze(0), sr) + metrics = _spectral_metrics(wav, sr) if metrics: - img = _save_spectrogram(wav.squeeze(0), sr, ckpt_dir / f"step_{step:05d}.png") + img = _save_spectrogram(wav, sr, ckpt_dir / f"step_{step:05d}.png") print(f"[TI Trainer] step {step} " f"centroid={metrics['spectral_centroid_hz']:.0f}Hz " f"flatness={metrics['spectral_flatness']:.4f} "