fix: raise segment_seconds max from 4s to 30s

Hardcoded max of 4.0 prevented using full 8s clips. Raised to 30s.
Also bumped default from 1.0 to 2.0 as a more sensible starting point.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

LORA_TRAINING.md

@@ -28,7 +28,24 @@ soundfile
 
 ## Step 1 — Prepare the dataset
 
-### 1.1 Extract visual features in ComfyUI
+### 1.1 Video format
+
+The feature extractor accepts any input but internally resamples frames to fixed square resolutions (384×384 for CLIP, 224×224 for Synchformer). Both encoders were trained on standard video datasets — predominantly landscape footage. This has two practical implications:
+
+**Aspect ratio** — use **16:9 landscape** whenever possible. Portrait clips (9:16) are mechanically supported but the bicubic stretch into square distorts the image relative to the encoders' training distribution, which can degrade sync feature quality. If your source is portrait, center-crop to square before extraction. Square (1:1) is also fine.
+
+**Resolution** — anything ≥ 480p is sufficient. The extractor downscales to 384px and 224px regardless of source resolution; higher resolution adds no benefit.
+
+**Frame rate** — any. Connect `VHS_VIDEOINFO` from VHS LoadVideo to the feature extractor so fps is read automatically from the file instead of being entered manually.
+
+| Format | Recommendation |
+|---|---|
+| Aspect ratio | 16:9 landscape (preferred) or 1:1 square |
+| Resolution | ≥ 480p (720p+ is fine, no upper limit that matters) |
+| Frame rate | Any — set via VHS_VIDEOINFO |
+| Portrait (9:16) | Center-crop to square before extraction |
+
+### 1.2 Extract visual features in ComfyUI
 
 For each video clip you want to train on:
 
@@ -71,7 +88,7 @@ If the video frame contains multiple moving objects, CLIP and sync features will
 - Leave `mask_strength` at `1.0` for clean isolation; lower it only if the masked region is very small and the model loses context.
 - Re-extract features with a mask even if you already have `.npz` files — better features directly reduce training noise.
 
-### 1.2 Collect clean audio
+### 1.3 Collect clean audio
 
 For each `.npz` file, place a matching audio file with the **same filename stem** in the same directory:
 
@@ -91,7 +108,7 @@ Supported audio formats: `.wav`, `.flac`, `.ogg`, `.aiff`, `.aif`
 
 The audio will be automatically resampled and trimmed/padded to match the model's expected duration. Use clean, isolated recordings — no background noise.
 
-### 1.3 Optional: prompts.txt
+### 1.4 Optional: prompts.txt
 
 If you want a different prompt at training time than the one embedded in the `.npz`, create a `prompts.txt` file in the dataset directory:
 
@@ -157,6 +174,11 @@ The script will:
 | `--resume` | `None` | Path to a step checkpoint to resume from (e.g. `lora_output/adapter_step04000.pt`) |
 | `--precision` | `bf16` | Mixed precision: `bf16`, `fp16`, `fp32` |
 | `--seed` | `42` | Random seed |
+| `--timestep_mode` | `uniform` | Timestep sampling: `uniform`, `logit_normal`, or `curriculum` |
+| `--logit_normal_sigma` | `1.0` | Spread of the logit-normal distribution. Only used with `logit_normal` / `curriculum` |
+| `--curriculum_switch` | `0.6` | Fraction of steps to use logit_normal before switching to uniform. Only with `curriculum` |
+| `--lora_dropout` | `0.0` | Dropout on the LoRA path only. `0.05`–`0.1` helps regularize on small datasets |
+| `--lora_plus_ratio` | `1.0` | LoRA+ LR ratio: `lr_B = lr × ratio`. `1.0` = standard LoRA, `16.0` = LoRA+ |
 
 ---
 
@@ -271,6 +293,51 @@ Add `linear1` to also adapt post-attention projections for large-scale domain sh
 
 Only add `linear1` once you have 150+ clips — it doubles the adapted parameter count and overfits faster on small datasets.
 
+### Timestep sampling mode
+
+Controls how training timesteps are sampled at each step.
+
+`uniform` (default) samples all timesteps equally — equivalent to original MMAudio training.
+
+`logit_normal` concentrates more steps near t=0.5 via `sigmoid(N(0, σ))`. This is the semantically rich mid-noise region. Consistently reaches a lower loss floor but the perceptual improvement on small datasets is marginal.
+
+`curriculum` uses logit_normal for the first `curriculum_switch` fraction of steps (default 60%), then switches to uniform for the remainder. The motivation: logit_normal accelerates early structure learning but undertrains the high-t boundary region; uniform then fills in the fine detail. A switch message is logged when the transition happens.
+
+| Mode | When to use |
+|---|---|
+| `uniform` (default) | Baseline — safe, equivalent to original training |
+| `logit_normal` | When you want a lower loss floor; marginal on small datasets |
+| `curriculum` | Experimental — may improve convergence quality on small datasets |
+
+The `logit_normal_sigma` parameter controls the width of the logit-normal distribution (used by both `logit_normal` and the first phase of `curriculum`):
+- σ=1.0: moderate peak at t=0.5, balanced coverage (default)
+- σ=0.5: sharper peak, less coverage of extremes
+- σ=2.0: broader, approaches uniform
+
+### LoRA dropout
+
+`lora_dropout` applies dropout to the input of the LoRA path (not the frozen base linear). It regularizes the low-rank update without disturbing pretrained weights — helpful on small datasets where the LoRA would otherwise overfit to the training clips.
+
+| Value | Use case |
+|---|---|
+| `0.0` (default) | No regularization — fine for 30+ clips |
+| `0.05` | Light regularization — recommended starting point on 10–20 clips |
+| `0.1` | Stronger regularization — use if loss plateaus but audio is still noisy |
+
+Dropout is not saved in the adapter file — it only affects training. Loading the adapter at inference does not require setting dropout.
+
+### LoRA+ (asymmetric learning rate)
+
+`lora_plus_ratio` splits the learning rate between LoRA A and B matrices: `lr_B = lr × ratio`. The B matrix is the output-side projection and benefits from a higher LR. Setting ratio to 16 enables the LoRA+ scheme from arXiv:2402.12354.
+
+| Ratio | Effect |
+|---|---|
+| `1.0` (default) | Standard LoRA — identical A and B learning rates |
+| `4.0` | Mild asymmetry |
+| `16.0` | LoRA+ — faster convergence, especially on early steps |
+
+LoRA+ is orthogonal to dropout and curriculum sampling — all three can be combined.
+
 ### Adapter strength at inference
 
 | Strength | Effect |

README.md

@@ -58,7 +58,7 @@ Generates audio from video features. Runs the rectified flow ODE with classifier
 
 | Input | Description |
 |-------|-------------|
-| `model` | From SelVA Model Loader |
+| `model` | From SelVA Model Loader (or any loader/loader chain) |
 | `features` | From SelVA Feature Extractor |
 | `prompt` | Text description — leave empty to use the prompt stored in features |
 | `negative_prompt` | What to suppress (e.g. `"speech, voice, talking"`) |
@@ -66,22 +66,261 @@ Generates audio from video features. Runs the rectified flow ODE with classifier
 | `steps` | Sampling steps (default: 25) |
 | `cfg_strength` | Classifier-free guidance scale (default: 4.5) |
 | `seed` | RNG seed |
-| `normalize` | Peak-normalize output to [-1, 1] (default: true) |
+| `normalize` | RMS-normalize output to `target_lufs` (default: true) |
+| `target_lufs` | *(optional)* Target RMS level in dBFS (default: -27) |
+| `steering_vectors` | *(optional)* From SelVA Activation Steering Loader |
+| `steering_strength` | *(optional)* Scale for steering vectors (default: 0.1) |
+| `textual_inversion` | *(optional)* From SelVA Textual Inversion Loader |
+| `ti_strength` | *(optional)* Blend strength for TI tokens (default: 1.0) |
 
 **Output:** `AUDIO`
 
 ---
 
-## Workflow
+### SelVA LoRA Loader
+
+Injects a trained LoRA adapter into the generator. Connect between Model Loader and Sampler.
+
+| Input | Description |
+|-------|-------------|
+| `model` | SELVA_MODEL from Model Loader |
+| `adapter_path` | Path to `adapter_final.pt` or any step checkpoint |
+| `strength` | 0.0 = disabled, 1.0 = full, >1.0 = exaggerated |
+
+**Output:** `model` (SELVA_MODEL with adapter injected)
+
+---
+
+### SelVA LoRA Trainer
+
+Fine-tunes LoRA adapters on a `.npz` feature dataset. See [LORA_TRAINING.md](LORA_TRAINING.md) for the full guide.
+
+**Output:** `adapter` (SELVA_LORA) and `summary_path` (STRING)
+
+---
+
+### SelVA LoRA Scheduler
+
+Runs a series of LoRA experiments from a JSON sweep file. The dataset is encoded once and reused across all runs. Results are collected in `experiment_summary.json` with overlaid loss curves.
+
+| Input | Description |
+|-------|-------------|
+| `model` | SELVA_MODEL |
+| `experiments_file` | Path to JSON sweep config |
+
+**Outputs:** `summary_path` (STRING), `comparison_curves` (IMAGE)
+
+---
+
+### SelVA Skip Experiment
+
+Signals a running SelVA LoRA Scheduler to skip the current experiment and move to the next. Queue this node while the scheduler is running.
+
+**Output:** `flag_path` (STRING)
+
+---
+
+### SelVA LoRA Evaluator
+
+Evaluates multiple LoRA adapters by generating audio from a fixed reference clip, then reports spectral metrics per adapter for comparison. Input is a JSON file listing adapter paths; an empty path means baseline (no LoRA).
+
+**Outputs:** `summary_path` (STRING), `comparison_image` (IMAGE)
+
+---
+
+### SelVA Dataset Browser
+
+Reads a `dataset.json` produced by the SelVA dataset preparation pipeline and exposes one entry at a time via an index. Useful for previewing and iterating through a prepared dataset.
+
+**Outputs:** video path, audio path, frames directory, label, total count
+
+---
+
+### SelVA VAE Roundtrip
+
+Encodes audio through the SelVA VAE then decodes it back. Use this to measure codec reconstruction quality in isolation — if the output sounds degraded relative to the input, the codec ceiling will limit any downstream fine-tuning approach.
+
+| Input | Description |
+|-------|-------------|
+| `model` | SELVA_MODEL |
+| `audio` | AUDIO to test |
+
+**Output:** `audio_reconstructed` (AUDIO)
+
+---
+
+### SelVA HF Smoother
+
+Attenuates high-frequency content that the SelVA codec handles poorly, by blending a low-pass filtered version of the audio with the original. Use before feature extraction to improve LoRA training targets.
+
+**Output:** `audio` (AUDIO)
+
+---
+
+### SelVA Spectral Matcher
+
+Applies a per-band gain correction to bring audio's spectral profile in line with the MMAudio VAE's expected distribution, derived from the normalization statistics baked into the VAE weights. Use on training audio to reduce codec mismatch.
+
+**Output:** `audio` (AUDIO)
+
+---
+
+### SelVA Textual Inversion Trainer
+
+Trains K learnable CLIP token embeddings against an audio dataset with all model weights frozen. The tokens are injected into the Sampler to guide generation toward a target style.
+
+> **Note:** Textual inversion via the text conditioning path has limited effectiveness for fine-grained timbral style transfer in SelVA due to mean-pooling in the text conditioning path. See [STYLE_TRANSFER.md](STYLE_TRANSFER.md) for the current recommended approach.
+
+**Outputs:** `embeddings_path` (STRING), `loss_curve` (IMAGE)
+
+---
+
+### SelVA Textual Inversion Loader
+
+Loads CLIP token embeddings from a `.pt` file produced by the Textual Inversion Trainer. Connect to the Sampler's `textual_inversion` input.
+
+**Output:** `textual_inversion` (TEXTUAL_INVERSION)
+
+---
+
+### SelVA TI Scheduler
+
+Runs a series of Textual Inversion experiments from a JSON sweep file, reusing the encoded dataset across runs.
+
+**Outputs:** `summary_path` (STRING), `comparison_curves` (IMAGE)
+
+---
+
+### SelVA Activation Steering Extractor
+
+Computes per-block activation steering vectors from a training dataset by comparing DiT hidden states under BJ conditioning vs. empty conditioning. The resulting vectors can nudge the denoising trajectory toward the target style at inference.
+
+| Input | Description |
+|-------|-------------|
+| `model` | SELVA_MODEL |
+| `data_dir` | Directory with `.npz` feature files |
+| `output_path` | Where to save `steering_vectors.pt` |
+| `n_samples` | Clips to average over (default: 16) |
+| `seed` | RNG seed |
+
+**Output:** `steering_path` (STRING)
+
+---
+
+### SelVA Activation Steering Loader
+
+Loads steering vectors from a `.pt` file produced by the Extractor. Connect to the Sampler's `steering_vectors` input.
+
+**Output:** `steering_vectors` (STEERING_VECTORS)
+
+---
+
+### SelVA BigVGAN Trainer
+
+Fine-tunes the BigVGAN vocoder (mel → waveform) on a set of target-style audio clips. Only the vocoder is modified — the DiT generator and VAE are completely untouched.
+
+Default mode (`snake_alpha_only`) tunes only the ~27K per-channel α parameters in Snake/SnakeBeta activations, which directly control harmonic periodicity. With 0.024% of parameters trainable the model cannot produce spectral averaging artifacts regardless of loss function. See [STYLE_TRANSFER.md](STYLE_TRANSFER.md) for the full rationale.
+
+| Input | Description |
+|-------|-------------|
+| `model` | SELVA_MODEL |
+| `data_dir` | Directory with target-style audio files (searched recursively) |
+| `output_path` | Where to save the fine-tuned vocoder `.pt` |
+| `train_mode` | `snake_alpha_only` (default) or `all_params` |
+| `steps` | Training steps (default: 2000) |
+| `lr` | Learning rate (default: 1e-4 for snake_alpha_only) |
+| `batch_size` | Clips per step (default: 4) |
+| `segment_seconds` | Audio segment length per training sample (default: 1.0 s) |
+| `lambda_l2sp` | L2-SP anchor regularization strength — penalizes drift from pretrained weights (default: 1e-3) |
+| `save_every` | Checkpoint interval in steps (default: 500) |
+| `seed` | RNG seed |
+| `discriminator_path` | *(optional)* Path to `bigvgan_discriminator_optimizer.pt` — when provided, frozen MPD+MRD feature matching replaces mel L1, directly penalizing harmonic smearing |
+
+**Output:** `checkpoint_path` (STRING) — load with SelVA BigVGAN Loader
+
+Saves eval samples and mel spectrogram PNGs at baseline, each checkpoint, and final.
+
+---
+
+### SelVA BigVGAN Loader
+
+Loads a fine-tuned BigVGAN vocoder checkpoint produced by SelVA BigVGAN Trainer and replaces the vocoder weights in a SELVA_MODEL in-place. Connect the output to SelVA Sampler instead of the base Model Loader.
+
+| Input | Description |
+|-------|-------------|
+| `model` | SELVA_MODEL from Model Loader |
+| `path` | Path to fine-tuned vocoder `.pt` (relative = ComfyUI output directory) |
+
+**Output:** `model` (SELVA_MODEL with fine-tuned vocoder)
+
+---
+
+### SelVA DITTO Optimizer
+
+Inference-time noise optimization ([arXiv:2401.12179](https://arxiv.org/abs/2401.12179), ICML 2024 Oral). Optimizes the initial noise latent x₀ to make the generated audio match a set of BJ reference clips, by backpropagating a mel style loss through the ODE solver. All model weights remain frozen — zero quality degradation risk.
+
+Style loss: mean spectrum + Gram matrix computed against reference mels. The Gram matrix captures covariance between frequency bands (timbral texture) without requiring temporal alignment with the reference clips. Optimization runs only through the DiT + VAE decoder; the vocoder is only invoked for the final output pass.
+
+| Input | Description |
+|-------|-------------|
+| `model` | SELVA_MODEL |
+| `features` | From SelVA Feature Extractor |
+| `prompt` | Sound description (leave empty to use features prompt) |
+| `negative_prompt` | Sounds to suppress |
+| `reference_dir` | Directory with BJ reference audio clips (.wav/.flac/.mp3) |
+| `n_opt_steps` | Gradient optimization steps on x₀ (default: 50) |
+| `opt_lr` | Adam LR for x₀ optimization (default: 0.1) |
+| `n_ode_steps` | ODE steps per optimization iteration (default: 10; lower = faster) |
+| `n_grad_steps` | ODE steps to differentiate through — truncated BPTT (default: 5) |
+| `style_weight` | Style loss weight (default: 1.0; increase for stronger BJ shift) |
+| `steps` | Euler steps for the final generation pass (default: 25) |
+| `cfg_strength` | CFG scale (default: 4.5) |
+| `seed` | RNG seed |
+| `normalize` | *(optional)* RMS normalize output (default: true) |
+| `target_lufs` | *(optional)* Target RMS level in dBFS (default: -27) |
+
+**Output:** `AUDIO`
+
+---
+
+## Workflows
+
+### Basic generation
 
 ```
-VHS LoadVideo ──► SelVA Feature Extractor ──────────────────────► SelVA Sampler ──► Save Audio
+VHS LoadVideo ──► SelVA Feature Extractor ─────────────────────► SelVA Sampler ──► Save Audio
-                      │ (video_info) ─► (fps auto)                      ▲
+                       │ (video_info)                                  ▲
-                      │ (features) ────────────────────────────────────►│
+                       │ (features) ──────────────────────────────────►│
-                      │ (prompt) ──────────────────────────────────────►│
+                       │ (prompt) ────────────────────────────────────►│
 ```
 
-Connect the `prompt` output of Feature Extractor directly to Sampler's `prompt` to keep them in sync. Leave Sampler's `prompt` empty and it will use whatever was stored during extraction.
+### DITTO style transfer (recommended first approach)
+
+```
+SelVA Model Loader ─────────────────────────────────────────────► SelVA DITTO Optimizer ──► Save Audio
+                                                                         ▲
+SelVA Feature Extractor ──(features)────────────────────────────────────►│
+                          (prompt) ──────────────────────────────────────►│
+BJ reference_dir ───────────────────────────────────────────────────────►│
+```
+
+No training required. Each run optimizes x₀ independently for the current video and reference set.
+
+### Vocoder fine-tuning
+
+```
+SelVA Model Loader ──► SelVA BigVGAN Trainer ──► (checkpoint .pt)
+                              ▲
+BJ audio clips ──(data_dir)──►│
+
+SelVA Model Loader ──► SelVA BigVGAN Loader ──► SelVA Sampler ──► Save Audio
+                              ▲                       ▲
+                        checkpoint .pt         SelVA Feature Extractor
+```
+
+### LoRA training
+
+See [LORA_TRAINING.md](LORA_TRAINING.md).
 
 ---
 
@@ -127,8 +366,15 @@ The `auto` offload strategy picks `keep_in_vram` if ≥ 16 GB VRAM is available,
 
 ---
 
+## Style Transfer
+
+For adapting SelVA to a specific audio style (e.g. BJ / Bladee / Jersey Club), see [STYLE_TRANSFER.md](STYLE_TRANSFER.md).
+
+---
+
 ## Credits
 
 - [SelVA](https://github.com/jnwnlee/selva) by Jaehwan Lee et al. — TextSynchformer and SelVA training
 - [MMAudio](https://github.com/hkchengrex/MMAudio) by Feng et al. — MM-DiT audio generator and flow matching framework
 - [BigVGAN](https://github.com/NVIDIA/BigVGAN) by NVIDIA — neural vocoder for 16 kHz synthesis
+- [DITTO](https://arxiv.org/abs/2401.12179) by Novack et al. — inference-time diffusion optimization

STYLE_TRANSFER.md

@@ -0,0 +1,158 @@
+# Style Transfer for SelVA
+
+This document covers approaches for adapting SelVA's audio output to a specific timbral style using a small reference dataset (~50 clips). The context here is BJ / Bladee / Jersey Club style — sharp metallic transients, saturated harmonics, 808 sub bass, glassy high-frequency content — but the methods apply to any style target.
+
+---
+
+## Why standard fine-tuning is hard
+
+SelVA's generation quality depends on the DiT (generator) outputting latents that fall in the high-density region of the VAE decoder's training distribution. BJ's audio maps to a sparse, tail region of that space — the VAE roundtrip already shows ~10–15 dB elevated HF noise floor on BJ material. Any training that pushes the generator toward exact BJ encoder outputs is training toward an already-degraded target.
+
+**LoRA** makes this worse: it introduces "intruder dimensions" — new high-rank singular vectors absent from the pretrained weight spectrum — that push DiT outputs further off-manifold. This mechanism is LR- and scale-independent. Reducing LoRA scale does not fix the direction, only the magnitude. Empirically: spectral flatness degrades to ~0.21–0.26 (vs. baseline 0.013) at every scale from 0.0625 to 1.0.
+
+**Textual inversion** via the text conditioning path suffers from mean-pooling: SelVA's text features are pooled into a single global vector before injection into the DiT. The optimizer finds a spectral bias (noise/buzz) as the cheapest way to reduce reconstruction loss — not a semantic style shift.
+
+The approaches below are ordered by expected quality and ease of use.
+
+---
+
+## Tier 1 — DITTO (recommended first try)
+
+**Node: SelVA DITTO Optimizer**
+
+Inference-time noise optimization. Keeps all model weights frozen and only optimizes the initial noise latent x₀ using a style loss computed against the reference clips. Since the weights never change, there is zero risk of quality degradation — the model still generates from its original manifold, just from a better starting point.
+
+**Style loss:** mean spectrum + Gram matrix of mel spectrograms. The Gram matrix captures covariance between frequency bands (timbral texture) without requiring temporal alignment with the reference. Optimization runs entirely before the vocoder — BigVGAN is only called for the final output pass.
+
+**How it works:**
+
+For each video clip you want to process:
+1. Run SelVA Feature Extractor as usual.
+2. Instead of SelVA Sampler, connect to **SelVA DITTO Optimizer** with your BJ `reference_dir`.
+3. The node runs N optimization steps, each backpropagating through the last few ODE Euler steps to compute `∂loss/∂x₀`.
+4. After optimization, one final full-ODE pass generates the output audio from the refined x₀.
+
+```
+SelVA Model Loader ────────────────────────────────► SelVA DITTO Optimizer ──► audio
+                                                             ▲
+SelVA Feature Extractor ──(features)────────────────────────►│
+                          (prompt) ──────────────────────────►│
+BJ clips ───────────────────────────(reference_dir) ─────────►│
+```
+
+**Tuning guide:**
+
+| Parameter | Starting value | When to adjust |
+|---|---|---|
+| `n_opt_steps` | 50 | Increase to 100–200 if style shift is too subtle |
+| `opt_lr` | 0.1 | Lower to 0.05 if coherence breaks; raise to 0.3 for stronger shift |
+| `n_ode_steps` | 10 | Lower = faster optimization, less accurate gradient |
+| `n_grad_steps` | 5 | Number of ODE steps to differentiate through — must be ≤ n_ode_steps |
+| `style_weight` | 1.0 | Increase to 2–5 for stronger BJ character; watch for incoherence |
+
+**Memory:** Each opt step stores activations for `n_grad_steps` DiT forward passes with gradient checkpointing. At n_grad_steps=5, expect ~4–6 GB additional VRAM over baseline inference.
+
+**Time per video clip:** ~50 opt steps × (10 ODE steps × 2 passes for checkpointing) + 25 final steps ≈ 5–15 minutes depending on GPU.
+
+**Limitations:** DITTO with mel Gram matrix loss shifts timbral statistics but cannot precisely match the BJ transient sharpness — the Gram matrix is a texture descriptor, not a transient detector. See Tier 2 (vocoder fine-tuning) for that.
+
+---
+
+## Tier 2 — Vocoder Fine-tuning
+
+**Nodes: SelVA BigVGAN Trainer → SelVA BigVGAN Loader**
+
+The BigVGAN vocoder (mel → waveform) is the component most responsible for the final timbral character of the output. Fine-tuning only the vocoder keeps the DiT completely untouched — latents stay on-manifold, only the waveform rendering changes.
+
+### Why plain mel L1 loss fails
+
+BigVGAN was trained with `L_G = Σ[L_adv + 2·L_fm] + 45·L_mel`. The adversarial and feature-matching terms do the perceptual heavy lifting — they prevent the generator from averaging over high-variance harmonic content. Dropping them for a plain mel L1 loss is a loss-function topology problem: the model minimizes expected reconstruction error by averaging over harmonic uncertainty, eroding the saturated 3–8 kHz harmonics visible as "green smear" in spectrograms. This happens regardless of LR or step count.
+
+### `snake_alpha_only` mode (default, recommended)
+
+BigVGAN's AMP blocks use Snake/SnakeBeta activations: `y = x + (1/α)·sin²(α·x)` where α is a per-channel learnable scalar. Alpha parameters directly control the harmonic periodicity of each layer's output — they are the "harmonic tuning knobs" of the vocoder.
+
+With `train_mode=snake_alpha_only`, only the ~27K alpha parameters (0.024% of the 112M parameter model) are trained. The conv weights encoding waveform structure remain frozen. With this few trainable parameters the model physically cannot reshape the spectrum significantly regardless of loss function — no capacity for the green smear.
+
+**Loss in snake_alpha_only mode:** mel L1 + multi-resolution STFT L1 are still used but can only shift harmonic emphasis, not spectral shape.
+
+### `all_params` mode with discriminator
+
+For a stronger shift — or to use proper perceptual losses — run with `train_mode=all_params` and provide a `discriminator_path` (the `bigvgan_discriminator_optimizer.pt` from the BigVGAN pretrained release):
+
+1. The frozen pretrained MPD and MRD discriminators are loaded and used as fixed perceptual feature extractors.
+2. Loss becomes `2·L_fm(frozen_D) + 0.1·L_mel` — feature matching directly penalizes harmonic smearing through the discriminator's learned perceptual space.
+3. `lambda_l2sp` (default 1e-3) anchors all parameters to their pretrained values — prevents catastrophic drift on 50 clips.
+
+This is the highest-quality vocoder fine-tuning path but requires the discriminator checkpoint.
+
+### Workflow
+
+```
+SelVA Model Loader ──► SelVA BigVGAN Trainer ──► bigvgan_bj.pt
+                              ▲
+BJ audio clips ──(data_dir)──►│
+
+SelVA Model Loader ──► SelVA BigVGAN Loader ──► SelVA Sampler
+                              ▲                       ▲
+                      bigvgan_bj.pt           SelVA Feature Extractor
+```
+
+### Tuning guide
+
+| Parameter | Default | Notes |
+|---|---|---|
+| `train_mode` | snake_alpha_only | Safe default; use all_params only with discriminator_path |
+| `steps` | 2000 | 1000–2000 for snake_alpha_only; 3000–5000 for all_params |
+| `lr` | 1e-4 | For snake_alpha_only; lower to 1e-5 for all_params |
+| `lambda_l2sp` | 1e-3 | Increase to 1e-2 for all_params to limit drift |
+| `batch_size` | 4 | 4–8 for stable gradients |
+| `segment_seconds` | 1.0 | 1–2 s segments recommended |
+
+**Eval samples:** The trainer saves `.wav` and mel spectrogram `.png` files at baseline, each checkpoint, and final. Compare the spectrograms — saturation (red values in high-frequency bands) should increase relative to baseline.
+
+---
+
+## Tier 3 — DITTO + Vocoder (combined)
+
+Stack both:
+
+```
+SelVA Model Loader ──► SelVA BigVGAN Loader ──► SelVA DITTO Optimizer ──► audio
+                              ▲                          ▲
+                       bigvgan_bj.pt         SelVA Feature Extractor + reference_dir
+```
+
+The fine-tuned vocoder handles waveform rendering; DITTO shifts the latent trajectory. Each addresses a different aspect of style transfer.
+
+---
+
+## What doesn't work (and why)
+
+### Standard LoRA
+
+LoRA introduces "intruder dimensions" — high-rank singular vectors absent from the pretrained weight spectrum — at initialization. These push DiT outputs into decoder-hostile latent regions regardless of scale or LR. The failure is direction-based, not magnitude-based, so reducing LoRA scale does not fix it.
+
+PiSSA initialization (`init_lora_weights="pissa"`) and rsLoRA scaling (`use_rslora=True`) reduce intruder dimension formation by starting in the pretrained weight subspace. These are planned as future improvements.
+
+### Textual inversion
+
+SelVA mean-pools all 77 CLIP tokens into a single AdaLN bias vector. Every token contributes equally to a scalar offset; the optimizer finds spectral buzz as the minimum-cost way to reduce flow-matching reconstruction loss. More tokens make it worse.
+
+### Activation steering (global mean difference)
+
+The raw mean difference between BJ and empty conditions is not a clean style basis — it carries noise from the diversity of the training clips and the many attention blocks that have nothing to do with timbral character. Global injection (all blocks at any strength) kills the sound. Targeted layer injection (only the 3–6 blocks most predictive of BJ style) is theoretically sound but requires per-layer delta magnitude ranking to identify the right layers first.
+
+---
+
+## Reference dataset preparation
+
+Use the same audio clips for both DITTO and vocoder fine-tuning:
+
+- **Minimum:** 20–30 clips. DITTO works from 5+; vocoder benefits from 40+.
+- **Format:** `.wav` or `.flac` at native sample rate. The trainer resamples automatically.
+- **Length:** Any length ≥ 1 s. Longer is fine — the trainer segments internally.
+- **Quality:** Clean, full-mix BJ clips. Avoid heavily compressed or streaming-ripped files. Use HF Smoother if HF content sounds brittle after VAE roundtrip.
+- **Diversity:** Vary tempo, key, vocal density. 20 diverse clips > 50 copies of the same 8-bar loop.
+
+Normalize all clips to consistent loudness (e.g. -14 LUFS) before training. Inconsistent levels increase loss variance and slow convergence.

docs/audio_dataset_pipeline.md

@@ -0,0 +1,170 @@
+# Audio Dataset Pipeline for Generative Model Training
+
+Research notes on audio cleaning, augmentation, and quality metrics for LoRA fine-tuning of MMAudio/SelVA. Based on papers and tooling survey (April 2026).
+
+---
+
+## Core Principle
+
+Augmentation for generative models ≠ augmentation for classifiers.
+The goal is **not invariance** — it is expanding the training manifold so the model learns the distribution of a sound rather than memorizing a fixed set of waveforms.
+
+With 10 clips, velocity field collapse (arXiv:2410.23594) is mathematically expected: the flow-matching model memorizes the training trajectories instead of generalizing. More diverse data is the only real fix.
+
+---
+
+## Recommended Pipeline
+
+### Step 1 — Quality Screening
+
+```python
+# Clipping check
+clip_ratio = np.sum(np.abs(audio) >= 0.99) / len(audio)  # flag if > 0.1%
+
+# DC offset check + removal
+dc = np.mean(audio)
+audio -= dc
+
+# LUFS normalization to -14 LUFS (essential for training consistency)
+# pip install pyloudnorm
+import pyloudnorm as pyln
+meter = pyln.Meter(sr)
+loudness = meter.integrated_loudness(audio)
+audio = pyln.normalize.loudness(audio, loudness, -14.0)
+# Or via ffmpeg: ffmpeg -af loudnorm=I=-14:LRA=7:TP=-1
+
+# DNSMOS quality gate (discard if OVRL < 3.5 for training; < 2.5 is unusable)
+# from Microsoft DNS-Challenge repo
+```
+
+### Step 2 — Cleaning
+
+| Tool | Install | Use |
+|---|---|---|
+| **AudioSep** | `pip install audiosep` | Isolate target sound from background — most impactful tool |
+| **noisereduce** | `pip install noisereduce` | Light stationary/non-stationary denoising, preserves character |
+| **librosa** | `pip install librosa` | Silence trimming: `librosa.effects.trim(audio, top_db=30)` |
+| **torchaudio.transforms.Fade** | (torchaudio) | Prevent click artifacts at clip edges |
+| **DeepFilterNet** | `pip install deepfilternet` | Heavy denoising — good for speech, may alter tonal sounds |
+
+**AudioSep usage:**
+```python
+from audiosep import AudioSep
+model = AudioSep.from_pretrained("audio-agi/audiosep")
+# ~1.5 GB checkpoint, ~4 GB VRAM
+model.inference(audio_path, "a dog barking loudly", output_path)
+```
+
+### Step 3 — Waveform Augmentation (10 clips → 50–100)
+
+Apply stochastically per clip:
+
+| Transform | Params | Notes |
+|---|---|---|
+| **PitchShift** | ±1–3 semitones | 3 variants per clip. Limit to ±1 st for tonal/pitched sounds |
+| **ApplyImpulseResponse** | 5 different RIRs | 5 variants per clip — EchoThief (~150 free IRs) or pyroomacoustics |
+| **LoudnessNormalization** | ±2 dB random | Subtle level variation |
+| **SevenBandParametricEQ** | ±3 dB | Gentle spectral variation |
+| **TimeStretch** | 0.9–1.1× only | Do NOT use 2× to pad short clips — breaks video sync |
+
+```python
+# pip install audiomentations pedalboard pyroomacoustics
+import audiomentations as A
+
+augment = A.Compose([
+    A.PitchShift(min_semitones=-2, max_semitones=2, p=0.5),
+    A.ApplyImpulseResponse(ir_paths="path/to/irs/", p=0.5),
+    A.SevenBandParametricEQ(min_gain_db=-3, max_gain_db=3, p=0.3),
+    A.LoudnessNormalization(min_lufs=-16, max_lufs=-12, p=0.5),
+    A.TimeStretch(min_rate=0.9, max_rate=1.1, p=0.3),
+])
+audio_aug = augment(samples=audio, sample_rate=sr)
+```
+
+**RIR sources:**
+- EchoThief: ~150 free real-world IRs (churches, caves, parking garages)
+- pyroomacoustics: synthetic room simulation, fully controllable
+
+### Step 4 — Latent Augmentation (at training time)
+
+After VAE encoding:
+
+**Latent mixup** between same-category pairs:
+```python
+# Mix latents BEFORE flow-matching noise is added
+# Only mix clips from the same sound category — cross-category mixing produces garbage
+lam = torch.distributions.Beta(0.4, 0.4).sample()
+z_mix = lam * z1 + (1 - lam) * z2
+```
+With 10 clips: C(10,2) = 45 possible pairs → significant expansion without new recordings.
+
+**Small Gaussian noise:**
+```python
+z_noised = z + torch.randn_like(z) * 0.02 * z.std()
+```
+Prevents trivial memorization of exact latent coordinates.
+
+MusicLDM (arXiv:2308.01546) shows latent mixup > waveform mixup for generative quality.
+
+---
+
+## Transforms to AVOID for Generative Training
+
+| Transform | Why |
+|---|---|
+| ClippingDistortion, BitCrush, TanhDistortion, Mp3Compression | Model learns the artifact |
+| Reverse | Breaks temporal structure for video-to-audio task |
+| TimeMask (creating silence gaps) | Unnatural — model learns to produce silence |
+| TimeStretch > 1.3× | Phase vocoder artifacts become part of the target distribution |
+| Heavy background noise (< 15 dB SNR) | Model learns to reproduce the noise |
+
+---
+
+## Quality Metrics
+
+| Metric | Tool | Threshold |
+|---|---|---|
+| DNSMOS P.835 (SIG/BAK/OVRL) | Microsoft DNS-Challenge | OVRL > 3.5 for training |
+| LUFS | pyloudnorm | Normalize all clips to -14 LUFS |
+| WADA-SNR | (standalone) | No-reference SNR estimate |
+| Clipping ratio | NumPy | Flag if > 0.1% of samples at ±0.99 |
+
+---
+
+## Tool Reference
+
+| Tool | Install | Purpose |
+|---|---|---|
+| audiomentations | `pip install audiomentations` | Primary augmentation library |
+| pedalboard | `pip install pedalboard` | Higher quality pitch shift, IR convolution |
+| AudioSep | `pip install audiosep` | Source separation / isolation |
+| noisereduce | `pip install noisereduce` | Non-stationary denoising |
+| DeepFilterNet | `pip install deepfilternet` | Heavy denoising (speech-optimized) |
+| pyloudnorm | `pip install pyloudnorm` | LUFS normalization |
+| Silero VAD | `pip install silero-vad` | Voice/silence detection |
+| pyroomacoustics | `pip install pyroomacoustics` | Synthetic RIR generation |
+
+---
+
+## Integration with PrismAudio / SelVA
+
+No established ComfyUI audio preprocessing ecosystem as of early 2026. Build thin wrapper nodes around the tools above. PrismAudio already has all required patterns (subprocess isolation, AUDIO type transport).
+
+**Target node set:**
+- `SelVA Dataset Cleaner` — wraps noisereduce + LUFS normalization + trim + DNSMOS gate
+- `SelVA Dataset Augmenter` — wraps audiomentations Compose pipeline
+
+Steps 1–3 are preprocessing (run once before feature extraction).
+Step 4 (latent mixup) is a training loop modification — integrate into `selva_lora_trainer.py`.
+
+---
+
+## Key Papers
+
+| Paper | ArXiv | Finding |
+|---|---|---|
+| MusicLDM | 2308.01546 | Latent mixup > waveform mixup for generative quality |
+| EDMSound | 2311.08667 | Memorization documented — same failure mode as 10-clip training |
+| Synthio | 2410.02056 | Synthetic audio as augmentation data (ICLR 2025) |
+| HunyuanVideo-Foley | 2508.16930 | V2A data pipeline at scale (100K hrs) |
+| FM memorization | 2410.23594 | Velocity field collapse theory — proves early overfitting on small datasets |

docs/audiox_evaluation.md

@@ -0,0 +1,184 @@
+# AudioX vs SelVA — Evaluation
+
+AudioX (arXiv:2503.10522, ICLR 2026) is a unified multimodal audio generation model from HKUST.
+This document compares it against SelVA/MMAudio and assesses the cost of adding it to PrismAudio.
+
+---
+
+## Quick Decision Guide
+
+| Situation | Use |
+|---|---|
+| Video → realistic sound effects | **SelVA** — faster, purpose-built, MIT license |
+| Music generation from video or text | **AudioX** — SelVA cannot do this |
+| Audio inpainting / music continuation | **AudioX** — SelVA cannot do this |
+| LoRA fine-tuning on a custom sound | **SelVA** — full training infrastructure already exists |
+| Variable output duration | **AudioX** — SelVA is fixed at 8 s |
+| Inference speed matters | **SelVA** — 25 steps vs 250 (10× faster) |
+| Non-commercial research | Either |
+| Any commercial use | **SelVA only** — AudioX is CC-BY-NC-4.0 |
+
+---
+
+## Architecture
+
+| Dimension | SelVA (MMAudio) | AudioX-MAF |
+|---|---|---|
+| Core paradigm | Flow matching | Diffusion (k-diffusion / DPM++) |
+| Inference steps | 25 ODE steps (Euler) | 250 diffusion steps (DPM++ 3M SDE) |
+| Sample rate | 44.1 kHz (large) / 16 kHz (small) | 48 kHz (fixed) |
+| Generator | MM-DiT, velocity prediction | ContinuousMMDiTTransformer |
+| Video encoder | Synchformer | Synchformer (AudioX custom re-impl, same concept) |
+| VAE / codec | DAC (descript-audio-codec) | DAC + AudioCraft options |
+| Text encoder | T5-large | T5 (configurable small → XXL) |
+| Video-audio fusion | Cross-attention in MM-DiT | MAF: dual-projection (dim alignment + seq length alignment) |
+| Output duration | Fixed 8 s | Configurable via `sample_size` (default ~44 s at 48kHz) |
+| Training data | ~2 M samples (MMAudio paper) | 7 M samples (IF-caps dataset, curated) |
+| License | MIT | CC-BY-NC-4.0 |
+
+**MAF (Multimodal Adaptive Fusion):** AudioX's key architectural contribution. Instead of directly
+concatenating multimodal tokens into the DiT's cross-attention, MAF projects each modality to
+match the latent's sequence length via a dedicated linear + transposed-conv stack, then applies
+`MMDitSingleBlock` layers for cross-modal fusion. The paper reports this improves cross-modal
+alignment particularly for video-to-audio tasks.
+
+**Flow matching vs diffusion:** Flow matching (SelVA) trains a single velocity field to move
+directly from noise to data along a straight trajectory — this is why 25 steps suffice. Standard
+diffusion (AudioX) approximates a longer stochastic path, requiring 250 steps for quality output.
+This is not a quality difference per se; flow matching is simply more efficient.
+
+---
+
+## Capabilities
+
+| Task | SelVA | AudioX |
+|---|---|---|
+| Video → sound effects | ✓ (primary use case) | ✓ |
+| Text → sound effects | Partial (T5 conditions quality but not primary) | ✓ (strong benchmark scores) |
+| Video → music | ✗ | ✓ |
+| Text → music | ✗ | ✓ |
+| Audio inpainting | ✗ | ✓ (mask_args parameter) |
+| Music continuation | ✗ | ✓ (init_audio parameter) |
+| Variable output duration | ✗ (fixed 8 s) | ✓ |
+| Multiple input modalities simultaneously | Partial | ✓ (text + video + audio at once) |
+
+AudioX benchmarks claim superior results on text-to-audio (AudioCaps) and text-to-music
+(MusicCaps) vs prior models. Video-to-audio comparison against MMAudio specifically is not
+prominently featured in the paper. Perceptual evaluation confirms this: AudioX does not sound
+noticeably better than SelVA on video-to-audio tasks. AudioX's advantage is **breadth**
+(music, inpainting, variable duration), not raw video-to-audio quality.
+
+---
+
+## Integration Cost
+
+Adding AudioX inference-only nodes to PrismAudio would require:
+
+### New nodes (3 files)
+
+```
+nodes/
+  audiox_model_loader.py    AUDIOX_MODEL loader — get_pretrained_model("HKUSTAudio/AudioX-MAF")
+  audiox_sampler.py         wraps generate_diffusion_cond(), inputs: model + text + video + audio
+  audiox_feature_extractor.py  optional — pre-extract Synchformer sync features (caching)
+```
+
+### Installation
+
+```bash
+pip install git+https://github.com/ZeyueT/AudioX.git
+```
+
+New dependencies not currently in PrismAudio:
+- `pytorch-lightning==2.4.0`
+- `k-diffusion==0.1.1`
+- `v-diffusion-pytorch==0.0.2`
+- `descript-audio-codec==1.0.0` (already used by SelVA — no conflict, same package)
+- `gradio==4.44.1` (optional — only for the upstream Gradio UI)
+
+Model weights: `HKUSTAudio/AudioX-MAF` on HuggingFace (~several GB).
+
+### Inference API surface
+
+```python
+from audiox import get_pretrained_model
+from audiox.inference.generation import generate_diffusion_cond
+
+model, config = get_pretrained_model("HKUSTAudio/AudioX-MAF")
+
+output = generate_diffusion_cond(
+    model,
+    steps=250,
+    cfg_scale=6.0,
+    conditioning={
+        "text_prompt": "a dog barking",
+        "video_prompt": {"video": frames_tensor, "sync_features": sync_feat},
+        "seconds_total": 8.0,
+    },
+    sample_size=384000,   # 8 s at 48kHz
+    sample_rate=48000,
+    device="cuda",
+)
+# output: torch.Tensor (batch, channels, num_samples) float32 [-1, 1]
+```
+
+---
+
+## LoRA Training
+
+Adding AudioX LoRA training to PrismAudio is **significantly harder** than the SelVA trainer:
+
+| Aspect | SelVA LoRA | AudioX LoRA |
+|---|---|---|
+| Loss function | Single MSE velocity loss | Diffusion loss over 250-step schedule |
+| Training steps needed | ~2000 steps practical | Unknown — likely much more |
+| Step cost | Fast (1 velocity prediction) | Slow (full diffusion forward pass per step) |
+| Existing infrastructure | Full trainer + scheduler + experiments | Nothing — would need to build from scratch |
+| Noise schedule | Trivial (linear interpolation) | Cosine alpha-sigma schedule |
+| Prior art for LoRA | LoRA on flow matching well-studied | Less explored; closer to Stable Diffusion LoRA |
+
+**Conclusion:** AudioX LoRA training is feasible (it would follow SD-style LoRA with the DPM++
+noise schedule) but would be a substantial new project. Not worth building until inference nodes
+are stable and there is a clear use case that SelVA cannot serve.
+
+---
+
+## License
+
+AudioX weights are released under **CC-BY-NC-4.0** (Creative Commons Non-Commercial).
+
+- Free for personal use, research, and non-commercial projects
+- **Cannot be used in commercial products or services** without a separate agreement
+- Attribution required
+- SelVA/MMAudio: MIT (no restrictions)
+
+If PrismAudio is ever distributed as part of a commercial tool, AudioX nodes must be clearly
+opt-in with a license warning, or excluded entirely.
+
+---
+
+## Recommendation
+
+**Short term:** AudioX is not a replacement for SelVA for the current use case (video → custom
+sound effects with LoRA fine-tuning). SelVA is faster, has full training infrastructure, and
+is MIT licensed.
+
+**When AudioX becomes worth integrating:**
+- If you need to generate background music synchronized to video
+- If you need audio inpainting (fill a gap in an existing audio track)
+- If you need text-to-audio generation without a video input
+- After verifying the CC-BY-NC-4.0 license is acceptable for your use
+
+**Estimated integration effort for inference nodes only:** 2–3 days of work (3 new node files,
+dependency management, testing). No changes to existing SelVA nodes required — they would
+coexist in the same package.
+
+---
+
+## References
+
+- Paper: arXiv:2503.10522 — *AudioX: Diffusion Transformer for Anything-to-Audio Generation*
+- GitHub: https://github.com/ZeyueT/AudioX
+- Model weights: https://huggingface.co/HKUSTAudio/AudioX-MAF
+- Demo: https://huggingface.co/spaces/Zeyue7/AudioX
+- Project page: https://zeyuet.github.io/AudioX/

experiments/alpha_scale_sweep.json

@@ -0,0 +1,77 @@
+{
+  "name": "alpha_scale_sweep",
+  "description": "Fix LoRA noise contamination (flatness 0.013→0.094 at alpha=rank). Root cause: alpha=rank (scale=1.0) at high rank drowns base model priors. Testing dramatically lower alpha to nudge rather than overwrite. All runs at lr=3e-4 (best stable LR from r128_sweet_spot).",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/alpha_scale_sweep",
+  "base": {
+    "steps": 6000,
+    "lr": 3e-4,
+    "batch_size": 16,
+    "warmup_steps": 200,
+    "grad_accum": 1,
+    "save_every": 2000,
+    "seed": 42,
+    "target": "attn.qkv",
+    "timestep_mode": "uniform",
+    "logit_normal_sigma": 1.0,
+    "curriculum_switch": 0.6,
+    "lora_dropout": 0.0,
+    "lora_plus_ratio": 1.0,
+    "lr_schedule": "constant"
+  },
+  "experiments": [
+
+    {
+      "id": "g1_r16_alpha4",
+      "group": "conservative",
+      "description": "Back to basics: rank=16 alpha=4 (scale=0.25). Small adapter, gentle scale — cleanest possible LoRA signal.",
+      "rank": 16,
+      "alpha": 4.0
+    },
+    {
+      "id": "g1_r16_alpha16",
+      "group": "conservative",
+      "description": "rank=16 alpha=16 (scale=1.0) — the original default. Reference point: is the noise issue rank-specific or universal?",
+      "rank": 16,
+      "alpha": 16.0
+    },
+
+    {
+      "id": "g2_r32_alpha8",
+      "group": "mid",
+      "description": "rank=32 alpha=8 (scale=0.25). More capacity than r16 but still gentle scale.",
+      "rank": 32,
+      "alpha": 8.0
+    },
+    {
+      "id": "g2_r32_alpha32",
+      "group": "mid",
+      "description": "rank=32 alpha=32 (scale=1.0). Same rank, full scale — isolates whether scale or rank is causing noise.",
+      "rank": 32,
+      "alpha": 32.0
+    },
+
+    {
+      "id": "g3_r128_alpha8",
+      "group": "high_rank_low_alpha",
+      "description": "rank=128 alpha=8 (scale=0.0625). High capacity, very gentle contribution — can r128 stay clean at low alpha?",
+      "rank": 128,
+      "alpha": 8.0
+    },
+    {
+      "id": "g3_r128_alpha16",
+      "group": "high_rank_low_alpha",
+      "description": "rank=128 alpha=16 (scale=0.125). Slightly more signal than alpha=8.",
+      "rank": 128,
+      "alpha": 16.0
+    },
+    {
+      "id": "g3_r128_alpha32",
+      "group": "high_rank_low_alpha",
+      "description": "rank=128 alpha=32 (scale=0.25). Same scale as r16_alpha4 and r32_alpha8 — comparable contribution across ranks.",
+      "rank": 128,
+      "alpha": 32.0
+    }
+
+  ]
+}

experiments/eval_r128_candidates.json

@@ -0,0 +1,39 @@
+{
+  "name": "eval_r128_candidates",
+  "description": "Top candidates from r128_sweet_spot. Comparing the two lowest-loss runs, the stable lr=3e-4, and the curriculum run that hit 0.161 before regressing. Baseline included as perceptual reference.",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_dir": "/media/unraid/davinci/Selva/BJ/evals/r128_candidates",
+  "steps": 25,
+  "seed": 42,
+  "adapters": [
+    {
+      "id": "baseline",
+      "description": "No LoRA — base model output for perceptual reference"
+    },
+    {
+      "id": "lr_5e4_r128",
+      "description": "Best loss overall (0.137), still descending at step 10k",
+      "path": "/media/unraid/davinci/Selva/BJ/experiment/r128_sweet_spot/g1_r128_lr_5e4/adapter_final.pt"
+    },
+    {
+      "id": "lr_3e4_r256",
+      "description": "Tied with lr_5e4 at 0.139, higher rank — does extra capacity help perceptually?",
+      "path": "/media/unraid/davinci/Selva/BJ/experiment/r128_sweet_spot/g4_r256_lr_3e4/adapter_final.pt"
+    },
+    {
+      "id": "lr_3e4_r128",
+      "description": "Stable plateau from step 4k to 10k (0.221) — visually confirmed clean spectrograms",
+      "path": "/media/unraid/davinci/Selva/BJ/experiment/r128_sweet_spot/g1_r128_lr_3e4/adapter_final.pt"
+    },
+    {
+      "id": "curriculum_lr_3e4",
+      "description": "Best min loss of all (0.161 at step 6k), regressed to 0.193 after curriculum switch — curious if the early checkpoint sounds better",
+      "path": "/media/unraid/davinci/Selva/BJ/experiment/r128_sweet_spot/g2_r128_lr_3e4_curriculum/adapter_final.pt"
+    },
+    {
+      "id": "curriculum_lr_3e4_step6000",
+      "description": "Same run at its actual best step (before regression) — compare against adapter_final to hear the regression",
+      "path": "/media/unraid/davinci/Selva/BJ/experiment/r128_sweet_spot/g2_r128_lr_3e4_curriculum/adapter_step06000.pt"
+    }
+  ]
+}

experiments/r128_sweet_spot.json

@@ -0,0 +1,103 @@
+{
+  "name": "r128_sweet_spot",
+  "description": "Find the noise-free sweet spot on rank 128. LoRA+ ratio=16 caused noise — testing higher base LR without LoRA+ as a cleaner alternative. Target loss range 0.25–0.35. Also probing rank 256 since 102GB VRAM allows it.",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/r128_sweet_spot",
+  "base": {
+    "steps": 10000,
+    "rank": 128,
+    "alpha": 0.0,
+    "lr": 1e-4,
+    "batch_size": 16,
+    "warmup_steps": 200,
+    "grad_accum": 1,
+    "save_every": 2000,
+    "seed": 42,
+    "target": "attn.qkv",
+    "timestep_mode": "uniform",
+    "logit_normal_sigma": 1.0,
+    "curriculum_switch": 0.6,
+    "lora_dropout": 0.0,
+    "lora_plus_ratio": 1.0
+  },
+  "experiments": [
+
+    {
+      "id": "g1_r128_lr_2e4",
+      "group": "lr",
+      "description": "LR=2e-4. Conservative 2× step up from baseline — noise-free descent toward sweet spot.",
+      "lr": 2e-4
+    },
+    {
+      "id": "g1_r128_lr_3e4",
+      "group": "lr",
+      "description": "LR=3e-4. 3× baseline — landed at 0.41 on r64, should reach 0.25–0.35 on r128.",
+      "lr": 3e-4
+    },
+    {
+      "id": "g1_r128_lr_5e4",
+      "group": "lr",
+      "description": "LR=5e-4. Aggressive but no LoRA+ B-matrix asymmetry — cleaner noise profile.",
+      "lr": 5e-4
+    },
+
+    {
+      "id": "g2_r128_curriculum",
+      "group": "curriculum",
+      "description": "Curriculum only at baseline LR. Clean slow descent — reference for what curriculum contributes alone.",
+      "timestep_mode": "curriculum"
+    },
+    {
+      "id": "g2_r128_lr_3e4_curriculum",
+      "group": "curriculum",
+      "description": "LR=3e-4 + curriculum. Speed of higher LR with coverage of curriculum — no LoRA+.",
+      "lr": 3e-4,
+      "timestep_mode": "curriculum"
+    },
+    {
+      "id": "g2_r128_lr_3e4_curriculum_dropout",
+      "group": "curriculum",
+      "description": "LR=3e-4 + curriculum + dropout=0.05. Full controlled stack without LoRA+.",
+      "lr": 3e-4,
+      "timestep_mode": "curriculum",
+      "lora_dropout": 0.05
+    },
+
+    {
+      "id": "g3_r128_lora_plus_4",
+      "group": "lora_plus",
+      "description": "LoRA+ ratio=4 (lr_B=4e-4). Much more conservative than ratio=16 — tests if noise came from ratio not the technique.",
+      "lora_plus_ratio": 4.0
+    },
+
+    {
+      "id": "g4_r256_baseline",
+      "group": "rank256",
+      "description": "Rank 256 at baseline LR. 102GB VRAM makes this viable — does more capacity keep helping?",
+      "rank": 256
+    },
+    {
+      "id": "g4_r256_lr_3e4",
+      "group": "rank256",
+      "description": "Rank 256 + LR=3e-4. Best rank + best LR candidate combined.",
+      "rank": 256,
+      "lr": 3e-4
+    },
+
+    {
+      "id": "g5_r128_lr_2e4_cosine",
+      "group": "cosine",
+      "description": "LR=2e-4 + cosine decay. Fixes the oscillation observed at step 6000–8000 by decaying LR to ~0 instead of staying flat.",
+      "lr": 2e-4,
+      "lr_schedule": "cosine"
+    },
+    {
+      "id": "g5_r128_lr_3e4_cosine",
+      "group": "cosine",
+      "description": "LR=3e-4 + cosine decay. Higher LR with decay — should reach lower loss faster then lock in.",
+      "lr": 3e-4,
+      "lr_schedule": "cosine"
+    }
+
+  ]
+}

experiments/r64_overnight.json

@@ -0,0 +1,130 @@
+{
+  "name": "r64_overnight",
+  "description": "Focused rank-64 overnight sweep. All experiments use rank 64 as base — confirmed best from tier1_thorough early results. 8000 steps to reach convergence (none converged at 4000).",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/r64_overnight",
+  "base": {
+    "steps": 8000,
+    "rank": 64,
+    "alpha": 0.0,
+    "lr": 1e-4,
+    "batch_size": 16,
+    "warmup_steps": 200,
+    "grad_accum": 1,
+    "save_every": 2000,
+    "seed": 42,
+    "target": "attn.qkv",
+    "timestep_mode": "uniform",
+    "logit_normal_sigma": 1.0,
+    "curriculum_switch": 0.6,
+    "lora_dropout": 0.0,
+    "lora_plus_ratio": 1.0
+  },
+  "experiments": [
+
+    {
+      "id": "g1_r64_baseline",
+      "group": "rank",
+      "description": "Rank 64 baseline — clean reference at 8000 steps."
+    },
+    {
+      "id": "g1_r128_baseline",
+      "group": "rank",
+      "description": "Rank 128 — 102GB VRAM makes this free. Does doubling rank from 64 help further?",
+      "rank": 128
+    },
+
+    {
+      "id": "g2_r64_alpha_32",
+      "group": "alpha",
+      "description": "Rank 64 alpha=32 (scale=0.5). Reduces intruder singular dimensions (arXiv:2410.21228).",
+      "alpha": 32.0
+    },
+    {
+      "id": "g2_r64_alpha_16",
+      "group": "alpha",
+      "description": "Rank 64 alpha=16 (scale=0.25). More aggressive scale reduction — may over-constrain.",
+      "alpha": 16.0
+    },
+
+    {
+      "id": "g3_r64_lora_plus",
+      "group": "regularisation",
+      "description": "LoRA+ ratio=16. lr_B = 16 × lr_A. Faster convergence at constant step budget.",
+      "lora_plus_ratio": 16.0
+    },
+    {
+      "id": "g3_r64_dropout_0.05",
+      "group": "regularisation",
+      "description": "Dropout=0.05. Light sparsity regularisation on LoRA path.",
+      "lora_dropout": 0.05
+    },
+    {
+      "id": "g3_r64_dropout_0.1",
+      "group": "regularisation",
+      "description": "Dropout=0.1. Stronger regularisation — tests if 49 clips needs heavier constraint.",
+      "lora_dropout": 0.1
+    },
+    {
+      "id": "g3_r64_curriculum",
+      "group": "regularisation",
+      "description": "Curriculum sampling: logit_normal for steps 1-4800, then uniform (arXiv:2603.12517).",
+      "timestep_mode": "curriculum"
+    },
+
+    {
+      "id": "g4_r64_lr_low",
+      "group": "lr",
+      "description": "LR=3e-5. 3× lower — checks if 1e-4 is overshooting at rank 64.",
+      "lr": 3e-5
+    },
+    {
+      "id": "g4_r64_lr_high",
+      "group": "lr",
+      "description": "LR=3e-4. 3× higher — may converge faster but risk instability.",
+      "lr": 3e-4
+    },
+
+    {
+      "id": "g5_r64_target_full",
+      "group": "target",
+      "description": "Rank 64 targeting attn.qkv + linear1 (FFN projections). Doubles LoRA coverage.",
+      "target": "attn.qkv linear1"
+    },
+    {
+      "id": "g5_r128_target_full",
+      "group": "target",
+      "description": "Rank 128 + full target. Maximum possible coverage with available VRAM.",
+      "rank": 128,
+      "target": "attn.qkv linear1"
+    },
+
+    {
+      "id": "g6_r64_full_tier1",
+      "group": "combined",
+      "description": "All Tier 1 at rank 64: LoRA+ 16 + dropout 0.05 + curriculum. Full stack at 8000 steps.",
+      "lora_plus_ratio": 16.0,
+      "lora_dropout": 0.05,
+      "timestep_mode": "curriculum"
+    },
+    {
+      "id": "g6_r64_alpha32_full",
+      "group": "combined",
+      "description": "Rank 64 alpha=32 + all Tier 1. Best alpha scaling + best regularisation stack.",
+      "alpha": 32.0,
+      "lora_plus_ratio": 16.0,
+      "lora_dropout": 0.05,
+      "timestep_mode": "curriculum"
+    },
+    {
+      "id": "g6_r128_full_tier1",
+      "group": "combined",
+      "description": "Rank 128 + all Tier 1. Tests if more capacity + regularisation beats rank 64 full.",
+      "rank": 128,
+      "lora_plus_ratio": 16.0,
+      "lora_dropout": 0.05,
+      "timestep_mode": "curriculum"
+    }
+
+  ]
+}

experiments/ti_sweep_1.json

@@ -0,0 +1,52 @@
+{
+  "name": "ti_sweep_1",
+  "description": "First TI sweep. n4_baseline (suffix, batch=16, lr=1e-3) completed — buzz artifact diagnosed as token norm drifting to 3.2x outside CLIP manifold. All new experiments use norm clamping (auto from dataset) + corrected lr/batch.",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/ti_sweep_1",
+  "base": {
+    "steps": 3000,
+    "batch_size": 4,
+    "warmup_steps": 100,
+    "save_every": 1000,
+    "seed": 42,
+    "init_text": "",
+    "lr": 2e-4,
+    "n_tokens": 4,
+    "inject_mode": "suffix"
+  },
+  "experiments": [
+
+    {
+      "id": "n4_baseline",
+      "group": "reference",
+      "description": "COMPLETED (old code, no norm clamp). batch=16, lr=1e-3. Token norm drifted to 3.2 → buzz artifact. Kept for loss curve comparison only."
+    },
+
+    {
+      "id": "n4_clamped",
+      "group": "norm_clamp",
+      "description": "Same as baseline but with norm clamping enabled. Primary diagnostic: does clamping alone fix the buzz? lr=2e-4, batch=4, suffix."
+    },
+    {
+      "id": "n4_prefix_clamped",
+      "group": "norm_clamp",
+      "description": "Prefix injection + norm clamping. Best of both: high-attention positions, tokens stay on CLIP manifold.",
+      "inject_mode": "prefix"
+    },
+    {
+      "id": "n8_prefix_clamped",
+      "group": "norm_clamp",
+      "description": "8 tokens, prefix, clamped. More capacity without the artifact.",
+      "n_tokens": 8,
+      "inject_mode": "prefix"
+    },
+    {
+      "id": "n4_prefix_warm_clamped",
+      "group": "norm_clamp",
+      "description": "4 tokens, prefix, warm init from 'mechanical impact sound design', clamped. Should converge fastest — starts in-manifold, stays in-manifold.",
+      "inject_mode": "prefix",
+      "init_text": "mechanical impact sound design"
+    }
+
+  ]
+}

experiments/tier1_sweep.json

@@ -0,0 +1,61 @@
+{
+  "name": "tier1_sweep",
+  "description": "Ablation of Tier 1 improvements: LoRA+, dropout, curriculum sampling. Baseline = uniform, no regularisation.",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_root": "lora_sweeps/tier1_sweep",
+  "base": {
+    "steps": 4000,
+    "rank": 16,
+    "alpha": 0.0,
+    "lr": 1e-4,
+    "batch_size": 16,
+    "warmup_steps": 100,
+    "grad_accum": 1,
+    "save_every": 500,
+    "seed": 42,
+    "target": "attn.qkv",
+    "timestep_mode": "uniform",
+    "logit_normal_sigma": 1.0,
+    "curriculum_switch": 0.6,
+    "lora_dropout": 0.0,
+    "lora_plus_ratio": 1.0
+  },
+  "experiments": [
+    {
+      "id": "baseline",
+      "description": "Standard LoRA — no Tier 1 changes. Reference point."
+    },
+    {
+      "id": "lora_plus_16",
+      "description": "LoRA+ only: lr_B = 16 * lr_A. Should converge faster in early steps.",
+      "lora_plus_ratio": 16.0
+    },
+    {
+      "id": "dropout_0.05",
+      "description": "LoRA dropout 0.05 only. Light regularisation for 49-clip dataset.",
+      "lora_dropout": 0.05
+    },
+    {
+      "id": "dropout_0.1",
+      "description": "LoRA dropout 0.1 only. Stronger regularisation — may prevent overfitting past step 2000.",
+      "lora_dropout": 0.1
+    },
+    {
+      "id": "curriculum",
+      "description": "Curriculum sampling only: logit_normal for steps 1-2400, then uniform. Should improve convergence vs pure uniform.",
+      "timestep_mode": "curriculum"
+    },
+    {
+      "id": "full_tier1",
+      "description": "All Tier 1 combined: LoRA+ + dropout 0.05 + curriculum.",
+      "lora_plus_ratio": 16.0,
+      "lora_dropout": 0.05,
+      "timestep_mode": "curriculum"
+    },
+    {
+      "id": "rank_64",
+      "description": "Rank 64 baseline — MMAudio LoRA guide default. More expressive adapter for 49-clip dataset.",
+      "rank": 64
+    }
+  ]
+}

experiments/tier1_thorough.json

@@ -0,0 +1,144 @@
+{
+  "name": "tier1_thorough",
+  "description": "Full overnight Tier 1 ablation on 49-clip BJ dataset. 4 groups: rank, alpha, regularisation, and best combinations. ~10-12h depending on GPU.",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/tier1_thorough",
+  "base": {
+    "steps": 4000,
+    "rank": 16,
+    "alpha": 0.0,
+    "lr": 1e-4,
+    "batch_size": 16,
+    "warmup_steps": 100,
+    "grad_accum": 1,
+    "save_every": 1000,
+    "seed": 42,
+    "target": "attn.qkv",
+    "timestep_mode": "uniform",
+    "logit_normal_sigma": 1.0,
+    "curriculum_switch": 0.6,
+    "lora_dropout": 0.0,
+    "lora_plus_ratio": 1.0
+  },
+  "experiments": [
+    {
+      "id": "g1_rank_16",
+      "group": "rank",
+      "description": "Rank 16 baseline — reference point for all groups."
+    },
+    {
+      "id": "g1_rank_32",
+      "group": "rank",
+      "description": "Rank 32 — midpoint. Does doubling rank improve quality without overfitting?",
+      "rank": 32
+    },
+    {
+      "id": "g1_rank_64",
+      "group": "rank",
+      "description": "Rank 64 — MMAudio LoRA guide default. Maximum expressiveness at 49 clips.",
+      "rank": 64
+    },
+    {
+      "id": "g2_alpha_half_r16",
+      "group": "alpha",
+      "description": "Alpha=8 with rank 16 (scale=0.5). Reduces intruder singular dimensions (arXiv:2410.21228).",
+      "alpha": 8.0
+    },
+    {
+      "id": "g2_alpha_half_r64",
+      "group": "alpha",
+      "description": "Alpha=32 with rank 64 (scale=0.5). Best-practice scaling for high-rank adapters.",
+      "rank": 64,
+      "alpha": 32.0
+    },
+    {
+      "id": "g3_lora_plus_4",
+      "group": "regularisation",
+      "description": "LoRA+ ratio=4 — conservative asymmetric LR. Lower bound for the technique.",
+      "lora_plus_ratio": 4.0
+    },
+    {
+      "id": "g3_lora_plus_16",
+      "group": "regularisation",
+      "description": "LoRA+ ratio=16 — standard from FLUX LoRA literature. Faster early convergence.",
+      "lora_plus_ratio": 16.0
+    },
+    {
+      "id": "g3_dropout_0.05",
+      "group": "regularisation",
+      "description": "LoRA dropout 0.05 only. Light sparsity regularisation (arXiv:2404.09610).",
+      "lora_dropout": 0.05
+    },
+    {
+      "id": "g3_dropout_0.1",
+      "group": "regularisation",
+      "description": "LoRA dropout 0.1 only. Stronger regularisation — may prevent overfitting past step 2000.",
+      "lora_dropout": 0.1
+    },
+    {
+      "id": "g3_curriculum",
+      "group": "regularisation",
+      "description": "Curriculum sampling only: logit_normal steps 1-2400, then uniform (arXiv:2603.12517).",
+      "timestep_mode": "curriculum"
+    },
+    {
+      "id": "g4_full_r16",
+      "group": "combined",
+      "description": "All Tier 1 at rank 16: LoRA+ 16 + dropout 0.05 + curriculum.",
+      "lora_plus_ratio": 16.0,
+      "lora_dropout": 0.05,
+      "timestep_mode": "curriculum"
+    },
+    {
+      "id": "g4_full_r64",
+      "group": "combined",
+      "description": "All Tier 1 at rank 64 + alpha=32. Best expressiveness + best regularisation.",
+      "rank": 64,
+      "alpha": 32.0,
+      "lora_plus_ratio": 16.0,
+      "lora_dropout": 0.05,
+      "timestep_mode": "curriculum"
+    },
+
+    {
+      "id": "g5_lr_low",
+      "group": "lr",
+      "description": "LR=3e-5 — 3× lower than baseline. Tests if 1e-4 is overshooting.",
+      "lr": 3e-5
+    },
+    {
+      "id": "g5_lr_high",
+      "group": "lr",
+      "description": "LR=3e-4 — 3× higher than baseline. Tests if 1e-4 is too conservative.",
+      "lr": 3e-4
+    },
+
+    {
+      "id": "g6_target_full_r16",
+      "group": "target",
+      "description": "Rank 16 targeting attn.qkv + linear1 (FFN projections). Doubles LoRA coverage.",
+      "target": "attn.qkv linear1"
+    },
+    {
+      "id": "g6_target_full_r64",
+      "group": "target",
+      "description": "Rank 64 + alpha=32 targeting attn.qkv + linear1. Maximum coverage + expressiveness.",
+      "rank": 64,
+      "alpha": 32.0,
+      "target": "attn.qkv linear1"
+    },
+
+    {
+      "id": "g4_full_r64_6k",
+      "group": "combined",
+      "description": "All Tier 1 at rank 64 + alpha=32, extended to 6000 steps. Checks if convergence is done at 4000.",
+      "rank": 64,
+      "alpha": 32.0,
+      "lora_plus_ratio": 16.0,
+      "lora_dropout": 0.05,
+      "timestep_mode": "curriculum",
+      "steps": 6000,
+      "save_every": 1000
+    }
+  ]
+}

nodes/__init__.py

@@ -7,6 +7,21 @@ _NODES = {
     "SelvaSampler":          (".selva_sampler",           "SelvaSampler",          "SelVA Sampler"),
     "SelvaLoraLoader":       (".selva_lora_loader",       "SelvaLoraLoader",       "SelVA LoRA Loader"),
     "SelvaLoraTrainer":      (".selva_lora_trainer",      "SelvaLoraTrainer",      "SelVA LoRA Trainer"),
+    "SelvaLoraScheduler":    (".selva_lora_scheduler",    "SelvaLoraScheduler",    "SelVA LoRA Scheduler"),
+    "SelvaDatasetBrowser":   (".selva_dataset_browser",   "SelvaDatasetBrowser",   "SelVA Dataset Browser"),
+    "SelvaSkipExperiment":   (".selva_skip_experiment",   "SelvaSkipExperiment",   "SelVA Skip Experiment"),
+    "SelvaLoraEvaluator":    (".selva_lora_evaluator",    "SelvaLoraEvaluator",    "SelVA LoRA Evaluator"),
+    "SelvaVaeRoundtrip":     (".selva_vae_roundtrip",     "SelvaVaeRoundtrip",     "SelVA VAE Roundtrip"),
+    "SelvaHfSmoother":       (".selva_audio_preprocessors", "SelvaHfSmoother",       "SelVA HF Smoother"),
+    "SelvaSpectralMatcher":  (".selva_audio_preprocessors", "SelvaSpectralMatcher",  "SelVA Spectral Matcher"),
+    "SelvaTextualInversionTrainer": (".selva_textual_inversion_trainer", "SelvaTextualInversionTrainer", "SelVA Textual Inversion Trainer"),
+    "SelvaTextualInversionLoader":  (".selva_textual_inversion_loader",  "SelvaTextualInversionLoader",  "SelVA Textual Inversion Loader"),
+    "SelvaTiScheduler":                      (".selva_ti_scheduler",                      "SelvaTiScheduler",                      "SelVA TI Scheduler"),
+    "SelvaActivationSteeringExtractor":      (".selva_activation_steering_extractor",      "SelvaActivationSteeringExtractor",      "SelVA Activation Steering Extractor"),
+    "SelvaActivationSteeringLoader":         (".selva_activation_steering_loader",         "SelvaActivationSteeringLoader",         "SelVA Activation Steering Loader"),
+    "SelvaBigvganTrainer":                   (".selva_bigvgan_trainer",                    "SelvaBigvganTrainer",                   "SelVA BigVGAN Trainer"),
+    "SelvaBigvganLoader":                    (".selva_bigvgan_loader",                     "SelvaBigvganLoader",                    "SelVA BigVGAN Loader"),
+    "SelvaDittoOptimizer":                   (".selva_ditto_optimizer",                    "SelvaDittoOptimizer",                   "SelVA DITTO Optimizer"),
 }
 
 for key, (module_path, class_name, display_name) in _NODES.items():

nodes/selva_activation_steering_extractor.py

@@ -0,0 +1,201 @@
+"""SelVA Activation Steering Extractor.
+
+Computes per-block steering vectors by running the frozen generator on the
+training dataset and recording how BJ's conditioning shifts the DiT hidden
+states vs. empty/unconditional conditioning.
+
+For each block i:
+    steering[i] = mean(latent_hidden | BJ conditions)
+                - mean(latent_hidden | empty conditions)
+
+The resulting vectors are injected at inference time (via SelVA Sampler's
+steering_strength input) to nudge the denoising trajectory toward BJ's
+activation patterns without modifying any model weights.
+"""
+
+import random
+from pathlib import Path
+
+import torch
+import comfy.utils
+import folder_paths
+
+from .utils import SELVA_CATEGORY, get_device, soft_empty_cache
+from .selva_lora_trainer import _prepare_dataset
+
+
+def _collect_activations(generator, conditions, latent, t_tensor):
+    """Run one predict_flow call, collecting latent hidden states per block.
+
+    Returns a list of [seq, hidden_dim] float32 CPU tensors,
+    one per block (joint_blocks first, then fused_blocks).
+    """
+    activations = []
+
+    def make_hook(is_joint):
+        def hook(module, input, output):
+            h = output[0] if is_joint else output
+            activations.append(h.detach().float().mean(0).cpu())  # [seq, hidden]
+        return hook
+
+    handles = []
+    for block in generator.joint_blocks:
+        handles.append(block.register_forward_hook(make_hook(is_joint=True)))
+    for block in generator.fused_blocks:
+        handles.append(block.register_forward_hook(make_hook(is_joint=False)))
+
+    try:
+        with torch.no_grad():
+            generator.predict_flow(latent, t_tensor, conditions)
+    finally:
+        for h in handles:
+            h.remove()
+
+    return activations  # list of n_blocks tensors [seq, hidden]
+
+
+class SelvaActivationSteeringExtractor:
+    """Computes activation steering vectors from a training dataset.
+
+    Runs the frozen generator on N clips at random timesteps with both
+    BJ-conditioned and empty-conditioned inputs, then saves the mean
+    difference per DiT block to a .pt file.
+    """
+
+    OUTPUT_NODE = True
+    CATEGORY    = SELVA_CATEGORY
+    FUNCTION    = "extract"
+    RETURN_TYPES  = ("STRING",)
+    RETURN_NAMES  = ("steering_path",)
+    OUTPUT_TOOLTIPS = ("Path to saved steering_vectors.pt — load with SelVA Activation Steering Loader.",)
+    DESCRIPTION = (
+        "Computes per-block activation steering vectors: mean(BJ activations) − "
+        "mean(empty activations) at each DiT block. Load the result with "
+        "SelVA Activation Steering Loader and connect to the Sampler."
+    )
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("SELVA_MODEL",),
+                "data_dir": ("STRING", {
+                    "default": "",
+                    "tooltip": "Directory containing .npz feature files (same as LoRA/TI trainer).",
+                }),
+                "output_path": ("STRING", {
+                    "default": "steering_vectors.pt",
+                    "tooltip": "Where to save the steering vectors. Relative paths resolve to ComfyUI output directory.",
+                }),
+                "n_samples": ("INT", {
+                    "default": 16, "min": 1, "max": 256,
+                    "tooltip": "Number of clips to average over. More = more stable vectors, slower extraction.",
+                }),
+                "seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFF}),
+            },
+        }
+
+    def extract(self, model, data_dir, output_path, n_samples, seed):
+        device = get_device()
+        dtype  = model["dtype"]
+        seq_cfg = model["seq_cfg"]
+
+        data_dir = Path(data_dir.strip())
+        if not data_dir.is_absolute():
+            data_dir = Path(folder_paths.models_dir) / data_dir
+        if not data_dir.exists():
+            raise FileNotFoundError(f"[Steering] data_dir not found: {data_dir}")
+
+        out_path = Path(output_path.strip())
+        if not out_path.is_absolute():
+            out_path = Path(folder_paths.get_output_directory()) / out_path
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+
+        print(f"\n[Steering] Extracting steering vectors  n_samples={n_samples}", flush=True)
+        print(f"[Steering] data_dir = {data_dir}", flush=True)
+        print(f"[Steering] output   = {out_path}\n", flush=True)
+
+        dataset   = _prepare_dataset(model, data_dir, device)
+        generator = model["generator"]
+        generator.eval()
+
+        torch.manual_seed(seed)
+        random.seed(seed)
+        indices = random.choices(range(len(dataset)), k=n_samples)
+
+        n_blocks     = len(generator.joint_blocks) + len(generator.fused_blocks)
+        bj_sums      = [None] * n_blocks
+        empty_sums   = [None] * n_blocks
+        counts       = [0] * n_blocks
+
+        pbar = comfy.utils.ProgressBar(n_samples)
+
+        for sample_i, clip_idx in enumerate(indices):
+            x1_cpu, clip_f_cpu, sync_f_cpu, text_clip_cpu = dataset[clip_idx]
+
+            clip_f    = clip_f_cpu.to(device, dtype)    # [1, T_clip, 1024]
+            sync_f    = sync_f_cpu.to(device, dtype)    # [1, T_sync, 768]
+            text_clip = text_clip_cpu.to(device, dtype) # [1, 77, 1024]
+
+            # x1 shape is [1, latent_seq_len, latent_dim] — dim 1 is the sequence length.
+            clip_latent_seq_len = x1_cpu.shape[1]
+
+            generator.update_seq_lengths(
+                latent_seq_len=clip_latent_seq_len,
+                clip_seq_len=clip_f.shape[1],
+                sync_seq_len=sync_f.shape[1],
+            )
+
+            conditions       = generator.preprocess_conditions(clip_f, sync_f, text_clip)
+            empty_conditions = generator.get_empty_conditions(bs=1)
+
+            # Random timestep and noise latent for this clip
+            t_val    = torch.rand(1).item()
+            t_tensor = torch.tensor([t_val], device=device, dtype=dtype)
+            latent   = torch.randn(
+                1, clip_latent_seq_len, generator.latent_dim,
+                device=device, dtype=dtype,
+            )
+
+            bj_acts    = _collect_activations(generator, conditions,       latent, t_tensor)
+            empty_acts = _collect_activations(generator, empty_conditions, latent, t_tensor)
+
+            for i, (bj, em) in enumerate(zip(bj_acts, empty_acts)):
+                if bj_sums[i] is None:
+                    bj_sums[i]    = bj.clone()
+                    empty_sums[i] = em.clone()
+                else:
+                    bj_sums[i]    += bj
+                    empty_sums[i] += em
+                counts[i] += 1
+
+            pbar.update(1)
+            if (sample_i + 1) % 4 == 0 or sample_i == n_samples - 1:
+                print(f"[Steering] Processed {sample_i + 1}/{n_samples} clips", flush=True)
+
+        # Steering vector per block: mean(BJ) - mean(empty)
+        steering_vectors = []
+        for i in range(n_blocks):
+            vec = (bj_sums[i] - empty_sums[i]) / counts[i]   # [hidden]
+            steering_vectors.append(vec)
+
+            norm = vec.norm().item()
+            print(f"[Steering] Block {i:2d}  steering_norm={norm:.4f}", flush=True)
+
+        n_joint = len(generator.joint_blocks)
+        payload = {
+            "steering_vectors": steering_vectors,   # list of [seq, hidden] tensors
+            "n_blocks":         n_blocks,
+            "n_joint":          n_joint,
+            "n_fused":          len(generator.fused_blocks),
+            "latent_seq_len":   seq_cfg.latent_seq_len,
+            "n_samples":        n_samples,
+            "seed":             seed,
+            "mode":             model["mode"],
+            "variant":          model["variant"],
+        }
+        torch.save(payload, str(out_path))
+        print(f"\n[Steering] Saved: {out_path}", flush=True)
+
+        soft_empty_cache()
+        return (str(out_path),)

nodes/selva_activation_steering_loader.py

@@ -0,0 +1,62 @@
+"""SelVA Activation Steering Loader.
+
+Loads a steering_vectors.pt bundle produced by SelVA Activation Steering Extractor
+and returns a STEERING_VECTORS dict for use by SelVA Sampler.
+"""
+
+from pathlib import Path
+
+import torch
+import folder_paths
+
+from .utils import SELVA_CATEGORY
+
+
+class SelvaActivationSteeringLoader:
+    CATEGORY    = SELVA_CATEGORY
+    FUNCTION    = "load"
+    RETURN_TYPES  = ("STEERING_VECTORS",)
+    RETURN_NAMES  = ("steering_vectors",)
+    OUTPUT_TOOLTIPS = ("Steering vectors bundle — connect to SelVA Sampler's steering_vectors input.",)
+    DESCRIPTION = (
+        "Loads activation steering vectors from a .pt file produced by "
+        "SelVA Activation Steering Extractor. Connect to SelVA Sampler to nudge "
+        "denoising toward the target activation patterns."
+    )
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "path": ("STRING", {
+                    "default": "steering_vectors.pt",
+                    "tooltip": "Path to steering_vectors.pt. Relative paths resolve to ComfyUI output directory.",
+                }),
+            },
+        }
+
+    def load(self, path):
+        p = Path(path.strip())
+        if not p.is_absolute():
+            p = Path(folder_paths.get_output_directory()) / p
+        if not p.exists():
+            raise FileNotFoundError(f"[Steering] File not found: {p}")
+
+        payload = torch.load(str(p), map_location="cpu", weights_only=False)
+
+        n_blocks = payload["n_blocks"]
+        n_joint  = payload["n_joint"]
+        n_fused  = payload["n_fused"]
+        n_vecs   = len(payload["steering_vectors"])
+
+        print(f"[Steering] Loaded: {p}", flush=True)
+        print(f"[Steering] blocks={n_blocks} (joint={n_joint} fused={n_fused})  "
+              f"latent_seq_len={payload['latent_seq_len']}  "
+              f"n_samples={payload['n_samples']}", flush=True)
+        print(f"[Steering] mode={payload.get('mode')}  variant={payload.get('variant')}", flush=True)
+
+        norms = [payload["steering_vectors"][i].norm().item() for i in range(n_vecs)]
+        mean_norm = sum(norms) / len(norms)
+        print(f"[Steering] Mean steering norm across {n_vecs} blocks: {mean_norm:.4f}", flush=True)
+
+        return (payload,)

nodes/selva_audio_preprocessors.py

@@ -0,0 +1,293 @@
+"""SelVA Audio Preprocessors — condition training clips for codec compatibility.
+
+Two nodes that reduce the domain mismatch between custom training audio and the
+MMAudio VAE's expected spectral distribution, improving LoRA training quality:
+
+  SelvaHfSmoother      — soft low-pass blend to attenuate extreme HF content
+  SelvaSpectralMatcher — adaptive per-band EQ toward the codec's training distribution
+
+Root cause they address: MMAudio was trained on natural sounds (speech, foley, env)
+with limited engineered HF content. The BigVGANv2 vocoder (frozen, pre-trained) handles
+the codec's HF reconstruction poorly for sound design / music training clips, because
+those clips land in a latent-space region the vocoder never saw during training.
+
+Recommended order: SpectralMatcher → HfSmoother → feature extraction → LoRA training.
+"""
+
+import numpy as np
+import torch
+import torchaudio.functional as AF
+
+from .utils import SELVA_CATEGORY
+
+
+# ── Mel filterbank (same algorithm as selva_core/ext/mel_converter.py) ────────
+
+def _mel_filterbank(sr: int, n_fft: int, n_mels: int,
+                    fmin: float, fmax: float) -> torch.Tensor:
+    """Returns mel filterbank matrix [n_mels, n_fft//2+1]."""
+    def hz_to_mel(f):
+        return 2595.0 * np.log10(1.0 + np.asarray(f) / 700.0)
+
+    def mel_to_hz(m):
+        return 700.0 * (10.0 ** (np.asarray(m) / 2595.0) - 1.0)
+
+    n_freqs   = n_fft // 2 + 1
+    fft_freqs = np.linspace(0.0, sr / 2.0, n_freqs)
+    mel_pts   = np.linspace(hz_to_mel(fmin), hz_to_mel(fmax), n_mels + 2)
+    hz_pts    = mel_to_hz(mel_pts)
+
+    fb = np.zeros((n_mels, n_freqs), dtype=np.float32)
+    for m in range(1, n_mels + 1):
+        lo, mid, hi = hz_pts[m - 1], hz_pts[m], hz_pts[m + 1]
+        up   = (fft_freqs - lo)  / (mid - lo  + 1e-12)
+        down = (hi - fft_freqs)  / (hi  - mid + 1e-12)
+        fb[m - 1] = np.maximum(0.0, np.minimum(up, down))
+    return torch.from_numpy(fb)
+
+
+# ── VAE target log-mel means (source: selva_core/ext/autoencoder/vae.py) ──────
+# These are the per-band expected log-mel energy means from MMAudio's training data.
+# Used as the spectral matching target: clips are EQ'd to match this profile.
+
+_TARGET_MEAN_80D = [
+    -1.6058, -1.3676, -1.2520, -1.2453, -1.2078, -1.2224, -1.2419, -1.2439,
+    -1.2922, -1.2927, -1.3170, -1.3543, -1.3401, -1.3836, -1.3907, -1.3912,
+    -1.4313, -1.4152, -1.4527, -1.4728, -1.4568, -1.5101, -1.5051, -1.5172,
+    -1.5623, -1.5373, -1.5746, -1.5687, -1.6032, -1.6131, -1.6081, -1.6331,
+    -1.6489, -1.6489, -1.6700, -1.6738, -1.6953, -1.6969, -1.7048, -1.7280,
+    -1.7361, -1.7495, -1.7658, -1.7814, -1.7889, -1.8064, -1.8221, -1.8377,
+    -1.8417, -1.8643, -1.8857, -1.8929, -1.9173, -1.9379, -1.9531, -1.9673,
+    -1.9824, -2.0042, -2.0215, -2.0436, -2.0766, -2.1064, -2.1418, -2.1855,
+    -2.2319, -2.2767, -2.3161, -2.3572, -2.3954, -2.4282, -2.4659, -2.5072,
+    -2.5552, -2.6074, -2.6584, -2.7107, -2.7634, -2.8266, -2.8981, -2.9673,
+]
+
+_TARGET_MEAN_128D = [
+    -3.3462, -2.6723, -2.4893, -2.3143, -2.2664, -2.3317, -2.1802, -2.4006,
+    -2.2357, -2.4597, -2.3717, -2.4690, -2.5142, -2.4919, -2.6610, -2.5047,
+    -2.7483, -2.5926, -2.7462, -2.7033, -2.7386, -2.8112, -2.7502, -2.9594,
+    -2.7473, -3.0035, -2.8891, -2.9922, -2.9856, -3.0157, -3.1191, -2.9893,
+    -3.1718, -3.0745, -3.1879, -3.2310, -3.1424, -3.2296, -3.2791, -3.2782,
+    -3.2756, -3.3134, -3.3509, -3.3750, -3.3951, -3.3698, -3.4505, -3.4509,
+    -3.5089, -3.4647, -3.5536, -3.5788, -3.5867, -3.6036, -3.6400, -3.6747,
+    -3.7072, -3.7279, -3.7283, -3.7795, -3.8259, -3.8447, -3.8663, -3.9182,
+    -3.9605, -3.9861, -4.0105, -4.0373, -4.0762, -4.1121, -4.1488, -4.1874,
+    -4.2461, -4.3170, -4.3639, -4.4452, -4.5282, -4.6297, -4.7019, -4.7960,
+    -4.8700, -4.9507, -5.0303, -5.0866, -5.1634, -5.2342, -5.3242, -5.4053,
+    -5.4927, -5.5712, -5.6464, -5.7052, -5.7619, -5.8410, -5.9188, -6.0103,
+    -6.0955, -6.1673, -6.2362, -6.3120, -6.3926, -6.4797, -6.5565, -6.6511,
+    -6.8130, -6.9961, -7.1275, -7.2457, -7.3576, -7.4663, -7.6136, -7.7469,
+    -7.8815, -8.0132, -8.1515, -8.3071, -8.4722, -8.7418, -9.3975, -9.6628,
+    -9.7671, -9.8863, -9.9992, -10.0860, -10.1709, -10.5418, -11.2795, -11.3861,
+]
+
+_MEL_CONFIGS = {
+    "16k": dict(sr=16_000, n_fft=1024, n_mels=80,  hop=256, fmin=0, fmax=8_000,
+                target=_TARGET_MEAN_80D,  log10=True),
+    "44k": dict(sr=44_100, n_fft=2048, n_mels=128, hop=512, fmin=0, fmax=22_050,
+                target=_TARGET_MEAN_128D, log10=False),
+}
+
+
+# ── Node 1: HF Smoother ───────────────────────────────────────────────────────
+
+class SelvaHfSmoother:
+    """Soft high-frequency attenuation for LoRA training clip preprocessing.
+
+    Blends a low-pass filtered copy of the audio with the original.  Attenuates
+    the extreme HF content common in engineered sound design that the BigVGANv2
+    vocoder handles poorly, bringing the clip closer to the spectral region the
+    MMAudio codec was trained on (natural sounds with limited HF energy).
+
+    A blend of 0.7 at 12 kHz is a transparent starting point — audible only on
+    close comparison. Increase blend or lower cutoff if roundtrip quality is still
+    poor after spectral matching.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio": ("AUDIO",),
+                "cutoff_hz": ("FLOAT", {
+                    "default": 12000.0, "min": 2000.0, "max": 20000.0, "step": 500.0,
+                    "tooltip": "Low-pass cutoff. 12 kHz is gentle; lower = more aggressive.",
+                }),
+                "blend": ("FLOAT", {
+                    "default": 0.7, "min": 0.0, "max": 1.0, "step": 0.05,
+                    "tooltip": "0 = original, 1 = fully filtered. 0.7 is a transparent starting point.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = ("AUDIO",)
+    RETURN_NAMES  = ("audio",)
+    FUNCTION      = "process"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Blends a low-pass filtered version of the audio with the original to gently attenuate "
+        "high-frequency content that the SelVA codec handles poorly. "
+        "Use before feature extraction to improve LoRA training targets. "
+        "Run after SelVA Spectral Matcher for best results."
+    )
+
+    def process(self, audio, cutoff_hz: float, blend: float):
+        waveform = audio["waveform"].float()   # [1, C, L]
+        sr       = audio["sample_rate"]
+
+        filtered = AF.lowpass_biquad(waveform, sr, cutoff_hz)
+        out = blend * filtered + (1.0 - blend) * waveform
+
+        # Preserve RMS level — LPF removes energy, keep the clip at its original loudness
+        rms_in  = waveform.pow(2).mean().sqrt().clamp(min=1e-8)
+        rms_out = out.pow(2).mean().sqrt().clamp(min=1e-8)
+        out     = out * (rms_in / rms_out)
+
+        peak = out.abs().max()
+        if peak > 1.0:
+            out = out / peak
+
+        print(f"[HF Smoother] cutoff={cutoff_hz:.0f} Hz  blend={blend:.2f}  "
+              f"rms={rms_in:.4f}→{out.pow(2).mean().sqrt():.4f}  "
+              f"peak={out.abs().max():.4f}", flush=True)
+
+        return ({"waveform": out, "sample_rate": sr},)
+
+
+# ── Node 2: Spectral Matcher ──────────────────────────────────────────────────
+
+class SelvaSpectralMatcher:
+    """Adaptive per-band EQ toward the SelVA VAE's expected spectral distribution.
+
+    Computes the log-mel energy profile of the clip and compares it to the per-band
+    means stored in the VAE's normalization buffers (the statistics MMAudio was trained
+    on).  Applies a smooth frequency-domain gain correction so the clip's spectral shape
+    matches what the codec expects, improving encode→decode roundtrip quality and
+    therefore LoRA training target quality.
+
+    The correction is additive in log space (multiplicative in linear), so it only
+    changes spectral balance — not the waveform's timing or phase structure.
+
+    max_gain_db clamps the correction to prevent extreme boosts on very quiet bands.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio": ("AUDIO",),
+                "mode": (["44k", "16k"], {
+                    "tooltip": "Must match the SelVA model you are training. "
+                               "44k = large model, 16k = small model.",
+                }),
+                "strength": ("FLOAT", {
+                    "default": 0.8, "min": 0.0, "max": 1.0, "step": 0.05,
+                    "tooltip": "0 = no correction, 1 = full match to VAE distribution. "
+                               "0.8 is a good starting point.",
+                }),
+                "max_gain_db": ("FLOAT", {
+                    "default": 12.0, "min": 1.0, "max": 30.0, "step": 1.0,
+                    "tooltip": "Clamps per-band gain to ±dB. Prevents extreme boosts on "
+                               "very quiet frequency bands. 12 dB is conservative.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = ("AUDIO",)
+    RETURN_NAMES  = ("audio",)
+    FUNCTION      = "process"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Applies a smooth per-band gain correction to bring the audio's spectral profile "
+        "in line with the MMAudio VAE's expected distribution, derived from the per-band "
+        "normalization statistics baked into the VAE weights. "
+        "Use before feature extraction to improve LoRA training target quality. "
+        "Run before SelVA HF Smoother."
+    )
+
+    def process(self, audio, mode: str, strength: float, max_gain_db: float):
+        cfg      = _MEL_CONFIGS[mode]
+        waveform = audio["waveform"].float()   # [1, C, L]
+        sr_in    = audio["sample_rate"]
+        sr_tgt   = cfg["sr"]
+        n_fft    = cfg["n_fft"]
+        hop      = cfg["hop"]
+
+        # ── flatten to mono and resample if needed ────────────────────────────
+        wav = waveform[0].mean(0)   # [L]
+        if sr_in != sr_tgt:
+            wav = AF.resample(wav.unsqueeze(0), sr_in, sr_tgt).squeeze(0)
+
+        device = wav.device
+        window = torch.hann_window(n_fft, device=device)
+
+        # ── STFT ──────────────────────────────────────────────────────────────
+        stft = torch.stft(wav, n_fft, hop_length=hop, win_length=n_fft,
+                          window=window, center=True, return_complex=True)  # [n_freqs, T]
+        mag  = stft.abs()                                                    # [n_freqs, T]
+
+        # ── current log-mel mean per band ─────────────────────────────────────
+        fb = _mel_filterbank(sr_tgt, n_fft, cfg["n_mels"],
+                             cfg["fmin"], cfg["fmax"]).to(device)   # [n_mels, n_freqs]
+
+        mel_mag = torch.matmul(fb, mag).clamp(min=1e-5)   # [n_mels, T]
+        if cfg["log10"]:
+            mel_log = torch.log10(mel_mag)
+        else:
+            mel_log = torch.log(mel_mag)
+
+        current_mean = mel_log.mean(dim=-1)                               # [n_mels]
+        target_mean  = torch.tensor(cfg["target"], device=device)         # [n_mels]
+
+        # ── per-mel-band gain (log space) ─────────────────────────────────────
+        mel_gain = (target_mean - current_mean) * strength                # [n_mels]
+
+        # Clamp to ±max_gain_db
+        if cfg["log10"]:
+            max_log = max_gain_db / 20.0        # log10: 20 log10 = dB
+        else:
+            max_log = max_gain_db / 8.6859      # ln: 20 * log10(e) ≈ 8.686
+        mel_gain = mel_gain.clamp(-max_log, max_log)
+
+        # ── map mel gains → STFT frequency bins (weighted average) ────────────
+        fb_sum     = fb.sum(0).clamp(min=1e-8)             # [n_freqs]
+        freq_gain  = (mel_gain @ fb) / fb_sum              # [n_freqs]
+
+        if cfg["log10"]:
+            linear_gain = 10.0 ** freq_gain                # [n_freqs]
+        else:
+            linear_gain = torch.exp(freq_gain)             # [n_freqs]
+
+        # ── apply gain in frequency domain and reconstruct ───────────────────
+        stft_out = stft * linear_gain.unsqueeze(-1)        # [n_freqs, T]
+        wav_out  = torch.istft(stft_out, n_fft, hop_length=hop, win_length=n_fft,
+                               window=window, center=True,
+                               length=wav.shape[0])         # [L]
+
+        # ── resample back to original sr ──────────────────────────────────────
+        if sr_in != sr_tgt:
+            wav_out = AF.resample(wav_out.unsqueeze(0), sr_tgt, sr_in).squeeze(0)
+
+        # ── preserve original RMS level ───────────────────────────────────────
+        rms_in  = wav.pow(2).mean().sqrt().clamp(min=1e-8)
+        rms_out = wav_out.pow(2).mean().sqrt().clamp(min=1e-8)
+        wav_out = wav_out * (rms_in / rms_out)
+
+        peak = wav_out.abs().max()
+        if peak > 1.0:
+            wav_out = wav_out / peak
+
+        # ── reshape to match input layout [1, C, L] ───────────────────────────
+        out = wav_out.unsqueeze(0).unsqueeze(0)
+        if waveform.shape[1] > 1:
+            out = out.expand(-1, waveform.shape[1], -1).clone()
+
+        gain_db_range = (
+            20.0 * torch.log10(linear_gain.clamp(min=1e-8))
+        )
+        print(f"[Spectral Matcher] mode={mode}  strength={strength:.2f}  "
+              f"gain [{gain_db_range.min():.1f}, {gain_db_range.max():.1f}] dB  "
+              f"rms={rms_in:.4f}→{out.pow(2).mean().sqrt():.4f}", flush=True)
+
+        return ({"waveform": out, "sample_rate": sr_in},)

nodes/selva_bigvgan_loader.py

@@ -0,0 +1,67 @@
+"""SelVA BigVGAN Loader.
+
+Loads a fine-tuned BigVGAN vocoder checkpoint produced by SelVA BigVGAN Trainer
+and replaces the vocoder weights in the loaded SELVA_MODEL in-place.
+
+The model is modified in-place so ComfyUI's model cache is updated — no need to
+reload the full SelVA model. Subsequent Sampler runs will use the fine-tuned vocoder.
+"""
+
+from pathlib import Path
+
+import torch
+import folder_paths
+
+from .utils import SELVA_CATEGORY
+
+
+class SelvaBigvganLoader:
+    CATEGORY    = SELVA_CATEGORY
+    FUNCTION    = "load"
+    RETURN_TYPES  = ("SELVA_MODEL",)
+    RETURN_NAMES  = ("model",)
+    OUTPUT_TOOLTIPS = ("SELVA_MODEL with the fine-tuned BigVGAN vocoder injected.",)
+    DESCRIPTION = (
+        "Loads a fine-tuned BigVGAN/BigVGANv2 vocoder checkpoint from SelVA BigVGAN Trainer "
+        "and replaces the vocoder weights in the SELVA_MODEL in-place. "
+        "Supports both 16k and 44k models. "
+        "Connect the output to SelVA Sampler instead of the base model loader."
+    )
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("SELVA_MODEL",),
+                "path": ("STRING", {
+                    "default": "bigvgan_bj.pt",
+                    "tooltip": "Path to fine-tuned vocoder checkpoint (.pt). "
+                               "Relative paths resolve to ComfyUI output directory.",
+                }),
+            },
+        }
+
+    def load(self, model, path):
+        p = Path(path.strip())
+        if not p.is_absolute():
+            p = Path(folder_paths.get_output_directory()) / p
+        if not p.exists():
+            raise FileNotFoundError(f"[BigVGAN] Checkpoint not found: {p}")
+
+        ckpt = torch.load(str(p), map_location="cpu", weights_only=False)
+        if "generator" not in ckpt:
+            raise ValueError(f"[BigVGAN] Expected {{'generator': ...}} in checkpoint, got keys: {list(ckpt.keys())}")
+
+        mode = model["mode"]
+        if mode == "16k":
+            vocoder = model["feature_utils"].tod.vocoder.vocoder  # BigVGANVocoder
+        elif mode == "44k":
+            vocoder = model["feature_utils"].tod.vocoder           # BigVGANv2 directly
+        else:
+            raise ValueError(f"[BigVGAN] Unknown mode: {mode}")
+
+        vocoder.load_state_dict(ckpt["generator"])
+        vocoder.eval()
+
+        print(f"[BigVGAN] Loaded fine-tuned vocoder from: {p}", flush=True)
+        return (model,)

nodes/selva_bigvgan_trainer.py

@@ -0,0 +1,720 @@
+"""SelVA BigVGAN Vocoder Fine-tuner.
+
+Tier-1 approach based on research: snake alpha fine-tuning + L2-SP anchor
+regularization + optional frozen discriminator feature matching.
+
+Root cause of harmonic smearing with plain mel/STFT losses:
+  Spectral L1 minimizes expected reconstruction error — averaging over
+  high-variance harmonics. This is a loss-function topology problem, not
+  an LR/step-count problem. The fix is either (a) restrict trainable params
+  so the model lacks capacity to smear, or (b) use a perceptual loss that
+  penalizes harmonic averaging.
+
+Tier-1 implementation:
+  1. snake_alpha_only mode — only tune ~5K per-channel α parameters in
+     Snake/SnakeBeta activations. These control harmonic periodicity per
+     channel. With only 5K trainable params, the model physically cannot
+     reshape the spectrum enough to cause the "green smear".
+  2. L2-SP anchor loss — penalizes parameter drift from pretrained values
+     (strictly better than weight decay, which anchors to zero).
+  3. Frozen discriminator feature matching — if a BigVGAN discriminator
+     checkpoint is provided, the pretrained MPD+MRD networks are used as
+     fixed perceptual feature extractors. Feature matching loss penalizes
+     harmonic smearing directly without any GAN instability.
+
+Save format: {'generator': vocoder.state_dict()} — same as the original
+BigVGAN checkpoint so it can be loaded with SelVA BigVGAN Loader.
+"""
+
+import random
+import threading
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+import comfy.utils
+import comfy.model_management
+import folder_paths
+
+from .utils import SELVA_CATEGORY, get_device, soft_empty_cache
+
+
+# ---------------------------------------------------------------------------
+# Minimal MPD + MRD discriminators matching BigVGAN pretrained checkpoint keys
+# ---------------------------------------------------------------------------
+
+def _get_pad(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class _DiscriminatorP(nn.Module):
+    """Multi-Period Discriminator sub-module (HiFi-GAN / BigVGAN style)."""
+    def __init__(self, period):
+        super().__init__()
+        self.period = period
+        from torch.nn.utils.parametrizations import weight_norm
+        norm = weight_norm
+        self.convs = nn.ModuleList([
+            norm(nn.Conv2d(1,   32,  (5, 1), (3, 1), (_get_pad(5, 1), 0))),
+            norm(nn.Conv2d(32,  128, (5, 1), (3, 1), (_get_pad(5, 1), 0))),
+            norm(nn.Conv2d(128, 512, (5, 1), (3, 1), (_get_pad(5, 1), 0))),
+            norm(nn.Conv2d(512, 1024,(5, 1), (3, 1), (_get_pad(5, 1), 0))),
+            norm(nn.Conv2d(1024,1024,(5, 1), 1,      (_get_pad(5, 1), 0))),
+        ])
+        self.conv_post = norm(nn.Conv2d(1024, 1, (3, 1), 1, (1, 0)))
+
+    def forward(self, x):
+        fmap = []
+        b, c, t = x.shape
+        if t % self.period != 0:
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, 0.1)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return fmap
+
+
+class _MultiPeriodDiscriminator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.discriminators = nn.ModuleList([
+            _DiscriminatorP(p) for p in [2, 3, 5, 7, 11]
+        ])
+
+    def forward(self, y):
+        fmaps = []
+        for d in self.discriminators:
+            fmaps.extend(d(y))
+        return fmaps
+
+
+class _DiscriminatorR(nn.Module):
+    """Multi-Resolution Discriminator sub-module."""
+    def __init__(self, fft_size, shift_size, win_length):
+        super().__init__()
+        self.fft_size   = fft_size
+        self.shift_size = shift_size
+        self.win_length = win_length
+        from torch.nn.utils.parametrizations import weight_norm
+        norm = weight_norm
+        self.convs = nn.ModuleList([
+            norm(nn.Conv2d(1,   128, (3, 9), padding=(1, 4))),
+            norm(nn.Conv2d(128, 128, (3, 9), stride=(1, 2), padding=(1, 4))),
+            norm(nn.Conv2d(128, 128, (3, 9), stride=(1, 2), padding=(1, 4))),
+            norm(nn.Conv2d(128, 128, (3, 9), stride=(1, 2), padding=(1, 4))),
+            norm(nn.Conv2d(128, 128, (3, 3), padding=(1, 1))),
+        ])
+        self.conv_post = norm(nn.Conv2d(128, 1, (3, 3), padding=(1, 1)))
+
+    def spectrogram(self, x):
+        """x: [B, 1, T] → [B, 1, freq, time]"""
+        n, hop, win = self.fft_size, self.shift_size, self.win_length
+        window = torch.hann_window(win, device=x.device)
+        x = x.squeeze(1)  # [B, T]
+        pad = (win - hop) // 2
+        x = F.pad(x, (pad, pad + (win - hop) % 2), mode="reflect")
+        x = torch.stft(x, n, hop, win, window, center=False, return_complex=True)
+        x = x.abs().unsqueeze(1)  # [B, 1, freq, time]
+        return x
+
+    def forward(self, x):
+        fmap = []
+        x = self.spectrogram(x)
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, 0.1)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return fmap
+
+
+class _MultiResolutionDiscriminator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        resolutions = [(1024, 120, 600), (2048, 240, 1200), (512, 50, 240)]
+        self.discriminators = nn.ModuleList([
+            _DiscriminatorR(*r) for r in resolutions
+        ])
+
+    def forward(self, y):
+        fmaps = []
+        for d in self.discriminators:
+            fmaps.extend(d(y))
+        return fmaps
+
+
+def _feature_matching_loss(fmaps_real, fmaps_gen):
+    """L1 between paired feature map lists (both already detach-safe for real)."""
+    loss = torch.zeros(1, device=fmaps_gen[0].device)
+    for fr, fg in zip(fmaps_real, fmaps_gen):
+        T = min(fr.shape[-1], fg.shape[-1])
+        loss = loss + F.l1_loss(fg[..., :T], fr[..., :T].detach())
+    return loss / len(fmaps_real)
+
+
+# ---------------------------------------------------------------------------
+# Utility helpers
+# ---------------------------------------------------------------------------
+
+def _save_spectrogram(path, mel_tensor):
+    """Save mel spectrogram [1, n_mels, T] as a PNG using PIL (no matplotlib dep)."""
+    try:
+        from PIL import Image
+        import numpy as np
+
+        mel = mel_tensor.squeeze(0).float().cpu().numpy()  # [n_mels, T]
+        mel = mel[::-1]                                    # low freq at bottom
+        lo, hi = mel.min(), mel.max()
+        if hi > lo:
+            mel = (mel - lo) / (hi - lo)
+        else:
+            mel = mel - lo
+        img_u8 = (mel * 255).clip(0, 255).astype(np.uint8)
+
+        # Simple blue→green→yellow colour map (viridis-ish) via LUT
+        lut_r = np.array([int(max(0, min(255, 255 * (v * 2 - 1)))) for v in np.linspace(0, 1, 256)], dtype=np.uint8)
+        lut_g = np.array([int(max(0, min(255, 255 * (1 - abs(v * 2 - 1))))) for v in np.linspace(0, 1, 256)], dtype=np.uint8)
+        lut_b = np.array([int(max(0, min(255, 255 * (1 - v * 2)))) for v in np.linspace(0, 1, 256)], dtype=np.uint8)
+        r = Image.fromarray(lut_r[img_u8])
+        g = Image.fromarray(lut_g[img_u8])
+        b = Image.fromarray(lut_b[img_u8])
+        Image.merge("RGB", (r, g, b)).save(str(path))
+    except Exception as e:
+        print(f"[BigVGAN] Spectrogram save failed: {e}", flush=True)
+
+
+def _save_wav(path, wav_tensor, sample_rate):
+    """Save [channels, samples] float32 tensor to .wav.
+
+    Tries torchaudio first; falls back to soundfile when the ffmpeg/torchcodec
+    backend is unavailable (same environment constraint as _load_wav).
+    """
+    try:
+        torchaudio.save(str(path), wav_tensor, sample_rate)
+        return
+    except Exception:
+        pass
+    import soundfile as sf
+    data = wav_tensor.numpy()
+    if data.ndim == 2:
+        data = data.T   # soundfile expects [samples, channels]
+    sf.write(str(path), data, sample_rate)
+
+
+def _load_wav(path):
+    """Load audio file to [channels, samples] float32 tensor.
+
+    Tries torchaudio first; falls back to soundfile for wav/flac when the
+    ffmpeg/torchcodec backend is unavailable (e.g. libavutil soname mismatch).
+    """
+    try:
+        return torchaudio.load(str(path))
+    except Exception:
+        pass
+    # soundfile fallback — handles wav, flac, ogg natively without ffmpeg
+    import soundfile as sf
+    data, sr = sf.read(str(path), dtype="float32", always_2d=True)
+    wav = torch.from_numpy(data.T)  # [channels, samples]
+    return wav, sr
+
+
+# Multi-resolution STFT windows — same three resolutions as BigVGAN discriminator config.
+_STFT_RESOLUTIONS = [
+    (1024, 120,  600),
+    (2048, 240, 1200),
+    (512,   50,  240),
+]
+
+
+def _stft_mag(wav, n_fft, hop_length, win_length, device):
+    """Magnitude STFT.  wav: [B, T]  →  [B, n_fft//2+1, T']"""
+    window = torch.hann_window(win_length, device=device)
+    spec = torch.stft(
+        wav, n_fft=n_fft, hop_length=hop_length, win_length=win_length,
+        window=window, center=True, return_complex=True,
+    )
+    return spec.abs()
+
+
+def _multi_resolution_stft_loss(pred_wav, target_wav, device):
+    """Average L1 mag loss across three STFT resolutions.  inputs: [B, 1, T]"""
+    pred   = pred_wav.squeeze(1)    # [B, T]
+    target = target_wav.squeeze(1)
+    loss = torch.zeros(1, device=device)
+    for n_fft, hop, win in _STFT_RESOLUTIONS:
+        pm = _stft_mag(pred,   n_fft, hop, win, device)
+        tm = _stft_mag(target, n_fft, hop, win, device)
+        T  = min(pm.shape[-1], tm.shape[-1])
+        loss = loss + F.l1_loss(pm[..., :T], tm[..., :T])
+    return loss / len(_STFT_RESOLUTIONS)
+
+
+# ---------------------------------------------------------------------------
+# Node
+# ---------------------------------------------------------------------------
+
+class SelvaBigvganTrainer:
+    OUTPUT_NODE = True
+    CATEGORY    = SELVA_CATEGORY
+    FUNCTION    = "train"
+    RETURN_TYPES  = ("STRING",)
+    RETURN_NAMES  = ("checkpoint_path",)
+    OUTPUT_TOOLTIPS = ("Path to saved vocoder checkpoint — load with SelVA BigVGAN Loader.",)
+    DESCRIPTION = (
+        "Fine-tunes the BigVGAN vocoder (mel→waveform) on BJ audio clips. "
+        "Default mode (snake_alpha_only) tunes only the ~5K Snake activation α "
+        "parameters — cannot cause harmonic smearing. Add a discriminator path "
+        "for perceptual feature matching loss. DiT and VAE stay frozen."
+    )
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("SELVA_MODEL",),
+                "data_dir": ("STRING", {
+                    "default": "",
+                    "tooltip": "Directory with BJ audio files (.wav/.flac/.mp3). Searched recursively.",
+                }),
+                "output_path": ("STRING", {
+                    "default": "bigvgan_bj.pt",
+                    "tooltip": "Where to save the fine-tuned vocoder. Relative paths → ComfyUI output dir.",
+                }),
+                "train_mode": (["snake_alpha_only", "all_params"], {
+                    "default": "snake_alpha_only",
+                    "tooltip": (
+                        "snake_alpha_only: only tune ~5K per-channel α parameters in Snake/SnakeBeta "
+                        "activations. These control harmonic periodicity. Cannot cause spectral smearing. "
+                        "all_params: tune all vocoder weights — set lambda_l2sp>0 to prevent drift."
+                    ),
+                }),
+                "steps": ("INT", {
+                    "default": 2000, "min": 100, "max": 50000,
+                    "tooltip": "Training steps. 1000–2000 is a good first experiment with snake_alpha_only.",
+                }),
+                "lr": ("FLOAT", {
+                    "default": 1e-4, "min": 1e-6, "max": 1e-2, "step": 1e-5,
+                    "tooltip": "Learning rate. 1e-4 for snake_alpha_only, 1e-5 for all_params.",
+                }),
+                "batch_size": ("INT", {"default": 4, "min": 1, "max": 32}),
+                "segment_seconds": ("FLOAT", {
+                    "default": 2.0, "min": 0.25, "max": 30.0, "step": 0.25,
+                    "tooltip": "Audio segment length per training sample in seconds.",
+                }),
+                "lambda_l2sp": ("FLOAT", {
+                    "default": 1e-3, "min": 0.0, "max": 0.1, "step": 1e-4,
+                    "tooltip": (
+                        "L2-SP anchor regularization: penalizes parameter drift from pretrained values. "
+                        "0 = disabled. 1e-3 is good for snake_alpha_only. "
+                        "Increase to 1e-2 for all_params to prevent catastrophic forgetting."
+                    ),
+                }),
+                "save_every": ("INT", {"default": 500, "min": 50, "max": 10000}),
+                "seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFF}),
+            },
+            "optional": {
+                "discriminator_path": ("STRING", {
+                    "default": "",
+                    "tooltip": (
+                        "Optional path to BigVGAN discriminator checkpoint "
+                        "(bigvgan_discriminator_optimizer.pt from the BigVGAN pretrained release). "
+                        "When provided, frozen MPD+MRD feature matching replaces mel L1 — "
+                        "the key fix for harmonic smearing. Leave empty to use mel+STFT losses only."
+                    ),
+                }),
+            },
+        }
+
+    def train(self, model, data_dir, output_path, train_mode, steps, lr, batch_size,
+              segment_seconds, lambda_l2sp, save_every, seed, discriminator_path=""):
+        import traceback
+
+        device        = get_device()
+        mode          = model["mode"]
+        dtype         = model["dtype"]
+        feature_utils = model["feature_utils"]
+        mel_converter = feature_utils.mel_converter
+        strategy      = model["strategy"]
+
+        if mode == "16k":
+            vocoder     = feature_utils.tod.vocoder.vocoder
+            sample_rate = 16_000
+        elif mode == "44k":
+            vocoder     = feature_utils.tod.vocoder
+            sample_rate = 44_100
+        else:
+            raise ValueError(f"[BigVGAN] Unknown mode: {mode}")
+
+        # Resolve paths
+        data_dir = Path(data_dir.strip())
+        if not data_dir.is_absolute():
+            data_dir = Path(folder_paths.models_dir) / data_dir
+        if not data_dir.exists():
+            raise FileNotFoundError(f"[BigVGAN] data_dir not found: {data_dir}")
+
+        out_path = Path(output_path.strip())
+        if not out_path.is_absolute():
+            out_path = Path(folder_paths.get_output_directory()) / out_path
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+
+        disc_path = None
+        if discriminator_path and discriminator_path.strip():
+            disc_path = Path(discriminator_path.strip())
+            if not disc_path.is_absolute():
+                disc_path = Path(folder_paths.get_output_directory()) / disc_path
+            if not disc_path.exists():
+                raise FileNotFoundError(f"[BigVGAN] Discriminator checkpoint not found: {disc_path}")
+
+        # Find and pre-load audio clips
+        segment_samples = int(segment_seconds * sample_rate)
+        audio_files = []
+        for ext in ("*.wav", "*.flac", "*.mp3", "*.ogg", "*.aac"):
+            audio_files.extend(data_dir.rglob(ext))
+        if not audio_files:
+            raise FileNotFoundError(f"[BigVGAN] No audio files found in {data_dir}")
+
+        print(f"[BigVGAN] Loading {len(audio_files)} audio files...", flush=True)
+        clips = []
+        for af in audio_files:
+            try:
+                wav, sr = _load_wav(af)
+                if wav.shape[0] > 1:
+                    wav = wav.mean(0, keepdim=True)
+                if sr != sample_rate:
+                    wav = torchaudio.functional.resample(wav, sr, sample_rate)
+                wav = wav.squeeze(0)  # [L]
+                if wav.shape[0] >= segment_samples:
+                    clips.append(wav.cpu())
+                else:
+                    print(f"  [BigVGAN] Skip {af.name}: shorter than {segment_seconds}s", flush=True)
+            except Exception as e:
+                print(f"  [BigVGAN] Failed {af.name}: {e}", flush=True)
+                traceback.print_exc()
+
+        if not clips:
+            raise RuntimeError(
+                f"[BigVGAN] No usable clips found (need audio >= {segment_seconds}s)"
+            )
+
+        print(f"[BigVGAN] {len(clips)} clips ready  mode={train_mode}  "
+              f"segment={segment_seconds}s  steps={steps}  lr={lr}  "
+              f"batch={batch_size}  lambda_l2sp={lambda_l2sp}\n", flush=True)
+
+        # Unload all other ComfyUI models (SelVA generator, etc.) to free VRAM
+        # before starting training. BigVGAN + discriminator need the headroom.
+        comfy.model_management.unload_all_models()
+        soft_empty_cache()
+
+        if strategy == "offload_to_cpu":
+            feature_utils.to(device)
+            soft_empty_cache()
+
+        mel_converter.to(device)
+
+        pbar = comfy.utils.ProgressBar(steps)
+
+        # -----------------------------------------------------------------------
+        # Run the entire training in a fresh thread.
+        #
+        # ComfyUI executes nodes inside torch.inference_mode(). Even with an inner
+        # inference_mode(False) context, factory functions and operations may still
+        # produce inference tensors in some environments (e.g. when the outer
+        # context is set via an async wrapper or a third-party hook).
+        #
+        # torch.inference_mode is THREAD-LOCAL. A new thread always starts with
+        # inference_mode disabled, so all tensor operations in the worker thread
+        # produce normal, autograd-compatible tensors — no flags to fight.
+        # -----------------------------------------------------------------------
+        _result = [None]
+        _exc    = [None]
+
+        def _worker():
+            try:
+                _result[0] = _do_train(
+                    vocoder, mel_converter, clips,
+                    device, dtype, strategy, feature_utils,
+                    segment_samples, sample_rate,
+                    train_mode, steps, lr, batch_size, lambda_l2sp,
+                    save_every, seed, out_path, disc_path, pbar,
+                )
+            except Exception as e:
+                _exc[0] = e
+                traceback.print_exc()
+
+        t = threading.Thread(target=_worker, daemon=True)
+        t.start()
+        t.join()
+
+        if _exc[0] is not None:
+            raise _exc[0]
+        return (_result[0],)
+
+
+# ---------------------------------------------------------------------------
+# Training worker
+# ---------------------------------------------------------------------------
+
+def _do_train(vocoder, mel_converter, clips,
+              device, dtype, strategy, feature_utils,
+              segment_samples, sample_rate,
+              train_mode, steps, lr, batch_size, lambda_l2sp,
+              save_every, seed, out_path, disc_path, pbar):
+    """Execute training. Called in a fresh thread — no inference_mode active.
+
+    Even though inference_mode is off here, tensors created in the calling
+    thread's inference_mode carry the inference flag on the object itself.
+    Operations on inference tensors produce inference tensors regardless of
+    the current context.  The ONLY way to strip the flag is to call .clone()
+    from outside inference_mode — which is exactly where we are now.
+    """
+    import torch.nn as nn_mod
+
+    # ── Strip inference flag from all inputs that came from the main thread ──
+    # 1. Audio clips (loaded in ComfyUI's inference_mode).
+    clips = [c.clone() for c in clips]
+
+    # 2. mel_converter buffers (mel_basis, hann_window) — same origin.
+    for name, buf in list(mel_converter._buffers.items()):
+        if buf is not None:
+            mel_converter._buffers[name] = buf.clone()
+
+    # 3. Vocoder parameters are handled below with clone().detach().
+    # ─────────────────────────────────────────────────────────────────────────
+
+    torch.manual_seed(seed)
+    random.seed(seed)
+
+    # Reference segment for eval samples — always clip 0, full length
+    ref_wav = clips[0].to(device, dtype)                      # full first clip [T]
+    ref_mel = mel_converter(ref_wav.unsqueeze(0))             # [1, n_mels, T_mel]
+
+    # Ground-truth spectrogram — saved once alongside baseline for comparison
+    gt_spec_path = out_path.parent / f"{out_path.stem}_gt_spec.png"
+    _save_spectrogram(gt_spec_path, ref_mel)
+    print(f"[BigVGAN] GT spectrogram: {gt_spec_path}", flush=True)
+
+    def _save_sample(label):
+        try:
+            voc_device = next(vocoder.parameters()).device
+            mel = ref_mel.to(voc_device)
+            with torch.no_grad():
+                wav = vocoder(mel)
+                if wav.dim() == 2:
+                    wav = wav.unsqueeze(1)
+                wav = wav.float().cpu().clamp(-1, 1)
+            wav_path  = out_path.parent / f"{out_path.stem}_{label}.wav"
+            spec_path = out_path.parent / f"{out_path.stem}_{label}_spec.png"
+            _save_wav(wav_path, wav.squeeze(0), sample_rate)
+            with torch.no_grad():
+                pred_mel = mel_converter(wav.squeeze(1).to(mel_converter.mel_basis.device))
+            _save_spectrogram(spec_path, pred_mel)
+            print(f"[BigVGAN] Sample: {wav_path}  spec: {spec_path}", flush=True)
+        except Exception as e:
+            print(f"[BigVGAN] Sample save failed ({label}): {e}", flush=True)
+
+    _save_sample("baseline")
+
+    # Sanitize all inference tensors in the vocoder.
+    # Three categories to handle (all loaded in ComfyUI's inference_mode):
+    #
+    # 1. Registered parameters (_parameters): covers bias, alpha, etc.
+    #
+    # 2. Plain tensor attributes in __dict__: torch.nn.utils.parametrize.
+    #    remove_parametrizations() calls setattr(module, name, tensor) with
+    #    a raw tensor, NOT nn.Parameter. Module.__setattr__ stores raw tensors
+    #    in __dict__ (not _parameters), so our parameter loop misses them.
+    #    This is how BigVGAN's conv.weight ends up invisible to _parameters.
+    #    Fix: re-register as Parameter, which also makes them trainable.
+    #
+    # 3. Registered buffers (_buffers): Activation1d's anti-aliasing filter
+    #    tensors. Not trainable, but operations on inference buffers produce
+    #    inference tensor outputs — which breaks the backward graph mid-network.
+    #    Fix: clone to strip the inference flag (not registered as parameters).
+    for module in vocoder.modules():
+        # Category 1: registered parameters
+        for pname, param in list(module._parameters.items()):
+            if param is not None:
+                module._parameters[pname] = nn_mod.Parameter(
+                    param.data.clone().detach(), requires_grad=True
+                )
+
+        # Category 2: plain tensor attributes (e.g. weight left by remove_parametrizations)
+        for name, val in list(module.__dict__.items()):
+            if (isinstance(val, torch.Tensor)
+                    and not isinstance(val, nn_mod.Parameter)
+                    and name not in module._buffers
+                    and name not in module._modules):
+                module.register_parameter(name, nn_mod.Parameter(val.clone()))
+
+        # Category 3: buffers (Activation1d filter, etc.) — clone, don't parametrize
+        for bname, buf in list(module._buffers.items()):
+            if buf is not None:
+                module._buffers[bname] = buf.clone()
+
+    # ── Training mode: select which parameters to train ──────────────────────
+    if train_mode == "snake_alpha_only":
+        alpha_params = []
+        for name, param in vocoder.named_parameters():
+            if "alpha" in name:
+                param.requires_grad_(True)
+                alpha_params.append(param)
+            else:
+                param.requires_grad_(False)
+        n_trainable = sum(p.numel() for p in alpha_params)
+        print(f"[BigVGAN] snake_alpha_only: {n_trainable} trainable params "
+              f"({len(alpha_params)} alpha tensors)", flush=True)
+        trainable_params = alpha_params
+    else:  # all_params
+        for param in vocoder.parameters():
+            param.requires_grad_(True)
+        n_trainable = sum(p.numel() for p in vocoder.parameters())
+        print(f"[BigVGAN] all_params: {n_trainable} trainable params", flush=True)
+        trainable_params = list(vocoder.parameters())
+
+    # ── L2-SP: cache reference parameter values (before any gradient steps) ──
+    ref_params = {}
+    if lambda_l2sp > 0.0:
+        for name, param in vocoder.named_parameters():
+            if param.requires_grad:
+                ref_params[name] = param.data.clone().detach()
+        print(f"[BigVGAN] L2-SP anchor: {len(ref_params)} params  λ={lambda_l2sp}", flush=True)
+
+    # ── Optional: load pretrained discriminator for feature matching ──────────
+    mpd = mrd = None
+    if disc_path is not None:
+        try:
+            ckpt_d = torch.load(str(disc_path), map_location="cpu", weights_only=False)
+            mpd = _MultiPeriodDiscriminator()
+            mrd = _MultiResolutionDiscriminator()
+            # Try common key names used by different BigVGAN releases
+            for mpd_key in ("mpd", "discriminator_mpd", "MPD"):
+                if mpd_key in ckpt_d:
+                    mpd.load_state_dict(ckpt_d[mpd_key], strict=False)
+                    print(f"[BigVGAN] Loaded MPD from key '{mpd_key}'", flush=True)
+                    break
+            for mrd_key in ("mrd", "discriminator_mrd", "MRD", "msd", "discriminator_msd"):
+                if mrd_key in ckpt_d:
+                    mrd.load_state_dict(ckpt_d[mrd_key], strict=False)
+                    print(f"[BigVGAN] Loaded MRD from key '{mrd_key}'", flush=True)
+                    break
+            mpd.to(device).eval()
+            mrd.to(device).eval()
+            for p in mpd.parameters():
+                p.requires_grad_(False)
+            for p in mrd.parameters():
+                p.requires_grad_(False)
+            print(f"[BigVGAN] Frozen discriminators ready for feature matching", flush=True)
+        except Exception as e:
+            print(f"[BigVGAN] WARNING: Could not load discriminator ({e}), "
+                  f"falling back to mel+STFT losses", flush=True)
+            mpd = mrd = None
+
+    optimizer = torch.optim.AdamW(trainable_params, lr=lr, betas=(0.8, 0.99))
+    vocoder.train()
+
+    try:
+        for step in range(steps):
+            # Sample random batch — clips are CPU floats, move to device
+            batch = []
+            for _ in range(batch_size):
+                clip  = random.choice(clips)
+                start = random.randint(0, clip.shape[0] - segment_samples)
+                batch.append(clip[start : start + segment_samples])
+
+            target_flat = torch.stack(batch).to(device, dtype)   # [B, T]
+            target_wav  = target_flat.unsqueeze(1)                # [B, 1, T]
+
+            with torch.no_grad():
+                target_mel = mel_converter(target_flat)           # [B, n_mels, T_mel]
+
+            # Gradient checkpointing: recompute BigVGAN activations during
+            # backward instead of storing them. The 512x upsampling stack
+            # produces enormous intermediate tensors — checkpointing trades
+            # ~2x compute for a large reduction in activation memory, allowing
+            # batch_size > 1 without OOM.
+            pred_wav = torch.utils.checkpoint.checkpoint(
+                vocoder, target_mel, use_reentrant=False
+            )                                                     # [B, 1, T_wav]
+
+            T = min(pred_wav.shape[-1], target_wav.shape[-1])
+            pred_t   = pred_wav[...,  :T]
+            target_t = target_wav[..., :T]
+
+            # ── Compute loss ─────────────────────────────────────────────────
+            if mpd is not None and mrd is not None:
+                # Perceptual feature matching via frozen discriminators
+                with torch.no_grad():
+                    fmaps_real_mpd = mpd(target_t)
+                    fmaps_real_mrd = mrd(target_t)
+                fmaps_gen_mpd = mpd(pred_t)
+                fmaps_gen_mrd = mrd(pred_t)
+                fm_loss = (
+                    _feature_matching_loss(fmaps_real_mpd, fmaps_gen_mpd) +
+                    _feature_matching_loss(fmaps_real_mrd, fmaps_gen_mrd)
+                )
+                # Keep a small mel loss for stable frequency alignment
+                pred_mel = mel_converter(pred_t.squeeze(1))
+                T_mel    = min(pred_mel.shape[-1], target_mel.shape[-1])
+                mel_loss = F.l1_loss(pred_mel[..., :T_mel], target_mel[..., :T_mel])
+                primary_loss = 2.0 * fm_loss + 0.1 * mel_loss
+                loss_desc = f"fm={fm_loss.item():.4f}  mel={mel_loss.item():.4f}"
+            else:
+                # Fallback: mel L1 + multi-resolution STFT L1
+                pred_mel = mel_converter(pred_t.squeeze(1))
+                T_mel    = min(pred_mel.shape[-1], target_mel.shape[-1])
+                mel_loss  = F.l1_loss(pred_mel[..., :T_mel], target_mel[..., :T_mel])
+                stft_loss = _multi_resolution_stft_loss(pred_t, target_t, device)
+                primary_loss = mel_loss + stft_loss
+                loss_desc = f"mel={mel_loss.item():.4f}  stft={stft_loss.item():.4f}"
+
+            # ── L2-SP regularization ─────────────────────────────────────────
+            l2sp_loss = torch.zeros(1, device=device)
+            if lambda_l2sp > 0.0 and ref_params:
+                for name, param in vocoder.named_parameters():
+                    if name in ref_params and param.requires_grad:
+                        l2sp_loss = l2sp_loss + F.mse_loss(
+                            param, ref_params[name], reduction="sum"
+                        )
+                l2sp_loss = l2sp_loss * lambda_l2sp
+
+            loss = primary_loss + l2sp_loss
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(trainable_params, 1.0)
+            optimizer.step()
+
+            pbar.update(1)
+
+            if (step + 1) % max(1, steps // 20) == 0 or step == steps - 1:
+                l2sp_str = f"  l2sp={l2sp_loss.item():.4e}" if lambda_l2sp > 0 else ""
+                print(f"[BigVGAN] {step+1}/{steps}  {loss_desc}"
+                      f"  total={loss.item():.4f}{l2sp_str}", flush=True)
+
+            if (step + 1) % save_every == 0 and (step + 1) < steps:
+                step_path = out_path.parent / f"{out_path.stem}_step{step+1}{out_path.suffix}"
+                torch.save({"generator": vocoder.state_dict()}, str(step_path))
+                print(f"[BigVGAN] Checkpoint: {step_path}", flush=True)
+                vocoder.eval()
+                _save_sample(f"step{step+1}")
+                vocoder.train()
+
+    finally:
+        vocoder.requires_grad_(False)
+        vocoder.eval()
+        if strategy == "offload_to_cpu":
+            feature_utils.to("cpu")
+            soft_empty_cache()
+
+    torch.save({"generator": vocoder.state_dict()}, str(out_path))
+    print(f"\n[BigVGAN] Saved: {out_path}", flush=True)
+    _save_sample("final")
+    return str(out_path)

nodes/selva_dataset_browser.py

@@ -0,0 +1,106 @@
+import json
+from pathlib import Path
+
+import folder_paths
+
+from .utils import SELVA_CATEGORY
+
+
+class SelvaDatasetBrowser:
+    """Browse a dataset.json file entry by entry using an integer index.
+
+    Each entry in the JSON is expected to have:
+      - "path"  : base path (no extension) — directory that holds frame images
+      - "label" : text description of the clip
+
+    Derived outputs:
+      - video_path  : path + ".mp4"
+      - audio_path  : path + ".wav"
+      - frames_dir  : path  (the directory itself, for image-sequence loaders)
+      - label       : entry["label"]
+      - count       : total number of entries in the file
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset_json": ("STRING", {
+                    "default": "",
+                    "tooltip": "Absolute or ComfyUI-relative path to a dataset.json file.",
+                }),
+                "index": ("INT", {
+                    "default": 0,
+                    "min": 0,
+                    "max": 9999,
+                    "step": 1,
+                    "tooltip": "Zero-based index of the entry to inspect.",
+                }),
+            },
+        }
+
+    RETURN_TYPES  = ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "INT")
+    RETURN_NAMES  = ("video_path", "audio_wav", "audio_flac", "features_path", "frames_dir", "mask_dir", "label", "max_index")
+    OUTPUT_TOOLTIPS = (
+        "path + '.mp4'",
+        "features/ + name + '.wav'",
+        "features/ + name + '.flac'",
+        "features/ + name + '.npz'  (pre-extracted SelVA features)",
+        "path  (image-sequence directory)",
+        "path + '_mask'  (mask image-sequence directory)",
+        "Text label for this clip",
+        "count - 1 — wire to a primitive INT's max to constrain the index widget",
+    )
+    FUNCTION  = "browse"
+    CATEGORY  = SELVA_CATEGORY
+    DESCRIPTION = (
+        "Reads a dataset.json produced by the SelVA dataset preparation pipeline "
+        "and exposes one entry at a time via an integer index. "
+        "Outputs the video path, audio path, frames directory, label, and total entry count."
+    )
+
+    # Re-read the file every call so edits are picked up without restarting ComfyUI.
+    IS_CHANGED = classmethod(lambda cls, **_: float("nan"))
+
+    def browse(self, dataset_json: str, index: int):
+        p = Path(dataset_json.strip())
+        if not p.is_absolute():
+            p = Path(folder_paths.base_path) / p
+        if not p.exists():
+            raise FileNotFoundError(f"[SelVA Dataset Browser] File not found: {p}")
+
+        with p.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        if not isinstance(data, list) or len(data) == 0:
+            raise ValueError(f"[SelVA Dataset Browser] Expected a non-empty JSON array in {p}")
+
+        count = len(data)
+        if index >= count:
+            raise IndexError(
+                f"[SelVA Dataset Browser] index {index} is out of range "
+                f"(dataset has {count} entries, last index is {count - 1})"
+            )
+        entry = data[index]
+
+        base  = entry["path"]
+        label = entry.get("label", "")
+
+        p_base    = Path(base)
+        feat_base = str(p_base.parent / "features" / p_base.name)
+
+        print(
+            f"[SelVA Dataset Browser] {index + 1}/{count}  label='{label}'  base={base}",
+            flush=True,
+        )
+
+        return (
+            base + ".mp4",
+            feat_base + ".wav",
+            feat_base + ".flac",
+            feat_base + ".npz",
+            base,
+            base + "_mask",
+            label,
+            count - 1,
+        )

nodes/selva_ditto_optimizer.py

@@ -0,0 +1,431 @@
+"""SelVA DITTO Optimizer.
+
+Inference-time noise optimization: optimizes the initial noise latent x_0
+using a style loss against BJ reference clips, backpropagating through the
+ODE solver. All model weights remain frozen — only x_0 changes.
+
+Based on DITTO: Diffusion Inference-Time T-Optimization (arXiv:2401.12179,
+ICML 2024 Oral). Adapted for SelVA's flow-matching Euler ODE.
+
+Style loss: mel-spectrogram statistics matching (mean spectrum + Gram matrix)
+against BJ reference clips. Runs entirely before the vocoder — optimization
+only requires the DiT + VAE decoder, not BigVGAN.
+
+Memory strategy: gradient checkpointing at each ODE step — stores O(1 DiT
+forward pass activations) instead of O(N steps). Backward recomputes each
+step's activations on demand.
+"""
+
+import dataclasses
+import threading
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+import comfy.utils
+import comfy.model_management
+import folder_paths
+
+from .utils import SELVA_CATEGORY, get_device, get_offload_device, soft_empty_cache
+
+
+def _load_wav(path):
+    """Load audio file to [channels, samples] float32 tensor."""
+    try:
+        return torchaudio.load(str(path))
+    except Exception:
+        pass
+    import soundfile as sf
+    data, sr = sf.read(str(path), dtype="float32", always_2d=True)
+    wav = torch.from_numpy(data.T)
+    return wav, sr
+
+
+def _mel_style_loss(mel_gen, ref_mean, ref_gram):
+    """Style loss between generated mel and precomputed reference statistics.
+
+    mel_gen:  [1, n_mels, T]  generated mel spectrogram (with grad)
+    ref_mean: [n_mels]         mean spectrum of BJ reference clips (detached)
+    ref_gram: [n_mels, n_mels] Gram matrix of BJ reference clips  (detached)
+
+    Mean spectrum loss captures the spectral envelope (which harmonics are
+    boosted). Gram matrix loss captures timbral texture — covariance between
+    frequency bands — without requiring temporal alignment.
+    """
+    m = mel_gen.squeeze(0)  # [n_mels, T]
+
+    # Mean spectrum loss
+    gen_mean = m.mean(dim=-1)  # [n_mels]
+    loss_mean = F.l1_loss(gen_mean, ref_mean)
+
+    # Gram matrix loss (texture, position-invariant)
+    gram_gen = (m @ m.T) / m.shape[-1]  # [n_mels, n_mels]
+    loss_gram = F.mse_loss(gram_gen, ref_gram)
+
+    return loss_mean + 0.1 * loss_gram
+
+
+class SelvaDittoOptimizer:
+    """DITTO inference-time noise optimization.
+
+    Freezes all model weights and optimizes only the initial noise latent x_0
+    to make the generated audio sound like the BJ reference clips.
+    No training data or gradient updates to the model — per-video per-run.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("SELVA_MODEL",),
+                "features": ("SELVA_FEATURES",),
+                "prompt": ("STRING", {
+                    "default": "", "multiline": True,
+                    "tooltip": "Sound description. Leave empty to use features prompt.",
+                }),
+                "negative_prompt": ("STRING", {
+                    "default": "", "multiline": False,
+                }),
+                "reference_dir": ("STRING", {
+                    "default": "",
+                    "tooltip": "Directory with BJ reference audio files (.wav/.flac/.mp3). "
+                               "Reference mel statistics are precomputed from these once.",
+                }),
+                "n_opt_steps": ("INT", {
+                    "default": 50, "min": 5, "max": 500,
+                    "tooltip": "Gradient optimization steps on x_0. 50 is a good start; "
+                               "each step requires ~2 DiT forward passes.",
+                }),
+                "opt_lr": ("FLOAT", {
+                    "default": 0.1, "min": 0.001, "max": 2.0, "step": 0.01,
+                    "tooltip": "Adam learning rate for x_0 optimization. "
+                               "0.1 is the DITTO paper default.",
+                }),
+                "n_ode_steps": ("INT", {
+                    "default": 10, "min": 5, "max": 50,
+                    "tooltip": "Euler ODE steps run during each optimization iteration. "
+                               "Lower = faster optimization (10–15 is a good trade-off). "
+                               "Final generation always uses the steps parameter below.",
+                }),
+                "n_grad_steps": ("INT", {
+                    "default": 5, "min": 1, "max": 50,
+                    "tooltip": "ODE steps to differentiate through (truncated BPTT). "
+                               "Higher = more accurate gradient, more VRAM. "
+                               "Must be ≤ n_ode_steps. 5 is a good default.",
+                }),
+                "style_weight": ("FLOAT", {
+                    "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1,
+                    "tooltip": "Weight of the BJ style loss. Increase to push harder toward "
+                               "BJ style at the cost of coherence with the video.",
+                }),
+                "steps": ("INT", {
+                    "default": 25, "min": 1, "max": 200,
+                    "tooltip": "Euler steps for the final generation pass (after optimization).",
+                }),
+                "cfg_strength": ("FLOAT", {
+                    "default": 4.5, "min": 1.0, "max": 20.0, "step": 0.1}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
+            },
+            "optional": {
+                "normalize": ("BOOLEAN", {"default": True}),
+                "target_lufs": ("FLOAT", {
+                    "default": -27.0, "min": -40.0, "max": -6.0, "step": 1.0}),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    OUTPUT_TOOLTIPS = ("DITTO-optimized audio — x_0 steered toward BJ style.",)
+    FUNCTION = "optimize"
+    CATEGORY = SELVA_CATEGORY
+    DESCRIPTION = (
+        "DITTO inference-time noise optimization (arXiv:2401.12179). "
+        "Optimizes the initial noise latent x_0 to match BJ reference clips "
+        "via mel statistics style loss, backpropagating through the ODE. "
+        "All model weights frozen — zero quality degradation risk."
+    )
+
+    def optimize(self, model, features, prompt, negative_prompt,
+                 reference_dir, n_opt_steps, opt_lr, n_ode_steps, n_grad_steps,
+                 style_weight, steps, cfg_strength, seed,
+                 normalize=True, target_lufs=-27.0):
+        import traceback
+
+        device        = get_device()
+        dtype         = model["dtype"]
+        strategy      = model["strategy"]
+        net_generator = model["generator"]
+        feature_utils = model["feature_utils"]
+        mel_converter = feature_utils.mel_converter
+
+        # Validate variant match
+        feat_variant = features.get("variant")
+        if feat_variant is not None and feat_variant != model["variant"]:
+            raise ValueError(
+                f"[DITTO] Variant mismatch: features='{feat_variant}' model='{model['variant']}'. "
+                f"Re-run Feature Extractor."
+            )
+
+        if not prompt or not prompt.strip():
+            prompt = features.get("prompt", "")
+
+        # Resolve duration and seq_cfg
+        duration = features.get("duration", 0)
+        if duration <= 0:
+            raise ValueError("[DITTO] Features contain no duration field.")
+        seq_cfg = dataclasses.replace(model["seq_cfg"], duration=duration)
+        sample_rate = seq_cfg.sampling_rate
+
+        # Load and precompute reference mel statistics
+        ref_dir = Path(reference_dir.strip())
+        if not ref_dir.is_absolute():
+            ref_dir = Path(folder_paths.models_dir) / ref_dir
+        if not ref_dir.exists():
+            raise FileNotFoundError(f"[DITTO] reference_dir not found: {ref_dir}")
+
+        ref_files = []
+        for ext in ("*.wav", "*.flac", "*.mp3", "*.ogg"):
+            ref_files.extend(ref_dir.rglob(ext))
+        if not ref_files:
+            raise FileNotFoundError(f"[DITTO] No audio files in reference_dir: {ref_dir}")
+
+        print(f"[DITTO] Loading {len(ref_files)} reference clips...", flush=True)
+        mel_converter.to(device)
+
+        ref_mels = []
+        with torch.no_grad():
+            for rf in ref_files[:32]:  # cap at 32 for speed
+                try:
+                    wav, sr = _load_wav(rf)
+                    if wav.shape[0] > 1:
+                        wav = wav.mean(0, keepdim=True)
+                    if sr != sample_rate:
+                        wav = torchaudio.functional.resample(wav, sr, sample_rate)
+                    wav = wav.squeeze(0).to(device, dtype)
+                    mel = mel_converter(wav.unsqueeze(0))  # [1, n_mels, T]
+                    ref_mels.append(mel)
+                except Exception as e:
+                    print(f"  [DITTO] Skip {rf.name}: {e}", flush=True)
+
+        if not ref_mels:
+            raise RuntimeError("[DITTO] No usable reference clips.")
+
+        # Precompute reference statistics (done once — detached, no grad)
+        with torch.no_grad():
+            all_means = torch.stack([m.squeeze(0).mean(dim=-1) for m in ref_mels])
+            ref_mean  = all_means.mean(0)  # [n_mels]
+
+            all_grams = []
+            for m in ref_mels:
+                M = m.squeeze(0)  # [n_mels, T]
+                all_grams.append((M @ M.T) / M.shape[-1])
+            ref_gram = torch.stack(all_grams).mean(0)  # [n_mels, n_mels]
+
+        print(f"[DITTO] Reference stats computed from {len(ref_mels)} clips  "
+              f"n_opt={n_opt_steps}  lr={opt_lr}  ode_steps={n_ode_steps}  "
+              f"grad_steps={n_grad_steps}", flush=True)
+
+        if strategy == "offload_to_cpu":
+            net_generator.to(device)
+            feature_utils.to(device)
+            soft_empty_cache()
+
+        pbar = comfy.utils.ProgressBar(n_opt_steps + steps)
+
+        _result = [None]
+        _exc    = [None]
+
+        def _worker():
+            try:
+                _result[0] = _do_optimize(
+                    net_generator, feature_utils, mel_converter,
+                    features, prompt, negative_prompt,
+                    ref_mean, ref_gram,
+                    seq_cfg, sample_rate, device, dtype,
+                    n_opt_steps, opt_lr, n_ode_steps, n_grad_steps,
+                    style_weight, steps, cfg_strength, seed,
+                    normalize, target_lufs, pbar,
+                )
+            except Exception as e:
+                _exc[0] = e
+                traceback.print_exc()
+
+        t = threading.Thread(target=_worker, daemon=True)
+        t.start()
+        t.join()
+
+        if strategy == "offload_to_cpu":
+            net_generator.to(get_offload_device())
+            feature_utils.to(get_offload_device())
+            soft_empty_cache()
+
+        if _exc[0] is not None:
+            raise _exc[0]
+        return (_result[0],)
+
+
+def _do_optimize(net_generator, feature_utils, mel_converter,
+                 features, prompt, negative_prompt,
+                 ref_mean, ref_gram,
+                 seq_cfg, sample_rate, device, dtype,
+                 n_opt_steps, opt_lr, n_ode_steps, n_grad_steps,
+                 style_weight, steps, cfg_strength, seed,
+                 normalize, target_lufs, pbar):
+    """Optimization loop — runs in a fresh thread (no inference_mode active)."""
+
+    # Strip inference flags from ref stats (came from main thread)
+    ref_mean = ref_mean.clone().detach()
+    ref_gram = ref_gram.clone().detach()
+
+    torch.manual_seed(seed)
+
+    clip_f = features["clip_features"].to(device, dtype).clone()
+    sync_f = features["sync_features"].to(device, dtype).clone()
+
+    net_generator.update_seq_lengths(
+        latent_seq_len=seq_cfg.latent_seq_len,
+        clip_seq_len=clip_f.shape[1],
+        sync_seq_len=sync_f.shape[1],
+    )
+
+    with torch.no_grad():
+        text_clip = feature_utils.encode_text_clip([prompt])
+
+        neg_text_clip = feature_utils.encode_text_clip([negative_prompt]) \
+            if negative_prompt.strip() else None
+
+        conditions       = net_generator.preprocess_conditions(clip_f, sync_f, text_clip)
+        empty_conditions = net_generator.get_empty_conditions(
+            bs=1, negative_text_features=neg_text_clip
+        )
+
+    # Initial noise — x_0 is the parameter we optimize
+    x0_init = torch.randn(
+        1, seq_cfg.latent_seq_len, net_generator.latent_dim,
+        device=device, dtype=dtype,
+    )
+    x0 = torch.nn.Parameter(x0_init.clone())
+    optimizer = torch.optim.Adam([x0], lr=opt_lr)
+
+    # n_grad_steps must not exceed n_ode_steps
+    n_grad_steps = min(n_grad_steps, n_ode_steps)
+    n_free_steps = n_ode_steps - n_grad_steps  # steps run without gradient
+
+    ts = torch.linspace(0.0, 1.0, n_ode_steps + 1, device=device, dtype=dtype)
+
+    print(f"[DITTO] Optimizing x_0  "
+          f"free_steps={n_free_steps}  grad_steps={n_grad_steps}", flush=True)
+
+    # Freeze all model weights (double-check — should already be frozen at inference)
+    net_generator.requires_grad_(False)
+    feature_utils.requires_grad_(False)
+    mel_converter.requires_grad_(False)
+
+    for opt_step in range(n_opt_steps):
+        comfy.model_management.throw_exception_if_processing_interrupted()
+
+        # ── Phase 1: run first (n_ode_steps - n_grad_steps) steps without grad ──
+        # Detach from x0 so Phase 1 does not build a computation graph.
+        with torch.no_grad():
+            x = x0.detach()
+            for i in range(n_free_steps):
+                t  = ts[i]
+                dt = ts[i + 1] - t
+                flow = net_generator.ode_wrapper(t, x, conditions, empty_conditions, cfg_strength)
+                x = x + dt * flow
+
+        # Straight-through estimator: reconnect x to x0's gradient path by
+        # adding the zero tensor (x0 - x0.detach()). This adds zero value but
+        # creates a grad_fn pointing back to x0, so loss.backward() will
+        # propagate ∂loss/∂x (at the Phase-1/2 boundary) directly to x0.grad.
+        # The approximation is ∂x_prefix/∂x0 ≈ I — the no-grad prefix is
+        # treated as identity for gradient purposes (truncated BPTT).
+        #
+        # x may carry an inference tensor flag from Phase 1 (derived from
+        # conditions which were built outside inference_mode but may have
+        # propagated the flag). .clone() strips it so the STE addition does
+        # not try to save an inference tensor for backward.
+        x = x.clone()
+        x = x + (x0 - x0.detach())
+
+        # ── Phase 2: run last n_grad_steps with gradient + checkpointing ──
+        for i in range(n_free_steps, n_ode_steps):
+            t  = ts[i]
+            dt = ts[i + 1] - t
+
+            # Gradient checkpointing: recompute forward during backward,
+            # avoiding storage of DiT activations for each step.
+            def _ode_step(x_in, t=t):
+                return net_generator.ode_wrapper(t, x_in, conditions, empty_conditions, cfg_strength)
+
+            flow = torch.utils.checkpoint.checkpoint(
+                _ode_step, x, use_reentrant=False
+            )
+            x = x + dt * flow
+
+        # ── Decode to mel (no vocoder — cheap) ──────────────────────────────
+        # Wrap unnormalize + decode in gradient checkpointing so PyTorch does
+        # not try to save model weights for backward. The VAE / generator
+        # weights are inference-flagged tensors (loaded in the main thread);
+        # saving them for backward would raise "Inference tensors cannot be
+        # saved for backward". checkpoint(use_reentrant=False) recomputes the
+        # forward during backward instead of storing activations.
+        def _unnorm_decode(x_in):
+            x_un = net_generator.unnormalize(x_in)
+            return feature_utils.decode(x_un)
+
+        mel_gen = torch.utils.checkpoint.checkpoint(
+            _unnorm_decode, x, use_reentrant=False
+        )
+
+        # ── Style loss ───────────────────────────────────────────────────────
+        loss = style_weight * _mel_style_loss(mel_gen, ref_mean, ref_gram)
+
+        optimizer.zero_grad()
+        loss.backward()  # gradient flows through Phase 2 + STE back to x0.grad
+        torch.nn.utils.clip_grad_norm_([x0], 1.0)
+        optimizer.step()
+
+        pbar.update(1)
+
+        if (opt_step + 1) % max(1, n_opt_steps // 10) == 0:
+            print(f"[DITTO] {opt_step+1}/{n_opt_steps}  loss={loss.item():.4f}", flush=True)
+
+    # ── Final generation with optimized x_0 ─────────────────────────────────
+    print(f"[DITTO] Optimization done. Final generation ({steps} steps)...", flush=True)
+
+    with torch.no_grad():
+        fm_ts = torch.linspace(0.0, 1.0, steps + 1, device=device, dtype=dtype)
+        x = x0.detach()
+        for i in range(steps):
+            comfy.model_management.throw_exception_if_processing_interrupted()
+            t  = fm_ts[i]
+            dt = fm_ts[i + 1] - t
+            flow = net_generator.ode_wrapper(t, x, conditions, empty_conditions, cfg_strength)
+            x = x + dt * flow
+            pbar.update(1)
+
+        x1_unnorm = net_generator.unnormalize(x)
+        spec  = feature_utils.decode(x1_unnorm)
+        audio = feature_utils.vocode(spec)
+
+    print(f"[DITTO] latent stats: mean={x.float().mean():.4f} std={x.float().std():.4f}",
+          flush=True)
+
+    audio = audio.float()
+    if audio.dim() == 2:
+        audio = audio.unsqueeze(1)
+    elif audio.dim() == 3 and audio.shape[1] != 1:
+        audio = audio.mean(dim=1, keepdim=True)
+
+    if normalize:
+        target_rms = 10 ** (target_lufs / 20.0)
+        rms = audio.pow(2).mean().sqrt().clamp(min=1e-8)
+        audio = audio * (target_rms / rms)
+        peak = audio.abs().max().clamp(min=1e-8)
+        if peak > 1.0:
+            audio = audio / peak
+
+    print(f"[DITTO] audio: shape={tuple(audio.shape)} sr={sample_rate}", flush=True)
+    return ({"waveform": audio.cpu(), "sample_rate": sample_rate},)

nodes/selva_lora_evaluator.py

@@ -0,0 +1,417 @@
+"""SelVA LoRA Evaluator — generates eval samples from multiple adapters for comparison.
+
+JSON format:
+    {
+      "name": "eval_batch_1",
+      "data_dir": "/path/to/features",
+      "output_dir": "/path/to/evals/batch1",
+      "steps": 25,
+      "seed": 42,
+      "adapters": [
+        {"id": "baseline"},
+        {"id": "lr_3e4_10k", "path": "/path/to/adapter_final.pt"},
+        {"id": "lr_5e4_10k", "path": "/path/to/adapter_final.pt"}
+      ]
+    }
+
+Empty / missing "path" = baseline (no LoRA applied).
+"""
+
+import copy
+import json
+import sys
+import traceback
+from datetime import datetime, timezone
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchaudio
+
+import comfy.utils
+import folder_paths
+
+from .utils import SELVA_CATEGORY, get_device, soft_empty_cache
+from .selva_lora_trainer import (
+    _prepare_dataset,
+    _eval_sample,
+    _spectral_metrics,
+    _save_spectrogram,
+    _pil_to_tensor,
+    _find_audio,
+    _load_audio,
+)
+from selva_core.model.lora import apply_lora, load_lora
+
+
+def _avg_metrics(metrics_list: list) -> dict:
+    """Average spectral metrics across multiple clips, ignoring None entries."""
+    keys = ["hf_energy_ratio", "spectral_centroid_hz", "spectral_rolloff_hz",
+            "spectral_flatness", "temporal_variance"]
+    valid = [m for m in metrics_list if m]
+    if not valid:
+        return {}
+    return {k: round(float(sum(m[k] for m in valid) / len(valid)), 4) for k in keys}
+
+
+def _resolve_path(raw: str) -> Path:
+    p = Path(raw.strip())
+    unix_style_on_windows = sys.platform == "win32" and p.is_absolute() and not p.drive
+    if not p.is_absolute() or unix_style_on_windows:
+        p = Path(folder_paths.get_output_directory()) / p.relative_to(p.anchor)
+    return p
+
+
+def _safe_stem(adapter_id: str) -> str:
+    """Replace characters illegal in filenames."""
+    for ch in r'/\:*?"<>|':
+        adapter_id = adapter_id.replace(ch, "_")
+    return adapter_id
+
+
+def _draw_metric_comparison(adapter_ids: list, metrics_list: list, output_path: Path):
+    """Draw a 2×2 grid of horizontal bar charts comparing spectral metrics.
+
+    Saves a PNG to output_path and returns a ComfyUI IMAGE tensor.
+    """
+    from matplotlib.figure import Figure
+    from matplotlib.backends.backend_agg import FigureCanvasAgg
+
+    METRICS = [
+        ("hf_energy_ratio",      "HF Energy Ratio (>4 kHz)"),
+        ("spectral_centroid_hz", "Spectral Centroid (Hz)"),
+        ("spectral_flatness",    "Spectral Flatness"),
+        ("temporal_variance",    "Temporal Variance"),
+    ]
+    COLORS = [
+        "#4285F4", "#EA4335", "#34A853", "#FBBC05",
+        "#9B59B6", "#1ABC9C", "#E67E22", "#95A5A6",
+    ]
+
+    fig = Figure(figsize=(12, max(4, len(adapter_ids) * 0.6 + 2)), dpi=110, tight_layout=True)
+    axes = [fig.add_subplot(2, 2, i + 1) for i in range(4)]
+
+    for ax, (key, title) in zip(axes, METRICS):
+        values = []
+        colors = []
+        for i, m in enumerate(metrics_list):
+            v = m.get(key, 0.0) if m else 0.0
+            values.append(v)
+            colors.append(COLORS[i % len(COLORS)])
+
+        bars = ax.barh(adapter_ids, values, color=colors, height=0.6)
+        ax.set_title(title, fontsize=9)
+        ax.set_xlabel(key, fontsize=8)
+        ax.tick_params(axis="y", labelsize=7)
+        ax.tick_params(axis="x", labelsize=7)
+
+        # Value labels on bars
+        for bar, val in zip(bars, values):
+            w = bar.get_width()
+            ax.text(w * 1.01, bar.get_y() + bar.get_height() / 2,
+                    f"{val:.3f}", va="center", ha="left", fontsize=6)
+
+    canvas = FigureCanvasAgg(fig)
+    canvas.draw()
+    canvas.print_figure(str(output_path), dpi=110)
+
+    buf = canvas.buffer_rgba()
+    w, h = canvas.get_width_height()
+    arr = np.frombuffer(buf, dtype=np.uint8).reshape(h, w, 4)[:, :, :3]
+    from PIL import Image
+    return _pil_to_tensor(Image.fromarray(arr))
+
+
+class SelvaLoraEvaluator:
+    """Evaluates a batch of LoRA adapters on a fixed reference clip.
+
+    Generates one audio sample per adapter, computes spectral metrics for each,
+    and produces a comparison chart. Use this after a sweep to compare candidates
+    before running the next round of training.
+    """
+
+    OUTPUT_NODE = True
+    CATEGORY    = SELVA_CATEGORY
+    FUNCTION    = "run"
+    RETURN_TYPES  = ("STRING", "IMAGE")
+    RETURN_NAMES  = ("summary_path", "comparison_image")
+    OUTPUT_TOOLTIPS = (
+        "Path to eval_summary.json — contains spectral metrics per adapter.",
+        "Bar chart comparing spectral metrics across all evaluated adapters.",
+    )
+    DESCRIPTION = (
+        "Evaluates multiple LoRA adapters by generating one audio sample per adapter "
+        "from a fixed reference clip, then collects spectral metrics for comparison. "
+        "Input is a JSON file listing adapter paths. Empty path = baseline (no LoRA)."
+    )
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("SELVA_MODEL",),
+                "eval_file": ("STRING", {
+                    "default": "eval_batch.json",
+                    "tooltip": (
+                        "Path to the JSON evaluation spec. Relative paths resolve "
+                        "to the ComfyUI output directory. "
+                        "Each adapter entry needs an 'id' and an optional 'path'. "
+                        "Omit 'path' for a no-LoRA baseline."
+                    ),
+                }),
+            }
+        }
+
+    def run(self, model, eval_file):
+        # ------------------------------------------------------------------
+        # 1. Resolve and parse the JSON file
+        # ------------------------------------------------------------------
+        eval_path = Path(eval_file.strip())
+        if not eval_path.is_absolute():
+            candidate = Path(folder_paths.models_dir) / eval_path
+            if not candidate.exists():
+                candidate = Path(folder_paths.get_output_directory()) / eval_path
+            eval_path = candidate
+        if not eval_path.exists():
+            raise FileNotFoundError(f"[LoRA Evaluator] Eval file not found: {eval_path}")
+
+        spec = json.loads(eval_path.read_text(encoding="utf-8"))
+
+        if "adapters" not in spec or not spec["adapters"]:
+            raise ValueError("[LoRA Evaluator] 'adapters' list is missing or empty.")
+        for i, a in enumerate(spec["adapters"]):
+            if "id" not in a:
+                raise ValueError(f"[LoRA Evaluator] Adapter at index {i} missing 'id'.")
+
+        if "data_dir" not in spec:
+            raise ValueError("[LoRA Evaluator] 'data_dir' is required.")
+        if "output_dir" not in spec:
+            raise ValueError("[LoRA Evaluator] 'output_dir' is required.")
+
+        name       = spec.get("name", eval_path.stem)
+        data_dir   = _resolve_path(spec["data_dir"])
+        output_dir = _resolve_path(spec["output_dir"])
+        steps      = int(spec.get("steps", 25))
+        seed       = int(spec.get("seed",  42))
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        print(f"\n[LoRA Evaluator] '{name}': {len(spec['adapters'])} adapter(s)", flush=True)
+        print(f"[LoRA Evaluator] data_dir   = {data_dir}", flush=True)
+        print(f"[LoRA Evaluator] output_dir = {output_dir}\n", flush=True)
+
+        # ------------------------------------------------------------------
+        # 2. Prepare dataset (VAE encode once)
+        # ------------------------------------------------------------------
+        device = get_device()
+        dtype  = model["dtype"]
+        dataset = _prepare_dataset(model, data_dir, device)
+
+        feature_utils_orig = model["feature_utils"]
+        seq_cfg            = model["seq_cfg"]
+
+        # ------------------------------------------------------------------
+        # 3. Collect reference metrics for all dataset clips
+        # ------------------------------------------------------------------
+        import shutil
+        npz_files  = sorted(data_dir.glob("*.npz"))
+        ref_dir    = output_dir / "reference"
+        ref_dir.mkdir(exist_ok=True)
+        ref_clips  = []   # list of {clip, wav_path, spectral_metrics}
+
+        print(f"[LoRA Evaluator] Computing reference metrics for {len(npz_files)} clip(s)...",
+              flush=True)
+        for npz_path in npz_files:
+            audio_path = _find_audio(npz_path)
+            if audio_path is None:
+                continue
+            try:
+                ref_wav = _load_audio(audio_path, seq_cfg.sampling_rate, seq_cfg.duration)
+                ref_wav = ref_wav.unsqueeze(0)   # [1, L]
+                ref_out = ref_dir / f"{npz_path.stem}{audio_path.suffix}"
+                shutil.copy2(str(audio_path), str(ref_out))
+                metrics = _spectral_metrics(ref_wav, seq_cfg.sampling_rate)
+                ref_clips.append({
+                    "clip":             npz_path.stem,
+                    "wav_path":         str(ref_out),
+                    "spectral_metrics": metrics,
+                })
+            except Exception as e:
+                print(f"[LoRA Evaluator] Reference {npz_path.name} failed: {e}", flush=True)
+
+        # Average reference metrics across all clips
+        ref_avg = _avg_metrics([c["spectral_metrics"] for c in ref_clips])
+        print(f"[LoRA Evaluator] Reference avg — "
+              f"centroid={ref_avg.get('spectral_centroid_hz', 0):.0f}Hz  "
+              f"hf={ref_avg.get('hf_energy_ratio', 0):.3f}  "
+              f"flatness={ref_avg.get('spectral_flatness', 0):.4f}", flush=True)
+
+        # ------------------------------------------------------------------
+        # 4. Build summary skeleton
+        # ------------------------------------------------------------------
+        summary = {
+            "name":               name,
+            "started_at":         datetime.now(timezone.utc).isoformat(),
+            "completed_at":       None,
+            "data_dir":           str(data_dir),
+            "output_dir":         str(output_dir),
+            "n_clips":            len(ref_clips),
+            "steps":              steps,
+            "seed":               seed,
+            "reference_avg":      ref_avg,
+            "reference_clips":    ref_clips,
+            "adapters":           [],
+        }
+        summary_path = output_dir / "eval_summary.json"
+
+        def _write_summary():
+            summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+        _write_summary()
+
+        # ------------------------------------------------------------------
+        # 5. Per-adapter evaluation loop (all clips)
+        # ------------------------------------------------------------------
+        n_clips = len(dataset)
+        pbar = comfy.utils.ProgressBar(len(spec["adapters"]) * n_clips)
+
+        for adapter_spec in spec["adapters"]:
+            adapter_id   = adapter_spec["id"]
+            adapter_path = (adapter_spec.get("path") or "").strip()
+            safe_id      = _safe_stem(adapter_id)
+            clip_dir     = output_dir / safe_id
+            clip_dir.mkdir(exist_ok=True)
+
+            record = {
+                "id":               adapter_id,
+                "path":             adapter_path or None,
+                "meta":             None,
+                "clips":            [],
+                "avg_metrics":      None,
+                "status":           "running",
+            }
+
+            print(f"\n[LoRA Evaluator] ── '{adapter_id}' ({n_clips} clips) ──", flush=True)
+
+            try:
+                with torch.inference_mode(False):
+                    generator = copy.deepcopy(model["generator"])
+
+                    if adapter_path:
+                        pt_path = Path(adapter_path)
+                        if not pt_path.is_absolute():
+                            pt_path = Path(folder_paths.base_path) / pt_path
+                        if not pt_path.exists():
+                            raise FileNotFoundError(f"Adapter not found: {pt_path}")
+
+                        ckpt = torch.load(str(pt_path), map_location="cpu",
+                                          weights_only=False)
+                        if isinstance(ckpt, dict) and "state_dict" in ckpt:
+                            state_dict = ckpt["state_dict"]
+                            meta       = ckpt.get("meta", {})
+                        else:
+                            state_dict = ckpt
+                            meta       = {}
+
+                        rank    = int(meta.get("rank",   16))
+                        alpha   = float(meta.get("alpha", float(rank)))
+                        target  = list(meta.get("target", ["attn.qkv"]))
+                        dropout = float(meta.get("lora_dropout", 0.0))
+                        record["meta"] = {"rank": rank, "alpha": alpha, "target": target}
+
+                        n = apply_lora(generator, rank=rank, alpha=alpha,
+                                       target_suffixes=tuple(target), dropout=dropout)
+                        if n == 0:
+                            raise RuntimeError(
+                                f"apply_lora matched 0 layers (target={target})"
+                            )
+                        load_lora(generator, state_dict)
+                        print(f"[LoRA Evaluator] Loaded {pt_path.name} "
+                              f"(rank={rank}, {n} layers)", flush=True)
+                    else:
+                        print("[LoRA Evaluator] Baseline (no LoRA)", flush=True)
+
+                    generator = generator.to(device, dtype)
+                    generator.update_seq_lengths(
+                        latent_seq_len=seq_cfg.latent_seq_len,
+                        clip_seq_len=seq_cfg.clip_seq_len,
+                        sync_seq_len=seq_cfg.sync_seq_len,
+                    )
+
+                    clip_metrics_list = []
+                    for clip_idx in range(n_clips):
+                        clip_stem = npz_files[clip_idx].stem
+                        wav, sr = _eval_sample(
+                            generator, feature_utils_orig, dataset,
+                            seq_cfg, device, dtype,
+                            num_steps=steps, seed=seed, clip_idx=clip_idx,
+                        )
+                        if wav is None:
+                            pbar.update(1)
+                            continue
+
+                        wav_path = clip_dir / f"{clip_stem}.wav"
+                        try:
+                            torchaudio.save(str(wav_path), wav, sr)
+                        except RuntimeError:
+                            import soundfile as sf
+                            sf.write(str(wav_path), wav.squeeze(0).numpy(), sr)
+
+                        metrics = _spectral_metrics(wav, sr)
+                        clip_metrics_list.append(metrics)
+                        record["clips"].append({
+                            "clip":             clip_stem,
+                            "wav_path":         str(wav_path),
+                            "spectral_metrics": metrics,
+                        })
+                        print(f"  [{clip_idx+1}/{n_clips}] {clip_stem}  "
+                              f"centroid={metrics['spectral_centroid_hz']:.0f}Hz  "
+                              f"hf={metrics['hf_energy_ratio']:.3f}", flush=True)
+                        pbar.update(1)
+
+                record["avg_metrics"] = _avg_metrics(clip_metrics_list)
+                record["status"] = "completed"
+                avg = record["avg_metrics"]
+                print(f"[LoRA Evaluator] '{adapter_id}' avg — "
+                      f"centroid={avg.get('spectral_centroid_hz', 0):.0f}Hz  "
+                      f"hf={avg.get('hf_energy_ratio', 0):.3f}  "
+                      f"flatness={avg.get('spectral_flatness', 0):.4f}", flush=True)
+
+            except Exception as e:
+                record["status"] = "failed"
+                record["error"]  = str(e)
+                print(f"[LoRA Evaluator] '{adapter_id}' failed: {e}", flush=True)
+                traceback.print_exc()
+                pbar.update(n_clips - len(record["clips"]))
+
+            finally:
+                try:
+                    del generator
+                except NameError:
+                    pass
+                soft_empty_cache()
+
+            summary["adapters"].append(record)
+            _write_summary()
+
+        # ------------------------------------------------------------------
+        # 5. Finalise summary
+        # ------------------------------------------------------------------
+        summary["completed_at"] = datetime.now(timezone.utc).isoformat()
+        _write_summary()
+        print(f"\n[LoRA Evaluator] Done. Summary: {summary_path}", flush=True)
+
+        # ------------------------------------------------------------------
+        # 6. Comparison chart
+        # ------------------------------------------------------------------
+        completed = [r for r in summary["adapters"] if r.get("status") == "completed"]
+        if completed:
+            ids          = ["reference"] + [r["id"] for r in completed]
+            metrics_list = [summary["reference_avg"]] + [r["avg_metrics"] for r in completed]
+            chart_path   = output_dir / "metric_comparison.png"
+            comparison   = _draw_metric_comparison(ids, metrics_list, chart_path)
+            print(f"[LoRA Evaluator] Comparison chart: {chart_path}", flush=True)
+        else:
+            from PIL import Image
+            comparison = _pil_to_tensor(Image.new("RGB", (400, 200), (255, 255, 255)))
+
+        return (str(summary_path), comparison)

nodes/selva_lora_scheduler.py

@@ -0,0 +1,532 @@
+"""SelVA LoRA Scheduler — runs a sweep of training experiments from a JSON file.
+
+Each experiment inherits from a shared `base` config and overrides specific keys.
+The dataset is loaded once and reused across all experiments. Results are written
+to `experiment_summary.json` (updated after each completed run) and a comparison
+loss-curve image showing all runs on the same axes.
+
+JSON format:
+    {
+      "name": "tier1_sweep",
+      "description": "optional human note",
+      "data_dir": "dataset/dog_bark",
+      "output_root": "lora_output/tier1_sweep",
+      "base": { "rank": 16, "lr": 1e-4, "steps": 2000, ... },
+      "experiments": [
+        {"id": "baseline", "description": "..."},
+        {"id": "lora_plus_16", "lora_plus_ratio": 16.0},
+        ...
+      ]
+    }
+"""
+
+import copy
+import json
+import sys
+import time
+import traceback
+from datetime import datetime, timezone
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image, ImageDraw
+
+import comfy.utils
+import folder_paths
+
+from .utils import SELVA_CATEGORY, get_device
+from .selva_lora_trainer import (
+    SelvaLoraTrainer,
+    SkipExperiment,
+    _prepare_dataset,
+    _smooth_losses,
+    _pil_to_tensor,
+)
+
+
+def _get_system_info() -> dict:
+    """Collect GPU / torch version info for the summary header."""
+    info: dict = {
+        "torch_version": torch.__version__,
+        "cuda_version":  torch.version.cuda or "N/A",
+        "gpu_name":      None,
+        "gpu_vram_gb":   None,
+    }
+    if torch.cuda.is_available():
+        try:
+            info["gpu_name"]    = torch.cuda.get_device_name(0)
+            props               = torch.cuda.get_device_properties(0)
+            info["gpu_vram_gb"] = round(props.total_memory / 1e9, 1)
+        except Exception:
+            pass
+    return info
+
+
+# Defaults mirror SelvaLoraTrainer INPUT_TYPES defaults
+_PARAM_DEFAULTS = {
+    "alpha":               0.0,
+    "target":              "attn.qkv",
+    "batch_size":          4,
+    "warmup_steps":        100,
+    "grad_accum":          1,
+    "save_every":          500,
+    "resume_path":         "",
+    "seed":                42,
+    "timestep_mode":       "uniform",
+    "logit_normal_sigma":  1.0,
+    "curriculum_switch":   0.6,
+    "lora_dropout":        0.0,
+    "lora_plus_ratio":     1.0,
+    "lr_schedule":         "constant",
+}
+
+# Palette for comparison chart: one color per experiment (cycles if > 8)
+_PALETTE = [
+    (66,  133, 244),   # blue
+    (234,  67,  53),   # red
+    (52,  168,  83),   # green
+    (251, 188,   5),   # yellow
+    (155,  89, 182),   # purple
+    (26,  188, 156),   # teal
+    (230, 126,  34),   # orange
+    (149, 165, 166),   # grey
+]
+
+
+def _resolve_path(raw: str) -> Path:
+    """Resolve path the same way SelvaLoraTrainer does (relative → ComfyUI output dir)."""
+    p = Path(raw.strip())
+    unix_style_on_windows = (
+        sys.platform == "win32" and p.is_absolute() and not p.drive
+    )
+    if not p.is_absolute() or unix_style_on_windows:
+        p = Path(folder_paths.get_output_directory()) / p.relative_to(p.anchor)
+    return p
+
+
+def _merge_config(base: dict, experiment: dict) -> dict:
+    """Merge base defaults + file base + experiment overrides."""
+    cfg = dict(_PARAM_DEFAULTS)
+    cfg.update(base)
+    # Don't carry id/description into the training params
+    cfg.update({k: v for k, v in experiment.items() if k not in ("id", "description")})
+    return cfg
+
+
+def _loss_at_steps(loss_history: list, log_interval: int, save_every: int,
+                   start_step: int, total_steps: int) -> dict:
+    """Build a dict of {step: loss} at each save_every boundary.
+
+    loss_history[i] = average loss over steps [start + i*log_interval + 1 …
+                                                 start + (i+1)*log_interval].
+    """
+    result = {}
+    targets = range(save_every, total_steps + 1, save_every)
+    for target in targets:
+        # index of the loss entry nearest to this step
+        idx = (target - start_step) // log_interval - 1
+        if 0 <= idx < len(loss_history):
+            result[str(target)] = round(loss_history[idx], 6)
+    return result
+
+
+def _draw_comparison_curves(
+    experiments_data: list,   # list of dicts: {id, loss_history, log_interval, start_step}
+) -> Image.Image:
+    """Draw all smoothed loss curves on the same axes, one color per experiment."""
+    W, H = 900, 420
+    pl, pr, pt, pb = 75, 160, 30, 50  # wider right margin for legend
+
+    img  = Image.new("RGB", (W, H), (255, 255, 255))
+    draw = ImageDraw.Draw(img)
+
+    pw = W - pl - pr
+    ph = H - pt - pb
+
+    # Collect all smoothed series
+    series = []
+    for i, ed in enumerate(experiments_data):
+        lh = ed.get("loss_history") or []
+        if len(lh) < 2:
+            continue
+        sm = _smooth_losses(lh)
+        series.append({
+            "id":         ed["id"],
+            "smoothed":   sm,
+            "log_interval": ed.get("log_interval", 50),
+            "start_step": ed.get("start_step", 0),
+            "color":      _PALETTE[i % len(_PALETTE)],
+        })
+
+    if not series:
+        draw.text((pl + 10, pt + 10), "No data to plot", fill=(80, 80, 80))
+        return img
+
+    all_vals = [v for s in series for v in s["smoothed"]]
+    lo, hi = min(all_vals), max(all_vals)
+    if hi == lo:
+        hi = lo + 1e-6
+    rng = hi - lo
+
+    # Horizontal grid + y-axis labels
+    for i in range(5):
+        y   = pt + int(i * ph / 4)
+        val = hi - i * rng / 4
+        draw.line([(pl, y), (W - pr, y)], fill=(220, 220, 220), width=1)
+        draw.text((2, y - 7), f"{val:.4f}", fill=(100, 100, 100))
+
+    # Draw each curve
+    for s in series:
+        n    = len(s["smoothed"])
+        pts  = []
+        for j, v in enumerate(s["smoothed"]):
+            x = pl + int(j * pw / max(n - 1, 1))
+            y = pt + int((1.0 - (v - lo) / rng) * ph)
+            pts.append((x, y))
+        draw.line(pts, fill=s["color"], width=2)
+
+    # Axes
+    draw.line([(pl, pt), (pl, H - pb)],         fill=(40, 40, 40), width=1)
+    draw.line([(pl, H - pb), (W - pr, H - pb)], fill=(40, 40, 40), width=1)
+    draw.text((pl + 4, 8), "Loss comparison (smoothed)", fill=(40, 40, 40))
+
+    # Legend (right side)
+    lx = W - pr + 10
+    ly = pt
+    for s in series:
+        draw.rectangle([(lx, ly + 3), (lx + 14, ly + 13)], fill=s["color"])
+        draw.text((lx + 18, ly), s["id"][:20], fill=(40, 40, 40))
+        ly += 20
+
+    return img
+
+
+class SelvaLoraScheduler:
+    """Runs a sweep of LoRA training experiments defined in a JSON file.
+
+    The dataset (VAE encoding + .npz loading) is performed once and shared
+    across all experiments. Each experiment deep-copies the generator and trains
+    independently. Results are written to `experiment_summary.json` after every
+    completed run so partial results are preserved if the sweep is interrupted.
+    """
+
+    OUTPUT_NODE = True
+    CATEGORY    = SELVA_CATEGORY
+    FUNCTION    = "run"
+    RETURN_TYPES  = ("STRING", "IMAGE")
+    RETURN_NAMES  = ("summary_path", "comparison_curves")
+    OUTPUT_TOOLTIPS = (
+        "Path to experiment_summary.json — share this file to compare runs.",
+        "All smoothed loss curves overlaid on the same axes.",
+    )
+    DESCRIPTION = (
+        "Runs a series of LoRA training experiments defined in a JSON sweep file. "
+        "The dataset is encoded once and reused across all experiments. "
+        "Results (loss, config, adapter paths) are collected in experiment_summary.json."
+    )
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("SELVA_MODEL",),
+                "experiments_file": ("STRING", {
+                    "default": "experiments.json",
+                    "tooltip": (
+                        "Path to JSON sweep file. Relative paths resolve to the ComfyUI "
+                        "models directory; absolute paths are used as-is. "
+                        "See LORA_TRAINING.md for the file format."
+                    ),
+                }),
+            }
+        }
+
+    def run(self, model, experiments_file):
+        # ------------------------------------------------------------------
+        # 1. Read + validate the JSON file
+        # ------------------------------------------------------------------
+        exp_path = Path(experiments_file.strip())
+        if not exp_path.is_absolute():
+            # Try relative to ComfyUI models dir first, then output dir
+            candidate = Path(folder_paths.models_dir) / exp_path
+            if not candidate.exists():
+                candidate = Path(folder_paths.get_output_directory()) / exp_path
+            exp_path = candidate
+        if not exp_path.exists():
+            raise FileNotFoundError(
+                f"[LoRA Scheduler] Experiment file not found: {exp_path}"
+            )
+
+        spec = json.loads(exp_path.read_text(encoding="utf-8"))
+
+        if "experiments" not in spec or not spec["experiments"]:
+            raise ValueError("[LoRA Scheduler] 'experiments' list is missing or empty.")
+        for i, exp in enumerate(spec["experiments"]):
+            if "id" not in exp:
+                raise ValueError(
+                    f"[LoRA Scheduler] Experiment at index {i} is missing required 'id' field."
+                )
+
+        sweep_name  = spec.get("name", exp_path.stem)
+        description = spec.get("description", "")
+        base_cfg    = spec.get("base", {})
+
+        # ------------------------------------------------------------------
+        # 2. Resolve data_dir and output_root
+        # ------------------------------------------------------------------
+        if "data_dir" not in spec:
+            raise ValueError("[LoRA Scheduler] 'data_dir' is required in the sweep file.")
+        data_dir    = _resolve_path(spec["data_dir"])
+        output_root = _resolve_path(spec.get("output_root", f"lora_sweeps/{sweep_name}"))
+        output_root.mkdir(parents=True, exist_ok=True)
+
+        device = get_device()
+        dtype  = model["dtype"]
+
+        print(f"\n[LoRA Scheduler] Sweep '{sweep_name}': "
+              f"{len(spec['experiments'])} experiment(s)", flush=True)
+        if description:
+            print(f"[LoRA Scheduler] {description}", flush=True)
+        print(f"[LoRA Scheduler] data_dir    = {data_dir}", flush=True)
+        print(f"[LoRA Scheduler] output_root = {output_root}\n", flush=True)
+
+        # ------------------------------------------------------------------
+        # 3. Load + encode dataset once
+        # ------------------------------------------------------------------
+        n_clips = len(list(data_dir.glob("*.npz")))
+        dataset = _prepare_dataset(model, data_dir, device)
+
+        # ------------------------------------------------------------------
+        # 4. Build or restore the summary (resume-aware)
+        # ------------------------------------------------------------------
+        summary_path   = output_root / "experiment_summary.json"
+        completed_ids  = set()
+        all_curve_data = []   # collected for comparison image
+
+        if summary_path.exists():
+            try:
+                existing = json.loads(summary_path.read_text(encoding="utf-8"))
+                for rec in existing.get("experiments", []):
+                    if rec.get("results", {}).get("status") == "completed":
+                        completed_ids.add(rec["id"])
+                        lh = rec["results"].get("loss_history", [])
+                        all_curve_data.append({
+                            "id":           rec["id"],
+                            "loss_history": lh,
+                            "log_interval": rec["results"].get("log_interval", 50),
+                            "start_step":   0,
+                        })
+                # Restore the original summary, clear completed_at so it gets set again
+                summary = existing
+                summary["completed_at"] = None
+                if completed_ids:
+                    print(f"[LoRA Scheduler] Resuming — skipping {len(completed_ids)} "
+                          f"completed experiment(s): {sorted(completed_ids)}", flush=True)
+            except Exception as e:
+                print(f"[LoRA Scheduler] Could not read existing summary ({e}) — starting fresh",
+                      flush=True)
+                completed_ids = set()
+                all_curve_data = []
+                summary = None
+
+        if not completed_ids:
+            summary = {
+                "sweep_name":   sweep_name,
+                "description":  description,
+                "sweep_file":   str(exp_path),
+                "started_at":   datetime.now(timezone.utc).isoformat(),
+                "completed_at": None,
+                "system":       _get_system_info(),
+                "data_dir":     str(data_dir),
+                "n_clips":      n_clips,
+                "experiments":  [],
+            }
+
+        def _write_summary():
+            summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+        _write_summary()
+
+        # ------------------------------------------------------------------
+        # 5. Run each experiment
+        # ------------------------------------------------------------------
+        trainer      = SelvaLoraTrainer()
+        pbar_outer   = comfy.utils.ProgressBar(len(spec["experiments"]))
+        log_interval = 50   # matches _train_inner
+
+        feature_utils_orig = model["feature_utils"]
+        seq_cfg            = model["seq_cfg"]
+        variant            = model["variant"]
+        mode               = model["mode"]
+
+        for exp in spec["experiments"]:
+            exp_id  = exp["id"]
+            exp_desc = exp.get("description", "")
+
+            if exp_id in completed_ids:
+                print(f"[LoRA Scheduler] Skipping '{exp_id}' (already completed)", flush=True)
+                pbar_outer.update(1)
+                continue
+            cfg     = _merge_config(base_cfg, exp)
+
+            # Required training params
+            steps       = int(cfg.get("steps",       2000))
+            rank        = int(cfg.get("rank",        16))
+            lr          = float(cfg.get("lr",         1e-4))
+            alpha       = float(cfg.get("alpha",      0.0))
+            target      = str(cfg.get("target",       "attn.qkv"))
+            batch_size  = int(cfg.get("batch_size",   4))
+            warmup      = int(cfg.get("warmup_steps", 100))
+            grad_accum  = int(cfg.get("grad_accum",   1))
+            save_every  = int(cfg.get("save_every",   500))
+            resume_path = str(cfg.get("resume_path",  ""))
+            seed        = int(cfg.get("seed",         42))
+            ts_mode     = str(cfg.get("timestep_mode",      "uniform"))
+            ln_sigma    = float(cfg.get("logit_normal_sigma", 1.0))
+            curr_switch = float(cfg.get("curriculum_switch",  0.6))
+            dropout     = float(cfg.get("lora_dropout",       0.0))
+            plus_ratio  = float(cfg.get("lora_plus_ratio",    1.0))
+            lr_schedule = str(cfg.get("lr_schedule",          "constant"))
+            alpha_val   = alpha if alpha > 0.0 else float(rank)
+            target_suffixes = tuple(target.strip().split())
+
+            output_dir = output_root / exp_id
+            output_dir.mkdir(parents=True, exist_ok=True)
+
+            print(f"\n[LoRA Scheduler] ── Experiment '{exp_id}' ──", flush=True)
+            if exp_desc:
+                print(f"[LoRA Scheduler] {exp_desc}", flush=True)
+
+            exp_record = {
+                "id":          exp_id,
+                "description": exp_desc,
+                "config": {
+                    "rank": rank, "alpha": alpha_val, "lr": lr, "steps": steps,
+                    "batch_size": batch_size, "warmup_steps": warmup,
+                    "grad_accum": grad_accum, "save_every": save_every,
+                    "seed": seed, "target": list(target_suffixes),
+                    "timestep_mode": ts_mode, "logit_normal_sigma": ln_sigma,
+                    "curriculum_switch": curr_switch,
+                    "lora_dropout": dropout, "lora_plus_ratio": plus_ratio,
+                    "lr_schedule": lr_schedule,
+                },
+                "results":     {"status": "running"},
+                "adapter_path": None,
+                "output_dir":   str(output_dir),
+            }
+            summary["experiments"].append(exp_record)
+            _write_summary()
+
+            t_start = time.monotonic()
+            try:
+                with torch.inference_mode(False), torch.enable_grad():
+                    r = trainer._train_inner(
+                        model, dataset, feature_utils_orig, seq_cfg,
+                        device, dtype, variant, mode,
+                        data_dir, output_dir, steps, rank, lr,
+                        alpha_val, target_suffixes, batch_size, warmup,
+                        grad_accum, save_every, resume_path, seed,
+                        ts_mode, ln_sigma, curr_switch, dropout, plus_ratio,
+                        lr_schedule,
+                    )
+
+                duration          = time.monotonic() - t_start
+                loss_history      = r["loss_history"]
+                grad_norm_history = r.get("grad_norm_history", [])
+                spectral_metrics  = r.get("spectral_metrics", {})
+                run_start_step    = r.get("start_step", 0)
+                smoothed          = _smooth_losses(loss_history) if loss_history else []
+
+                # Scalar summary metrics
+                final_loss    = round(smoothed[-1], 6) if smoothed else None
+                min_loss      = round(min(smoothed), 6) if smoothed else None
+                min_idx       = smoothed.index(min(smoothed)) if smoothed else None
+                min_loss_step = (
+                    run_start_step + (min_idx + 1) * log_interval
+                    if min_idx is not None else None
+                )
+
+                # Stability: std-dev of raw loss over last 25% of steps
+                if loss_history:
+                    quarter    = max(1, len(loss_history) // 4)
+                    last_q     = loss_history[-quarter:]
+                    loss_std_last_quarter = round(float(np.std(last_q)), 6)
+                else:
+                    loss_std_last_quarter = None
+
+                exp_record["results"] = {
+                    "status":               "completed",
+                    "final_loss":           final_loss,
+                    "min_loss":             min_loss,
+                    "min_loss_step":        min_loss_step,
+                    "loss_std_last_quarter": loss_std_last_quarter,
+                    "loss_at_steps":        _loss_at_steps(
+                        loss_history, log_interval, save_every, run_start_step, steps
+                    ),
+                    "loss_history":         [round(v, 6) for v in loss_history],
+                    "grad_norm_history":    grad_norm_history,
+                    "spectral_metrics":     {str(k): v for k, v in spectral_metrics.items()},
+                    "log_interval":         log_interval,
+                    "duration_seconds":     round(duration, 1),
+                }
+                exp_record["adapter_path"] = r["adapter_path"]
+
+                all_curve_data.append({
+                    "id":           exp_id,
+                    "loss_history": loss_history,
+                    "log_interval": log_interval,
+                    "start_step":   0,
+                })
+
+            except SkipExperiment as e:
+                duration = time.monotonic() - t_start
+                print(f"[LoRA Scheduler] Experiment '{exp_id}' skipped: {e}", flush=True)
+                partial   = getattr(e, "partial", {})
+                lh        = partial.get("loss_history", [])
+                smoothed  = _smooth_losses(lh) if lh else []
+                exp_record["results"] = {
+                    "status":            "skipped",
+                    "stopped_at_step":   partial.get("stopped_at_step"),
+                    "final_loss":        round(smoothed[-1], 6) if smoothed else None,
+                    "loss_history":      [round(v, 6) for v in lh],
+                    "grad_norm_history": partial.get("grad_norm_history", []),
+                    "spectral_metrics":  {str(k): v for k, v in partial.get("spectral_metrics", {}).items()},
+                    "duration_seconds":  round(duration, 1),
+                }
+                _write_summary()
+                pbar_outer.update(1)
+                continue
+
+            except Exception as e:
+                duration = time.monotonic() - t_start
+                print(f"[LoRA Scheduler] Experiment '{exp_id}' failed: {e}", flush=True)
+                traceback.print_exc()
+                exp_record["results"] = {
+                    "status":           "failed",
+                    "error":            str(e),
+                    "duration_seconds": round(duration, 1),
+                }
+                _write_summary()
+                pbar_outer.update(1)
+                # Continue to next experiment rather than aborting the whole sweep
+                continue
+
+            _write_summary()
+            pbar_outer.update(1)
+
+        # ------------------------------------------------------------------
+        # 6. Finalise summary
+        # ------------------------------------------------------------------
+        summary["completed_at"] = datetime.now(timezone.utc).isoformat()
+        _write_summary()
+        print(f"\n[LoRA Scheduler] Sweep complete. Summary: {summary_path}", flush=True)
+
+        # ------------------------------------------------------------------
+        # 7. Comparison image
+        # ------------------------------------------------------------------
+        comparison_img = _draw_comparison_curves(all_curve_data)
+        comparison_img.save(str(output_root / "loss_comparison.png"))
+        comparison_tensor = _pil_to_tensor(comparison_img)
+
+        return (str(summary_path), comparison_tensor)

nodes/selva_lora_trainer.py

@@ -1,9 +1,14 @@
 import copy
 import json
+import math
 import random
 import traceback
 from pathlib import Path
 
+
+class SkipExperiment(Exception):
+    """Raised when skip_current.flag is found — signals the scheduler to move to the next experiment."""
+
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -88,21 +93,23 @@ def _load_npz(path: Path) -> dict:
 # ---------------------------------------------------------------------------
 
 def _eval_sample(generator, feature_utils_orig, dataset, seq_cfg, device, dtype,
-                 num_steps: int = 8):
+                 num_steps: int = 25, seed: int = 42, clip_idx: int = 0):
-    """Run a quick no-CFG inference pass on a random training clip.
+    """Run a quick no-CFG inference pass on a training clip.
 
+    Uses dataset[clip_idx] and a fixed noise seed so samples across checkpoints
+    are directly comparable — you can hear the model improve step by step.
     Returns (waveform [1, L] float32 cpu, sample_rate) or (None, None) on failure.
-    Uses fewer ODE steps than inference (8 vs 25) for speed.
     """
     generator.eval()
     try:
-        _, clip_f_cpu, sync_f_cpu, text_clip_cpu = random.choice(dataset)
+        _, clip_f_cpu, sync_f_cpu, text_clip_cpu = dataset[clip_idx]
         clip_f    = clip_f_cpu.to(device, dtype)
         sync_f    = sync_f_cpu.to(device, dtype)
         text_clip = text_clip_cpu.to(device, dtype)
 
+        rng = torch.Generator(device=device).manual_seed(seed)
         x0 = torch.randn(1, seq_cfg.latent_seq_len, generator.latent_dim,
-                         device=device, dtype=dtype)
+                         device=device, dtype=dtype, generator=rng)
 
         eval_fm = FlowMatching(min_sigma=0, inference_mode="euler", num_steps=num_steps)
 
@@ -131,8 +138,12 @@ def _eval_sample(generator, feature_utils_orig, dataset, seq_cfg, device, dtype,
         elif audio.dim() == 3 and audio.shape[1] != 1:
             audio = audio.mean(dim=1, keepdim=True)
 
+        target_rms = 10 ** (-27.0 / 20.0)   # -27 dBFS matches measured RMS of training clips
+        rms = audio.pow(2).mean().sqrt().clamp(min=1e-8)
+        audio = audio * (target_rms / rms)
         peak = audio.abs().max().clamp(min=1e-8)
-        audio = (audio / peak).clamp(-1, 1)
+        if peak > 1.0:
+            audio = audio / peak
         return audio.squeeze(0), seq_cfg.sampling_rate   # [1, L]
 
     except Exception as e:
@@ -142,6 +153,122 @@ def _eval_sample(generator, feature_utils_orig, dataset, seq_cfg, device, dtype,
         generator.train()
 
 
+# ---------------------------------------------------------------------------
+# Eval spectrogram rendering
+# ---------------------------------------------------------------------------
+
+_SPEC_N_FFT    = 2048
+_SPEC_HOP      = 512
+_SPEC_DB_FLOOR = -80.0
+_SPEC_LOG_BINS = 256
+
+
+def _spectral_metrics(wav: torch.Tensor, sr: int) -> dict:
+    """Compute spectral quality metrics for a mono [1, L] float32 CPU tensor.
+
+    Returns:
+      hf_energy_ratio   — energy above 4kHz / total energy  (low bitrate → low value)
+      spectral_centroid_hz — energy-weighted mean frequency
+      spectral_rolloff_hz  — frequency below which 85% of energy sits
+    """
+    import numpy as np
+    wav_np  = wav.squeeze(0).numpy()
+    hop     = min(_SPEC_HOP, _SPEC_N_FFT)
+    window  = torch.hann_window(_SPEC_N_FFT)
+    stft    = torch.stft(torch.from_numpy(wav_np), n_fft=_SPEC_N_FFT, hop_length=hop,
+                         window=window, return_complex=True)
+    power   = stft.abs().pow(2).mean(dim=1).numpy()   # [n_freqs] averaged over time
+    freqs   = np.linspace(0, sr / 2, len(power))
+
+    total   = power.sum() + 1e-12
+    hf_mask = freqs >= 4000
+    hf_ratio = float(power[hf_mask].sum() / total)
+
+    centroid = float((freqs * power).sum() / total)
+
+    cumsum  = np.cumsum(power)
+    rolloff_idx = np.searchsorted(cumsum, 0.85 * cumsum[-1])
+    rolloff = float(freqs[min(rolloff_idx, len(freqs) - 1)])
+
+    # Spectral flatness (Wiener entropy): geometric_mean / arithmetic_mean of power
+    # 0.0 = pure tone, 1.0 = white noise — rising value = noise contamination
+    log_power   = np.log(power + 1e-12)
+    flatness    = float(np.exp(log_power.mean()) / (power.mean() + 1e-12))
+
+    # Temporal energy variance — how dynamic the audio is
+    # Compute RMS per frame, take std. Low value = compressed/lifeless
+    hop       = min(_SPEC_HOP, _SPEC_N_FFT)
+    window    = torch.hann_window(_SPEC_N_FFT)
+    stft_full = torch.stft(torch.from_numpy(wav_np), n_fft=_SPEC_N_FFT, hop_length=hop,
+                           window=window, return_complex=True)
+    frame_rms = stft_full.abs().pow(2).mean(dim=0).sqrt().numpy()   # [n_frames]
+    temporal_variance = float(frame_rms.std() / (frame_rms.mean() + 1e-12))
+
+    return {
+        "hf_energy_ratio":      round(hf_ratio, 4),
+        "spectral_centroid_hz": round(centroid, 1),
+        "spectral_rolloff_hz":  round(rolloff, 1),
+        "spectral_flatness":    round(flatness, 4),
+        "temporal_variance":    round(temporal_variance, 4),
+    }
+
+
+def _save_spectrogram(wav: torch.Tensor, sr: int, path: Path) -> None:
+    """Save a log-frequency dB spectrogram PNG for an eval sample.
+
+    wav: [1, L] float32 CPU tensor (mono).
+    """
+    import numpy as np
+    from matplotlib.figure import Figure
+    from matplotlib.backends.backend_agg import FigureCanvasAgg
+
+    wav_np = wav.squeeze(0).numpy()
+    hop    = min(_SPEC_HOP, _SPEC_N_FFT)
+    window = torch.hann_window(_SPEC_N_FFT)
+    stft   = torch.stft(torch.from_numpy(wav_np), n_fft=_SPEC_N_FFT, hop_length=hop,
+                        window=window, return_complex=True)
+    mag    = stft.abs().numpy()
+    db     = 20.0 * np.log10(np.maximum(mag, 1e-8))
+    db     = np.maximum(db, db.max() + _SPEC_DB_FLOOR).astype(np.float32)
+
+    # Log-frequency resampling
+    n_freqs = db.shape[0]
+    src_idx = np.logspace(0, np.log10(max(n_freqs - 1, 2)), _SPEC_LOG_BINS)
+    lo   = np.floor(src_idx).astype(int).clip(0, n_freqs - 2)
+    frac = (src_idx - lo)[:, None]
+    spec = ((1 - frac) * db[lo] + frac * db[lo + 1]).astype(np.float32)
+    spec = spec[::-1]   # low freq at bottom
+
+    # Y-tick positions (Hz labels)
+    tgt_hz = [100, 500, 1000, 2000, 4000, 8000, 16000]
+    tpos, tlbl = [], []
+    for hz in tgt_hz:
+        bin_f = hz * _SPEC_N_FFT / sr
+        if bin_f < 1 or bin_f >= n_freqs:
+            continue
+        pos = int(np.searchsorted(src_idx, bin_f))
+        tpos.append(_SPEC_LOG_BINS - 1 - min(pos, _SPEC_LOG_BINS - 1))
+        tlbl.append(f"{hz // 1000}k" if hz >= 1000 else str(hz))
+
+    vmin = float(np.percentile(spec, 2.0))
+    vmax = float(np.percentile(spec, 99.5))
+
+    fig = Figure(figsize=(12, 3), dpi=120, tight_layout=True)
+    ax  = fig.add_subplot(1, 1, 1)
+    im  = ax.imshow(spec, aspect="auto", cmap="inferno", origin="upper",
+                    vmin=vmin, vmax=vmax, interpolation="antialiased")
+    ax.set_yticks(tpos)
+    ax.set_yticklabels(tlbl, fontsize=8)
+    ax.set_ylabel("Hz", fontsize=9)
+    ax.set_xlabel("Time frames", fontsize=9)
+    ax.set_title(path.stem, fontsize=9)
+    fig.colorbar(im, ax=ax, label="dB", fraction=0.02, pad=0.01)
+
+    canvas = FigureCanvasAgg(fig)
+    canvas.draw()
+    canvas.print_figure(str(path.with_suffix(".png")), dpi=120)
+
+
 # ---------------------------------------------------------------------------
 # Loss curve rendering
 # ---------------------------------------------------------------------------
@@ -220,115 +347,17 @@ def _pil_to_tensor(img: Image.Image) -> torch.Tensor:
     return torch.from_numpy(arr).unsqueeze(0)
 
 
-# ---------------------------------------------------------------------------
+def _prepare_dataset(model: dict, data_dir: Path, device) -> list:
-# Node
+    """Load VAE, encode audio clips, load .npz features.
-# ---------------------------------------------------------------------------
 
-class SelvaLoraTrainer:
+    Returns a list of (latents, clip_features, sync_features, text_clip) CPU tensors.
-    @classmethod
+    The VAE is freed after encoding. Call this once and reuse the dataset across
-    def INPUT_TYPES(cls):
+    multiple training jobs (e.g. in the scheduler).
-        return {
+    """
-            "required": {
-                "model":      ("SELVA_MODEL",),
-                "data_dir":   ("STRING", {
-                    "default": "",
-                    "tooltip": "Directory containing .npz feature files and paired audio files.",
-                }),
-                "output_dir": ("STRING", {
-                    "default": "lora_output",
-                    "tooltip": "Where to save adapter checkpoints.",
-                }),
-                "steps": ("INT", {
-                    "default": 2000, "min": 100, "max": 100000,
-                    "tooltip": "Total training steps.",
-                }),
-                "rank": ("INT", {
-                    "default": 16, "min": 1, "max": 128,
-                    "tooltip": "LoRA rank. Higher = more capacity, more VRAM. 16 is a safe default.",
-                }),
-                "lr": ("FLOAT", {
-                    "default": 1e-4, "min": 1e-6, "max": 1e-2, "step": 1e-6,
-                    "tooltip": "Learning rate.",
-                }),
-            },
-            "optional": {
-                "alpha": ("FLOAT", {
-                    "default": 0.0, "min": 0.0, "max": 256.0, "step": 0.5,
-                    "tooltip": "LoRA alpha. 0 = use rank value (scale = 1.0).",
-                }),
-                "target": ("STRING", {
-                    "default": "attn.qkv",
-                    "tooltip": "Space-separated layer name suffixes to wrap. Default targets all QKV projections. Add 'linear1' for post-attention projections.",
-                }),
-                "batch_size":   ("INT", {"default": 4,   "min": 1, "max": 32,
-                                         "tooltip": "Number of clips per training step. Higher = more stable gradients, more VRAM."}),
-                "warmup_steps": ("INT", {"default": 100, "min": 0, "max": 5000}),
-                "grad_accum":   ("INT", {"default": 1,   "min": 1, "max": 32,
-                                         "tooltip": "Gradient accumulation steps. Usually 1 when batch_size > 1."}),
-                "save_every":   ("INT", {"default": 500, "min": 50, "max": 10000}),
-                "resume_path":  ("STRING", {
-                    "default": "",
-                    "tooltip": "Path to a step checkpoint (.pt) to resume training from.",
-                }),
-                "seed": ("INT", {"default": 42}),
-            },
-        }
-
-    RETURN_TYPES  = ("SELVA_MODEL", "STRING", "IMAGE")
-    RETURN_NAMES  = ("model", "adapter_path", "loss_curve")
-    OUTPUT_TOOLTIPS = (
-        "Model with trained LoRA adapter applied — connect directly to Sampler.",
-        "Path to adapter_final.pt — use with SelVA LoRA Loader in future sessions.",
-        "Training loss curve.",
-    )
-    FUNCTION  = "train"
-    CATEGORY  = SELVA_CATEGORY
-    DESCRIPTION = (
-        "Trains a LoRA adapter on a dataset of .npz feature files + paired audio files. "
-        "Blocks the queue for the duration of training. "
-        "Prepare the dataset with SelVA Feature Extractor (set a name to get numbered .npz files) "
-        "and pair each .npz with a clean audio file of the same stem."
-    )
-
-    def train(self, model, data_dir, output_dir, steps, rank, lr,
-              alpha=0.0, target="attn.qkv", batch_size=4, warmup_steps=100,
-              grad_accum=1, save_every=500, resume_path="", seed=42):
-
-        torch.manual_seed(seed)
-        random.seed(seed)
-
-        device   = get_device()
-        dtype    = model["dtype"]
-        variant  = model["variant"]
     mode               = model["mode"]
     seq_cfg            = model["seq_cfg"]
     feature_utils_orig = model["feature_utils"]
 
-        data_dir   = Path(data_dir.strip())
-
-        _out_str = output_dir.strip()
-        _out_p   = Path(_out_str)
-        # On Windows a Unix-style path like "/lora_output" is technically absolute
-        # (drive-relative) but the user almost certainly meant a subfolder of the
-        # ComfyUI output directory. Treat any non-absolute path AND any path whose
-        # only "absolute" anchor is a leading slash (no drive letter) as relative to
-        # the ComfyUI output folder.
-        import sys as _sys
-        _unix_style_on_windows = (
-            _sys.platform == "win32"
-            and _out_p.is_absolute()
-            and not _out_p.drive  # e.g. Path("/foo").drive == "" on Windows
-        )
-        if not _out_p.is_absolute() or _unix_style_on_windows:
-            _out_p = Path(folder_paths.get_output_directory()) / _out_p.relative_to(_out_p.anchor)
-            print(f"[LoRA Trainer] output_dir resolved to: {_out_p}", flush=True)
-        output_dir = _out_p
-        output_dir.mkdir(parents=True, exist_ok=True)
-
-        alpha_val      = float(alpha) if alpha > 0.0 else float(rank)
-        target_suffixes = tuple(target.strip().split())
-
-        # --- Load VAE encoder (not present in inference model) ---
     vae_name = "v1-16.pth" if mode == "16k" else "v1-44.pth"
     vae_path = _SELVA_DIR / "ext" / vae_name
     if not vae_path.exists():
@@ -345,7 +374,6 @@ class SelvaLoraTrainer:
         need_vae_encoder=True,
     ).to(device).eval()
 
-        # --- Pre-load dataset ---
     npz_files = sorted(data_dir.glob("*.npz"))
     if not npz_files:
         raise ValueError(f"[LoRA Trainer] No .npz files found in {data_dir}")
@@ -418,18 +446,171 @@ class SelvaLoraTrainer:
         raise ValueError("[LoRA Trainer] No clips could be loaded.")
     print(f"[LoRA Trainer] {len(dataset)} clip(s) ready.", flush=True)
 
+    return dataset
+
+
+# ---------------------------------------------------------------------------
+# Node
+# ---------------------------------------------------------------------------
+
+class SelvaLoraTrainer:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model":      ("SELVA_MODEL",),
+                "data_dir":   ("STRING", {
+                    "default": "",
+                    "tooltip": "Directory containing .npz feature files and paired audio files.",
+                }),
+                "output_dir": ("STRING", {
+                    "default": "lora_output",
+                    "tooltip": "Where to save adapter checkpoints.",
+                }),
+                "steps": ("INT", {
+                    "default": 2000, "min": 100, "max": 100000,
+                    "tooltip": "Total training steps.",
+                }),
+                "rank": ("INT", {
+                    "default": 16, "min": 1, "max": 128,
+                    "tooltip": "LoRA rank. Higher = more capacity, more VRAM. 16 is a safe default.",
+                }),
+                "lr": ("FLOAT", {
+                    "default": 1e-4, "min": 1e-6, "max": 1e-2, "step": 1e-6,
+                    "tooltip": "Learning rate.",
+                }),
+            },
+            "optional": {
+                "alpha": ("FLOAT", {
+                    "default": 0.0, "min": 0.0, "max": 256.0, "step": 0.5,
+                    "tooltip": "LoRA alpha. 0 = use rank value (scale = 1.0).",
+                }),
+                "target": ("STRING", {
+                    "default": "attn.qkv",
+                    "tooltip": "Space-separated layer name suffixes to wrap. Default targets all QKV projections. Add 'linear1' for post-attention projections.",
+                }),
+                "batch_size":   ("INT", {"default": 4,   "min": 1, "max": 32,
+                                         "tooltip": "Number of clips per training step. Higher = more stable gradients, more VRAM."}),
+                "warmup_steps": ("INT", {"default": 100, "min": 0, "max": 5000}),
+                "grad_accum":   ("INT", {"default": 1,   "min": 1, "max": 32,
+                                         "tooltip": "Gradient accumulation steps. Usually 1 when batch_size > 1."}),
+                "save_every":   ("INT", {"default": 500, "min": 50, "max": 10000}),
+                "resume_path":  ("STRING", {
+                    "default": "",
+                    "tooltip": "Path to a step checkpoint (.pt) to resume training from.",
+                }),
+                "seed": ("INT", {"default": 42}),
+                "timestep_mode": (["uniform", "logit_normal", "curriculum"], {
+                    "default": "uniform",
+                    "tooltip": "How to sample training timesteps. "
+                               "uniform: all timesteps equally (matches original MMAudio). "
+                               "logit_normal: concentrates near t=0.5. "
+                               "curriculum: logit_normal for first curriculum_switch% of steps then uniform (recommended for small datasets).",
+                }),
+                "logit_normal_sigma": ("FLOAT", {
+                    "default": 1.0, "min": 0.1, "max": 3.0, "step": 0.1,
+                    "tooltip": "Spread of the logit-normal distribution. "
+                               "1.0 = moderate peak at t=0.5. Higher approaches uniform. "
+                               "Used with logit_normal and curriculum modes.",
+                }),
+                "curriculum_switch": ("FLOAT", {
+                    "default": 0.6, "min": 0.1, "max": 0.9, "step": 0.05,
+                    "tooltip": "Fraction of steps to run logit_normal before switching to uniform. "
+                               "0.6 = switch at 60% of total steps. Only used with timestep_mode=curriculum.",
+                }),
+                "lora_dropout": ("FLOAT", {
+                    "default": 0.0, "min": 0.0, "max": 0.3, "step": 0.01,
+                    "tooltip": "Dropout applied to the LoRA path only (not the frozen base weights). "
+                               "0=disabled. 0.05–0.1 helps regularize on small datasets (arXiv:2404.09610).",
+                }),
+                "lora_plus_ratio": ("FLOAT", {
+                    "default": 1.0, "min": 1.0, "max": 32.0, "step": 1.0,
+                    "tooltip": "LoRA+ LR ratio: lr_B = lr × ratio. "
+                               "1.0 = standard LoRA. 16.0 = LoRA+ (arXiv:2402.12354).",
+                }),
+                "lr_schedule": (["constant", "cosine"], {
+                    "default": "constant",
+                    "tooltip": "LR schedule after warmup. "
+                               "constant: flat LR for all steps. "
+                               "cosine: decay from lr to ~0 following a cosine curve — "
+                               "prevents oscillation when LR is slightly too high.",
+                }),
+            },
+        }
+
+    RETURN_TYPES  = ("SELVA_MODEL", "STRING", "IMAGE")
+    RETURN_NAMES  = ("model", "adapter_path", "loss_curve")
+    OUTPUT_TOOLTIPS = (
+        "Model with trained LoRA adapter applied — connect directly to Sampler.",
+        "Path to adapter_final.pt — use with SelVA LoRA Loader in future sessions.",
+        "Training loss curve.",
+    )
+    FUNCTION  = "train"
+    CATEGORY  = SELVA_CATEGORY
+    DESCRIPTION = (
+        "Trains a LoRA adapter on a dataset of .npz feature files + paired audio files. "
+        "Blocks the queue for the duration of training. "
+        "Prepare the dataset with SelVA Feature Extractor (set a name to get numbered .npz files) "
+        "and pair each .npz with a clean audio file of the same stem."
+    )
+
+    def train(self, model, data_dir, output_dir, steps, rank, lr,
+              alpha=0.0, target="attn.qkv", batch_size=4, warmup_steps=100,
+              grad_accum=1, save_every=500, resume_path="", seed=42,
+              timestep_mode="uniform", logit_normal_sigma=1.0, curriculum_switch=0.6,
+              lora_dropout=0.0, lora_plus_ratio=1.0, lr_schedule="constant"):
+
+        torch.manual_seed(seed)
+        random.seed(seed)
+
+        device   = get_device()
+        dtype    = model["dtype"]
+        variant  = model["variant"]
+        mode     = model["mode"]
+        seq_cfg  = model["seq_cfg"]
+        feature_utils_orig = model["feature_utils"]
+
+        data_dir   = Path(data_dir.strip())
+
+        _out_str = output_dir.strip()
+        _out_p   = Path(_out_str)
+        # On Windows a Unix-style path like "/lora_output" is technically absolute
+        # (drive-relative) but the user almost certainly meant a subfolder of the
+        # ComfyUI output directory. Treat any non-absolute path AND any path whose
+        # only "absolute" anchor is a leading slash (no drive letter) as relative to
+        # the ComfyUI output folder.
+        import sys as _sys
+        _unix_style_on_windows = (
+            _sys.platform == "win32"
+            and _out_p.is_absolute()
+            and not _out_p.drive  # e.g. Path("/foo").drive == "" on Windows
+        )
+        if not _out_p.is_absolute() or _unix_style_on_windows:
+            _out_p = Path(folder_paths.get_output_directory()) / _out_p.relative_to(_out_p.anchor)
+            print(f"[LoRA Trainer] output_dir resolved to: {_out_p}", flush=True)
+        output_dir = _out_p
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        alpha_val      = float(alpha) if alpha > 0.0 else float(rank)
+        target_suffixes = tuple(target.strip().split())
+
+        dataset = _prepare_dataset(model, data_dir, device)
+
         # ComfyUI executes nodes inside torch.inference_mode(). Inference tensors
         # can't participate in autograd even with enable_grad — disable inference
         # mode entirely so deepcopy, apply_lora, and the training loop all run
         # with a clean autograd context.
         with torch.inference_mode(False), torch.enable_grad():
-            return self._train_inner(
+            r = self._train_inner(
                 model, dataset, feature_utils_orig, seq_cfg,
                 device, dtype, variant, mode,
                 data_dir, output_dir, steps, rank, lr,
                 alpha_val, target_suffixes, batch_size, warmup_steps,
                 grad_accum, save_every, resume_path, seed,
+                timestep_mode, logit_normal_sigma, curriculum_switch,
+                lora_dropout, lora_plus_ratio, lr_schedule,
             )
+            return (r["patched_model"], r["adapter_path"], r["loss_curve"])
 
     def _train_inner(
         self, model, dataset, feature_utils_orig, seq_cfg,
@@ -437,18 +618,21 @@ class SelvaLoraTrainer:
         data_dir, output_dir, steps, rank, lr,
         alpha_val, target_suffixes, batch_size, warmup_steps,
         grad_accum, save_every, resume_path, seed,
+        timestep_mode="uniform", logit_normal_sigma=1.0, curriculum_switch=0.6,
+        lora_dropout=0.0, lora_plus_ratio=1.0, lr_schedule="constant",
     ):
         # --- Prepare generator copy with LoRA ---
         generator = copy.deepcopy(model["generator"]).to(device, dtype)
 
         n_lora = apply_lora(generator, rank=rank, alpha=alpha_val,
-                            target_suffixes=target_suffixes)
+                            target_suffixes=target_suffixes, dropout=lora_dropout)
         if n_lora == 0:
             raise RuntimeError(
                 f"[LoRA Trainer] No layers matched target={target_suffixes}. "
                 "Check the 'target' field."
             )
-        print(f"[LoRA Trainer] Wrapped {n_lora} layers (rank={rank}, alpha={alpha_val})", flush=True)
+        print(f"[LoRA Trainer] Wrapped {n_lora} layers "
+              f"(rank={rank}, alpha={alpha_val}, dropout={lora_dropout})", flush=True)
 
         for name, p in generator.named_parameters():
             p.requires_grad_("lora_" in name)
@@ -460,9 +644,25 @@ class SelvaLoraTrainer:
         )
 
         # --- Optimizer + scheduler ---
-        lora_params = [p for p in generator.parameters() if p.requires_grad]
+        # LoRA+: split A and B into separate param groups with different LRs.
-        optimizer   = torch.optim.AdamW(lora_params, lr=lr, weight_decay=1e-2)
+        # ratio=1.0 = standard LoRA (same LR for both). ratio=16 = LoRA+.
+        lora_A_params = [p for n, p in generator.named_parameters() if "lora_A" in n and p.requires_grad]
+        lora_B_params = [p for n, p in generator.named_parameters() if "lora_B" in n and p.requires_grad]
+        optimizer = torch.optim.AdamW([
+            {"params": lora_A_params, "lr": lr},
+            {"params": lora_B_params, "lr": lr * lora_plus_ratio},
+        ], weight_decay=1e-2)
+        if lora_plus_ratio != 1.0:
+            print(f"[LoRA Trainer] LoRA+: lr_A={lr:.2e}  lr_B={lr * lora_plus_ratio:.2e}", flush=True)
 
+        if lr_schedule == "cosine":
+            def lr_lambda(s):
+                if s < warmup_steps:
+                    return s / max(1, warmup_steps)
+                progress = (s - warmup_steps) / max(1, steps - warmup_steps)
+                return max(1e-6 / lr, 0.5 * (1.0 + math.cos(math.pi * progress)))
+            print(f"[LoRA Trainer] LR schedule: cosine decay {lr:.2e} → 0", flush=True)
+        else:
             def lr_lambda(s):
                 return s / max(1, warmup_steps) if s < warmup_steps else 1.0
 
@@ -491,9 +691,19 @@ class SelvaLoraTrainer:
 
         log_interval = 50
         remaining    = steps - start_step
+        if remaining < log_interval:
+            raise ValueError(
+                f"[LoRA Trainer] Only {remaining} steps remaining (steps={steps}, "
+                f"start_step={start_step}). Need at least {log_interval} steps to "
+                "record any loss — increase 'steps' or lower the resume checkpoint."
+            )
         pbar_train   = comfy.utils.ProgressBar(remaining)
         loss_history      = []
         running_loss      = 0.0
+        grad_norm_history = []
+        spectral_metrics  = {}   # {step: {hf_energy_ratio, spectral_centroid_hz, spectral_rolloff_hz}}
+        running_grad_norm = 0.0
+        grad_norm_count   = 0
 
         meta = {
             "variant":             variant,
@@ -501,10 +711,21 @@ class SelvaLoraTrainer:
             "alpha":               alpha_val,
             "target":              list(target_suffixes),
             "steps":               steps,
+            "timestep_mode":       timestep_mode,
+            "logit_normal_sigma":  logit_normal_sigma,
+            "curriculum_switch":   curriculum_switch,
+            "lora_dropout":        lora_dropout,
+            "lora_plus_ratio":     lora_plus_ratio,
+            "lr_schedule":         lr_schedule,
         }
 
+        # For curriculum mode: compute the step at which we switch from logit_normal to uniform
+        curriculum_switch_step = start_step + int((steps - start_step) * curriculum_switch)
+        _curriculum_switched = False
+
         print(f"\n[LoRA Trainer] Training {remaining} steps "
-              f"(step {start_step + 1} → {steps}, batch_size={batch_size})\n", flush=True)
+              f"(step {start_step + 1} → {steps}, batch_size={batch_size}, "
+              f"timestep_mode={timestep_mode})\n", flush=True)
 
         last_step = start_step
         completed = False
@@ -518,9 +739,19 @@ class SelvaLoraTrainer:
                 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)
 
-                generator.normalize(x1)
+                x1 = generator.normalize(x1)
 
+                if timestep_mode == "logit_normal" or (
+                    timestep_mode == "curriculum" and step <= curriculum_switch_step
+                ):
+                    u = torch.randn(batch_size, device=device, dtype=dtype) * logit_normal_sigma
+                    t = torch.sigmoid(u)
+                else:
                     t = torch.rand(batch_size, device=device, dtype=dtype)
+
+                if timestep_mode == "curriculum" and step == curriculum_switch_step + 1 and not _curriculum_switched:
+                    print(f"[LoRA Trainer] Curriculum switch: logit_normal → uniform at step {step}", flush=True)
+                    _curriculum_switched = True
                 x0 = torch.randn_like(x1)
                 xt = fm.get_conditional_flow(x0, x1, t)
 
@@ -530,18 +761,46 @@ class SelvaLoraTrainer:
                 running_loss += loss.item() * grad_accum
 
                 if step % grad_accum == 0:
-                    torch.nn.utils.clip_grad_norm_(lora_params, max_norm=1.0)
+                    grad_norm = torch.nn.utils.clip_grad_norm_(
+                        lora_A_params + lora_B_params, max_norm=1.0
+                    ).item()
+                    running_grad_norm += grad_norm
+                    grad_norm_count   += 1
                     optimizer.step()
                     scheduler.step()
                     optimizer.zero_grad()
 
                 if step % log_interval == 0:
+                    skip_flag = output_dir.parent / "skip_current.flag"
+                    if skip_flag.exists():
+                        skip_flag.unlink()
+                        exc = SkipExperiment(f"skip_current.flag detected at step {step} — skipping to next experiment")
+                        exc.partial = {
+                            "loss_history":      list(loss_history),
+                            "grad_norm_history": list(grad_norm_history),
+                            "spectral_metrics":  dict(spectral_metrics),
+                            "stopped_at_step":   step,
+                        }
+                        raise exc
+
                     avg = running_loss / log_interval
                     loss_history.append(avg)
+                    # grad_norm_count can be 0 when grad_accum > log_interval
+                    # (no optimizer step fired in this interval yet)
+                    if grad_norm_count > 0:
+                        avg_gnorm = running_grad_norm / grad_norm_count
+                        grad_norm_history.append(round(avg_gnorm, 6))
+                        gnorm_str = f"  grad_norm={avg_gnorm:.4f}"
+                    else:
+                        grad_norm_history.append(None)
+                        gnorm_str = ""
                     lr_now = scheduler.get_last_lr()[0]
                     print(f"[LoRA Trainer] step {step:5d}/{steps}  "
-                          f"loss={avg:.4f}  lr={lr_now:.2e}  bs={batch_size}", flush=True)
+                          f"loss={avg:.4f}{gnorm_str}  "
+                          f"lr={lr_now:.2e}  bs={batch_size}", flush=True)
                     running_loss      = 0.0
+                    running_grad_norm = 0.0
+                    grad_norm_count   = 0
 
                     # Live preview: send updated loss curve to ComfyUI frontend
                     preview_img = _draw_loss_curve(loss_history, log_interval, start_step,
@@ -561,17 +820,30 @@ class SelvaLoraTrainer:
                     }, ckpt_path)
                     print(f"[LoRA Trainer] Saved {ckpt_path}", flush=True)
 
-                    # Save a quick eval sample next to the checkpoint
+                    # Save a quick eval sample in samples/ subfolder
+                    samples_dir = output_dir / "samples"
+                    samples_dir.mkdir(exist_ok=True)
                     wav, sr = _eval_sample(generator, feature_utils_orig,
-                                           dataset, seq_cfg, device, dtype)
+                                           dataset, seq_cfg, device, dtype, seed=seed)
                     if wav is not None:
-                        wav_path = output_dir / f"sample_step{step:05d}.wav"
+                        wav_path = samples_dir / f"step_{step:05d}.wav"
                         try:
                             torchaudio.save(str(wav_path), wav, sr)
                         except RuntimeError:
                             import soundfile as sf
                             sf.write(str(wav_path), wav.squeeze(0).numpy(), sr)
                         print(f"[LoRA Trainer] Sample saved: {wav_path}", flush=True)
+                        try:
+                            metrics = _spectral_metrics(wav, sr)
+                            spectral_metrics[step] = metrics
+                            print(f"[LoRA Trainer] Spectral: hf_ratio={metrics['hf_energy_ratio']:.3f}  "
+                                  f"centroid={metrics['spectral_centroid_hz']:.0f}Hz  "
+                                  f"rolloff={metrics['spectral_rolloff_hz']:.0f}Hz  "
+                                  f"flatness={metrics['spectral_flatness']:.3f}  "
+                                  f"temporal_var={metrics['temporal_variance']:.3f}", flush=True)
+                            _save_spectrogram(wav, sr, wav_path)
+                        except Exception as e:
+                            print(f"[LoRA Trainer] Spectral/spectrogram failed: {e}", flush=True)
 
                 last_step = step
                 pbar_train.update(1)
@@ -614,4 +886,14 @@ class SelvaLoraTrainer:
         patched = {**model, "generator": generator}
 
         loss_curve = _pil_to_tensor(smoothed_img)
-        return (patched, str(final_path), loss_curve)
+        return {
+            "patched_model":     patched,
+            "adapter_path":      str(final_path),
+            "loss_curve":        loss_curve,
+            "loss_history":      loss_history,
+            "grad_norm_history": grad_norm_history,
+            "spectral_metrics":  spectral_metrics,
+            "start_step":        start_step,
+            "meta":              meta,
+            "completed":         True,
+        }

nodes/selva_sampler.py

@@ -3,6 +3,7 @@ import comfy.utils
 import comfy.model_management
 
 from .utils import SELVA_CATEGORY, get_device, get_offload_device, soft_empty_cache
+from .selva_textual_inversion_trainer import _inject_tokens
 
 
 class SelvaSampler:
@@ -31,9 +32,31 @@ class SelvaSampler:
                 "seed":     ("INT",   {"default": 0,   "min": 0,   "max": 0xFFFFFFFF}),
             },
             "optional": {
+                "steering_vectors": ("STEERING_VECTORS", {
+                    "tooltip": "Activation steering bundle from SelVA Activation Steering Loader. "
+                               "Nudges each DiT block's hidden state toward the extracted pattern.",
+                }),
+                "steering_strength": ("FLOAT", {
+                    "default": 0.1, "min": 0.0, "max": 2.0, "step": 0.05,
+                    "tooltip": "Scale applied to each steering vector before adding to block output. "
+                               "Start around 0.1–0.3; higher values risk destabilizing the ODE.",
+                }),
                 "normalize": ("BOOLEAN", {
                     "default": True,
-                    "tooltip": "Peak-normalize output to [-1, 1]. Disable to preserve the raw decoder output level.",
+                    "tooltip": "Normalize output level. Uses RMS normalization to target_lufs rather than peak normalization, so level matches typical audio content.",
+                }),
+                "target_lufs": ("FLOAT", {
+                    "default": -27.0, "min": -40.0, "max": -6.0, "step": 1.0,
+                    "tooltip": "Target RMS level in dBFS when normalize=True. -27 matches the measured RMS of LUFS-normalized training clips. Increase toward -20 for louder output.",
+                }),
+                "textual_inversion": ("TEXTUAL_INVERSION", {
+                    "tooltip": "Learned token embeddings from SelVA Textual Inversion Loader. "
+                               "Injects style tokens into CLIP conditioning without modifying model weights.",
+                }),
+                "ti_strength": ("FLOAT", {
+                    "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05,
+                    "tooltip": "Blends between original CLIP conditioning (0.0) and full TI injection (1.0). "
+                               "Reduce toward 0.3–0.5 if TI produces buzz artifacts.",
                 }),
             },
         }
@@ -45,7 +68,7 @@ class SelvaSampler:
     CATEGORY = SELVA_CATEGORY
     DESCRIPTION = "Generates audio from video features using SelVA's flow matching ODE. Supports text prompts and negative prompts via classifier-free guidance."
 
-    def generate(self, model, features, prompt, negative_prompt, duration, steps, cfg_strength, seed, normalize=True):
+    def generate(self, model, features, prompt, negative_prompt, duration, steps, cfg_strength, seed, steering_vectors=None, steering_strength=0.1, normalize=True, target_lufs=-27.0, textual_inversion=None, ti_strength=1.0):
         import dataclasses
         from selva_core.model.flow_matching import FlowMatching
 
@@ -110,6 +133,19 @@ class SelvaSampler:
                 neg_text_clip = feature_utils.encode_text_clip([negative_prompt]) \
                     if negative_prompt.strip() else None
 
+                # Inject textual inversion tokens into CLIP conditioning
+                if textual_inversion is not None:
+                    emb         = textual_inversion["embeddings"].to(device, dtype)  # [K, 1024]
+                    K           = emb.shape[0]
+                    inject_mode = textual_inversion.get("inject_mode", "suffix")
+                    ti_text     = _inject_tokens(text_clip, emb, K, inject_mode)
+                    text_clip   = torch.lerp(text_clip, ti_text, ti_strength)
+                    if neg_text_clip is not None:
+                        ti_neg    = _inject_tokens(neg_text_clip, emb, K, inject_mode)
+                        neg_text_clip = torch.lerp(neg_text_clip, ti_neg, ti_strength)
+                    print(f"[SelVA] Textual inversion: {K} tokens  mode={inject_mode}  strength={ti_strength}",
+                          flush=True)
+
                 conditions = net_generator.preprocess_conditions(clip_f, sync_f, text_clip)
                 empty_conditions = net_generator.get_empty_conditions(
                     bs=1, negative_text_features=neg_text_clip
@@ -123,6 +159,63 @@ class SelvaSampler:
                     device=gen_device, dtype=dtype, generator=rng,
                 ).to(device)
 
+                # Activation steering: apply only during the conditional predict_flow pass
+                # so steering gets amplified by cfg_strength rather than canceling out.
+                steering_handles = []
+                _orig_predict_flow = None
+                if steering_vectors is not None and steering_strength > 0.0:
+                    vecs = steering_vectors["steering_vectors"]
+                    n_joint = steering_vectors["n_joint"]
+
+                    # Patch predict_flow to flag which pass is conditional.
+                    # ode_wrapper calls predict_flow(conditions) and predict_flow(empty_conditions);
+                    # identity check tells us which is which.
+                    _is_cond_pass = [False]
+                    _orig_predict_flow = net_generator.predict_flow
+
+                    def _tracked_predict_flow(latent, t, cond):
+                        _is_cond_pass[0] = (cond is conditions)
+                        return _orig_predict_flow(latent, t, cond)
+
+                    net_generator.predict_flow = _tracked_predict_flow
+
+                    def _make_steering_hook(vec_cpu, is_joint, strength, dev, dt):
+                        vec = vec_cpu.to(dev, dt)  # [seq, hidden]
+                        def hook(module, input, output):
+                            if not _is_cond_pass[0]:
+                                return  # skip unconditional pass; steering amplified by cfg_strength
+                            # Interpolate steering vec to match actual output seq length
+                            # (handles generation at different duration than extraction)
+                            if is_joint:
+                                out_seq = output[0].shape[1]
+                            else:
+                                out_seq = output.shape[1]
+                            v = vec
+                            if v.shape[0] != out_seq:
+                                v = torch.nn.functional.interpolate(
+                                    v.T.unsqueeze(0),   # [1, hidden, seq_orig]
+                                    size=out_seq,
+                                    mode="linear",
+                                    align_corners=False,
+                                ).squeeze(0).T          # [seq_new, hidden]
+                            if is_joint:
+                                latent_out = output[0] + strength * v
+                                return (latent_out,) + output[1:]
+                            else:
+                                return output + strength * v
+                        return hook
+
+                    blocks = list(net_generator.joint_blocks) + list(net_generator.fused_blocks)
+                    for i, block in enumerate(blocks):
+                        is_joint = i < n_joint
+                        if i < len(vecs):
+                            h = block.register_forward_hook(
+                                _make_steering_hook(vecs[i], is_joint, steering_strength, device, dtype)
+                            )
+                            steering_handles.append(h)
+                    print(f"[SelVA] Activation steering: {len(steering_handles)} blocks  "
+                          f"strength={steering_strength}  (conditional pass only)", flush=True)
+
                 # Flow matching ODE (Euler)
                 fm = FlowMatching(min_sigma=0, inference_mode="euler", num_steps=steps)
                 pbar = comfy.utils.ProgressBar(steps)
@@ -139,6 +232,11 @@ class SelvaSampler:
                         "[SelVA] CUDA out of memory during generation. Try switching offload_strategy "
                         "to 'offload_to_cpu', using a smaller variant, or reducing duration."
                     )
+                finally:
+                    if _orig_predict_flow is not None:
+                        net_generator.predict_flow = _orig_predict_flow
+                    for h in steering_handles:
+                        h.remove()
 
             print(f"[SelVA] latent stats: mean={x1.float().mean():.4f} std={x1.float().std():.4f}", flush=True)
 
@@ -168,8 +266,14 @@ class SelvaSampler:
             audio = audio.mean(dim=1, keepdim=True)  # stereo → mono
 
         if normalize:
+            target_rms = 10 ** (target_lufs / 20.0)
+            rms = audio.pow(2).mean().sqrt().clamp(min=1e-8)
+            audio = audio * (target_rms / rms)
+            # If RMS normalization pushes peaks into clipping, scale back to
+            # preserve dynamics rather than hard-clipping (no saturation)
             peak = audio.abs().max().clamp(min=1e-8)
-            audio = (audio / peak).clamp(-1, 1)
+            if peak > 1.0:
+                audio = audio / peak
         print(f"[SelVA] audio: shape={tuple(audio.shape)} sr={sample_rate}", flush=True)
 
         return ({"waveform": audio.cpu(), "sample_rate": sample_rate},)

nodes/selva_skip_experiment.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+import folder_paths
+
+from .utils import SELVA_CATEGORY
+
+
+class SelvaSkipExperiment:
+    """Writes skip_current.flag into a sweep output_root.
+
+    Queue this node while a SelVA LoRA Scheduler sweep is running to skip
+    the current experiment and move to the next one. The trainer picks up
+    the flag within 50 steps (~a few seconds).
+    """
+
+    OUTPUT_NODE = True
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "output_root": ("STRING", {
+                    "default": "",
+                    "tooltip": "output_root of the running sweep — same value as in your experiments JSON.",
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("flag_path",)
+    OUTPUT_TOOLTIPS = ("Path where the flag was written.",)
+    FUNCTION = "skip"
+    CATEGORY = SELVA_CATEGORY
+    DESCRIPTION = (
+        "Signals the running SelVA LoRA Scheduler to skip the current experiment "
+        "and move to the next one. Queue this node while the scheduler is running. "
+        "Partial scalars collected so far are saved in the summary."
+    )
+
+    def skip(self, output_root: str):
+        p = Path(output_root.strip())
+        if not p.is_absolute():
+            p = Path(folder_paths.get_output_directory()) / p
+        if not p.exists():
+            raise FileNotFoundError(f"[SelVA Skip] output_root not found: {p}")
+
+        flag = p / "skip_current.flag"
+        flag.touch()
+        print(f"[SelVA Skip] Flag written: {flag}", flush=True)
+        return (str(flag),)

nodes/selva_textual_inversion_loader.py

@@ -0,0 +1,70 @@
+"""SelVA Textual Inversion Loader.
+
+Loads a .pt file produced by SelvaTextualInversionTrainer and returns a
+TEXTUAL_INVERSION bundle that the SelVA Sampler can inject into text conditioning.
+"""
+
+from pathlib import Path
+
+import torch
+import folder_paths
+
+from .utils import SELVA_CATEGORY
+
+
+class SelvaTextualInversionLoader:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "path": ("STRING", {
+                    "default": "textual_inversion.pt",
+                    "tooltip": "Path to a .pt file produced by SelVA Textual Inversion Trainer. "
+                               "Relative paths resolve to the ComfyUI output directory.",
+                }),
+            },
+        }
+
+    RETURN_TYPES  = ("TEXTUAL_INVERSION",)
+    RETURN_NAMES  = ("textual_inversion",)
+    OUTPUT_TOOLTIPS = ("Learned token embeddings — connect to SelVA Sampler's textual_inversion input.",)
+    FUNCTION  = "load"
+    CATEGORY  = SELVA_CATEGORY
+    DESCRIPTION = (
+        "Loads learned CLIP token embeddings produced by SelVA Textual Inversion Trainer. "
+        "Connect the output to the SelVA Sampler's optional textual_inversion input to guide "
+        "generation toward the training data style without degrading audio quality."
+    )
+
+    def load(self, path: str) -> tuple:
+        p = Path(path.strip())
+        if not p.is_absolute():
+            p = Path(folder_paths.get_output_directory()) / p
+        if not p.exists():
+            raise FileNotFoundError(f"[TI Loader] File not found: {p}")
+
+        data = torch.load(str(p), map_location="cpu", weights_only=False)
+
+        embeddings = data["embeddings"]   # [K, 1024]
+        n_tokens   = int(data.get("n_tokens", embeddings.shape[0]))
+
+        print(f"[TI Loader] Loaded '{p.name}'  n_tokens={n_tokens}  "
+              f"shape={tuple(embeddings.shape)}", flush=True)
+        if data.get("init_text"):
+            print(f"[TI Loader]   init_text='{data['init_text']}'", flush=True)
+        if data.get("step"):
+            print(f"[TI Loader]   trained {data['step']} / {data.get('steps', '?')} steps  "
+                  f"lr={data.get('lr', '?')}", flush=True)
+
+        inject_mode = data.get("inject_mode", "suffix")
+        print(f"[TI Loader]   inject_mode='{inject_mode}'", flush=True)
+
+        bundle = {
+            "embeddings":  embeddings,      # [K, 1024] float32 on CPU
+            "n_tokens":    n_tokens,
+            "inject_mode": inject_mode,
+            "path":        str(p),
+            "init_text":   data.get("init_text", ""),
+        }
+        return (bundle,)

nodes/selva_textual_inversion_trainer.py

@@ -0,0 +1,451 @@
+"""SelVA Textual Inversion Trainer.
+
+Learns K token embedding vectors in CLIP space that guide the base model
+to generate audio in the style of the training clips — without modifying
+any model weights.
+
+Key difference from LoRA:
+- ALL generator parameters are frozen (requires_grad=False)
+- Only K×1024 token embeddings receive gradients
+- Latents stay on the decoder's natural manifold → no quality degradation
+- The learned tokens shift WHICH latents are generated, not HOW
+
+Usage:
+  1. Train on your .npz audio features
+  2. Load result with SelVA Textual Inversion Loader
+  3. Connect to SelVA Sampler optional input
+"""
+
+import copy
+import random
+import traceback
+from pathlib import Path
+
+import torch
+import torchaudio
+import comfy.utils
+import folder_paths
+
+from .utils import SELVA_CATEGORY, get_device, soft_empty_cache
+from selva_core.model.flow_matching import FlowMatching
+from .selva_lora_trainer import (
+    _prepare_dataset,
+    _eval_sample,
+    _spectral_metrics,
+    _save_spectrogram,
+    _smooth_losses,
+    _draw_loss_curve,
+    _pil_to_tensor,
+)
+
+
+# ---------------------------------------------------------------------------
+# Eval helper with token injection
+# ---------------------------------------------------------------------------
+
+def _inject_tokens(text_clip: torch.Tensor, tokens: torch.Tensor,
+                   n_tokens: int, inject_mode: str) -> torch.Tensor:
+    """Build a text_clip tensor with learned tokens injected.
+
+    inject_mode:
+      "suffix" — replace last n_tokens positions (EOS/padding zone)
+      "prefix" — replace positions 1:1+n_tokens (after BOS, before content)
+
+    Always uses torch.cat so gradient flows to `tokens` when tokens.requires_grad.
+    Works for both training (tokens is a Parameter) and eval (tokens is detached).
+    """
+    if inject_mode == "prefix":
+        bos  = text_clip[:, :1, :].detach()                      # [B, 1, D]
+        toks = tokens.unsqueeze(0).expand(text_clip.shape[0], -1, -1)  # [B, K, D]
+        rest = text_clip[:, 1 + n_tokens:, :].detach()           # [B, 75-K, D]
+        return torch.cat([bos, toks, rest], dim=1)                # [B, 77, D]
+    else:  # suffix (default)
+        front = text_clip[:, :-n_tokens, :].detach()             # [B, 77-K, D]
+        toks  = tokens.unsqueeze(0).expand(text_clip.shape[0], -1, -1)  # [B, K, D]
+        return torch.cat([front, toks], dim=1)                   # [B, 77, D]
+
+
+def _eval_sample_ti(generator, learned_tokens, n_tokens, inject_mode,
+                    feature_utils_orig, dataset, seq_cfg,
+                    device, dtype, num_steps=25, seed=42, clip_idx=0):
+    """Inference pass with learned tokens injected into text conditioning."""
+    generator.eval()
+    try:
+        _, clip_f_cpu, sync_f_cpu, text_clip_cpu = dataset[clip_idx]
+        clip_f    = clip_f_cpu.to(device, dtype)
+        sync_f    = sync_f_cpu.to(device, dtype)
+        text_clip = text_clip_cpu.to(device, dtype).clone()
+
+        emb = learned_tokens.detach().to(device, dtype)
+        text_input = _inject_tokens(text_clip, emb, n_tokens, inject_mode)
+
+        rng = torch.Generator(device=device).manual_seed(seed)
+        x0  = torch.randn(1, seq_cfg.latent_seq_len, generator.latent_dim,
+                          device=device, dtype=dtype, generator=rng)
+
+        eval_fm = FlowMatching(min_sigma=0, inference_mode="euler", num_steps=num_steps)
+
+        def velocity_fn(t, x):
+            return generator.forward(x, clip_f, sync_f, text_input,
+                                     t.reshape(1).to(device, dtype))
+
+        with torch.no_grad():
+            x1_pred   = eval_fm.to_data(velocity_fn, x0)
+            x1_unnorm = generator.unnormalize(x1_pred)
+
+            orig_dev = next(feature_utils_orig.parameters()).device
+            if orig_dev != device:
+                feature_utils_orig.to(device)
+            try:
+                spec  = feature_utils_orig.decode(x1_unnorm)
+                audio = feature_utils_orig.vocode(spec)
+            finally:
+                if orig_dev != device:
+                    feature_utils_orig.to(orig_dev)
+
+        audio = audio.float().cpu()
+        if audio.dim() == 2:
+            audio = audio.unsqueeze(1)
+        elif audio.dim() == 3 and audio.shape[1] != 1:
+            audio = audio.mean(dim=1, keepdim=True)
+
+        target_rms = 10 ** (-27.0 / 20.0)
+        rms = audio.pow(2).mean().sqrt().clamp(min=1e-8)
+        audio = (audio * (target_rms / rms))
+        peak = audio.abs().max().clamp(min=1e-8)
+        if peak > 1.0:
+            audio = audio / peak
+        return audio.squeeze(0), seq_cfg.sampling_rate
+
+    except Exception as e:
+        print(f"[TI Trainer] Eval sample failed: {e}", flush=True)
+        traceback.print_exc()
+        return None, None
+    finally:
+        generator.train()
+
+
+# ---------------------------------------------------------------------------
+# Node
+# ---------------------------------------------------------------------------
+
+class SelvaTextualInversionTrainer:
+    """Learns K CLIP token embeddings that steer SelVA toward a target audio style.
+
+    Unlike LoRA, all model weights are frozen. Only the K×1024 embedding tensor
+    receives gradients, keeping generated latents on the decoder's natural manifold
+    and preserving base model audio quality while shifting generation style.
+    """
+
+    OUTPUT_NODE = True
+    CATEGORY    = SELVA_CATEGORY
+    FUNCTION    = "train"
+    RETURN_TYPES  = ("STRING", "IMAGE")
+    RETURN_NAMES  = ("embeddings_path", "loss_curve")
+    OUTPUT_TOOLTIPS = (
+        "Path to saved .pt embeddings — load with SelVA Textual Inversion Loader.",
+        "Smoothed training loss curve.",
+    )
+    DESCRIPTION = (
+        "Trains K learnable CLIP token embeddings against your audio dataset "
+        "with all model weights frozen. The tokens are then injected into the "
+        "sampler to guide generation toward the training data style without "
+        "degrading audio quality."
+    )
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("SELVA_MODEL",),
+                "data_dir": ("STRING", {
+                    "default": "",
+                    "tooltip": "Directory containing .npz feature files and paired audio files (same as LoRA trainer).",
+                }),
+                "output_path": ("STRING", {
+                    "default": "textual_inversion.pt",
+                    "tooltip": "Where to save the learned embeddings. Relative paths resolve to ComfyUI output directory.",
+                }),
+                "n_tokens": ("INT", {
+                    "default": 4, "min": 1, "max": 16,
+                    "tooltip": "Number of learnable token vectors. More tokens = more expressive but slower to train. 4 is a good default.",
+                }),
+                "steps": ("INT", {
+                    "default": 3000, "min": 100, "max": 50000,
+                    "tooltip": "Training steps. 3000 is a reasonable starting point.",
+                }),
+                "lr": ("FLOAT", {
+                    "default": 2e-4, "min": 1e-5, "max": 1e-1, "step": 1e-5,
+                    "tooltip": "Learning rate. 2e-4 matches the LoRA working regime. Higher LR (1e-3) causes token norm to drift without plateauing on small datasets.",
+                }),
+                "batch_size": ("INT", {
+                    "default": 4, "min": 1, "max": 64,
+                    "tooltip": "Clips sampled per training step. Smaller batch (4–8) gives more diverse gradients and helps token norm saturate rather than drift.",
+                }),
+                "seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFF}),
+                "save_every": ("INT", {
+                    "default": 1000, "min": 100, "max": 10000,
+                    "tooltip": "Save a checkpoint and generate an eval sample every N steps.",
+                }),
+            },
+            "optional": {
+                "inject_mode": (["suffix", "prefix"], {
+                    "default": "suffix",
+                    "tooltip": (
+                        "Where to inject the learned tokens in the 77-token CLIP sequence. "
+                        "'suffix' replaces the last K positions (EOS/padding — may be ignored by the model). "
+                        "'prefix' replaces positions 1:1+K right after BOS — higher attention weight, stronger style signal."
+                    ),
+                }),
+                "init_text": ("STRING", {
+                    "default": "",
+                    "tooltip": "Optional text phrase to warm-start token values via CLIP. Leave empty for random init (N(0, 0.02)). Example: 'industrial sound design'.",
+                }),
+                "warmup_steps": ("INT", {
+                    "default": 100, "min": 0, "max": 1000,
+                    "tooltip": "Linear LR warmup steps.",
+                }),
+            },
+        }
+
+    def train(self, model, data_dir, output_path, n_tokens, steps, lr,
+              batch_size, seed, save_every,
+              inject_mode="suffix", init_text="", warmup_steps=100):
+
+        device = get_device()
+        dtype  = model["dtype"]
+        mode   = model["mode"]
+        seq_cfg = model["seq_cfg"]
+        feature_utils_orig = model["feature_utils"]
+
+        # --- Resolve paths ---
+        data_dir = Path(data_dir.strip())
+        if not data_dir.is_absolute():
+            data_dir = Path(folder_paths.models_dir) / data_dir
+        if not data_dir.exists():
+            raise FileNotFoundError(f"[TI Trainer] data_dir not found: {data_dir}")
+
+        out_path = Path(output_path.strip())
+        if not out_path.is_absolute():
+            out_path = Path(folder_paths.get_output_directory()) / out_path
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+
+        print(f"\n[TI Trainer] n_tokens={n_tokens}  steps={steps}  lr={lr:.2e}", flush=True)
+        print(f"[TI Trainer] data_dir   = {data_dir}", flush=True)
+        print(f"[TI Trainer] output     = {out_path}\n", flush=True)
+
+        # --- Load dataset (reuse LoRA trainer helper) ---
+        dataset = _prepare_dataset(model, data_dir, device)
+
+        # Training must run outside inference_mode so autograd works
+        with torch.inference_mode(False), torch.enable_grad():
+            r = self._train_inner(
+                model, dataset, feature_utils_orig, seq_cfg,
+                device, dtype, mode,
+                data_dir, out_path,
+                n_tokens, steps, lr, batch_size,
+                warmup_steps, seed, save_every, init_text, inject_mode,
+            )
+        smoothed = _smooth_losses(r["loss_history"]) if r["loss_history"] else []
+        curve_img = _draw_loss_curve(r["loss_history"], log_interval=50, smoothed=smoothed)
+        return (r["embeddings_path"], _pil_to_tensor(curve_img))
+
+    def _train_inner(
+        self, model, dataset, feature_utils_orig, seq_cfg,
+        device, dtype, mode,
+        data_dir, out_path,
+        n_tokens, steps, lr, batch_size,
+        warmup_steps, seed, save_every, init_text, inject_mode="suffix",
+    ):
+        torch.manual_seed(seed)
+
+        # --- Generator (frozen) ---
+        generator = copy.deepcopy(model["generator"]).to(device, dtype)
+        generator.requires_grad_(False)
+        generator.update_seq_lengths(
+            latent_seq_len=seq_cfg.latent_seq_len,
+            clip_seq_len=seq_cfg.clip_seq_len,
+            sync_seq_len=seq_cfg.sync_seq_len,
+        )
+
+        # --- Init learned tokens ---
+        # Call encode_text_clip outside the grad context (it has @inference_mode),
+        # grab values only (no grad needed), then wrap as nn.Parameter.
+        if init_text.strip():
+            with torch.no_grad():
+                init_embed = feature_utils_orig.encode_text_clip([init_text.strip()])
+            # Positions 1:1+n_tokens — after BOS, before EOS — have actual content
+            init_vals = init_embed[0, 1:1 + n_tokens, :].detach().clone().float()
+            if init_vals.shape[0] < n_tokens:
+                # Prompt was very short; pad remaining with small noise
+                pad = torch.randn(n_tokens - init_vals.shape[0], init_vals.shape[1]) * 0.02
+                init_vals = torch.cat([init_vals, pad], dim=0)
+            learned_tokens = torch.nn.Parameter(init_vals.to(device, dtype))
+            print(f"[TI Trainer] Init from '{init_text.strip()}' (positions 1–{n_tokens})", flush=True)
+        else:
+            learned_tokens = torch.nn.Parameter(
+                torch.randn(n_tokens, 1024, device=device, dtype=dtype) * 0.02
+            )
+            print(f"[TI Trainer] Init: random N(0, 0.02)", flush=True)
+
+        # --- Measure CLIP token norm from the dataset (content positions 1–20) ---
+        # Learned tokens must stay within this range or the model treats them as
+        # out-of-distribution and produces buzz artifacts instead of style shift.
+        with torch.no_grad():
+            sample_norms = []
+            for item in dataset[:min(len(dataset), 20)]:
+                tc = item[3].squeeze(0)          # [77, 1024]
+                sample_norms.append(tc[1:20].norm(dim=-1))  # skip BOS/EOS
+            clip_norm_ref   = torch.cat(sample_norms).mean().item()
+            clip_norm_limit = clip_norm_ref * 1.5   # 50% headroom above real tokens
+        print(f"[TI Trainer] CLIP token norm ref={clip_norm_ref:.4f}  "
+              f"limit={clip_norm_limit:.4f}", flush=True)
+
+        # --- Optimizer + scheduler ---
+        optimizer = torch.optim.AdamW([learned_tokens], lr=lr, weight_decay=1e-2)
+
+        def lr_lambda(s):
+            return s / max(1, warmup_steps) if s < warmup_steps else 1.0
+
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+        fm = FlowMatching(min_sigma=0, inference_mode="euler", num_steps=25)
+
+        # --- Checkpoint dir ---
+        ckpt_dir = out_path.parent / out_path.stem
+        ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+        # --- Baseline sample (once, before any training) ---
+        print(f"[TI Trainer] Generating baseline sample...", flush=True)
+        baseline_wav, baseline_sr = _eval_sample(
+            generator, feature_utils_orig, dataset, seq_cfg, device, dtype, seed=seed,
+        )
+        if baseline_wav is not None:
+            baseline_path = ckpt_dir / "baseline.wav"
+            try:
+                torchaudio.save(str(baseline_path), baseline_wav, baseline_sr)
+            except RuntimeError:
+                import soundfile as sf
+                sf.write(str(baseline_path), baseline_wav.squeeze(0).numpy(), baseline_sr)
+            try:
+                _save_spectrogram(baseline_wav, baseline_sr, ckpt_dir / "baseline.png")
+            except Exception:
+                pass
+            print(f"[TI Trainer] Baseline saved: {baseline_path}", flush=True)
+
+        # --- Training loop ---
+        generator.train()
+        optimizer.zero_grad()
+
+        log_interval = 50
+        pbar         = comfy.utils.ProgressBar(steps)
+        loss_history = []
+        running_loss = 0.0
+
+        print(f"[TI Trainer] Training {steps} steps  batch_size={batch_size}\n", flush=True)
+
+        for step in range(1, steps + 1):
+            batch = random.choices(dataset, k=batch_size)
+            x1_list, clip_list, sync_list, text_list = zip(*batch)
+
+            x1        = torch.stack([x.squeeze(0) for x in x1_list]).to(device, dtype)
+            clip_f    = torch.stack([x.squeeze(0) for x in clip_list]).to(device, dtype)
+            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 — gradient flows via torch.cat (not in-place assignment).
+            text_input = _inject_tokens(text_clip, learned_tokens, n_tokens, inject_mode)
+
+            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_input, t)
+            loss   = fm.loss(v_pred, x0, x1).mean()
+            loss.backward()
+
+            torch.nn.utils.clip_grad_norm_([learned_tokens], max_norm=1.0)
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+
+            # Clamp token norm to CLIP manifold — prevents out-of-distribution
+            # embeddings that cause buzz artifacts instead of style shift.
+            with torch.no_grad():
+                norms = learned_tokens.norm(dim=-1, keepdim=True).clamp(min=1e-8)
+                scale = (clip_norm_limit / norms).clamp(max=1.0)
+                learned_tokens.data.mul_(scale)
+
+            running_loss += loss.item()
+            pbar.update(1)
+
+            if step % log_interval == 0:
+                avg = running_loss / log_interval
+                loss_history.append(round(avg, 6))
+                running_loss = 0.0
+                lr_now = scheduler.get_last_lr()[0]
+                norm = learned_tokens.norm(dim=-1).mean().item()
+                print(f"[TI Trainer] step {step:5d}/{steps}  "
+                      f"loss={avg:.4f}  lr={lr_now:.2e}  "
+                      f"token_norm={norm:.4f}/{clip_norm_limit:.4f}", flush=True)
+
+            if step % save_every == 0 or step == steps:
+                # Save checkpoint
+                ckpt = {
+                    "embeddings":  learned_tokens.detach().cpu(),
+                    "n_tokens":    n_tokens,
+                    "inject_mode": inject_mode,
+                    "step":        step,
+                    "init_text":   init_text,
+                    "lr":          lr,
+                    "steps":       steps,
+                    "loss_history": loss_history,
+                }
+                ckpt_path = ckpt_dir / f"step_{step:05d}.pt"
+                torch.save(ckpt, str(ckpt_path))
+
+                # Eval sample
+                wav, sr = _eval_sample_ti(
+                    generator, learned_tokens, n_tokens, inject_mode,
+                    feature_utils_orig, dataset, seq_cfg,
+                    device, dtype, seed=seed,
+                )
+                if wav is not None:
+                    wav_path = ckpt_dir / f"step_{step:05d}.wav"
+                    try:
+                        torchaudio.save(str(wav_path), wav, sr)
+                    except RuntimeError:
+                        import soundfile as sf
+                        sf.write(str(wav_path), wav.squeeze(0).numpy(), sr)
+
+                    try:
+                        metrics = _spectral_metrics(wav, sr)
+                        _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}  "
+                              f"hf={metrics['hf_energy_ratio']:.3f}", flush=True)
+                    except Exception as e:
+                        print(f"[TI Trainer] Spectral/spectrogram failed: {e}", flush=True)
+
+                print(f"[TI Trainer] Checkpoint: {ckpt_path}", flush=True)
+
+        # --- Final save ---
+        final = {
+            "embeddings":  learned_tokens.detach().cpu(),
+            "n_tokens":    n_tokens,
+            "inject_mode": inject_mode,
+            "step":        steps,
+            "init_text":   init_text,
+            "lr":          lr,
+            "steps":       steps,
+            "loss_history": loss_history,
+        }
+        torch.save(final, str(out_path))
+        print(f"\n[TI Trainer] Done. Saved: {out_path}", flush=True)
+
+        soft_empty_cache()
+        return {
+            "embeddings_path": str(out_path),
+            "loss_history":    loss_history,
+        }

nodes/selva_ti_scheduler.py

@@ -0,0 +1,479 @@
+"""SelVA Textual Inversion Scheduler — sweeps TI training experiments from a JSON file.
+
+Each experiment inherits from a shared `base` config and overrides specific keys.
+The dataset is loaded once and reused across all experiments. Results are written
+to `experiment_summary.json` (updated after each completed run) and a comparison
+loss-curve image showing all runs on the same axes.
+
+JSON format:
+    {
+      "name": "ti_sweep_1",
+      "description": "optional human note",
+      "data_dir": "dataset/bj_sounds",
+      "output_root": "ti_output/sweep_1",
+      "base": {
+        "n_tokens": 4,
+        "lr": 1e-3,
+        "steps": 3000,
+        "batch_size": 16,
+        "warmup_steps": 100,
+        "seed": 42,
+        "save_every": 1000
+      },
+      "experiments": [
+        {"id": "baseline",       "description": "default 4 tokens"},
+        {"id": "n8_tokens",      "n_tokens": 8},
+        {"id": "lr_5e4",         "lr": 5e-4},
+        {"id": "warm_init",      "init_text": "industrial sound design"},
+        {"id": "n4_more_steps",  "steps": 5000}
+      ]
+    }
+"""
+
+import json
+import sys
+import time
+import traceback
+from datetime import datetime, timezone
+from pathlib import Path
+
+import numpy as np
+import torch
+
+import comfy.utils
+import folder_paths
+
+from .utils import SELVA_CATEGORY, get_device
+from .selva_lora_trainer import (
+    _prepare_dataset,
+    _smooth_losses,
+    _pil_to_tensor,
+)
+from .selva_textual_inversion_trainer import SelvaTextualInversionTrainer
+
+
+# ---------------------------------------------------------------------------
+# Helpers (shared with LoRA scheduler, inlined to keep modules independent)
+# ---------------------------------------------------------------------------
+
+def _get_system_info() -> dict:
+    info: dict = {
+        "torch_version": torch.__version__,
+        "cuda_version":  torch.version.cuda or "N/A",
+        "gpu_name":      None,
+        "gpu_vram_gb":   None,
+    }
+    if torch.cuda.is_available():
+        try:
+            info["gpu_name"]    = torch.cuda.get_device_name(0)
+            props               = torch.cuda.get_device_properties(0)
+            info["gpu_vram_gb"] = round(props.total_memory / 1e9, 1)
+        except Exception:
+            pass
+    return info
+
+
+_PARAM_DEFAULTS = {
+    "n_tokens":     4,
+    "lr":           2e-4,
+    "steps":        3000,
+    "batch_size":   4,
+    "warmup_steps": 100,
+    "seed":         42,
+    "save_every":   1000,
+    "init_text":    "",
+    "inject_mode":  "suffix",
+}
+
+_PALETTE = [
+    (66,  133, 244),
+    (234,  67,  53),
+    (52,  168,  83),
+    (251, 188,   5),
+    (155,  89, 182),
+    (26,  188, 156),
+    (230, 126,  34),
+    (149, 165, 166),
+]
+
+
+def _resolve_path(raw: str) -> Path:
+    p = Path(raw.strip())
+    unix_style_on_windows = (
+        sys.platform == "win32" and p.is_absolute() and not p.drive
+    )
+    if not p.is_absolute() or unix_style_on_windows:
+        p = Path(folder_paths.get_output_directory()) / p.relative_to(p.anchor)
+    return p
+
+
+def _merge_config(base: dict, experiment: dict) -> dict:
+    cfg = dict(_PARAM_DEFAULTS)
+    cfg.update(base)
+    cfg.update({k: v for k, v in experiment.items() if k not in ("id", "description")})
+    return cfg
+
+
+def _loss_at_steps(loss_history: list, log_interval: int, save_every: int,
+                   total_steps: int) -> dict:
+    result = {}
+    for target in range(save_every, total_steps + 1, save_every):
+        idx = target // log_interval - 1
+        if 0 <= idx < len(loss_history):
+            result[str(target)] = round(loss_history[idx], 6)
+    return result
+
+
+def _draw_comparison_curves(experiments_data: list) -> "Image.Image":
+    from PIL import Image, ImageDraw
+
+    W, H = 900, 420
+    pl, pr, pt, pb = 75, 160, 30, 50
+
+    img  = Image.new("RGB", (W, H), (255, 255, 255))
+    draw = ImageDraw.Draw(img)
+    pw   = W - pl - pr
+    ph   = H - pt - pb
+
+    series = []
+    for i, ed in enumerate(experiments_data):
+        lh = ed.get("loss_history") or []
+        if len(lh) < 2:
+            continue
+        sm = _smooth_losses(lh)
+        series.append({
+            "id":       ed["id"],
+            "smoothed": sm,
+            "color":    _PALETTE[i % len(_PALETTE)],
+        })
+
+    if not series:
+        draw.text((pl + 10, pt + 10), "No data to plot", fill=(80, 80, 80))
+        return img
+
+    all_vals = [v for s in series for v in s["smoothed"]]
+    lo, hi = min(all_vals), max(all_vals)
+    if hi == lo:
+        hi = lo + 1e-6
+    rng = hi - lo
+
+    for i in range(5):
+        y   = pt + int(i * ph / 4)
+        val = hi - i * rng / 4
+        draw.line([(pl, y), (W - pr, y)], fill=(220, 220, 220), width=1)
+        draw.text((2, y - 7), f"{val:.4f}", fill=(100, 100, 100))
+
+    for s in series:
+        n   = len(s["smoothed"])
+        pts = []
+        for j, v in enumerate(s["smoothed"]):
+            x = pl + int(j * pw / max(n - 1, 1))
+            y = pt + int((1.0 - (v - lo) / rng) * ph)
+            pts.append((x, y))
+        draw.line(pts, fill=s["color"], width=2)
+
+    draw.line([(pl, pt), (pl, H - pb)],         fill=(40, 40, 40), width=1)
+    draw.line([(pl, H - pb), (W - pr, H - pb)], fill=(40, 40, 40), width=1)
+    draw.text((pl + 4, 8), "TI loss comparison (smoothed)", fill=(40, 40, 40))
+
+    lx, ly = W - pr + 10, pt
+    for s in series:
+        draw.rectangle([(lx, ly + 3), (lx + 14, ly + 13)], fill=s["color"])
+        draw.text((lx + 18, ly), s["id"][:20], fill=(40, 40, 40))
+        ly += 20
+
+    return img
+
+
+# ---------------------------------------------------------------------------
+# Node
+# ---------------------------------------------------------------------------
+
+class SelvaTiScheduler:
+    """Runs a sweep of Textual Inversion experiments defined in a JSON file.
+
+    The dataset is loaded once and reused. Each experiment calls
+    SelvaTextualInversionTrainer._train_inner() with its own config.
+    Results are written to experiment_summary.json after every completed run.
+    """
+
+    OUTPUT_NODE = True
+    CATEGORY    = SELVA_CATEGORY
+    FUNCTION    = "run"
+    RETURN_TYPES  = ("STRING", "IMAGE")
+    RETURN_NAMES  = ("summary_path", "comparison_curves")
+    OUTPUT_TOOLTIPS = (
+        "Path to experiment_summary.json — compare runs across sweeps.",
+        "All smoothed loss curves overlaid on the same axes.",
+    )
+    DESCRIPTION = (
+        "Runs a series of Textual Inversion experiments from a JSON sweep file. "
+        "The dataset is encoded once and reused. Results (loss, config, embeddings "
+        "paths) are collected in experiment_summary.json after each run."
+    )
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("SELVA_MODEL",),
+                "experiments_file": ("STRING", {
+                    "default": "ti_experiments.json",
+                    "tooltip": (
+                        "Path to JSON sweep file. Relative paths resolve to the ComfyUI "
+                        "output directory. See node description for the file format."
+                    ),
+                }),
+            }
+        }
+
+    def run(self, model, experiments_file):
+        # ------------------------------------------------------------------
+        # 1. Read + validate JSON
+        # ------------------------------------------------------------------
+        exp_path = Path(experiments_file.strip())
+        if not exp_path.is_absolute():
+            candidate = Path(folder_paths.models_dir) / exp_path
+            if not candidate.exists():
+                candidate = Path(folder_paths.get_output_directory()) / exp_path
+            exp_path = candidate
+        if not exp_path.exists():
+            raise FileNotFoundError(
+                f"[TI Scheduler] Experiment file not found: {exp_path}"
+            )
+
+        spec = json.loads(exp_path.read_text(encoding="utf-8"))
+
+        if "experiments" not in spec or not spec["experiments"]:
+            raise ValueError("[TI Scheduler] 'experiments' list is missing or empty.")
+        for i, exp in enumerate(spec["experiments"]):
+            if "id" not in exp:
+                raise ValueError(
+                    f"[TI Scheduler] Experiment at index {i} is missing required 'id' field."
+                )
+
+        sweep_name  = spec.get("name", exp_path.stem)
+        description = spec.get("description", "")
+        base_cfg    = spec.get("base", {})
+
+        # ------------------------------------------------------------------
+        # 2. Resolve data_dir and output_root
+        # ------------------------------------------------------------------
+        if "data_dir" not in spec:
+            raise ValueError("[TI Scheduler] 'data_dir' is required in the sweep file.")
+        data_dir    = _resolve_path(spec["data_dir"])
+        output_root = _resolve_path(spec.get("output_root", f"ti_sweeps/{sweep_name}"))
+        output_root.mkdir(parents=True, exist_ok=True)
+
+        device = get_device()
+        dtype  = model["dtype"]
+        mode   = model["mode"]
+        seq_cfg            = model["seq_cfg"]
+        feature_utils_orig = model["feature_utils"]
+
+        print(f"\n[TI Scheduler] Sweep '{sweep_name}': "
+              f"{len(spec['experiments'])} experiment(s)", flush=True)
+        if description:
+            print(f"[TI Scheduler] {description}", flush=True)
+        print(f"[TI Scheduler] data_dir    = {data_dir}", flush=True)
+        print(f"[TI Scheduler] output_root = {output_root}\n", flush=True)
+
+        # ------------------------------------------------------------------
+        # 3. Load dataset once
+        # ------------------------------------------------------------------
+        n_clips = len(list(data_dir.glob("*.npz")))
+        dataset = _prepare_dataset(model, data_dir, device)
+
+        # ------------------------------------------------------------------
+        # 4. Build or restore summary (resume-aware)
+        # ------------------------------------------------------------------
+        summary_path   = output_root / "experiment_summary.json"
+        completed_ids  = set()
+        all_curve_data = []
+
+        if summary_path.exists():
+            try:
+                existing = json.loads(summary_path.read_text(encoding="utf-8"))
+                for rec in existing.get("experiments", []):
+                    if rec.get("results", {}).get("status") == "completed":
+                        completed_ids.add(rec["id"])
+                        all_curve_data.append({
+                            "id":           rec["id"],
+                            "loss_history": rec["results"].get("loss_history", []),
+                        })
+                summary = existing
+                summary["completed_at"] = None
+                if completed_ids:
+                    print(f"[TI Scheduler] Resuming — skipping {len(completed_ids)} "
+                          f"completed experiment(s): {sorted(completed_ids)}", flush=True)
+            except Exception as e:
+                print(f"[TI Scheduler] Could not read existing summary ({e}) — starting fresh",
+                      flush=True)
+                completed_ids  = set()
+                all_curve_data = []
+                summary        = None
+
+        if not completed_ids:
+            summary = {
+                "sweep_name":   sweep_name,
+                "description":  description,
+                "sweep_file":   str(exp_path),
+                "started_at":   datetime.now(timezone.utc).isoformat(),
+                "completed_at": None,
+                "system":       _get_system_info(),
+                "data_dir":     str(data_dir),
+                "n_clips":      n_clips,
+                "experiments":  [],
+            }
+
+        comparison_img_path = output_root / "loss_comparison.png"
+
+        def _write_summary():
+            summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+        def _save_comparison():
+            try:
+                img = _draw_comparison_curves(all_curve_data)
+                img.save(str(comparison_img_path))
+            except Exception as e:
+                print(f"[TI Scheduler] Comparison image failed: {e}", flush=True)
+
+        _write_summary()
+
+        # ------------------------------------------------------------------
+        # 5. Run each experiment
+        # ------------------------------------------------------------------
+        trainer    = SelvaTextualInversionTrainer()
+        pbar_outer = comfy.utils.ProgressBar(len(spec["experiments"]))
+        log_interval = 50   # matches _train_inner
+
+        for exp in spec["experiments"]:
+            exp_id   = exp["id"]
+            exp_desc = exp.get("description", "")
+
+            if exp_id in completed_ids:
+                print(f"[TI Scheduler] Skipping '{exp_id}' (already completed)", flush=True)
+                pbar_outer.update(1)
+                continue
+
+            cfg = _merge_config(base_cfg, exp)
+
+            n_tokens    = int(cfg["n_tokens"])
+            lr          = float(cfg["lr"])
+            steps       = int(cfg["steps"])
+            batch_size  = int(cfg["batch_size"])
+            warmup      = int(cfg["warmup_steps"])
+            seed        = int(cfg["seed"])
+            save_every  = int(cfg["save_every"])
+            init_text   = str(cfg["init_text"])
+            inject_mode = str(cfg["inject_mode"])
+
+            output_dir = output_root / exp_id
+            output_dir.mkdir(parents=True, exist_ok=True)
+            out_path = output_dir / "embeddings.pt"
+
+            print(f"\n[TI Scheduler] ── Experiment '{exp_id}' ──", flush=True)
+            if exp_desc:
+                print(f"[TI Scheduler] {exp_desc}", flush=True)
+            print(f"[TI Scheduler] n_tokens={n_tokens}  lr={lr:.2e}  steps={steps}  "
+                  f"batch_size={batch_size}  warmup={warmup}  seed={seed}  "
+                  f"inject_mode={inject_mode}", flush=True)
+            if init_text:
+                print(f"[TI Scheduler] init_text='{init_text}'", flush=True)
+
+            exp_record = {
+                "id":          exp_id,
+                "description": exp_desc,
+                "config": {
+                    "n_tokens":     n_tokens,
+                    "lr":           lr,
+                    "steps":        steps,
+                    "batch_size":   batch_size,
+                    "warmup_steps": warmup,
+                    "seed":         seed,
+                    "save_every":   save_every,
+                    "init_text":    init_text,
+                    "inject_mode":  inject_mode,
+                },
+                "results":        {"status": "running"},
+                "embeddings_path": None,
+                "output_dir":      str(output_dir),
+            }
+            summary["experiments"].append(exp_record)
+            _write_summary()
+
+            t_start = time.monotonic()
+            try:
+                with torch.inference_mode(False), torch.enable_grad():
+                    r = trainer._train_inner(
+                        model, dataset, feature_utils_orig, seq_cfg,
+                        device, dtype, mode,
+                        data_dir, out_path,
+                        n_tokens, steps, lr, batch_size,
+                        warmup, seed, save_every, init_text, inject_mode,
+                    )
+
+                duration     = time.monotonic() - t_start
+                loss_history = r["loss_history"]
+                smoothed     = _smooth_losses(loss_history) if loss_history else []
+
+                final_loss    = round(smoothed[-1], 6) if smoothed else None
+                min_loss      = round(min(smoothed), 6) if smoothed else None
+                min_idx       = smoothed.index(min(smoothed)) if smoothed else None
+                min_loss_step = (min_idx + 1) * log_interval if min_idx is not None else None
+
+                loss_std_last_quarter = None
+                if loss_history:
+                    quarter = max(1, len(loss_history) // 4)
+                    loss_std_last_quarter = round(float(np.std(loss_history[-quarter:])), 6)
+
+                exp_record["results"] = {
+                    "status":                "completed",
+                    "final_loss":            final_loss,
+                    "min_loss":              min_loss,
+                    "min_loss_step":         min_loss_step,
+                    "loss_std_last_quarter": loss_std_last_quarter,
+                    "loss_at_steps":         _loss_at_steps(
+                        loss_history, log_interval, save_every, steps
+                    ),
+                    "loss_history":          [round(v, 6) for v in loss_history],
+                    "log_interval":          log_interval,
+                    "duration_seconds":      round(duration, 1),
+                }
+                exp_record["embeddings_path"] = r["embeddings_path"]
+
+                all_curve_data.append({
+                    "id":           exp_id,
+                    "loss_history": loss_history,
+                })
+
+            except Exception as e:
+                duration = time.monotonic() - t_start
+                print(f"[TI Scheduler] Experiment '{exp_id}' failed: {e}", flush=True)
+                traceback.print_exc()
+                exp_record["results"] = {
+                    "status":           "failed",
+                    "error":            str(e),
+                    "duration_seconds": round(duration, 1),
+                }
+                _write_summary()
+                pbar_outer.update(1)
+                continue
+
+            _write_summary()
+            _save_comparison()
+            pbar_outer.update(1)
+
+        # ------------------------------------------------------------------
+        # 6. Finalise
+        # ------------------------------------------------------------------
+        summary["completed_at"] = datetime.now(timezone.utc).isoformat()
+        _write_summary()
+        print(f"\n[TI Scheduler] Sweep complete. Summary: {summary_path}", flush=True)
+
+        # ------------------------------------------------------------------
+        # 7. Comparison image (final update, then return to ComfyUI)
+        # ------------------------------------------------------------------
+        _save_comparison()
+        comparison_img = _draw_comparison_curves(all_curve_data)
+        return (str(summary_path), _pil_to_tensor(comparison_img))

nodes/selva_vae_roundtrip.py

@@ -0,0 +1,158 @@
+"""SelVA VAE Roundtrip — encode audio through the VAE then decode straight back.
+
+Useful for diagnosing codec reconstruction quality: if the output sounds
+saturated/degraded compared to the input, the VAE/DAC is the bottleneck,
+not the diffusion model or LoRA.
+"""
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from pathlib import Path
+
+import folder_paths
+
+from .utils import SELVA_CATEGORY, get_device, soft_empty_cache
+
+
+_SELVA_DIR = Path(folder_paths.models_dir) / "selva"
+
+
+class SelvaVaeRoundtrip:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("SELVA_MODEL",),
+                "audio": ("AUDIO",),
+            },
+        }
+
+    RETURN_TYPES  = ("AUDIO",)
+    RETURN_NAMES  = ("audio_reconstructed",)
+    OUTPUT_TOOLTIPS = (
+        "Audio after VAE encode → decode roundtrip. "
+        "Compare to the input to hear codec reconstruction quality.",
+    )
+    FUNCTION  = "roundtrip"
+    CATEGORY  = SELVA_CATEGORY
+    DESCRIPTION = (
+        "Encodes the input audio through the SelVA VAE then decodes it straight back. "
+        "Use this to isolate codec reconstruction quality from generation quality. "
+        "If the output sounds degraded compared to the input, the VAE/DAC is the "
+        "bottleneck — not the model or LoRA."
+    )
+
+    def roundtrip(self, model, audio):
+        from selva_core.model.utils.features_utils import FeaturesUtils
+
+        mode          = model["mode"]
+        seq_cfg       = model["seq_cfg"]
+        dtype         = model["dtype"]
+        device        = get_device()
+        generator     = model["generator"]
+        feature_utils = model["feature_utils"]
+
+        vae_name = "v1-16.pth" if mode == "16k" else "v1-44.pth"
+        vae_path = _SELVA_DIR / "ext" / vae_name
+        if not vae_path.exists():
+            raise FileNotFoundError(
+                f"[VAE Roundtrip] VAE weight not found: {vae_path}. "
+                "Run SelVA Model Loader first to auto-download weights."
+            )
+
+        # Load encoder only — decoder/vocoder come from model["feature_utils"]
+        # to mirror exactly what the sampler uses.
+        # AutoEncoderModule requires vocoder_ckpt_path even when only encoding,
+        # so pass the BigVGAN path (weights won't actually be used for decode here).
+        bigvgan_path = _SELVA_DIR / "ext" / "best_netG.pt"
+        print("[VAE Roundtrip] Loading VAE encoder...", flush=True)
+        vae_enc = FeaturesUtils(
+            tod_vae_ckpt=str(vae_path),
+            enable_conditions=False,
+            mode=mode,
+            need_vae_encoder=True,
+            bigvgan_vocoder_ckpt=str(bigvgan_path) if bigvgan_path.exists() else None,
+        ).to(device).eval()
+
+        try:
+            # Prepare input audio
+            waveform = audio["waveform"]   # [1, C, L]
+            sr_in    = audio["sample_rate"]
+
+            wav = waveform[0].mean(0)  # mono [L]
+
+            if sr_in != seq_cfg.sampling_rate:
+                wav = torchaudio.functional.resample(
+                    wav.unsqueeze(0), sr_in, seq_cfg.sampling_rate
+                ).squeeze(0)
+                print(f"[VAE Roundtrip] Resampled {sr_in} → {seq_cfg.sampling_rate} Hz",
+                      flush=True)
+
+            target_len = int(seq_cfg.duration * seq_cfg.sampling_rate)
+            if wav.shape[0] > target_len:
+                wav = wav[:target_len]
+            elif wav.shape[0] < target_len:
+                wav = F.pad(wav, (0, target_len - wav.shape[0]))
+
+            wav_b = wav.unsqueeze(0).to(device).float()  # [1, L]
+
+            with torch.no_grad():
+                # Encode: audio → raw latent [1, latent_dim, T]
+                dist   = vae_enc.encode_audio(wav_b)
+                latent = dist.mode().clone()
+
+                # Trim/pad to exact model sequence length (same as _prepare_dataset)
+                tgt = seq_cfg.latent_seq_len
+                if latent.shape[2] < tgt:
+                    latent = F.pad(latent, (0, tgt - latent.shape[2]))
+                elif latent.shape[2] > tgt:
+                    latent = latent[:, :, :tgt]
+
+                # To [B, T, latent_dim] — layout the generator uses
+                latent_t = latent.transpose(1, 2).to(dtype)
+                print(f"[VAE Roundtrip] Encoded:    mean={latent_t.mean():.4f}  std={latent_t.std():.4f}",
+                      flush=True)
+
+                # Normalize → unnormalize mirrors the training/inference pipeline:
+                # training normalizes encoded latents; sampler unnormalizes before decode.
+                # This ensures the latent is in the same space the decoder expects.
+                latent_norm   = generator.normalize(latent_t.clone())
+                latent_unnorm = generator.unnormalize(latent_norm)
+                print(f"[VAE Roundtrip] Norm→unnorm: mean={latent_unnorm.mean():.4f}  std={latent_unnorm.std():.4f}",
+                      flush=True)
+
+                # Decode using model's feature_utils — same path as the sampler
+                orig_device = next(feature_utils.parameters()).device
+                if orig_device != device:
+                    feature_utils.to(device)
+                try:
+                    spec = feature_utils.decode(latent_unnorm)
+                    out  = feature_utils.vocode(spec)
+                finally:
+                    if orig_device != device:
+                        feature_utils.to(orig_device)
+
+            out = out.float().cpu()
+            if out.dim() == 1:
+                out = out.unsqueeze(0).unsqueeze(0)
+            elif out.dim() == 2:
+                out = out.unsqueeze(1)
+            elif out.dim() == 3 and out.shape[1] != 1:
+                out = out.mean(dim=1, keepdim=True)
+
+            rms = out.pow(2).mean().sqrt().clamp(min=1e-8)
+            target_rms = 10 ** (-27.0 / 20.0)
+            out = out * (target_rms / rms)
+            out = out.clamp(-1.0, 1.0)
+
+            print(f"[VAE Roundtrip] Output: shape={tuple(out.shape)}  "
+                  f"peak={out.abs().max():.4f}  rms={out.pow(2).mean().sqrt():.4f}",
+                  flush=True)
+
+        finally:
+            del vae_enc
+            soft_empty_cache()
+
+        return ({"waveform": out, "sample_rate": seq_cfg.sampling_rate},)

selva_core/model/lora.py

@@ -25,13 +25,14 @@ import torch.nn as nn
 class LoRALinear(nn.Module):
     """nn.Linear with a frozen base weight and trainable low-rank A/B matrices.
 
-    Output: base(x) + (x @ A.T @ B.T) * (alpha / rank)
+    Output: base(x) + (dropout(x) @ A.T @ B.T) * (alpha / rank)
 
     A is initialised with Kaiming uniform; B is initialised to zero so the
     adapter contribution starts at zero and does not disturb pretrained behaviour.
+    Dropout is applied only to the LoRA path, not the base linear.
     """
 
-    def __init__(self, linear: nn.Linear, rank: int, alpha: float):
+    def __init__(self, linear: nn.Linear, rank: int, alpha: float, dropout: float = 0.0):
         super().__init__()
         in_f  = linear.in_features
         out_f = linear.out_features
@@ -46,16 +47,18 @@ class LoRALinear(nn.Module):
         self.lora_A = nn.Parameter(torch.empty(rank, in_f, dtype=ref_dtype, device=ref_device))
         self.lora_B = nn.Parameter(torch.zeros(out_f, rank, dtype=ref_dtype, device=ref_device))
         self.scale  = alpha / rank
+        self.dropout = nn.Dropout(p=dropout) if dropout > 0.0 else nn.Identity()
 
         nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.linear(x) + (x @ self.lora_A.T @ self.lora_B.T) * self.scale
+        return self.linear(x) + (self.dropout(x) @ self.lora_A.T @ self.lora_B.T) * self.scale
 
     def extra_repr(self) -> str:
         rank = self.lora_A.shape[0]
+        p = self.dropout.p if isinstance(self.dropout, nn.Dropout) else 0.0
         return (f"in={self.linear.in_features}, out={self.linear.out_features}, "
-                f"rank={rank}, scale={self.scale:.4f}")
+                f"rank={rank}, scale={self.scale:.4f}, dropout={p}")
 
 
 def apply_lora(
@@ -63,6 +66,7 @@ def apply_lora(
     rank: int = 16,
     alpha: float = None,
     target_suffixes: tuple = ("attn.qkv",),
+    dropout: float = 0.0,
 ) -> int:
     """Replace matching nn.Linear layers with LoRALinear in-place.
 
@@ -74,6 +78,8 @@ def apply_lora(
                          ("attn.qkv",) which targets all SelfAttention QKV
                          projections in the MM-DiT generator.
                          Add "linear1" to also wrap post-attention output projections.
+        dropout:         Dropout probability on the LoRA path (not the base linear).
+                         0.05–0.1 helps regularize on small datasets.
 
     Returns:
         Number of linear layers wrapped.
@@ -92,7 +98,7 @@ def apply_lora(
         parent = model
         for part in parts[:-1]:
             parent = getattr(parent, part)
-        setattr(parent, parts[-1], LoRALinear(module, rank, alpha))
+        setattr(parent, parts[-1], LoRALinear(module, rank, alpha, dropout=dropout))
         count += 1
 
     return count

train_lora.py

@@ -167,6 +167,16 @@ def main():
                         help="Path to a step checkpoint (.pt) to resume training from.")
     parser.add_argument("--precision",          default="bf16", choices=["bf16", "fp16", "fp32"])
     parser.add_argument("--seed",               type=int,   default=42)
+    parser.add_argument("--timestep_mode",      default="uniform", choices=["uniform", "logit_normal", "curriculum"],
+                        help="Timestep sampling. uniform=original MMAudio, logit_normal=concentrated near t=0.5, curriculum=logit_normal then uniform.")
+    parser.add_argument("--logit_normal_sigma", type=float, default=1.0,
+                        help="Spread of logit-normal distribution.")
+    parser.add_argument("--curriculum_switch",  type=float, default=0.6,
+                        help="Fraction of steps to use logit_normal before switching to uniform (curriculum mode only).")
+    parser.add_argument("--lora_dropout",       type=float, default=0.0,
+                        help="Dropout on the LoRA path only. 0.05–0.1 helps on small datasets.")
+    parser.add_argument("--lora_plus_ratio",    type=float, default=1.0,
+                        help="LoRA+ LR ratio: lr_B = lr * ratio. 1.0=standard, 16.0=LoRA+.")
     args = parser.parse_args()
 
     torch.manual_seed(args.seed)
@@ -230,8 +240,9 @@ def main():
         rank=args.rank,
         alpha=args.alpha,
         target_suffixes=tuple(args.target),
+        dropout=args.lora_dropout,
     )
-    print(f"[LoRA] Wrapped {n_lora} linear layers (rank={args.rank}, target={args.target})")
+    print(f"[LoRA] Wrapped {n_lora} linear layers (rank={args.rank}, target={args.target}, dropout={args.lora_dropout})")
     if n_lora == 0:
         print("[LoRA] ERROR: no layers were wrapped — check --target names.")
         sys.exit(1)
@@ -311,8 +322,16 @@ def main():
     print(f"[LoRA] {len(dataset)} clip(s) ready.")
 
     # --- Optimizer + LR scheduler ---
-    lora_params = [p for p in net_generator.parameters() if p.requires_grad]
+    # LoRA+: separate param groups for A and B with different LRs.
-    optimizer   = torch.optim.AdamW(lora_params, lr=args.lr, weight_decay=1e-2)
+    # ratio=1.0 = standard LoRA. ratio=16 = LoRA+ (arXiv:2402.12354).
+    lora_A_params = [p for n, p in net_generator.named_parameters() if "lora_A" in n and p.requires_grad]
+    lora_B_params = [p for n, p in net_generator.named_parameters() if "lora_B" in n and p.requires_grad]
+    optimizer = torch.optim.AdamW([
+        {"params": lora_A_params, "lr": args.lr},
+        {"params": lora_B_params, "lr": args.lr * args.lora_plus_ratio},
+    ], weight_decay=1e-2)
+    if args.lora_plus_ratio != 1.0:
+        print(f"[LoRA] LoRA+: lr_A={args.lr:.2e}  lr_B={args.lr * args.lora_plus_ratio:.2e}")
 
     def lr_lambda(step):
         if step < args.warmup_steps:
@@ -347,6 +366,9 @@ def main():
           f"batch_size={args.batch_size}, lr={args.lr}, grad_accum={args.grad_accum}")
     print(f"[LoRA] Checkpoints every {args.save_every} steps → {output_dir}\n")
 
+    curriculum_switch_step = start_step + int((args.steps - start_step) * args.curriculum_switch)
+    _curriculum_switched = False
+
     total_loss = 0.0
     for step in range(start_step + 1, args.steps + 1):
         batch = random.choices(dataset, k=args.batch_size)
@@ -359,7 +381,18 @@ def main():
 
         net_generator.normalize(x1)
 
+        if args.timestep_mode == "logit_normal" or (
+            args.timestep_mode == "curriculum" and step <= curriculum_switch_step
+        ):
+            u = torch.randn(args.batch_size, device=device, dtype=dtype) * args.logit_normal_sigma
+            t = torch.sigmoid(u)
+        else:
             t = torch.rand(args.batch_size, device=device, dtype=dtype)
+
+        if args.timestep_mode == "curriculum" and step == curriculum_switch_step + 1 and not _curriculum_switched:
+            print(f"[LoRA] Curriculum switch: logit_normal → uniform at step {step}")
+            _curriculum_switched = True
+
         x0 = torch.randn_like(x1)
         xt = fm.get_conditional_flow(x0, x1, t)
 
@@ -370,7 +403,7 @@ def main():
         total_loss += loss.item() * args.grad_accum
 
         if step % args.grad_accum == 0:
-            torch.nn.utils.clip_grad_norm_(lora_params, max_norm=1.0)
+            torch.nn.utils.clip_grad_norm_(lora_A_params + lora_B_params, max_norm=1.0)
             optimizer.step()
             scheduler.step()
             optimizer.zero_grad()
@@ -394,6 +427,11 @@ def main():
                     "alpha":              args.alpha if args.alpha is not None else float(args.rank),
                     "target":             args.target,
                     "steps":             args.steps,
+                    "timestep_mode":      args.timestep_mode,
+                    "logit_normal_sigma": args.logit_normal_sigma,
+                    "curriculum_switch":  args.curriculum_switch,
+                    "lora_dropout":       args.lora_dropout,
+                    "lora_plus_ratio":    args.lora_plus_ratio,
                 },
             }, ckpt_path)
             print(f"[LoRA] Saved {ckpt_path}")
@@ -412,6 +450,11 @@ def main():
         "alpha":              args.alpha if args.alpha is not None else float(args.rank),
         "target":             args.target,
         "steps":              args.steps,
+        "timestep_mode":      args.timestep_mode,
+        "logit_normal_sigma": args.logit_normal_sigma,
+        "curriculum_switch":  args.curriculum_switch,
+        "lora_dropout":       args.lora_dropout,
+        "lora_plus_ratio":    args.lora_plus_ratio,
     }
     torch.save({"state_dict": get_lora_state_dict(net_generator), "meta": meta}, final)
     (output_dir / "meta.json").write_text(json.dumps(meta, indent=2))