feat: add experiment configs for logit+cosine combo and BigVGAN decoder fine-tuning

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
feat: add reference audio comparison metrics to LoRA trainer eval
2026-04-10 16:48:21 +02:00 · 2026-04-10 15:04:07 +02:00 · 2026-04-10 13:44:37 +02:00 · 2026-04-10 12:13:55 +02:00 · 2026-04-10 11:36:02 +02:00 · 2026-04-10 11:32:00 +02:00
24 changed files with 3726 additions and 175 deletions
@@ -0,0 +1,606 @@
+# Audio Dataset Pipeline Implementation Plan
+
+> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
+
+**Goal:** Add 5 chainable ComfyUI nodes for in-memory audio dataset preprocessing: load → resample → LUFS normalize → inspect/filter → extract single item.
+
+**Architecture:** Single new file `nodes/selva_dataset_pipeline.py` defines a custom `AUDIO_DATASET` type (list of dicts) and all 5 node classes. Nodes are stateless transforms — each takes `AUDIO_DATASET` and returns `AUDIO_DATASET`. No disk I/O except in the Loader. Register all nodes in `nodes/__init__.py`.
+
+**Tech Stack:** `pyloudnorm` (BS.1770-4 LUFS), `soxr` (VHQ resampling), `torchaudio`, `torch`. Both confirmed present in the ComfyUI environment at `/media/p5/miniforge3/envs/latestcomfyui`.
+
+---
+
+## The `AUDIO_DATASET` type
+
+Used as the ComfyUI type string `"AUDIO_DATASET"`. At runtime it is a Python list of dicts:
+
+```python
+[
+    {
+        "waveform": torch.Tensor,  # shape [1, C, L], float32, range [-1, 1]
+        "sample_rate": int,
+        "name": str,               # original filename stem, for reporting
+    },
+    ...
+]
+```
+
+---
+
+### Task 1: Create the file skeleton and AUDIO_DATASET constant
+
+**Files:**
+- Create: `nodes/selva_dataset_pipeline.py`
+
+**Step 1: Write the file with imports and type constant only**
+
+```python
+"""SelVA Audio Dataset Pipeline — chainable in-memory preprocessing nodes.
+
+Typical chain:
+  SelvaDatasetLoader
+      ↓ AUDIO_DATASET
+  SelvaDatasetResampler       (optional)
+      ↓ AUDIO_DATASET
+  SelvaDatasetLUFSNormalizer  (optional)
+      ↓ AUDIO_DATASET
+  SelvaDatasetInspector       (optional)
+      ↓ AUDIO_DATASET  +  STRING report
+  SelvaDatasetItemExtractor   → AUDIO (bridges to save/preview nodes)
+"""
+
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchaudio
+
+from .utils import SELVA_CATEGORY
+
+# ComfyUI custom type name — passed between all dataset pipeline nodes
+AUDIO_DATASET = "AUDIO_DATASET"
+
+_AUDIO_EXTS = {".wav", ".flac", ".mp3", ".ogg", ".aac", ".m4a"}
+```
+
+**Step 2: Verify import works (no test framework needed — just a quick smoke check)**
+
+```bash
+cd /media/p5/Comfyui-Prismaudio
+python3 -c "from nodes.selva_dataset_pipeline import AUDIO_DATASET; print(AUDIO_DATASET)"
+```
+Expected output: `AUDIO_DATASET`
+
+**Step 3: Commit**
+
+```bash
+git add nodes/selva_dataset_pipeline.py
+git commit -m "feat: add audio dataset pipeline skeleton"
+```
+
+---
+
+### Task 2: SelvaDatasetLoader
+
+**Files:**
+- Modify: `nodes/selva_dataset_pipeline.py`
+
+**Step 1: Add the Loader class**
+
+```python
+class SelvaDatasetLoader:
+    """Load all audio files in a folder into an in-memory AUDIO_DATASET."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "folder": ("STRING", {
+                    "default": "",
+                    "tooltip": "Absolute path to folder containing audio files. Searched recursively.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET,)
+    RETURN_NAMES  = ("dataset",)
+    FUNCTION      = "load"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = "Load all audio files from a folder into memory as an AUDIO_DATASET."
+
+    def load(self, folder: str):
+        folder = Path(folder.strip())
+        if not folder.exists():
+            raise FileNotFoundError(f"[DatasetLoader] Folder not found: {folder}")
+
+        files = [f for f in folder.rglob("*") if f.suffix.lower() in _AUDIO_EXTS]
+        if not files:
+            raise RuntimeError(f"[DatasetLoader] No audio files found in {folder}")
+
+        dataset = []
+        for f in sorted(files):
+            try:
+                wav, sr = torchaudio.load(str(f))          # [C, L]
+                wav = wav.unsqueeze(0).float()              # [1, C, L]
+                dataset.append({"waveform": wav, "sample_rate": sr, "name": f.stem})
+            except Exception as e:
+                print(f"[DatasetLoader] Skipping {f.name}: {e}", flush=True)
+
+        print(f"[DatasetLoader] Loaded {len(dataset)} clips from {folder}", flush=True)
+        return (dataset,)
+```
+
+**Step 2: Smoke test**
+
+```bash
+python3 -c "
+from nodes.selva_dataset_pipeline import SelvaDatasetLoader
+node = SelvaDatasetLoader()
+ds, = node.load('/media/unraid/davinci/Selva/BJ')
+print(len(ds), 'clips', ds[0]['name'], ds[0]['waveform'].shape, ds[0]['sample_rate'])
+"
+```
+Expected: prints clip count, first clip name, shape like `torch.Size([1, 2, 352800])`, sample rate.
+
+**Step 3: Commit**
+
+```bash
+git add nodes/selva_dataset_pipeline.py
+git commit -m "feat: add SelvaDatasetLoader node"
+```
+
+---
+
+### Task 3: SelvaDatasetResampler
+
+**Files:**
+- Modify: `nodes/selva_dataset_pipeline.py`
+
+**Step 1: Add the Resampler class**
+
+Uses `soxr` directly for VHQ quality. `soxr.resample` operates on numpy arrays, shape `[L, C]` (time-first).
+
+```python
+class SelvaDatasetResampler:
+    """Resample all clips in a dataset to a target sample rate using soxr VHQ."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "target_sr": ("INT", {
+                    "default": 44100, "min": 8000, "max": 192000,
+                    "tooltip": "Target sample rate. 44100 for large SelVA model, 16000 for small.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET,)
+    RETURN_NAMES  = ("dataset",)
+    FUNCTION      = "resample"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = "Resample all clips to target_sr using soxr VHQ. Skips clips already at target rate."
+
+    def resample(self, dataset, target_sr: int):
+        import soxr
+
+        out = []
+        changed = 0
+        for item in dataset:
+            sr = item["sample_rate"]
+            if sr == target_sr:
+                out.append(item)
+                continue
+
+            wav = item["waveform"][0]               # [C, L]
+            # soxr expects [L, C] (time-first), float64
+            wav_np = wav.permute(1, 0).double().numpy()   # [L, C]
+            wav_rs = soxr.resample(wav_np, sr, target_sr, quality="VHQ")
+            wav_t  = torch.from_numpy(wav_rs).float().permute(1, 0).unsqueeze(0)  # [1, C, L]
+            out.append({"waveform": wav_t, "sample_rate": target_sr, "name": item["name"]})
+            changed += 1
+
+        print(f"[DatasetResampler] {changed}/{len(dataset)} clips resampled → {target_sr} Hz", flush=True)
+        return (out,)
+```
+
+**Step 2: Smoke test**
+
+```bash
+python3 -c "
+from nodes.selva_dataset_pipeline import SelvaDatasetLoader, SelvaDatasetResampler
+ds, = SelvaDatasetLoader().load('/media/unraid/davinci/Selva/BJ')
+ds2, = SelvaDatasetResampler().resample(ds, 44100)
+print('ok', ds2[0]['sample_rate'], ds2[0]['waveform'].shape)
+"
+```
+
+**Step 3: Commit**
+
+```bash
+git add nodes/selva_dataset_pipeline.py
+git commit -m "feat: add SelvaDatasetResampler node (soxr VHQ)"
+```
+
+---
+
+### Task 4: SelvaDatasetLUFSNormalizer
+
+**Files:**
+- Modify: `nodes/selva_dataset_pipeline.py`
+
+**Step 1: Add the LUFS normalizer class**
+
+`pyloudnorm.Meter` requires numpy float64 array shape `[L]` (mono) or `[L, C]` (multichannel, channels last). True peak limit applied after gain.
+
+```python
+class SelvaDatasetLUFSNormalizer:
+    """Normalize each clip to a target integrated LUFS level + true peak limit."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "target_lufs": ("FLOAT", {
+                    "default": -23.0, "min": -40.0, "max": -6.0, "step": 0.5,
+                    "tooltip": "Target integrated loudness in LUFS. -23 is EBU R128 standard.",
+                }),
+                "true_peak_dbtp": ("FLOAT", {
+                    "default": -1.0, "min": -6.0, "max": 0.0, "step": 0.5,
+                    "tooltip": "True peak ceiling in dBTP. Applied after LUFS gain.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET,)
+    RETURN_NAMES  = ("dataset",)
+    FUNCTION      = "normalize"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Normalize each clip to target_lufs (BS.1770-4) then apply a true peak ceiling. "
+        "Skips clips that are too short for LUFS measurement (< 0.4 s)."
+    )
+
+    def normalize(self, dataset, target_lufs: float, true_peak_dbtp: float):
+        import pyloudnorm as pyln
+
+        tp_linear = 10.0 ** (true_peak_dbtp / 20.0)
+        out = []
+        skipped = 0
+
+        for item in dataset:
+            wav = item["waveform"][0]           # [C, L]
+            sr  = item["sample_rate"]
+
+            # pyloudnorm wants [L] mono or [L, C] multichannel, float64
+            wav_np = wav.permute(1, 0).double().numpy()  # [L, C]
+            if wav_np.shape[1] == 1:
+                wav_np = wav_np[:, 0]           # [L] mono
+
+            meter = pyln.Meter(sr)
+            try:
+                loudness = meter.integrated_loudness(wav_np)
+            except Exception:
+                skipped += 1
+                out.append(item)
+                continue
+
+            if not np.isfinite(loudness):
+                skipped += 1
+                out.append(item)
+                continue
+
+            gain_db     = target_lufs - loudness
+            gain_linear = 10.0 ** (gain_db / 20.0)
+
+            wav_norm = wav * gain_linear
+
+            # True peak limit
+            peak = wav_norm.abs().max().item()
+            if peak > tp_linear:
+                wav_norm = wav_norm * (tp_linear / peak)
+
+            out.append({"waveform": wav_norm.unsqueeze(0), "sample_rate": sr, "name": item["name"]})
+
+        print(
+            f"[LUFSNormalizer] {len(dataset) - skipped}/{len(dataset)} clips normalized  "
+            f"target={target_lufs} LUFS  TP={true_peak_dbtp} dBTP  skipped={skipped}",
+            flush=True,
+        )
+        return (out,)
+```
+
+**Step 2: Smoke test**
+
+```bash
+python3 -c "
+from nodes.selva_dataset_pipeline import SelvaDatasetLoader, SelvaDatasetLUFSNormalizer
+ds, = SelvaDatasetLoader().load('/media/unraid/davinci/Selva/BJ')
+ds2, = SelvaDatasetLUFSNormalizer().normalize(ds, -23.0, -1.0)
+print('ok', ds2[0]['name'], ds2[0]['waveform'].abs().max().item())
+"
+```
+Expected: peak ≤ ~0.89 (≈ -1 dBTP).
+
+**Step 3: Commit**
+
+```bash
+git add nodes/selva_dataset_pipeline.py
+git commit -m "feat: add SelvaDatasetLUFSNormalizer node (pyloudnorm BS.1770-4)"
+```
+
+---
+
+### Task 5: SelvaDatasetInspector
+
+**Files:**
+- Modify: `nodes/selva_dataset_pipeline.py`
+
+**Step 1: Add helper functions for artifact detection**
+
+```python
+def _check_hf_shelf(wav: torch.Tensor, sr: int) -> bool:
+    """Return True if clip looks codec-compressed (hard HF shelf above 15 kHz).
+
+    Method: compare mean energy in 1–5 kHz band vs 15–20 kHz band via STFT.
+    A ratio > 40 dB (i.e. near-silence above 15 kHz) flags codec artifacts.
+    """
+    if sr < 32000:
+        return False  # can't assess HF at low sample rates
+
+    n_fft  = 2048
+    hop    = 512
+    window = torch.hann_window(n_fft)
+    mono   = wav[0].mean(0)  # [L]
+    stft   = torch.stft(mono, n_fft, hop, n_fft, window, return_complex=True)
+    mag_sq = stft.abs().pow(2).mean(-1)  # [n_freqs]
+
+    freqs     = torch.linspace(0, sr / 2, n_fft // 2 + 1)
+    band_lo   = (freqs >= 1000) & (freqs < 5000)
+    band_hi   = (freqs >= 15000) & (freqs < 20000)
+
+    if band_hi.sum() == 0:
+        return False
+
+    energy_lo = mag_sq[band_lo].mean().clamp(min=1e-12)
+    energy_hi = mag_sq[band_hi].mean().clamp(min=1e-12)
+    ratio_db  = 10.0 * torch.log10(energy_lo / energy_hi).item()
+    return ratio_db > 40.0
+
+
+def _estimate_snr(wav: torch.Tensor) -> float:
+    """Rough SNR estimate: ratio of 95th-percentile frame RMS to 5th-percentile frame RMS."""
+    mono   = wav[0].mean(0)  # [L]
+    frames = mono.unfold(0, 2048, 512)          # [N, 2048]
+    rms    = frames.pow(2).mean(-1).sqrt()      # [N]
+    p95    = torch.quantile(rms, 0.95).item()
+    p05    = torch.quantile(rms, 0.05).clamp(min=1e-8).item()
+    return 20.0 * np.log10(p95 / p05 + 1e-8)
+```
+
+**Step 2: Add the Inspector class**
+
+```python
+class SelvaDatasetInspector:
+    """Analyze each clip for quality issues and optionally filter out flagged clips."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "skip_rejected": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "If True, flagged clips are removed from the output dataset. "
+                               "If False, all clips pass through but the report still lists issues.",
+                }),
+                "min_snr_db": ("FLOAT", {
+                    "default": 15.0, "min": 0.0, "max": 60.0, "step": 1.0,
+                    "tooltip": "Clips with estimated SNR below this value are flagged.",
+                }),
+                "check_codec_artifacts": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "Flag clips with a hard HF shelf above 15 kHz (MP3/codec artifact signature).",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET, "STRING")
+    RETURN_NAMES  = ("dataset", "report")
+    FUNCTION      = "inspect"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Analyze each clip for clipping, low SNR, and codec artifacts. "
+        "Outputs a filtered AUDIO_DATASET and a text report. "
+        "Connect report to a ShowText node to preview in the UI."
+    )
+
+    def inspect(self, dataset, skip_rejected: bool, min_snr_db: float, check_codec_artifacts: bool):
+        clean   = []
+        flagged = []
+        lines   = ["SelVA Dataset Inspector Report", "=" * 40]
+
+        for item in dataset:
+            wav    = item["waveform"]
+            sr     = item["sample_rate"]
+            name   = item["name"]
+            issues = []
+
+            # Clipping
+            peak = wav.abs().max().item()
+            if peak > 0.99:
+                issues.append(f"clipping (peak={peak:.3f})")
+
+            # Low SNR
+            snr = _estimate_snr(wav)
+            if snr < min_snr_db:
+                issues.append(f"low SNR ({snr:.1f} dB < {min_snr_db} dB)")
+
+            # Codec artifacts
+            if check_codec_artifacts and _check_hf_shelf(wav, sr):
+                issues.append("codec artifact (HF shelf > 15 kHz)")
+
+            if issues:
+                flagged.append(name)
+                lines.append(f"  FLAGGED  {name}: {', '.join(issues)}")
+                if not skip_rejected:
+                    clean.append(item)
+            else:
+                clean.append(item)
+                lines.append(f"  OK       {name}")
+
+        lines.append("=" * 40)
+        lines.append(
+            f"Total: {len(dataset)}  Clean: {len(clean)}  Flagged: {len(flagged)}"
+            + (" (removed)" if skip_rejected else " (kept)")
+        )
+        report = "\n".join(lines)
+        print(f"[DatasetInspector]\n{report}", flush=True)
+        return (clean, report)
+```
+
+**Step 3: Smoke test**
+
+```bash
+python3 -c "
+from nodes.selva_dataset_pipeline import SelvaDatasetLoader, SelvaDatasetInspector
+ds, = SelvaDatasetLoader().load('/media/unraid/davinci/Selva/BJ')
+clean, report = SelvaDatasetInspector().inspect(ds, skip_rejected=False, min_snr_db=15.0, check_codec_artifacts=True)
+print(report)
+"
+```
+Expected: report with per-clip OK/FLAGGED lines and summary counts.
+
+**Step 4: Commit**
+
+```bash
+git add nodes/selva_dataset_pipeline.py
+git commit -m "feat: add SelvaDatasetInspector node (codec artifacts, SNR, clipping)"
+```
+
+---
+
+### Task 6: SelvaDatasetItemExtractor
+
+**Files:**
+- Modify: `nodes/selva_dataset_pipeline.py`
+
+**Step 1: Add the extractor class**
+
+```python
+class SelvaDatasetItemExtractor:
+    """Extract a single AUDIO item from an AUDIO_DATASET by index.
+
+    Bridges the dataset pipeline to any node that accepts a standard AUDIO
+    input — save audio, HF Smoother, Spectral Matcher, etc.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "index":   ("INT", {
+                    "default": 0, "min": 0, "max": 9999,
+                    "tooltip": "0-based index. Wraps around if index >= dataset length.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = ("AUDIO", "STRING", "INT")
+    RETURN_NAMES  = ("audio",  "name",   "total")
+    FUNCTION      = "extract"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Extract one clip from an AUDIO_DATASET by index. "
+        "Returns standard AUDIO (compatible with all audio nodes), "
+        "the clip name, and the total dataset length."
+    )
+
+    def extract(self, dataset, index: int):
+        if not dataset:
+            raise RuntimeError("[DatasetItemExtractor] Dataset is empty.")
+        idx  = index % len(dataset)
+        item = dataset[idx]
+        audio = {"waveform": item["waveform"], "sample_rate": item["sample_rate"]}
+        print(
+            f"[DatasetItemExtractor] [{idx}/{len(dataset)-1}] {item['name']}  "
+            f"sr={item['sample_rate']}  shape={tuple(item['waveform'].shape)}",
+            flush=True,
+        )
+        return (audio, item["name"], len(dataset))
+```
+
+**Step 2: Smoke test**
+
+```bash
+python3 -c "
+from nodes.selva_dataset_pipeline import SelvaDatasetLoader, SelvaDatasetItemExtractor
+ds, = SelvaDatasetLoader().load('/media/unraid/davinci/Selva/BJ')
+audio, name, total = SelvaDatasetItemExtractor().extract(ds, 0)
+print(name, total, audio['waveform'].shape, audio['sample_rate'])
+"
+```
+
+**Step 3: Commit**
+
+```bash
+git add nodes/selva_dataset_pipeline.py
+git commit -m "feat: add SelvaDatasetItemExtractor node"
+```
+
+---
+
+### Task 7: Register all nodes in __init__.py
+
+**Files:**
+- Modify: `nodes/__init__.py:4-25`
+
+**Step 1: Add the 5 new entries to `_NODES`**
+
+Add inside the `_NODES` dict, after `"SelvaDittoOptimizer"`:
+
+```python
+    "SelvaDatasetLoader":          (".selva_dataset_pipeline", "SelvaDatasetLoader",          "SelVA Dataset Loader"),
+    "SelvaDatasetResampler":       (".selva_dataset_pipeline", "SelvaDatasetResampler",       "SelVA Dataset Resampler"),
+    "SelvaDatasetLUFSNormalizer":  (".selva_dataset_pipeline", "SelvaDatasetLUFSNormalizer",  "SelVA Dataset LUFS Normalizer"),
+    "SelvaDatasetInspector":       (".selva_dataset_pipeline", "SelvaDatasetInspector",       "SelVA Dataset Inspector"),
+    "SelvaDatasetItemExtractor":   (".selva_dataset_pipeline", "SelvaDatasetItemExtractor",   "SelVA Dataset Item Extractor"),
+```
+
+**Step 2: Verify registration**
+
+```bash
+python3 -c "
+import sys; sys.path.insert(0, '/media/p5/Comfyui-Prismaudio')
+from nodes import NODE_CLASS_MAPPINGS
+keys = [k for k in NODE_CLASS_MAPPINGS if 'Dataset' in k]
+print(keys)
+"
+```
+Expected: list of 5 dataset node keys.
+
+**Step 3: Final commit**
+
+```bash
+git add nodes/__init__.py
+git commit -m "feat: register audio dataset pipeline nodes in __init__.py"
+```
+
+---
+
+## Summary
+
+5 nodes in `nodes/selva_dataset_pipeline.py`, all registered in `__init__.py`:
+
+| Node | In | Out |
+|------|----|-----|
+| SelvaDatasetLoader | folder path | AUDIO_DATASET |
+| SelvaDatasetResampler | AUDIO_DATASET | AUDIO_DATASET |
+| SelvaDatasetLUFSNormalizer | AUDIO_DATASET | AUDIO_DATASET |
+| SelvaDatasetInspector | AUDIO_DATASET | AUDIO_DATASET + STRING |
+| SelvaDatasetItemExtractor | AUDIO_DATASET + index | AUDIO + name + total |
+
+Dependencies: `pyloudnorm`, `soxr` — both confirmed present in the ComfyUI env.
@@ -0,0 +1,31 @@
+{
+  "name": "bigvgan_disc_fm_retest",
+  "description": "Retest discriminator feature matching after bfloat16 dtype fix. Uses optimal config from overnight sweep (snake_alpha, GAFilter, lr=1e-4, phase=1.0, L2-SP=1e-3, 5000 steps).",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/bigvgan_disc_fm_retest",
+  "base": {
+    "train_mode": "snake_alpha_only",
+    "steps": 5000,
+    "lr": 1e-4,
+    "batch_size": 8,
+    "segment_seconds": 0.5,
+    "lambda_l2sp": 1e-3,
+    "use_gafilter": true,
+    "gafilter_kernel_size": 9,
+    "lambda_phase": 1.0,
+    "save_every": 1000,
+    "seed": 42,
+    "lora_adapter": "/media/unraid/davinci/Selva/BJ/experiment/pissa_sweep/standard_baseline/adapter_final.pt"
+  },
+  "experiments": [
+    {
+      "id": "snake_5k_control",
+      "description": "Control: best config from overnight sweep without discriminator. Baseline for A/B comparison."
+    },
+    {
+      "id": "disc_fm_5k",
+      "description": "Discriminator feature matching at 5k steps. Tests if perceptual FM loss improves over mel+phase alone.",
+      "discriminator_path": "/media/unraid/davinci/Selva/BJ/experiment/bigvgan_discriminator_optimizer.pt"
+    }
+  ]
+}
@@ -0,0 +1,35 @@
+{
+  "name": "bigvgan_optimized_dataset",
+  "description": "BigVGAN fine-tuning on optimized dataset (134 clips, 44.1kHz, LUFS-normalized). Standard mode (no LoRA) — trains decoder to faithfully reconstruct target domain audio from mel spectrograms. Uses optimal config from prior sweeps.",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features_v2_improved/",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/bigvgan_optimized_dataset",
+  "base": {
+    "train_mode": "snake_alpha_only",
+    "steps": 5000,
+    "lr": 1e-4,
+    "batch_size": 8,
+    "segment_seconds": 0.5,
+    "lambda_l2sp": 1e-3,
+    "use_gafilter": true,
+    "gafilter_kernel_size": 9,
+    "lambda_phase": 1.0,
+    "save_every": 1000,
+    "seed": 42
+  },
+  "experiments": [
+    {
+      "id": "standard_5k",
+      "description": "Standard mode: mel from clean FLAC → BigVGAN → reconstruct FLAC. No LoRA. Directly improves VAE roundtrip quality."
+    },
+    {
+      "id": "disc_fm_5k",
+      "description": "Standard mode + discriminator feature matching. Tests if perceptual loss helps on clean audio reconstruction.",
+      "discriminator_path": "/media/unraid/davinci/Selva/BJ/experiment/bigvgan_discriminator_optimizer.pt"
+    },
+    {
+      "id": "standard_10k",
+      "description": "Extended 10k steps. More data passes on 134 clips may extract more from the optimized dataset.",
+      "steps": 10000
+    }
+  ]
+}
@@ -0,0 +1,65 @@
+{
+  "name": "bigvgan_overnight",
+  "description": "BigVGAN vocoder quality sweep. Axes: snake_alpha steps, all_params short run, GAFilter on/off, discriminator FM, phase loss weight. All use LoRA-distorted mels as input so vocoder learns to fix LoRA artifacts.",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/bigvgan_overnight",
+  "base": {
+    "train_mode": "snake_alpha_only",
+    "steps": 3000,
+    "lr": 1e-4,
+    "batch_size": 8,
+    "segment_seconds": 0.5,
+    "lambda_l2sp": 1e-3,
+    "use_gafilter": true,
+    "gafilter_kernel_size": 9,
+    "lambda_phase": 1.0,
+    "save_every": 1000,
+    "seed": 42,
+    "lora_adapter": "/media/unraid/davinci/Selva/BJ/experiment/pissa_sweep/standard_baseline/adapter_final.pt"
+  },
+  "experiments": [
+    {
+      "id": "snake_3k_baseline",
+      "description": "Snake alpha + GAFilter baseline. 3000 steps, same as first successful run but longer."
+    },
+    {
+      "id": "snake_5k",
+      "description": "Snake alpha + GAFilter, 5000 steps. Test if longer training improves further.",
+      "steps": 5000
+    },
+    {
+      "id": "snake_no_gafilter",
+      "description": "Snake alpha only, no GAFilter. Isolate GAFilter contribution.",
+      "use_gafilter": false
+    },
+    {
+      "id": "snake_no_phase",
+      "description": "Snake alpha + GAFilter, no phase loss. Isolate phase loss contribution.",
+      "lambda_phase": 0.0
+    },
+    {
+      "id": "snake_phase_2",
+      "description": "Snake alpha + GAFilter, phase weight 2.0. Stronger phase penalty.",
+      "lambda_phase": 2.0
+    },
+    {
+      "id": "snake_lr5e-5",
+      "description": "Snake alpha + GAFilter, lower LR 5e-5. Test if slower converges better.",
+      "lr": 5e-5,
+      "steps": 5000
+    },
+    {
+      "id": "snake_disc_fm",
+      "description": "Snake alpha + GAFilter + discriminator feature matching. Perceptual loss should directly penalize harmonic smearing.",
+      "discriminator_path": "/media/unraid/davinci/Selva/BJ/experiment/bigvgan_discriminator_optimizer.pt"
+    },
+    {
+      "id": "all_2k_l2sp1e-2",
+      "description": "All params, 2000 steps, strong L2-SP (1e-2). Test if full param tuning with heavy anchor beats snake-only.",
+      "train_mode": "all_params",
+      "steps": 2000,
+      "lr": 1e-5,
+      "lambda_l2sp": 1e-2
+    }
+  ]
+}
@@ -0,0 +1,33 @@
+{
+  "name": "lora_logit_cosine_combo",
+  "description": "Combine the two best findings from optimized dataset sweep: logit-normal timestep sampling + cosine LR schedule. Both individually outperformed baseline by large margins (56% and 68% lower loss). Tests if gains stack.",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features_v2_improved/",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/lora_logit_cosine_combo",
+  "base": {
+    "rank": 128,
+    "lr": 3e-4,
+    "steps": 5000,
+    "batch_size": 16,
+    "warmup_steps": 100,
+    "save_every": 1000,
+    "seed": 42,
+    "init_mode": "pissa",
+    "use_rslora": true,
+    "target": "attn.qkv",
+    "timestep_mode": "uniform",
+    "lr_schedule": "constant"
+  },
+  "experiments": [
+    {
+      "id": "logit_normal_cosine",
+      "description": "Logit-normal timesteps + cosine LR decay. Combining the two best individual improvements.",
+      "timestep_mode": "logit_normal",
+      "lr_schedule": "cosine"
+    },
+    {
+      "id": "logit_normal_control",
+      "description": "Control: logit-normal only (constant LR). Reproduces previous winner for direct comparison.",
+      "timestep_mode": "logit_normal"
+    }
+  ]
+}
@@ -0,0 +1,64 @@
+{
+  "name": "lora_optimized_dataset",
+  "description": "LoRA training on optimized dataset (134 clips: resampled 44.1kHz, LUFS-normalized, spectral matched, HF smoothed, gain-augmented). Tests latent augmentation and schedule variants on top of known-best config (PiSSA, rank=128, lr=3e-4).",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features_v2_improved/",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/lora_optimized_dataset",
+  "base": {
+    "rank": 128,
+    "lr": 3e-4,
+    "steps": 5000,
+    "batch_size": 16,
+    "warmup_steps": 100,
+    "save_every": 1000,
+    "seed": 42,
+    "init_mode": "pissa",
+    "use_rslora": true,
+    "target": "attn.qkv",
+    "timestep_mode": "uniform",
+    "lr_schedule": "constant"
+  },
+  "experiments": [
+    {
+      "id": "baseline",
+      "description": "Control: known-best config (PiSSA r128 lr=3e-4) on the optimized dataset. No latent augmentation."
+    },
+    {
+      "id": "latent_mixup",
+      "description": "Latent mixup alpha=0.4 (MusicLDM). Tests if mixing training latents reduces memorization on 134 clips.",
+      "latent_mixup_alpha": 0.4
+    },
+    {
+      "id": "latent_noise",
+      "description": "Latent noise sigma=0.02. Mild Gaussian noise on training latents for regularization.",
+      "latent_noise_sigma": 0.02
+    },
+    {
+      "id": "mixup_and_noise",
+      "description": "Both latent mixup (0.4) and noise (0.02). Combined regularization.",
+      "latent_mixup_alpha": 0.4,
+      "latent_noise_sigma": 0.02
+    },
+    {
+      "id": "cosine_schedule",
+      "description": "Cosine LR decay. lr=3e-4 was stable with constant, but cosine may extract more from 5k steps.",
+      "lr_schedule": "cosine"
+    },
+    {
+      "id": "cosine_mixup",
+      "description": "Cosine LR + latent mixup. Best regularization combo candidate.",
+      "lr_schedule": "cosine",
+      "latent_mixup_alpha": 0.4
+    },
+    {
+      "id": "logit_normal",
+      "description": "Logit-normal timestep sampling (sigma=1.0). Concentrates training near t=0.5 where flow matching is hardest.",
+      "timestep_mode": "logit_normal"
+    },
+    {
+      "id": "curriculum_mixup",
+      "description": "Curriculum timesteps (logit_normal first 60%, then uniform) + latent mixup. Full regularization stack.",
+      "timestep_mode": "curriculum",
+      "latent_mixup_alpha": 0.4
+    }
+  ]
+}
@@ -0,0 +1,62 @@
+{
+  "name": "pissa_sweep",
+  "description": "PiSSA vs standard init ablation at rank 128. Best prior config (lr=3e-4, bs=16, 10k steps) as baseline. PiSSA starts on-manifold via SVD init — should eliminate intruder dimensions. rsLoRA stabilises scaling at high rank.",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/pissa_sweep",
+  "base": {
+    "steps": 10000,
+    "rank": 128,
+    "alpha": 0.0,
+    "lr": 3e-4,
+    "batch_size": 16,
+    "warmup_steps": 200,
+    "grad_accum": 1,
+    "save_every": 2000,
+    "seed": 42,
+    "target": "attn.qkv",
+    "timestep_mode": "uniform",
+    "lora_dropout": 0.0,
+    "lora_plus_ratio": 1.0,
+    "lr_schedule": "constant",
+    "init_mode": "pissa",
+    "use_rslora": true
+  },
+  "experiments": [
+    {
+      "id": "standard_baseline",
+      "description": "Standard Kaiming init + classic alpha/rank scaling. Replicates best prior config for A/B comparison.",
+      "init_mode": "standard",
+      "use_rslora": false
+    },
+    {
+      "id": "pissa_rslora",
+      "description": "PiSSA init + rsLoRA scaling. Full Tier-S config. Should start on-manifold and avoid intruder dimensions."
+    },
+    {
+      "id": "pissa_classic_scale",
+      "description": "PiSSA init + classic alpha/rank scaling. Isolates PiSSA contribution from rsLoRA.",
+      "use_rslora": false
+    },
+    {
+      "id": "standard_rslora",
+      "description": "Standard init + rsLoRA only. Isolates rsLoRA contribution from PiSSA.",
+      "init_mode": "standard"
+    },
+    {
+      "id": "pissa_rslora_lr1e-4",
+      "description": "PiSSA+rsLoRA at lower lr=1e-4. PiSSA starts closer to optimum — may need less aggressive lr.",
+      "lr": 1e-4
+    },
+    {
+      "id": "pissa_rslora_lr5e-4",
+      "description": "PiSSA+rsLoRA at higher lr=5e-4. Test if on-manifold start tolerates faster learning.",
+      "lr": 5e-4
+    },
+    {
+      "id": "pissa_rslora_dropout",
+      "description": "PiSSA+rsLoRA with dropout 0.05. Note: PiSSA forces dropout=0 (principal components should not be dropped) — this tests standard init with rsLoRA + dropout.",
+      "init_mode": "standard",
+      "lora_dropout": 0.05
+    }
+  ]
+}
@@ -0,0 +1,30 @@
+{
+  "name": "vocoder_finetune",
+  "description": "Single run with fine-tuned BJ BigVGAN vocoder injected. Validates vocoder integration with LoRA training. Best known config: lr=3e-4, rank=128.",
+  "data_dir": "/media/unraid/davinci/Selva/BJ/features",
+  "output_root": "/media/unraid/davinci/Selva/BJ/experiment/vocoder_finetune",
+  "base": {
+    "steps": 10000,
+    "rank": 128,
+    "alpha": 0.0,
+    "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": "r128_lr_3e4_bj_vocoder",
+      "description": "lr=3e-4 rank=128 with fine-tuned BJ BigVGAN vocoder. Direct comparison baseline against previous best g1_r128_lr_3e4."
+    }
+  ]
+}
@@ -21,7 +21,20 @@ _NODES = {
    "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"),
+    "SelvaBigvganScheduler":                 (".selva_bigvgan_scheduler",                  "SelvaBigvganScheduler",                 "SelVA BigVGAN Scheduler"),
    "SelvaDittoOptimizer":                   (".selva_ditto_optimizer",                    "SelvaDittoOptimizer",                   "SelVA DITTO Optimizer"),
+    "SelvaDatasetLoader":          (".selva_dataset_pipeline", "SelvaDatasetLoader",          "SelVA Dataset Loader"),
+    "SelvaDatasetResampler":       (".selva_dataset_pipeline", "SelvaDatasetResampler",       "SelVA Dataset Resampler"),
+    "SelvaDatasetLUFSNormalizer":  (".selva_dataset_pipeline", "SelvaDatasetLUFSNormalizer",  "SelVA Dataset LUFS Normalizer"),
+    "SelvaDatasetCompressor":      (".selva_dataset_pipeline", "SelvaDatasetCompressor",      "SelVA Dataset Compressor"),
+    "SelvaDatasetInspector":       (".selva_dataset_pipeline", "SelvaDatasetInspector",       "SelVA Dataset Inspector"),
+    "SelvaDatasetItemExtractor":   (".selva_dataset_pipeline", "SelvaDatasetItemExtractor",   "SelVA Dataset Item Extractor"),
+    "SelvaDatasetSaver":           (".selva_dataset_pipeline", "SelvaDatasetSaver",           "SelVA Dataset Saver"),
+    "SelvaHarmonicExciter":        (".selva_audio_postprocess", "SelvaHarmonicExciter",        "SelVA Harmonic Exciter"),
+    "SelvaOutputNormalizer":       (".selva_audio_postprocess", "SelvaOutputNormalizer",       "SelVA Output Normalizer"),
+    "SelvaDatasetSpectralMatcher": (".selva_dataset_pipeline",  "SelvaDatasetSpectralMatcher", "SelVA Dataset Spectral Matcher"),
+    "SelvaDatasetHfSmoother":      (".selva_dataset_pipeline",  "SelvaDatasetHfSmoother",      "SelVA Dataset HF Smoother"),
+    "SelvaDatasetAugmenter":       (".selva_dataset_pipeline",  "SelvaDatasetAugmenter",       "SelVA Dataset Augmenter"),
 }

 for key, (module_path, class_name, display_name) in _NODES.items():
@@ -1,15 +1,15 @@
 """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
+training dataset and recording how target style's conditioning shifts the DiT hidden
 states vs. empty/unconditional conditioning.

 For each block i:
-    steering[i] = mean(latent_hidden | BJ conditions)
+    steering[i] = mean(latent_hidden | target style 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
+steering_strength input) to nudge the denoising trajectory toward target style's
 activation patterns without modifying any model weights.
 """

@@ -58,7 +58,7 @@ 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
+    target style-conditioned and empty-conditioned inputs, then saves the mean
    difference per DiT block to a .pt file.
    """

@@ -69,7 +69,7 @@ class SelvaActivationSteeringExtractor:
    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) − "
+        "Computes per-block activation steering vectors: mean(target style activations) − "
        "mean(empty activations) at each DiT block. Load the result with "
        "SelVA Activation Steering Loader and connect to the Sampler."
    )
@@ -124,7 +124,7 @@ class SelvaActivationSteeringExtractor:
        indices = random.choices(range(len(dataset)), k=n_samples)

        n_blocks     = len(generator.joint_blocks) + len(generator.fused_blocks)
-        bj_sums      = [None] * n_blocks
+        style_sums      = [None] * n_blocks
        empty_sums   = [None] * n_blocks
        counts       = [0] * n_blocks

@@ -157,15 +157,15 @@ class SelvaActivationSteeringExtractor:
                device=device, dtype=dtype,
            )

-            bj_acts    = _collect_activations(generator, conditions,       latent, t_tensor)
+            style_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()
+            for i, (st, em) in enumerate(zip(style_acts, empty_acts)):
+                if style_sums[i] is None:
+                    style_sums[i]    = st.clone()
                    empty_sums[i] = em.clone()
                else:
-                    bj_sums[i]    += bj
+                    style_sums[i]    += st
                    empty_sums[i] += em
                counts[i] += 1

@@ -173,10 +173,10 @@ class SelvaActivationSteeringExtractor:
            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 vector per block: mean(target style) - mean(empty)
        steering_vectors = []
        for i in range(n_blocks):
-            vec = (bj_sums[i] - empty_sums[i]) / counts[i]   # [hidden]
+            vec = (style_sums[i] - empty_sums[i]) / counts[i]   # [hidden]
            steering_vectors.append(vec)

            norm = vec.norm().item()
@@ -0,0 +1,153 @@
+"""SelVA Audio Post-Processing nodes.
+
+Post-generation enhancement applied to standard AUDIO outputs:
+  SelvaHarmonicExciter    — multi-band harmonic exciter (HPF → tanh → mix)
+  SelvaOutputNormalizer   — LUFS normalization + true peak limiting
+"""
+
+import numpy as np
+import torch
+
+from .utils import SELVA_CATEGORY
+
+
+class SelvaHarmonicExciter:
+    """Multi-band harmonic exciter for post-generation enhancement.
+
+    Isolates high-frequency content above a cutoff, applies tanh saturation
+    to generate 2nd/3rd harmonics, then mixes back with the dry signal.
+    Restores harmonic richness lost during BigVGAN vocoder reconstruction.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio": ("AUDIO",),
+                "cutoff_hz": ("FLOAT", {
+                    "default": 3000.0, "min": 500.0, "max": 16000.0, "step": 100.0,
+                    "tooltip": "Highpass cutoff frequency in Hz. Only content above this is excited. "
+                               "3000 Hz targets the upper harmonics BigVGAN tends to smear.",
+                }),
+                "drive": ("FLOAT", {
+                    "default": 2.0, "min": 1.0, "max": 10.0, "step": 0.5,
+                    "tooltip": "Saturation drive. Higher = more harmonics generated. "
+                               "2-3 is subtle, 5+ is aggressive.",
+                }),
+                "mix": ("FLOAT", {
+                    "default": 0.15, "min": 0.0, "max": 1.0, "step": 0.05,
+                    "tooltip": "Wet/dry blend. 0.1-0.2 is subtle enhancement, "
+                               "0.5+ is aggressive harmonic addition.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = ("AUDIO",)
+    RETURN_NAMES  = ("audio",)
+    FUNCTION      = "excite"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Multi-band harmonic exciter. Applies tanh saturation to the high-frequency band "
+        "to restore harmonics lost during BigVGAN vocoder reconstruction. "
+        "Uses pedalboard.HighpassFilter for band isolation."
+    )
+
+    def excite(self, audio, cutoff_hz: float, drive: float, mix: float):
+        from pedalboard import Pedalboard, HighpassFilter
+
+        wav = audio["waveform"][0]   # [C, T]
+        sr  = audio["sample_rate"]
+
+        wav_np = wav.float().numpy()   # [C, T]
+
+        # Isolate HF band
+        board = Pedalboard([HighpassFilter(cutoff_frequency_hz=cutoff_hz)])
+        hf = board(wav_np, sr)         # [C, T]
+
+        # Tanh saturation — normalize by drive so output stays in [-1, 1]
+        excited = np.tanh(hf * drive) / max(drive, 1.0)
+
+        # Mix back with dry
+        mixed = wav_np + mix * excited
+
+        # Soft clip to prevent going over
+        mixed = np.tanh(mixed)
+
+        wav_out = torch.from_numpy(mixed).unsqueeze(0)  # [1, C, T]
+        print(
+            f"[HarmonicExciter] cutoff={cutoff_hz}Hz  drive={drive}  mix={mix:.0%}",
+            flush=True,
+        )
+        return ({"waveform": wav_out, "sample_rate": sr},)
+
+
+
+class SelvaOutputNormalizer:
+    """Normalize generated audio to a target LUFS level with true peak limiting.
+
+    Apply as the final node before saving — brings generated audio to a
+    consistent loudness target regardless of input video loudness variation.
+    Uses pyloudnorm (BS.1770-4).
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio": ("AUDIO",),
+                "target_lufs": ("FLOAT", {
+                    "default": -14.0, "min": -40.0, "max": -6.0, "step": 0.5,
+                    "tooltip": "Target integrated loudness in LUFS. "
+                               "-14 LUFS for streaming (Spotify/YouTube), "
+                               "-9 to -7 for production masters.",
+                }),
+                "true_peak_dbtp": ("FLOAT", {
+                    "default": -1.0, "min": -6.0, "max": 0.0, "step": 0.5,
+                    "tooltip": "True peak ceiling in dBTP applied after LUFS gain.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = ("AUDIO",)
+    RETURN_NAMES  = ("audio",)
+    FUNCTION      = "normalize"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Normalize output audio to a target LUFS level (BS.1770-4) with true peak limiting. "
+        "Apply as the last node before saving. Uses pyloudnorm."
+    )
+
+    def normalize(self, audio, target_lufs: float, true_peak_dbtp: float):
+        import pyloudnorm as pyln
+
+        wav = audio["waveform"][0]   # [C, T]
+        sr  = audio["sample_rate"]
+
+        tp_linear = 10.0 ** (true_peak_dbtp / 20.0)
+
+        wav_np = wav.permute(1, 0).double().numpy()   # [T, C]
+        if wav_np.shape[1] == 1:
+            wav_np = wav_np[:, 0]                     # [T] mono
+
+        meter    = pyln.Meter(sr)
+        loudness = meter.integrated_loudness(wav_np)
+
+        if not np.isfinite(loudness):
+            print("[OutputNormalizer] Could not measure loudness — clip too short or silent. Passing through.", flush=True)
+            return (audio,)
+
+        gain_db     = target_lufs - loudness
+        gain_linear = 10.0 ** (gain_db / 20.0)
+
+        wav_out = wav * gain_linear
+
+        peak = wav_out.abs().max().item()
+        if peak > tp_linear:
+            wav_out = wav_out * (tp_linear / peak)
+
+        print(
+            f"[OutputNormalizer] {loudness:.1f} LUFS → {target_lufs} LUFS  "
+            f"gain={gain_db:+.1f}dB  TP={true_peak_dbtp}dBTP",
+            flush=True,
+        )
+        return ({"waveform": wav_out.unsqueeze(0), "sample_rate": sr},)
@@ -13,6 +13,7 @@ import torch
 import folder_paths

 from .utils import SELVA_CATEGORY
+from .selva_bigvgan_trainer import inject_gafilters


 class SelvaBigvganLoader:
@@ -60,7 +61,16 @@ class SelvaBigvganLoader:
        else:
            raise ValueError(f"[BigVGAN] Unknown mode: {mode}")

+        # Remember device before injecting new modules (which default to CPU)
+        target_device = next(vocoder.parameters()).device
+
+        if ckpt.get("has_gafilter", False):
+            kernel_size = ckpt.get("gafilter_kernel_size", 9)
+            n_gaf = inject_gafilters(vocoder, kernel_size)
+            print(f"[BigVGAN] GAFilter injected: {n_gaf} filters  kernel={kernel_size}", flush=True)
+
        vocoder.load_state_dict(ckpt["generator"])
+        vocoder.to(target_device)
        vocoder.eval()

        print(f"[BigVGAN] Loaded fine-tuned vocoder from: {p}", flush=True)
@@ -0,0 +1,625 @@
+"""SelVA BigVGAN Vocoder Scheduler — runs a sweep of vocoder fine-tuning experiments.
+
+Each experiment inherits from a shared `base` config and overrides specific keys.
+Audio clips are 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.
+
+JSON format:
+    {
+      "name": "bigvgan_sweep",
+      "description": "optional note",
+      "data_dir": "/path/to/audio/clips",
+      "output_root": "/path/to/output",
+      "base": { "train_mode": "snake_alpha_only", "steps": 2000, "lr": 1e-4, ... },
+      "experiments": [
+        {"id": "baseline", "description": "..."},
+        {"id": "all_5k", "train_mode": "all_params", "steps": 5000, "lr": 1e-5},
+        ...
+      ]
+    }
+"""
+
+import copy
+import csv
+import json
+import threading
+import time
+import traceback
+from datetime import datetime, timezone
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchaudio
+
+import comfy.utils
+import comfy.model_management
+import folder_paths
+
+from .utils import SELVA_CATEGORY, get_device, soft_empty_cache
+from .selva_bigvgan_trainer import (
+    _do_train,
+    _pregenerate_lora_mels,
+    _load_wav,
+)
+from .selva_lora_trainer import _smooth_losses, _pil_to_tensor
+from .selva_lora_scheduler import (
+    _get_system_info,
+    _resolve_path,
+    _draw_comparison_curves,
+)
+
+
+# Defaults mirror SelvaBigvganTrainer INPUT_TYPES defaults
+_PARAM_DEFAULTS = {
+    "train_mode":           "snake_alpha_only",
+    "steps":                2000,
+    "lr":                   1e-4,
+    "batch_size":           4,
+    "segment_seconds":      2.0,
+    "lambda_l2sp":          1e-3,
+    "use_gafilter":         True,
+    "gafilter_kernel_size": 9,
+    "lambda_phase":         1.0,
+    "save_every":           500,
+    "seed":                 42,
+    "discriminator_path":   "",
+    "lora_adapter":         "",
+}
+
+
+def _merge_config(base: dict, experiment: dict) -> dict:
+    """Merge param defaults + file base + experiment overrides."""
+    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 _parse_training_log(log_path: Path) -> list:
+    """Parse BigVGAN training CSV → list of total_loss values."""
+    losses = []
+    if not log_path.exists():
+        return losses
+    try:
+        with open(log_path) as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                losses.append(float(row["total_loss"]))
+    except Exception:
+        pass
+    return losses
+
+
+def _loss_at_steps(loss_history: list, log_interval: int, save_every: int,
+                   total_steps: int) -> dict:
+    """Build {step: loss} at each save_every boundary.
+
+    Uses round-to-nearest to handle log_interval that doesn't divide
+    save_every evenly (e.g. steps=3000 → log_interval=150, save_every=1000).
+    """
+    result = {}
+    for target in range(save_every, total_steps + 1, save_every):
+        # loss_history[i] = loss at step (i+1)*log_interval
+        idx = round(target / log_interval) - 1
+        if 0 <= idx < len(loss_history):
+            result[str(target)] = round(loss_history[idx], 6)
+    return result
+
+
+class SelvaBigvganScheduler:
+    """Runs a sweep of BigVGAN vocoder fine-tuning experiments from a JSON file.
+
+    Audio clips are loaded once and reused across all experiments. Each experiment
+    deep-copies the vocoder 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 BigVGAN vocoder fine-tuning experiments defined in a JSON sweep file. "
+        "Audio clips are loaded once and reused across all experiments. "
+        "Results (loss, config, checkpoint paths) are collected in experiment_summary.json."
+    )
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("SELVA_MODEL",),
+                "experiments_file": ("STRING", {
+                    "default": "bigvgan_experiments.json",
+                    "tooltip": (
+                        "Path to JSON sweep file. Relative paths resolve to the ComfyUI "
+                        "models directory; absolute paths are used as-is."
+                    ),
+                }),
+            }
+        }
+
+    def run(self, model, experiments_file):
+        # ------------------------------------------------------------------
+        # 1. Read + validate the JSON file
+        # ------------------------------------------------------------------
+        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"[BigVGAN 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("[BigVGAN Scheduler] 'experiments' list is missing or empty.")
+        for i, exp in enumerate(spec["experiments"]):
+            if "id" not in exp:
+                raise ValueError(
+                    f"[BigVGAN 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("[BigVGAN 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"bigvgan_sweeps/{sweep_name}"))
+        output_root.mkdir(parents=True, exist_ok=True)
+
+        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":
+            original_vocoder = feature_utils.tod.vocoder.vocoder
+            sample_rate = 16_000
+        elif mode == "44k":
+            original_vocoder = feature_utils.tod.vocoder
+            sample_rate = 44_100
+        else:
+            raise ValueError(f"[BigVGAN Scheduler] Unknown mode: {mode}")
+
+        print(f"\n[BigVGAN Scheduler] Sweep '{sweep_name}': "
+              f"{len(spec['experiments'])} experiment(s)", flush=True)
+        if description:
+            print(f"[BigVGAN Scheduler] {description}", flush=True)
+        print(f"[BigVGAN Scheduler] data_dir    = {data_dir}", flush=True)
+        print(f"[BigVGAN Scheduler] output_root = {output_root}\n", flush=True)
+
+        # ------------------------------------------------------------------
+        # 3. Load audio clips once
+        # ------------------------------------------------------------------
+        # Find minimum segment length across all experiments so we load enough
+        min_segment_seconds = float("inf")
+        for exp in spec["experiments"]:
+            cfg = _merge_config(base_cfg, exp)
+            min_segment_seconds = min(min_segment_seconds, float(cfg.get("segment_seconds", 2.0)))
+        min_segment_samples = int(min_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 Scheduler] No audio files in {data_dir}")
+
+        print(f"[BigVGAN Scheduler] 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] >= min_segment_samples:
+                    clips.append(wav.cpu())
+                else:
+                    print(f"  [BigVGAN Scheduler] Skip {af.name}: "
+                          f"shorter than {min_segment_seconds}s", flush=True)
+            except Exception as e:
+                print(f"  [BigVGAN Scheduler] Failed {af.name}: {e}", flush=True)
+
+        if not clips:
+            raise RuntimeError(
+                f"[BigVGAN Scheduler] No usable clips (need audio >= {min_segment_seconds}s)"
+            )
+        print(f"[BigVGAN Scheduler] {len(clips)} clips ready\n", flush=True)
+
+        # ------------------------------------------------------------------
+        # 4. Offload unused components to free VRAM
+        # ------------------------------------------------------------------
+        comfy.model_management.unload_all_models()
+        feature_utils.to("cpu")
+        if "generator" in model:
+            model["generator"].to("cpu")
+        if "video_enc" in model:
+            model["video_enc"].to("cpu")
+        soft_empty_cache()
+
+        # ------------------------------------------------------------------
+        # 5. Pre-compute text CLIP embeddings if any experiment uses LoRA
+        # ------------------------------------------------------------------
+        text_clip_cache = {}
+        any_lora = any(
+            _merge_config(base_cfg, exp).get("lora_adapter", "")
+            for exp in spec["experiments"]
+        )
+        if any_lora:
+            npz_files = sorted(data_dir.glob("*.npz"))
+            if npz_files:
+                prompt_map = {}
+                prompts_file = data_dir / "prompts.txt"
+                if prompts_file.exists():
+                    for line in prompts_file.read_text(encoding="utf-8").splitlines():
+                        line = line.strip()
+                        if not line or line.startswith("#"):
+                            continue
+                        if "|" in line:
+                            fname, prompt = line.split("|", 1)
+                            prompt_map[fname.strip()] = prompt.strip()
+                default_prompt = data_dir.name
+
+                clip_model = feature_utils.clip_model
+                if clip_model is not None:
+                    clip_model.to(device)
+                try:
+                    for npz_path in npz_files:
+                        data = dict(np.load(str(npz_path), allow_pickle=False))
+                        prompt = prompt_map.get(
+                            npz_path.name, data.get("prompt", default_prompt)
+                        )
+                        if isinstance(prompt, np.ndarray):
+                            prompt = str(prompt)
+                        tc = feature_utils.encode_text_clip([prompt])
+                        text_clip_cache[npz_path.name] = tc.clone().detach().cpu()
+                finally:
+                    if clip_model is not None:
+                        clip_model.to("cpu")
+                    soft_empty_cache()
+                    if device.type == "cuda":
+                        torch.cuda.empty_cache()
+                print(f"[BigVGAN Scheduler] Pre-encoded {len(text_clip_cache)} "
+                      f"CLIP embeddings", flush=True)
+
+        # ------------------------------------------------------------------
+        # 6. Build or restore the 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"])
+                        lh = rec["results"].get("loss_history", [])
+                        all_curve_data.append({
+                            "id":           rec["id"],
+                            "loss_history": lh,
+                            "log_interval": rec["results"].get("log_interval", 100),
+                            "start_step":   0,
+                        })
+                summary = existing
+                summary["completed_at"] = None
+                if completed_ids:
+                    print(f"[BigVGAN Scheduler] Resuming — skipping "
+                          f"{len(completed_ids)} completed: "
+                          f"{sorted(completed_ids)}", flush=True)
+            except Exception as e:
+                print(f"[BigVGAN Scheduler] Could not read existing summary "
+                      f"({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":      len(clips),
+                "experiments":  [],
+            }
+
+        def _write_summary():
+            summary_path.write_text(
+                json.dumps(summary, indent=2), encoding="utf-8"
+            )
+
+        _write_summary()
+
+        # ------------------------------------------------------------------
+        # 7. Compute total steps for progress bar
+        # ------------------------------------------------------------------
+        total_steps = 0
+        for exp in spec["experiments"]:
+            if exp["id"] not in completed_ids:
+                cfg = _merge_config(base_cfg, exp)
+                total_steps += int(cfg.get("steps", 2000))
+        pbar = comfy.utils.ProgressBar(max(total_steps, 1))
+
+        # ------------------------------------------------------------------
+        # 8. Run experiments in a worker thread
+        # ------------------------------------------------------------------
+        # BigVGAN training requires a fresh thread because ComfyUI runs nodes
+        # inside torch.inference_mode(). inference_mode is thread-local — a
+        # new thread starts with it OFF, so all tensor operations produce
+        # normal autograd-compatible tensors.
+        _exc = [None]
+
+        def _worker():
+            try:
+                for exp in spec["experiments"]:
+                    exp_id   = exp["id"]
+                    exp_desc = exp.get("description", "")
+
+                    if exp_id in completed_ids:
+                        print(f"[BigVGAN Scheduler] Skipping '{exp_id}' "
+                              f"(already completed)", flush=True)
+                        continue
+
+                    cfg = _merge_config(base_cfg, exp)
+
+                    # ── Extract experiment parameters ────────────────────
+                    train_mode   = str(cfg.get("train_mode", "snake_alpha_only"))
+                    exp_steps    = int(cfg.get("steps", 2000))
+                    exp_lr       = float(cfg.get("lr", 1e-4))
+                    exp_bs       = int(cfg.get("batch_size", 4))
+                    exp_seg_s    = float(cfg.get("segment_seconds", 2.0))
+                    exp_l2sp     = float(cfg.get("lambda_l2sp", 1e-3))
+                    exp_gafilter = bool(cfg.get("use_gafilter", True))
+                    exp_gaf_ks   = int(cfg.get("gafilter_kernel_size", 9))
+                    exp_phase    = float(cfg.get("lambda_phase", 1.0))
+                    exp_save     = int(cfg.get("save_every", 500))
+                    exp_seed     = int(cfg.get("seed", 42))
+                    exp_disc     = str(cfg.get("discriminator_path", ""))
+                    exp_lora     = str(cfg.get("lora_adapter", ""))
+
+                    segment_samples = int(exp_seg_s * sample_rate)
+
+                    # Filter clips long enough for this experiment
+                    exp_clips = [c for c in clips if c.shape[0] >= segment_samples]
+                    if not exp_clips:
+                        print(f"[BigVGAN Scheduler] '{exp_id}' skipped: "
+                              f"no clips >= {exp_seg_s}s", flush=True)
+                        summary["experiments"].append({
+                            "id": exp_id, "description": exp_desc,
+                            "config": dict(cfg),
+                            "results": {
+                                "status": "failed",
+                                "error": f"No clips >= {exp_seg_s}s",
+                                "duration_seconds": 0,
+                            },
+                            "checkpoint_path": None,
+                            "output_dir": str(output_root / exp_id),
+                        })
+                        _write_summary()
+                        continue
+
+                    # ── Resolve discriminator path ───────────────────────
+                    disc_path = None
+                    if exp_disc:
+                        disc_path = Path(exp_disc.strip())
+                        if not disc_path.is_absolute():
+                            disc_path = (
+                                Path(folder_paths.get_output_directory()) / disc_path
+                            )
+                        if not disc_path.exists():
+                            print(f"[BigVGAN Scheduler] '{exp_id}': "
+                                  f"discriminator not found: {disc_path}",
+                                  flush=True)
+                            disc_path = None
+
+                    # ── Pre-generate LoRA mels (disk-cached) ─────────────
+                    lora_mel_pairs = None
+                    if exp_lora:
+                        lora_path = Path(exp_lora.strip())
+                        if not lora_path.is_absolute():
+                            lora_path = Path(folder_paths.base_path) / lora_path
+                        if lora_path.exists():
+                            seq_cfg = model["seq_cfg"]
+                            lora_mel_pairs = _pregenerate_lora_mels(
+                                model, data_dir, str(lora_path),
+                                device, dtype, sample_rate,
+                                seq_cfg.duration, seed=exp_seed,
+                                cache_dir=str(output_root),
+                                text_clip_cache=text_clip_cache,
+                            )
+                            if not lora_mel_pairs:
+                                print(f"[BigVGAN Scheduler] '{exp_id}': "
+                                      f"no LoRA mel pairs generated",
+                                      flush=True)
+                                lora_mel_pairs = None
+                            if device.type == "cuda":
+                                torch.cuda.empty_cache()
+                        else:
+                            print(f"[BigVGAN Scheduler] '{exp_id}': "
+                                  f"LoRA adapter not found: {lora_path}",
+                                  flush=True)
+
+                    # ── Output dir ───────────────────────────────────────
+                    exp_dir = output_root / exp_id
+                    exp_dir.mkdir(parents=True, exist_ok=True)
+                    out_path = exp_dir / f"bigvgan_{exp_id}.pt"
+
+                    print(f"\n[BigVGAN Scheduler] ── Experiment '{exp_id}' ──",
+                          flush=True)
+                    if exp_desc:
+                        print(f"[BigVGAN Scheduler] {exp_desc}", flush=True)
+                    print(f"[BigVGAN Scheduler] mode={train_mode}  "
+                          f"steps={exp_steps}  lr={exp_lr}  bs={exp_bs}  "
+                          f"seg={exp_seg_s}s  gafilter={exp_gafilter}  "
+                          f"phase={exp_phase}  l2sp={exp_l2sp}", flush=True)
+
+                    exp_record = {
+                        "id":          exp_id,
+                        "description": exp_desc,
+                        "config": {
+                            "train_mode": train_mode, "steps": exp_steps,
+                            "lr": exp_lr, "batch_size": exp_bs,
+                            "segment_seconds": exp_seg_s,
+                            "lambda_l2sp": exp_l2sp,
+                            "use_gafilter": exp_gafilter,
+                            "gafilter_kernel_size": exp_gaf_ks,
+                            "lambda_phase": exp_phase,
+                            "save_every": exp_save, "seed": exp_seed,
+                            "discriminator_path": exp_disc,
+                            "lora_adapter": exp_lora,
+                        },
+                        "results":         {"status": "running"},
+                        "checkpoint_path": None,
+                        "output_dir":      str(exp_dir),
+                    }
+                    summary["experiments"].append(exp_record)
+                    _write_summary()
+
+                    t_start = time.monotonic()
+                    try:
+                        # Ensure mel_converter is on device for this experiment
+                        mel_converter.to(device)
+
+                        # Fresh vocoder copy — _do_train modifies it in-place
+                        vocoder_copy = copy.deepcopy(original_vocoder)
+
+                        checkpoint_path = _do_train(
+                            vocoder_copy, mel_converter, exp_clips,
+                            device, dtype, strategy, feature_utils,
+                            segment_samples, sample_rate,
+                            train_mode, exp_steps, exp_lr, exp_bs,
+                            exp_l2sp, exp_gafilter, exp_gaf_ks,
+                            exp_phase, exp_save, exp_seed,
+                            out_path, disc_path, pbar,
+                            lora_mel_pairs,
+                        )
+
+                        duration = time.monotonic() - t_start
+
+                        # Parse training CSV for loss history
+                        log_path = exp_dir / f"bigvgan_{exp_id}_training_log.csv"
+                        loss_history = _parse_training_log(log_path)
+                        log_interval = max(1, exp_steps // 20)
+                        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
+                        )
+
+                        if loss_history:
+                            quarter = max(1, len(loss_history) // 4)
+                            loss_std = round(
+                                float(np.std(loss_history[-quarter:])), 6
+                            )
+                        else:
+                            loss_std = 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,
+                            "loss_at_steps":        _loss_at_steps(
+                                loss_history, log_interval,
+                                exp_save, exp_steps,
+                            ),
+                            "loss_history": [
+                                round(v, 6) for v in loss_history
+                            ],
+                            "log_interval":     log_interval,
+                            "duration_seconds": round(duration, 1),
+                        }
+                        exp_record["checkpoint_path"] = checkpoint_path
+
+                        all_curve_data.append({
+                            "id":           exp_id,
+                            "loss_history": loss_history,
+                            "log_interval": log_interval,
+                            "start_step":   0,
+                        })
+
+                    except Exception as e:
+                        duration = time.monotonic() - t_start
+                        print(f"[BigVGAN Scheduler] Experiment '{exp_id}' "
+                              f"failed: {e}", flush=True)
+                        traceback.print_exc()
+                        exp_record["results"] = {
+                            "status":           "failed",
+                            "error":            str(e),
+                            "duration_seconds": round(duration, 1),
+                        }
+                    finally:
+                        # Clean up vocoder copy to free VRAM
+                        soft_empty_cache()
+
+                    _write_summary()
+
+            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]
+
+        # ------------------------------------------------------------------
+        # 9. Finalise summary
+        # ------------------------------------------------------------------
+        summary["completed_at"] = datetime.now(timezone.utc).isoformat()
+        _write_summary()
+        print(f"\n[BigVGAN Scheduler] Sweep complete. "
+              f"Summary: {summary_path}", flush=True)
+
+        # ------------------------------------------------------------------
+        # 10. 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)
@@ -26,10 +26,14 @@ Save format: {'generator': vocoder.state_dict()} — same as the original
 BigVGAN checkpoint so it can be loaded with SelVA BigVGAN Loader.
 """

+import copy
+import hashlib
+import json as _json
 import random
 import threading
 from pathlib import Path

+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -121,8 +125,9 @@ class _DiscriminatorR(nn.Module):
        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]
+        orig_dtype = x.dtype
+        x = torch.stft(x.float(), n, hop, win, window, center=False, return_complex=True)
+        x = x.abs().to(orig_dtype).unsqueeze(1)  # [B, 1, freq, time]
        return x

    def forward(self, x):
@@ -161,6 +166,78 @@ def _feature_matching_loss(fmaps_real, fmaps_gen):
    return loss / len(fmaps_real)


+# ---------------------------------------------------------------------------
+# AF-Vocoder GAFilter (Interspeech 2025)
+# ---------------------------------------------------------------------------
+
+class GAFilter(nn.Module):
+    """Learnable per-channel depthwise FIR filter inserted after Snake activations.
+
+    Initialized as identity (delta at center) so training starts from the
+    pretrained vocoder's behaviour. Learns to shape the per-channel frequency
+    response to fix harmonic artifacts.
+    """
+    def __init__(self, channels: int, kernel_size: int = 9):
+        super().__init__()
+        self.conv = nn.Conv1d(
+            channels, channels, kernel_size,
+            padding=kernel_size // 2, groups=channels, bias=False,
+        )
+        nn.init.zeros_(self.conv.weight)
+        self.conv.weight.data[:, 0, kernel_size // 2] = 1.0  # identity
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class _ActivationWithGAFilter(nn.Module):
+    def __init__(self, activation: nn.Module, gafilter: GAFilter):
+        super().__init__()
+        self.activation = activation
+        self.gafilter   = gafilter
+
+    def forward(self, x):
+        T = x.shape[-1]
+        out = self.gafilter(self.activation(x))
+        # Guarantee exact length match — Activation1d's anti-alias resampling
+        # (Kaiser sinc filter with asymmetric pad_left/pad_right) can produce
+        # ±1-2 sample rounding in edge cases that break the resblock residual add.
+        if out.shape[-1] != T:
+            if out.shape[-1] > T:
+                out = out[..., :T]
+            else:
+                out = torch.nn.functional.pad(out, (0, T - out.shape[-1]))
+        return out
+
+
+def inject_gafilters(vocoder: nn.Module, kernel_size: int = 9) -> int:
+    """Inject GAFilter after each Activation1d in BigVGAN residual blocks.
+
+    Modifies vocoder in-place. GAFilter weights appear in vocoder.state_dict()
+    under resblocks.{i}.activations.{j}.gafilter.conv.weight — so a normal
+    load_state_dict call after injection will populate them correctly.
+
+    Returns the number of injected filters.
+    """
+    count = 0
+    for resblock in getattr(vocoder, "resblocks", []):
+        activations = getattr(resblock, "activations", None)
+        if activations is None:
+            continue
+        for j in range(len(activations)):
+            act1d = activations[j]
+            act   = getattr(act1d, "act", None)
+            if act is None:
+                continue
+            alpha = getattr(act, "alpha", None)
+            if alpha is None:
+                continue
+            channels = alpha.shape[0]
+            activations[j] = _ActivationWithGAFilter(act1d, GAFilter(channels, kernel_size))
+            count += 1
+    return count
+
+
 # ---------------------------------------------------------------------------
 # Utility helpers
 # ---------------------------------------------------------------------------
@@ -247,8 +324,9 @@ def _stft_mag(wav, n_fft, hop_length, win_length, device):

 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)
+    # cuFFT requires float32 regardless of model dtype
+    pred   = pred_wav.squeeze(1).float()    # [B, T]
+    target = target_wav.squeeze(1).float()
    loss = torch.zeros(1, device=device)
    for n_fft, hop, win in _STFT_RESOLUTIONS:
        pm = _stft_mag(pred,   n_fft, hop, win, device)
@@ -258,6 +336,251 @@ def _multi_resolution_stft_loss(pred_wav, target_wav, device):
    return loss / len(_STFT_RESOLUTIONS)


+def _phase_aware_stft_loss(pred_wav, target_wav, device):
+    """FA-GAN complex STFT loss: L1 on real, imaginary, and magnitude components.
+
+    Penalizes phase smearing directly — plain magnitude loss cannot distinguish
+    a correct spectrum with wrong phase from a smeared spectrum with random phase.
+    Based on FA-GAN (arXiv:2407.04575), applied across three STFT resolutions.
+    inputs: [B, 1, T]
+    """
+    # cuFFT requires float32 regardless of model dtype
+    pred   = pred_wav.squeeze(1).float()    # [B, T]
+    target = target_wav.squeeze(1).float()
+    loss = torch.zeros(1, device=device)
+    for n_fft, hop, win in _STFT_RESOLUTIONS:
+        window = torch.hann_window(win, device=device)
+        ps = torch.stft(pred,   n_fft, hop, win, window, center=True, return_complex=True)
+        ts = torch.stft(target, n_fft, hop, win, window, center=True, return_complex=True)
+        T  = min(ps.shape[-1], ts.shape[-1])
+        ps, ts = ps[..., :T], ts[..., :T]
+        loss = loss + F.l1_loss(ps.real, ts.real)
+        loss = loss + F.l1_loss(ps.imag, ts.imag)
+        loss = loss + F.l1_loss(ps.abs(), ts.abs())
+    return loss / (len(_STFT_RESOLUTIONS) * 3)
+
+
+# ---------------------------------------------------------------------------
+# LoRA mel pre-generation
+# ---------------------------------------------------------------------------
+
+_AUDIO_EXTS = (".wav", ".flac", ".mp3", ".ogg", ".aac")
+
+
+def _find_audio_for_npz(npz_path: Path):
+    """Find audio file matching an .npz stem (same as LoRA trainer _find_audio)."""
+    for ext in _AUDIO_EXTS:
+        c = npz_path.with_suffix(ext)
+        if c.exists():
+            return c
+    return None
+
+
+def _lora_mel_cache_key(lora_adapter_path, data_dir, seed, num_steps,
+                        cfg_strength, duration, sample_rate):
+    """Build a deterministic hash from all parameters that affect LoRA mel generation."""
+    # Hash the LoRA adapter file content (not path — same file moved = same cache)
+    h = hashlib.sha256()
+    with open(lora_adapter_path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    lora_hash = h.hexdigest()[:16]
+
+    # Hash the sorted .npz file list (names only — content is deterministic per name)
+    npz_names = sorted(p.name for p in Path(data_dir).glob("*.npz"))
+
+    key_data = _json.dumps({
+        "lora_hash": lora_hash,
+        "npz_files": npz_names,
+        "seed": seed,
+        "num_steps": num_steps,
+        "cfg_strength": cfg_strength,
+        "duration": duration,
+        "sample_rate": sample_rate,
+    }, sort_keys=True)
+    return hashlib.sha256(key_data.encode()).hexdigest()[:20]
+
+
+def _pregenerate_lora_mels(model, data_dir, lora_adapter_path, device, dtype,
+                           sample_rate, duration, seed=42, num_steps=25,
+                           cfg_strength=4.5, cache_dir=None,
+                           text_clip_cache=None):
+    """Generate LoRA mels for all clips with matching audio in data_dir.
+
+    Uses the LoRA adapter to run full ODE generation with CFG → VAE decode →
+    mel for each clip's conditioning features.  CFG matches the sampler's
+    default (4.5) so the degraded mels the vocoder trains on are representative
+    of what it will see at inference time.
+
+    If cache_dir is provided, results are cached to disk and reused when
+    generation parameters haven't changed.
+
+    text_clip_cache: dict mapping npz filename → pre-computed text CLIP
+    embedding tensor [1, seq, dim]. Pre-computed in the main thread where
+    inference_mode is active (CLIP weights are inference tensors).
+
+    Returns list of (mel [n_mels, T_mel], audio [L]) CPU tensors.
+    """
+    # ── Check cache ──────────────────────────────────────────────────────────
+    cache_path = None
+    if cache_dir is not None:
+        cache_key = _lora_mel_cache_key(
+            lora_adapter_path, data_dir, seed, num_steps,
+            cfg_strength, duration, sample_rate,
+        )
+        cache_path = Path(cache_dir) / f"lora_mels_{cache_key}.pt"
+        if cache_path.exists():
+            print(f"[BigVGAN] Loading cached LoRA mels: {cache_path.name}", flush=True)
+            cached = torch.load(str(cache_path), map_location="cpu", weights_only=True)
+            pairs = [(m, a) for m, a in zip(cached["mels"], cached["audios"])]
+            print(f"[BigVGAN] Loaded {len(pairs)} cached mel/audio pairs", flush=True)
+            return pairs
+
+    from selva_core.model.lora import apply_lora, load_lora
+    from selva_core.model.flow_matching import FlowMatching
+
+    seq_cfg = model["seq_cfg"]
+    feature_utils = model["feature_utils"]
+
+    # Load LoRA checkpoint
+    ckpt = torch.load(str(lora_adapter_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"]))
+    use_rslora = meta.get("use_rslora", False)
+
+    # Apply LoRA to a temporary generator copy
+    generator = copy.deepcopy(model["generator"]).to(device, dtype)
+    n = apply_lora(generator, rank=rank, alpha=alpha,
+                   target_suffixes=tuple(target),
+                   init_mode="standard", use_rslora=use_rslora)
+    load_lora(generator, state_dict)
+    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,
+    )
+    generator.eval()
+    print(f"[BigVGAN] LoRA loaded: {Path(lora_adapter_path).name} "
+          f"(rank={rank}, {n} layers)", flush=True)
+
+    # Load .npz features + matching audio
+    npz_files = sorted(data_dir.glob("*.npz"))
+    if not npz_files:
+        raise ValueError(f"[BigVGAN] No .npz files in {data_dir} — "
+                         "point data_dir to your LoRA training features directory")
+
+    if text_clip_cache is None:
+        text_clip_cache = {}
+
+    fm = FlowMatching(min_sigma=0, inference_mode="euler", num_steps=num_steps)
+    rng = torch.Generator(device=device).manual_seed(seed)
+
+    # Move only tod (VAE+vocoder) to GPU for decode.
+    # CLIP is NOT needed here — text embeddings are pre-computed in the main
+    # thread and passed via text_clip_cache.
+    tod = feature_utils.tod
+    tod_orig_dev = next(tod.parameters()).device
+    tod.to(device)
+
+    pairs = []
+    try:
+        with torch.no_grad():
+            for npz_path in npz_files:
+                audio_path = _find_audio_for_npz(npz_path)
+                if audio_path is None:
+                    print(f"  [BigVGAN] No audio for {npz_path.name}, skipping", flush=True)
+                    continue
+
+                # Load .npz conditioning features
+                data = dict(np.load(str(npz_path), allow_pickle=False))
+                clip_f = torch.from_numpy(data["clip_features"]).to(device, dtype)
+                sync_f = torch.from_numpy(data["sync_features"]).to(device, dtype)
+
+                # Pad/trim to expected sequence lengths
+                c_tgt = seq_cfg.clip_seq_len
+                if clip_f.shape[1] < c_tgt:
+                    clip_f = F.pad(clip_f, (0, 0, 0, c_tgt - clip_f.shape[1]))
+                elif clip_f.shape[1] > c_tgt:
+                    clip_f = clip_f[:, :c_tgt, :]
+
+                s_tgt = seq_cfg.sync_seq_len
+                if sync_f.shape[1] < s_tgt:
+                    sync_f = F.pad(sync_f, (0, 0, 0, s_tgt - sync_f.shape[1]))
+                elif sync_f.shape[1] > s_tgt:
+                    sync_f = sync_f[:, :s_tgt, :]
+
+                # Text CLIP embedding (pre-computed in main thread)
+                if npz_path.name in text_clip_cache:
+                    text_clip = text_clip_cache[npz_path.name].to(device, dtype)
+                else:
+                    print(f"  [BigVGAN] No text embedding for {npz_path.name}, skipping", flush=True)
+                    continue
+
+                # Load clean audio
+                try:
+                    wav, sr = _load_wav(audio_path)
+                    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)
+                    target_len = int(duration * sample_rate)
+                    if wav.shape[0] >= target_len:
+                        wav = wav[:target_len]
+                    else:
+                        wav = F.pad(wav, (0, target_len - wav.shape[0]))
+                except Exception as e:
+                    print(f"  [BigVGAN] Failed loading {audio_path.name}: {e}", flush=True)
+                    continue
+
+                # Generate LoRA latent via ODE with CFG (matches sampler)
+                conditions = generator.preprocess_conditions(clip_f, sync_f, text_clip)
+                empty_conditions = generator.get_empty_conditions(bs=1)
+
+                x0 = torch.randn(1, seq_cfg.latent_seq_len, generator.latent_dim,
+                                 device=device, dtype=dtype, generator=rng)
+
+                def velocity_fn(t, x, _cond=conditions, _empty=empty_conditions,
+                                _cfg=cfg_strength):
+                    return generator.ode_wrapper(t, x, _cond, _empty, _cfg)
+
+                x1_pred = fm.to_data(velocity_fn, x0)
+                x1_unnorm = generator.unnormalize(x1_pred.clone())
+
+                # VAE decode → mel
+                mel = feature_utils.decode(x1_unnorm)  # [1, n_mels, T_mel]
+
+                pairs.append((mel.squeeze(0).float().cpu(), wav.float().cpu()))
+                del x0, x1_pred, x1_unnorm, mel
+                print(f"  [BigVGAN] Generated: {npz_path.stem}", flush=True)
+
+    finally:
+        tod.to(tod_orig_dev)
+        del generator
+        soft_empty_cache()
+
+    print(f"[BigVGAN] Pre-generated {len(pairs)} LoRA mel / clean audio pairs", flush=True)
+
+    # ── Save cache ───────────────────────────────────────────────────────────
+    if cache_path is not None and pairs:
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({
+            "mels":   [m for m, _ in pairs],
+            "audios": [a for _, a in pairs],
+        }, str(cache_path))
+        print(f"[BigVGAN] Cached LoRA mels: {cache_path.name}", flush=True)
+
+    return pairs
+
+
 # ---------------------------------------------------------------------------
 # Node
 # ---------------------------------------------------------------------------
@@ -270,7 +593,7 @@ class SelvaBigvganTrainer:
    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. "
+        "Fine-tunes the BigVGAN vocoder (mel→waveform) on target 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."
@@ -283,7 +606,7 @@ class SelvaBigvganTrainer:
                "model": ("SELVA_MODEL",),
                "data_dir": ("STRING", {
                    "default": "",
-                    "tooltip": "Directory with BJ audio files (.wav/.flac/.mp3). Searched recursively.",
+                    "tooltip": "Directory with target audio files (.wav/.flac/.mp3). Searched recursively.",
                }),
                "output_path": ("STRING", {
                    "default": "bigvgan_bj.pt",
@@ -318,6 +641,26 @@ class SelvaBigvganTrainer:
                        "Increase to 1e-2 for all_params to prevent catastrophic forgetting."
                    ),
                }),
+                "use_gafilter": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": (
+                        "Inject AF-Vocoder GAFilter (Interspeech 2025) after each Snake activation. "
+                        "Adds a learnable depthwise FIR filter per channel, initialized as identity. "
+                        "Trained alongside Snake alphas. Saved into the checkpoint for inference."
+                    ),
+                }),
+                "gafilter_kernel_size": ("INT", {
+                    "default": 9, "min": 3, "max": 31, "step": 2,
+                    "tooltip": "FIR filter length for GAFilter. Must be odd. Larger = wider frequency response control.",
+                }),
+                "lambda_phase": ("FLOAT", {
+                    "default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1,
+                    "tooltip": (
+                        "FA-GAN phase-aware loss weight. Adds L1 loss on real + imaginary + magnitude "
+                        "STFT components, penalizing phase smearing directly. 0 = disabled. "
+                        "1.0 is a good starting point alongside other losses."
+                    ),
+                }),
                "save_every": ("INT", {"default": 500, "min": 50, "max": 10000}),
                "seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFF}),
            },
@@ -331,11 +674,22 @@ class SelvaBigvganTrainer:
                        "the key fix for harmonic smearing. Leave empty to use mel+STFT losses only."
                    ),
                }),
+                "lora_adapter": ("STRING", {
+                    "default": "",
+                    "tooltip": (
+                        "Optional path to a LoRA adapter .pt file. When provided, the trainer "
+                        "pre-generates LoRA-distorted mels for each training clip (using the full "
+                        "generation pipeline) and trains the vocoder to produce clean audio from them. "
+                        "data_dir must contain .npz feature files alongside audio files "
+                        "(same directory used for LoRA training)."
+                    ),
+                }),
            },
        }

    def train(self, model, data_dir, output_path, train_mode, steps, lr, batch_size,
-              segment_seconds, lambda_l2sp, save_every, seed, discriminator_path=""):
+              segment_seconds, lambda_l2sp, use_gafilter, gafilter_kernel_size, lambda_phase,
+              save_every, seed, discriminator_path="", lora_adapter=""):
        import traceback

        device        = get_device()
@@ -374,6 +728,14 @@ class SelvaBigvganTrainer:
            if not disc_path.exists():
                raise FileNotFoundError(f"[BigVGAN] Discriminator checkpoint not found: {disc_path}")

+        lora_path = None
+        if lora_adapter and lora_adapter.strip():
+            lora_path = Path(lora_adapter.strip())
+            if not lora_path.is_absolute():
+                lora_path = Path(folder_paths.base_path) / lora_path
+            if not lora_path.exists():
+                raise FileNotFoundError(f"[BigVGAN] LoRA adapter not found: {lora_path}")
+
        # Find and pre-load audio clips
        segment_samples = int(segment_seconds * sample_rate)
        audio_files = []
@@ -412,14 +774,62 @@ class SelvaBigvganTrainer:
        # 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()
+
+        # Move EVERYTHING to CPU first, then bring back only what we need.
+        # ComfyUI may have loaded the full model to GPU; unload_all_models
+        # doesn't always free model dicts passed between nodes.
+        feature_utils.to("cpu")
+        if "generator" in model:
+            model["generator"].to("cpu")
+        if "video_enc" in model:
+            model["video_enc"].to("cpu")
        soft_empty_cache()

-        if strategy == "offload_to_cpu":
-            feature_utils.to(device)
-            soft_empty_cache()
-
+        # Only move mel_converter to GPU — it's tiny and needed for training.
+        # _pregenerate_lora_mels handles its own device management for CLIP/tod.
        mel_converter.to(device)

+        # Pre-compute text CLIP embeddings in the main thread.
+        # CLIP weights are inference tensors from ComfyUI loading — they only
+        # work in inference_mode, which is thread-local and active here but NOT
+        # in the worker thread. Pre-computing avoids needing CLIP on GPU in the
+        # worker. Results are cloned+detached so they're normal tensors.
+        text_clip_cache = {}
+        if lora_path is not None:
+            npz_files = sorted(data_dir.glob("*.npz"))
+            if npz_files:
+                prompt_map = {}
+                prompts_file = data_dir / "prompts.txt"
+                if prompts_file.exists():
+                    for line in prompts_file.read_text(encoding="utf-8").splitlines():
+                        line = line.strip()
+                        if not line or line.startswith("#"):
+                            continue
+                        if "|" in line:
+                            fname, prompt = line.split("|", 1)
+                            prompt_map[fname.strip()] = prompt.strip()
+                default_prompt = data_dir.name
+
+                # Temporarily move CLIP to GPU for encoding
+                clip_model = feature_utils.clip_model
+                if clip_model is not None:
+                    clip_model.to(device)
+                try:
+                    for npz_path in npz_files:
+                        data = dict(np.load(str(npz_path), allow_pickle=False))
+                        prompt = prompt_map.get(npz_path.name, data.get("prompt", default_prompt))
+                        if isinstance(prompt, np.ndarray):
+                            prompt = str(prompt)
+                        tc = feature_utils.encode_text_clip([prompt])
+                        text_clip_cache[npz_path.name] = tc.clone().detach().cpu()
+                finally:
+                    if clip_model is not None:
+                        clip_model.to("cpu")
+                    soft_empty_cache()
+                    if device.type == "cuda":
+                        torch.cuda.empty_cache()
+                print(f"[BigVGAN] Pre-encoded {len(text_clip_cache)} text CLIP embeddings", flush=True)
+
        pbar = comfy.utils.ProgressBar(steps)

        # -----------------------------------------------------------------------
@@ -439,12 +849,37 @@ class SelvaBigvganTrainer:

        def _worker():
            try:
+                # Pre-generate LoRA mels in the worker thread (inference_mode is
+                # thread-local — off here) so deep-copied generator tensors are clean.
+                lora_mel_pairs = None
+                if lora_path is not None:
+                    seq_cfg = model["seq_cfg"]
+                    lora_mel_pairs = _pregenerate_lora_mels(
+                        model, data_dir, str(lora_path),
+                        device, dtype, sample_rate,
+                        seq_cfg.duration, seed=seed,
+                        cache_dir=out_path.parent,
+                        text_clip_cache=text_clip_cache,
+                    )
+                    if not lora_mel_pairs:
+                        raise RuntimeError(
+                            "[BigVGAN] LoRA adapter provided but no mel/audio pairs "
+                            "could be generated. Check that data_dir contains .npz "
+                            "files with matching audio files."
+                        )
+                    # Force-release the CUDA memory pool from pre-generation.
+                    # soft_empty_cache may not call torch.cuda.empty_cache().
+                    if device.type == "cuda":
+                        torch.cuda.empty_cache()
+
                _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,
+                    use_gafilter, gafilter_kernel_size, lambda_phase,
                    save_every, seed, out_path, disc_path, pbar,
+                    lora_mel_pairs,
                )
            except Exception as e:
                _exc[0] = e
@@ -467,7 +902,9 @@ 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):
+              use_gafilter, gafilter_kernel_size, lambda_phase,
+              save_every, seed, out_path, disc_path, pbar,
+              lora_mel_pairs=None):
    """Execute training. Called in a fresh thread — no inference_mode active.

    Even though inference_mode is off here, tensors created in the calling
@@ -483,9 +920,11 @@ def _do_train(vocoder, mel_converter, clips,
    clips = [c.clone() for c in clips]

    # 2. mel_converter buffers (mel_basis, hann_window) — same origin.
+    #    Also cast to float32: mel_converter receives float32 audio (cuFFT
+    #    requirement) so all internal buffers must match.
    for name, buf in list(mel_converter._buffers.items()):
        if buf is not None:
-            mel_converter._buffers[name] = buf.clone()
+            mel_converter._buffers[name] = buf.clone().float()

    # 3. Vocoder parameters are handled below with clone().detach().
    # ─────────────────────────────────────────────────────────────────────────
@@ -494,8 +933,8 @@ def _do_train(vocoder, mel_converter, clips,
    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]
+    ref_wav = clips[0].to(device)                             # full first clip [T]
+    ref_mel = mel_converter(ref_wav.float().unsqueeze(0))     # [1, n_mels, T_mel] (cuFFT needs float32)

    # Ground-truth spectrogram — saved once alongside baseline for comparison
    gt_spec_path = out_path.parent / f"{out_path.stem}_gt_spec.png"
@@ -504,8 +943,8 @@ def _do_train(vocoder, mel_converter, clips,

    def _save_sample(label):
        try:
-            voc_device = next(vocoder.parameters()).device
-            mel = ref_mel.to(voc_device)
+            voc_p = next(vocoder.parameters())
+            mel = ref_mel.to(voc_p.device, voc_p.dtype)
            with torch.no_grad():
                wav = vocoder(mel)
                if wav.dim() == 2:
@@ -521,8 +960,6 @@ def _do_train(vocoder, mel_converter, clips,
        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):
    #
@@ -560,18 +997,32 @@ def _do_train(vocoder, mel_converter, clips,
            if buf is not None:
                module._buffers[bname] = buf.clone()

+    # ── Move vocoder to training device/dtype ────────────────────────────────
+    # After cloning, vocoder may be on CPU (offloaded before training).
+    vocoder.to(device, dtype)
+
+    # Baseline sample — after sanitization so vocoder has normal tensors.
+    _save_sample("baseline")
+
+    # ── GAFilter injection ─────────────────────────────────────────────────
+    # GAFilter params are fresh tensors — no inference flag to strip.
+    if use_gafilter:
+        n_gaf = inject_gafilters(vocoder, gafilter_kernel_size)
+        vocoder.to(device, dtype)  # ensure new GAFilter params match
+        print(f"[BigVGAN] GAFilter injected: {n_gaf} filters  kernel={gafilter_kernel_size}", flush=True)
+
    # ── 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:
+            if "alpha" in name or (use_gafilter and "gafilter" 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)
+              f"({len(alpha_params)} tensors, gafilter={'yes' if use_gafilter else 'no'})", flush=True)
        trainable_params = alpha_params
    else:  # all_params
        for param in vocoder.parameters():
@@ -596,18 +1047,28 @@ def _do_train(vocoder, mel_converter, clips,
            mpd = _MultiPeriodDiscriminator()
            mrd = _MultiResolutionDiscriminator()
            # Try common key names used by different BigVGAN releases
+            mpd_loaded = False
            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)
+                    mpd_loaded = True
                    break
+            mrd_loaded = False
            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)
+                    mrd_loaded = True
                    break
-            mpd.to(device).eval()
-            mrd.to(device).eval()
+            if not (mpd_loaded and mrd_loaded):
+                raise RuntimeError(
+                    f"[BigVGAN] Could not find discriminator keys in checkpoint. "
+                    f"MPD loaded={mpd_loaded}, MRD loaded={mrd_loaded}. "
+                    f"Available keys: {list(ckpt_d.keys())}"
+                )
+            mpd.to(device, dtype).eval()
+            mrd.to(device, dtype).eval()
            for p in mpd.parameters():
                p.requires_grad_(False)
            for p in mrd.parameters():
@@ -621,9 +1082,51 @@ def _do_train(vocoder, mel_converter, clips,
    optimizer = torch.optim.AdamW(trainable_params, lr=lr, betas=(0.8, 0.99))
    vocoder.train()

+    log_path = out_path.parent / f"{out_path.stem}_training_log.csv"
+    log_file = open(log_path, "w", buffering=1)   # line-buffered
+    log_file.write("step,total_loss,fm_loss,mel_loss,stft_loss,phase_loss,l2sp_loss\n")
+
+    # ── Pre-compute mel segment sizes for LoRA mel cropping ───────────────
+    # LoRA mels have shape [n_mels, T_mel_full] for the full clip duration.
+    # We need to crop segment_seconds from both mel and audio at same position.
+    if lora_mel_pairs:
+        _example_mel = lora_mel_pairs[0][0]  # [n_mels, T_mel_full]
+        _example_audio = lora_mel_pairs[0][1]  # [L]
+        _mel_frames_full = _example_mel.shape[-1]
+        _audio_samples_full = _example_audio.shape[0]
+        # mel frames per audio sample
+        _mel_per_sample = _mel_frames_full / _audio_samples_full
+        _mel_segment = int(segment_samples * _mel_per_sample)
+        print(f"[BigVGAN] LoRA mel cropping: {_mel_segment} mel frames "
+              f"per {segment_samples} audio samples", flush=True)
+
    try:
        for step in range(steps):
-            # Sample random batch — clips are CPU floats, move to device
+            if lora_mel_pairs:
+                # LoRA mode: sample LoRA mel + matching clean audio from same pair.
+                # Crop both from the same time position for alignment.
+                audio_batch = []
+                mel_batch = []
+                for _ in range(batch_size):
+                    lora_mel, lora_audio = random.choice(lora_mel_pairs)
+                    max_start = lora_audio.shape[0] - segment_samples
+                    if max_start > 0:
+                        audio_start = random.randint(0, max_start)
+                    else:
+                        audio_start = 0
+                    audio_batch.append(lora_audio[audio_start : audio_start + segment_samples])
+                    mel_start = int(audio_start * _mel_per_sample)
+                    mel_crop = lora_mel[:, mel_start : mel_start + _mel_segment]
+                    # Pad if crop goes past edge
+                    if mel_crop.shape[-1] < _mel_segment:
+                        mel_crop = F.pad(mel_crop, (0, _mel_segment - mel_crop.shape[-1]))
+                    mel_batch.append(mel_crop)
+
+                target_flat = torch.stack(audio_batch).to(device, dtype)  # [B, T]
+                target_wav  = target_flat.unsqueeze(1)                     # [B, 1, T]
+                input_mel   = torch.stack(mel_batch).to(device, dtype)     # [B, n_mels, T_seg]
+            else:
+                # Standard mode: sample random crops from clean audio clips
                batch = []
                for _ in range(batch_size):
                    clip  = random.choice(clips)
@@ -634,7 +1137,12 @@ def _do_train(vocoder, mel_converter, clips,
                target_wav  = target_flat.unsqueeze(1)                # [B, 1, T]

                with torch.no_grad():
-                target_mel = mel_converter(target_flat)           # [B, n_mels, T_mel]
+                    input_mel = mel_converter(target_flat.float())    # [B, n_mels, T_mel] (cuFFT needs float32)
+
+            # Clean target mel for mel loss (always from clean audio)
+            with torch.no_grad():
+                target_mel = mel_converter(target_flat.float())       # [B, n_mels, T_mel]
+

            # Gradient checkpointing: recompute BigVGAN activations during
            # backward instead of storing them. The 512x upsampling stack
@@ -642,53 +1150,67 @@ def _do_train(vocoder, mel_converter, clips,
            # ~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
+                vocoder, input_mel.to(dtype), 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
+                # Discriminators are in model dtype (bfloat16); waveforms are float32
+                disc_dtype = next(mpd.parameters()).dtype
                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)
+                    fmaps_real_mpd = mpd(target_t.to(disc_dtype))
+                    fmaps_real_mrd = mrd(target_t.to(disc_dtype))
+                fmaps_gen_mpd = mpd(pred_t.to(disc_dtype))
+                fmaps_gen_mrd = mrd(pred_t.to(disc_dtype))
                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))
+                pred_mel = mel_converter(pred_t.squeeze(1).float())
                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))
+                pred_mel = mel_converter(pred_t.squeeze(1).float())
                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}"

+            # ── FA-GAN phase-aware loss (real + imag + mag STFT) ────────────
+            if lambda_phase > 0.0:
+                phase_loss    = _phase_aware_stft_loss(pred_t, target_t, device)
+                primary_loss  = primary_loss + lambda_phase * phase_loss
+                loss_desc    += f"  phase={phase_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:
+                    # Skip GAFilter — newly initialized, not pretrained; L2-SP
+                    # anchoring to identity would fight against learning.
+                    if name in ref_params and param.requires_grad and "gafilter" not in name:
                        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()

@@ -698,23 +1220,58 @@ def _do_train(vocoder, mel_converter, clips,
                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)
+                # CSV row
+                _fm    = fm_loss.item()    if mpd is not None else ""
+                _mel   = mel_loss.item()
+                _stft  = stft_loss.item()  if mpd is None else ""
+                _phase = phase_loss.item() if lambda_phase > 0.0 else ""
+                _l2sp  = l2sp_loss.item()
+                log_file.write(f"{step+1},{loss.item():.6f},{_fm},{_mel},{_stft},{_phase},{_l2sp}\n")

            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))
+                torch.save({
+                    "generator":            vocoder.state_dict(),
+                    "has_gafilter":         use_gafilter,
+                    "gafilter_kernel_size": gafilter_kernel_size if use_gafilter else 9,
+                }, str(step_path))
                print(f"[BigVGAN] Checkpoint: {step_path}", flush=True)
                vocoder.eval()
                _save_sample(f"step{step+1}")
                vocoder.train()

    finally:
+        log_file.close()
        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_dict = {
+        "generator":           vocoder.state_dict(),
+        "has_gafilter":        use_gafilter,
+        "gafilter_kernel_size": gafilter_kernel_size if use_gafilter else 9,
+    }
+    torch.save(save_dict, str(out_path))
+    print(f"\n[BigVGAN] Saved: {out_path}  gafilter={use_gafilter}", flush=True)
    _save_sample("final")
+
+    # Generate a LoRA mel → vocoder sample so the user can hear the improvement
+    if lora_mel_pairs:
+        try:
+            lora_mel_full = lora_mel_pairs[0][0]  # [n_mels, T_mel]
+            voc_device = next(vocoder.parameters()).device
+            voc_dtype = next(vocoder.parameters()).dtype
+            with torch.no_grad():
+                wav_lora = vocoder(lora_mel_full.unsqueeze(0).to(voc_device, voc_dtype))
+                if wav_lora.dim() == 2:
+                    wav_lora = wav_lora.unsqueeze(1)
+                wav_lora = wav_lora.float().cpu().clamp(-1, 1)
+            lora_wav_path = out_path.parent / f"{out_path.stem}_lora_sample.wav"
+            _save_wav(lora_wav_path, wav_lora.squeeze(0), sample_rate)
+            print(f"[BigVGAN] LoRA mel sample: {lora_wav_path}", flush=True)
+        except Exception as e:
+            print(f"[BigVGAN] LoRA sample failed: {e}", flush=True)
+
    return str(out_path)
@@ -0,0 +1,788 @@
+"""SelVA Audio Dataset Pipeline — chainable in-memory preprocessing nodes.
+
+Typical chain:
+  SelvaDatasetLoader
+      ↓ AUDIO_DATASET
+  SelvaDatasetResampler       (optional)
+      ↓ AUDIO_DATASET
+  SelvaDatasetLUFSNormalizer  (optional)
+      ↓ AUDIO_DATASET
+  SelvaDatasetCompressor      (optional)
+      ↓ AUDIO_DATASET
+  SelvaDatasetSpectralMatcher (optional — batch spectral EQ)
+      ↓ AUDIO_DATASET
+  SelvaDatasetHfSmoother      (optional — batch HF attenuation)
+      ↓ AUDIO_DATASET
+  SelvaDatasetAugmenter       (optional — gain/pitch/stretch variants)
+      ↓ AUDIO_DATASET
+  SelvaDatasetInspector       (optional)
+      ↓ AUDIO_DATASET  +  STRING report
+  SelvaDatasetItemExtractor   → AUDIO (bridges to save/preview nodes)
+"""
+
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchaudio
+
+from .utils import SELVA_CATEGORY
+
+# ComfyUI custom type name — passed between all dataset pipeline nodes
+AUDIO_DATASET = "AUDIO_DATASET"
+
+_AUDIO_EXTS = {".wav", ".flac", ".mp3", ".ogg", ".aac", ".m4a"}
+_SOUNDFILE_EXTS = {".wav", ".flac", ".ogg"}  # handled natively without FFmpeg
+
+
+def _load_audio(path: Path):
+    """Load audio file. Uses soundfile for WAV/FLAC/OGG to avoid torchcodec/FFmpeg issues."""
+    if path.suffix.lower() in _SOUNDFILE_EXTS:
+        import soundfile as sf
+        wav_np, sr = sf.read(str(path), dtype="float32", always_2d=True)  # [L, C]
+        wav = torch.from_numpy(wav_np).T.unsqueeze(0)                      # [1, C, L]
+    else:
+        wav, sr = torchaudio.load(str(path))   # [C, L]
+        wav = wav.unsqueeze(0).float()          # [1, C, L]
+    return wav, sr
+
+
+class SelvaDatasetLoader:
+    """Load all audio files in a folder into an in-memory AUDIO_DATASET."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "folder": ("STRING", {
+                    "default": "",
+                    "tooltip": "Absolute path to folder containing audio files. Searched recursively.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET,)
+    RETURN_NAMES  = ("dataset",)
+    FUNCTION      = "load"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = "Load all audio files from a folder into memory as an AUDIO_DATASET."
+
+    def load(self, folder: str):
+        folder = Path(folder.strip())
+        if not folder.exists():
+            raise FileNotFoundError(f"[DatasetLoader] Folder not found: {folder}")
+
+        files = [f for f in folder.rglob("*") if f.suffix.lower() in _AUDIO_EXTS]
+        if not files:
+            raise RuntimeError(f"[DatasetLoader] No audio files found in {folder}")
+
+        dataset = []
+        for f in sorted(files):
+            try:
+                wav, sr = _load_audio(f)
+                dataset.append({"waveform": wav, "sample_rate": sr, "name": f.stem})
+            except Exception as e:
+                print(f"[DatasetLoader] Skipping {f.name}: {e}", flush=True)
+
+        print(f"[DatasetLoader] Loaded {len(dataset)} clips from {folder}", flush=True)
+        return (dataset,)
+
+
+class SelvaDatasetResampler:
+    """Resample all clips in a dataset to a target sample rate using soxr VHQ."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "target_sr": ("INT", {
+                    "default": 44100, "min": 8000, "max": 192000,
+                    "tooltip": "Target sample rate. 44100 for large SelVA model, 16000 for small.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET,)
+    RETURN_NAMES  = ("dataset",)
+    FUNCTION      = "resample"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = "Resample all clips to target_sr using soxr VHQ. Skips clips already at target rate."
+
+    def resample(self, dataset, target_sr: int):
+        import soxr
+
+        out = []
+        changed = 0
+        for item in dataset:
+            sr = item["sample_rate"]
+            if sr == target_sr:
+                out.append(item)
+                continue
+
+            wav = item["waveform"][0]               # [C, L]
+            # soxr expects [L, C] (time-first), float64
+            wav_np = wav.permute(1, 0).double().numpy()   # [L, C]
+            wav_rs = soxr.resample(wav_np, sr, target_sr, quality="VHQ")
+            wav_t  = torch.from_numpy(wav_rs).float().permute(1, 0).unsqueeze(0)  # [1, C, L]
+            out.append({"waveform": wav_t, "sample_rate": target_sr, "name": item["name"]})
+            changed += 1
+
+        print(f"[DatasetResampler] {changed}/{len(dataset)} clips resampled → {target_sr} Hz", flush=True)
+        return (out,)
+
+
+class SelvaDatasetLUFSNormalizer:
+    """Normalize each clip to a target integrated LUFS level + true peak limit."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "target_lufs": ("FLOAT", {
+                    "default": -23.0, "min": -40.0, "max": -6.0, "step": 0.5,
+                    "tooltip": "Target integrated loudness in LUFS. -23 is EBU R128 standard.",
+                }),
+                "true_peak_dbtp": ("FLOAT", {
+                    "default": -1.0, "min": -6.0, "max": 0.0, "step": 0.5,
+                    "tooltip": "True peak ceiling in dBTP. Applied after LUFS gain.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET,)
+    RETURN_NAMES  = ("dataset",)
+    FUNCTION      = "normalize"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Normalize each clip to target_lufs (BS.1770-4) then apply a true peak ceiling. "
+        "Skips clips that are too short for LUFS measurement (< 0.4 s)."
+    )
+
+    def normalize(self, dataset, target_lufs: float, true_peak_dbtp: float):
+        import pyloudnorm as pyln
+
+        tp_linear = 10.0 ** (true_peak_dbtp / 20.0)
+        out = []
+        skipped = 0
+
+        for item in dataset:
+            wav = item["waveform"][0]           # [C, L]
+            sr  = item["sample_rate"]
+
+            # pyloudnorm wants [L] mono or [L, C] multichannel, float64
+            wav_np = wav.permute(1, 0).double().numpy()  # [L, C]
+            if wav_np.shape[1] == 1:
+                wav_np = wav_np[:, 0]           # [L] mono
+
+            meter = pyln.Meter(sr)
+            try:
+                loudness = meter.integrated_loudness(wav_np)
+            except Exception:
+                skipped += 1
+                out.append(item)
+                continue
+
+            if not np.isfinite(loudness):
+                skipped += 1
+                out.append(item)
+                continue
+
+            gain_db     = target_lufs - loudness
+            gain_linear = 10.0 ** (gain_db / 20.0)
+
+            wav_norm = wav * gain_linear
+
+            # True peak limit
+            peak = wav_norm.abs().max().item()
+            if peak > tp_linear:
+                wav_norm = wav_norm * (tp_linear / peak)
+
+            out.append({"waveform": wav_norm.unsqueeze(0), "sample_rate": sr, "name": item["name"]})
+
+        print(
+            f"[LUFSNormalizer] {len(dataset) - skipped}/{len(dataset)} clips normalized  "
+            f"target={target_lufs} LUFS  TP={true_peak_dbtp} dBTP  skipped={skipped}",
+            flush=True,
+        )
+        return (out,)
+
+
+class SelvaDatasetCompressor:
+    """Apply mild parallel compression to reduce within-clip loudness variance.
+
+    Uses pedalboard.Compressor (2:1–3:1 ratio). Parallel (New York) style:
+    blends compressed signal with dry so transients are preserved while
+    the dynamic range is gently tightened. Apply after LUFS normalization.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset":        (AUDIO_DATASET,),
+                "threshold_db":   ("FLOAT", {
+                    "default": -18.0, "min": -40.0, "max": -6.0, "step": 1.0,
+                    "tooltip": "Compression kicks in above this level. -18 dB is a safe starting point after LUFS normalization.",
+                }),
+                "ratio":          ("FLOAT", {
+                    "default": 2.5, "min": 1.5, "max": 4.0, "step": 0.5,
+                    "tooltip": "Compression ratio. 2:1–3:1 is mild; stay below 4:1 to avoid pumping.",
+                }),
+                "attack_ms":      ("FLOAT", {
+                    "default": 10.0, "min": 1.0, "max": 100.0, "step": 1.0,
+                    "tooltip": "Attack time in ms. Slower attack preserves transients.",
+                }),
+                "release_ms":     ("FLOAT", {
+                    "default": 100.0, "min": 20.0, "max": 500.0, "step": 10.0,
+                    "tooltip": "Release time in ms.",
+                }),
+                "mix":            ("FLOAT", {
+                    "default": 0.4, "min": 0.0, "max": 1.0, "step": 0.05,
+                    "tooltip": "Parallel blend: 0.0 = dry only, 1.0 = fully compressed. 0.3–0.5 is typical.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET,)
+    RETURN_NAMES  = ("dataset",)
+    FUNCTION      = "compress"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Mild parallel compression to reduce within-clip dynamic range. "
+        "Blends compressed signal with dry at the given mix ratio. "
+        "Apply after LUFS normalization."
+    )
+
+    def compress(self, dataset, threshold_db: float, ratio: float,
+                 attack_ms: float, release_ms: float, mix: float):
+        from pedalboard import Compressor, Pedalboard
+
+        board = Pedalboard([Compressor(
+            threshold_db=threshold_db,
+            ratio=ratio,
+            attack_ms=attack_ms,
+            release_ms=release_ms,
+        )])
+
+        out = []
+        for item in dataset:
+            wav = item["waveform"][0]   # [C, L]
+            sr  = item["sample_rate"]
+
+            # pedalboard expects [C, L] float32 numpy
+            wav_np = wav.float().numpy()                        # [C, L]
+            compressed = board(wav_np, sr)                      # [C, L]
+            mixed = (1.0 - mix) * wav_np + mix * compressed
+            wav_out = torch.from_numpy(mixed).unsqueeze(0)      # [1, C, L]
+            out.append({"waveform": wav_out, "sample_rate": sr, "name": item["name"]})
+
+        print(
+            f"[DatasetCompressor] {len(out)} clips compressed  "
+            f"thr={threshold_db}dB  ratio={ratio}:1  mix={mix:.0%}",
+            flush=True,
+        )
+        return (out,)
+
+
+def _check_hf_shelf(wav: torch.Tensor, sr: int) -> bool:
+    """Return True if clip looks codec-compressed (hard HF shelf above 15 kHz).
+
+    Method: compare mean energy in 1–5 kHz band vs 15–20 kHz band via STFT.
+    A ratio > 40 dB (i.e. near-silence above 15 kHz) flags codec artifacts.
+    """
+    if sr < 32000:
+        return False  # can't assess HF at low sample rates
+
+    n_fft  = 2048
+    hop    = 512
+    mono   = wav[0].mean(0)  # [L]
+    window = torch.hann_window(n_fft, device=mono.device)
+    stft   = torch.stft(mono, n_fft, hop, n_fft, window, return_complex=True)
+    mag_sq = stft.abs().pow(2).mean(-1)  # [n_freqs]
+
+    freqs     = torch.linspace(0, sr / 2, n_fft // 2 + 1, device=mono.device)
+    band_lo   = (freqs >= 1000) & (freqs < 5000)
+    band_hi   = (freqs >= 15000) & (freqs < 20000)
+
+    if band_hi.sum() == 0:
+        return False
+
+    energy_lo = mag_sq[band_lo].mean().clamp(min=1e-12)
+    energy_hi = mag_sq[band_hi].mean().clamp(min=1e-12)
+    ratio_db  = 10.0 * torch.log10(energy_lo / energy_hi).item()
+    return ratio_db > 40.0
+
+
+def _estimate_snr(wav: torch.Tensor) -> float:
+    """Rough SNR estimate: ratio of 95th-percentile frame RMS to 5th-percentile frame RMS."""
+    mono   = wav[0].mean(0)  # [L]
+    if mono.shape[0] < 2048:
+        return 60.0  # clip too short to frame — assume clean
+    frames = mono.unfold(0, 2048, 512)          # [N, 2048]
+    rms    = frames.pow(2).mean(-1).sqrt()      # [N]
+    p95    = torch.quantile(rms, 0.95).item()
+    p05    = torch.quantile(rms, 0.05).clamp(min=1e-8).item()
+    return 20.0 * np.log10(p95 / p05 + 1e-8)
+
+
+class SelvaDatasetInspector:
+    """Analyze each clip for quality issues and optionally filter out flagged clips."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "skip_rejected": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "If True, flagged clips are removed from the output dataset. "
+                               "If False, all clips pass through but the report still lists issues.",
+                }),
+                "min_snr_db": ("FLOAT", {
+                    "default": 15.0, "min": 0.0, "max": 60.0, "step": 1.0,
+                    "tooltip": "Clips with estimated SNR below this value are flagged.",
+                }),
+                "check_codec_artifacts": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "Flag clips with a hard HF shelf above 15 kHz (MP3/codec artifact signature).",
+                }),
+                "max_silence_fraction": ("FLOAT", {
+                    "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
+                    "tooltip": "Flag clips where more than this fraction of frames are near-silent "
+                               "(< -60 dBFS). Set to 0 to disable silence detection.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET, "STRING")
+    RETURN_NAMES  = ("dataset", "report")
+    FUNCTION      = "inspect"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Analyze each clip for clipping, low SNR, and codec artifacts. "
+        "Outputs a filtered AUDIO_DATASET and a text report. "
+        "Connect report to a ShowText node to preview in the UI."
+    )
+
+    def inspect(self, dataset, skip_rejected: bool, min_snr_db: float,
+                check_codec_artifacts: bool, max_silence_fraction: float = 0.5):
+        clean   = []
+        flagged = []
+        lines   = ["SelVA Dataset Inspector Report", "=" * 40]
+
+        for item in dataset:
+            wav    = item["waveform"]
+            sr     = item["sample_rate"]
+            name   = item["name"]
+            issues = []
+
+            # Clipping
+            peak = wav.abs().max().item()
+            if peak > 0.99:
+                issues.append(f"clipping (peak={peak:.3f})")
+
+            # Low SNR
+            snr = _estimate_snr(wav)
+            if snr < min_snr_db:
+                issues.append(f"low SNR ({snr:.1f} dB < {min_snr_db} dB)")
+
+            # Codec artifacts
+            if check_codec_artifacts and _check_hf_shelf(wav, sr):
+                issues.append("codec artifact (HF shelf > 15 kHz)")
+
+            # Silence detection
+            if max_silence_fraction > 0:
+                mono = wav[0].mean(0)
+                if mono.shape[0] >= 2048:
+                    frames = mono.unfold(0, 2048, 512)
+                    rms = frames.pow(2).mean(-1).sqrt()
+                    silent_frac = (rms < 1e-3).float().mean().item()
+                    if silent_frac > max_silence_fraction:
+                        issues.append(f"mostly silent ({silent_frac:.0%} < -60 dBFS)")
+
+            if issues:
+                flagged.append(name)
+                lines.append(f"  FLAGGED  {name}: {', '.join(issues)}")
+                if not skip_rejected:
+                    clean.append(item)
+            else:
+                clean.append(item)
+                lines.append(f"  OK       {name}")
+
+        lines.append("=" * 40)
+        lines.append(
+            f"Total: {len(dataset)}  Clean: {len(clean)}  Flagged: {len(flagged)}"
+            + (" (removed)" if skip_rejected else " (kept)")
+        )
+        report = "\n".join(lines)
+        print(f"[DatasetInspector]\n{report}", flush=True)
+        return (clean, report)
+
+
+class SelvaDatasetItemExtractor:
+    """Extract a single AUDIO item from an AUDIO_DATASET by index.
+
+    Bridges the dataset pipeline to any node that accepts a standard AUDIO
+    input — save audio, HF Smoother, Spectral Matcher, etc.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "index":   ("INT", {
+                    "default": 0, "min": 0, "max": 9999,
+                    "tooltip": "0-based index. Wraps around if index >= dataset length.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = ("AUDIO", "STRING", "INT")
+    RETURN_NAMES  = ("audio",  "name",   "total")
+    FUNCTION      = "extract"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Extract one clip from an AUDIO_DATASET by index. "
+        "Returns standard AUDIO (compatible with all audio nodes), "
+        "the clip name, and the total dataset length."
+    )
+
+    def extract(self, dataset, index: int):
+        if not dataset:
+            raise RuntimeError("[DatasetItemExtractor] Dataset is empty.")
+        idx  = index % len(dataset)
+        item = dataset[idx]
+        audio = {"waveform": item["waveform"], "sample_rate": item["sample_rate"]}
+        print(
+            f"[DatasetItemExtractor] [{idx}/{len(dataset)-1}] {item['name']}  "
+            f"sr={item['sample_rate']}  shape={tuple(item['waveform'].shape)}",
+            flush=True,
+        )
+        return (audio, item["name"], len(dataset))
+
+
+class SelvaDatasetSaver:
+    """Save all clips in an AUDIO_DATASET to disk as FLAC files.
+
+    Optionally copies matching .npz feature files from a source directory,
+    keeping FLAC/NPZ pairs in sync after the inspector has filtered clips.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "output_dir": ("STRING", {
+                    "default": "",
+                    "tooltip": "Absolute path to output folder. Created if it does not exist.",
+                }),
+            },
+            "optional": {
+                "npz_source_dir": ("STRING", {
+                    "default": "",
+                    "tooltip": "If set, copies {name}.npz from this folder alongside each saved FLAC. "
+                               "Missing NPZs are warned but do not abort the save.",
+                }),
+            },
+        }
+
+    RETURN_TYPES  = ("STRING",)
+    RETURN_NAMES  = ("report",)
+    OUTPUT_NODE   = True
+    FUNCTION      = "save"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Save every clip in an AUDIO_DATASET to output_dir as FLAC. "
+        "If npz_source_dir is provided, copies the matching .npz file for each clip — "
+        "so rejected clips never get their NPZ copied."
+    )
+
+    def save(self, dataset, output_dir: str, npz_source_dir: str = ""):
+        import shutil
+        import soundfile as sf
+
+        out = Path(output_dir.strip())
+        out.mkdir(parents=True, exist_ok=True)
+
+        npz_src = Path(npz_source_dir.strip()) if npz_source_dir.strip() else None
+
+        saved = 0
+        npz_copied = 0
+        npz_missing = []
+
+        for item in dataset:
+            name   = item["name"]
+            wav    = item["waveform"][0]          # [C, L]
+            sr     = item["sample_rate"]
+
+            # soundfile wants [L] mono or [L, C] multichannel, float32
+            wav_np = wav.permute(1, 0).float().numpy()   # [L, C]
+            if wav_np.shape[1] == 1:
+                wav_np = wav_np[:, 0]                    # [L] mono
+
+            flac_path = out / f"{name}.flac"
+            sf.write(str(flac_path), wav_np, sr, subtype="PCM_24")
+            saved += 1
+
+            if npz_src is not None:
+                # Augmented clips store their origin name — use it to find the .npz
+                lookup = item.get("origin_name", name)
+                npz_path = npz_src / f"{lookup}.npz"
+                if npz_path.exists():
+                    shutil.copy2(str(npz_path), str(out / f"{name}.npz"))
+                    npz_copied += 1
+                else:
+                    npz_missing.append(name)
+
+        lines = [
+            f"[DatasetSaver] Saved {saved} clips → {out}",
+        ]
+        if npz_src is not None:
+            lines.append(f"  NPZ copied: {npz_copied}  missing: {len(npz_missing)}")
+            for n in npz_missing:
+                lines.append(f"    MISSING NPZ: {n}")
+
+        report = "\n".join(lines)
+        print(report, flush=True)
+        return (report,)
+
+
+# ── Batch wrappers for audio preprocessors ───────────────────────────────────
+
+class SelvaDatasetSpectralMatcher:
+    """Apply SelVA Spectral Matcher to every clip in an AUDIO_DATASET.
+
+    Wraps SelvaSpectralMatcher so it works on batch datasets instead of
+    individual AUDIO items. Same parameters — see that node for details.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "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.",
+                }),
+                "max_gain_db": ("FLOAT", {
+                    "default": 12.0, "min": 1.0, "max": 30.0, "step": 1.0,
+                    "tooltip": "Clamps per-band gain to ±dB.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET,)
+    RETURN_NAMES  = ("dataset",)
+    FUNCTION      = "process"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Apply adaptive spectral matching to every clip in a dataset. "
+        "Batch version of SelVA Spectral Matcher — same per-band EQ toward the "
+        "VAE's expected distribution."
+    )
+
+    def process(self, dataset, mode: str, strength: float, max_gain_db: float):
+        from .selva_audio_preprocessors import SelvaSpectralMatcher
+        matcher = SelvaSpectralMatcher()
+        out = []
+        for item in dataset:
+            audio = {"waveform": item["waveform"], "sample_rate": item["sample_rate"]}
+            (result,) = matcher.process(audio, mode, strength, max_gain_db)
+            new_item = dict(item)  # preserve origin_name and any extra keys
+            new_item["waveform"] = result["waveform"]
+            new_item["sample_rate"] = result["sample_rate"]
+            out.append(new_item)
+        print(f"[DatasetSpectralMatcher] {len(out)} clips processed  "
+              f"mode={mode}  strength={strength}", flush=True)
+        return (out,)
+
+
+class SelvaDatasetHfSmoother:
+    """Apply SelVA HF Smoother to every clip in an AUDIO_DATASET.
+
+    Wraps SelvaHfSmoother so it works on batch datasets instead of
+    individual AUDIO items. Same parameters — see that node for details.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "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.",
+                }),
+            }
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET,)
+    RETURN_NAMES  = ("dataset",)
+    FUNCTION      = "process"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Apply soft HF attenuation to every clip in a dataset. "
+        "Batch version of SelVA HF Smoother — blends a low-pass filtered copy "
+        "with the original to tame extreme HF content."
+    )
+
+    def process(self, dataset, cutoff_hz: float, blend: float):
+        from .selva_audio_preprocessors import SelvaHfSmoother
+        smoother = SelvaHfSmoother()
+        out = []
+        for item in dataset:
+            audio = {"waveform": item["waveform"], "sample_rate": item["sample_rate"]}
+            (result,) = smoother.process(audio, cutoff_hz, blend)
+            new_item = dict(item)  # preserve origin_name and any extra keys
+            new_item["waveform"] = result["waveform"]
+            new_item["sample_rate"] = result["sample_rate"]
+            out.append(new_item)
+        print(f"[DatasetHfSmoother] {len(out)} clips processed  "
+              f"cutoff={cutoff_hz:.0f}Hz  blend={blend:.2f}", flush=True)
+        return (out,)
+
+
+# ── Dataset augmenter ────────────────────────────────────────────────────────
+
+class SelvaDatasetAugmenter:
+    """Create augmented variants of each clip to expand a small dataset.
+
+    Supports gain variation (always available) and optionally pitch shift
+    and time stretch via audiomentations. Install audiomentations for the
+    full feature set: ``pip install audiomentations``
+
+    Each original clip produces ``variants_per_clip`` augmented copies.
+    Originals are kept by default (toggle ``keep_originals``).
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dataset": (AUDIO_DATASET,),
+                "variants_per_clip": ("INT", {
+                    "default": 2, "min": 1, "max": 20,
+                    "tooltip": "Number of augmented copies per original clip.",
+                }),
+                "gain_range_db": ("FLOAT", {
+                    "default": 3.0, "min": 0.0, "max": 12.0, "step": 0.5,
+                    "tooltip": "Random gain ±dB applied to each variant. 3 dB is subtle.",
+                }),
+                "seed": ("INT", {"default": 42}),
+            },
+            "optional": {
+                "pitch_range_semitones": ("FLOAT", {
+                    "default": 0.0, "min": 0.0, "max": 4.0, "step": 0.25,
+                    "tooltip": "Random pitch shift ±semitones. Requires audiomentations. 0 = disabled.",
+                }),
+                "time_stretch_range": ("FLOAT", {
+                    "default": 0.0, "min": 0.0, "max": 0.3, "step": 0.05,
+                    "tooltip": "Random time stretch ±fraction (0.1 = 90%–110% speed). "
+                               "Requires audiomentations. 0 = disabled.",
+                }),
+                "keep_originals": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "Include the original unaugmented clips in the output.",
+                }),
+            },
+        }
+
+    RETURN_TYPES  = (AUDIO_DATASET,)
+    RETURN_NAMES  = ("dataset",)
+    FUNCTION      = "augment"
+    CATEGORY      = SELVA_CATEGORY
+    DESCRIPTION   = (
+        "Create augmented variants of each clip (gain, pitch, time stretch) "
+        "to expand small training datasets. Gain is always available; pitch and "
+        "time stretch require audiomentations (pip install audiomentations)."
+    )
+
+    def augment(self, dataset, variants_per_clip: int, gain_range_db: float,
+                seed: int, pitch_range_semitones: float = 0.0,
+                time_stretch_range: float = 0.0, keep_originals: bool = True):
+        rng = np.random.RandomState(seed)
+
+        # Try audiomentations for pitch/stretch
+        use_am = False
+        am_compose = None
+        needs_am = pitch_range_semitones > 0 or time_stretch_range > 0
+        if needs_am:
+            try:
+                import audiomentations as am
+                transforms = []
+                if pitch_range_semitones > 0:
+                    transforms.append(am.PitchShift(
+                        min_semitones=-pitch_range_semitones,
+                        max_semitones=pitch_range_semitones,
+                        p=0.5,
+                    ))
+                if time_stretch_range > 0:
+                    transforms.append(am.TimeStretch(
+                        min_rate=1.0 - time_stretch_range,
+                        max_rate=1.0 + time_stretch_range,
+                        leave_length_unchanged=True,
+                        p=0.5,
+                    ))
+                am_compose = am.Compose(transforms)
+                use_am = True
+            except ImportError:
+                print("[DatasetAugmenter] audiomentations not installed — "
+                      "pitch_shift and time_stretch disabled. "
+                      "Install: pip install audiomentations", flush=True)
+
+        out = []
+        if keep_originals:
+            out.extend(dataset)
+
+        for item in dataset:
+            wav = item["waveform"]       # [1, C, L]
+            sr  = item["sample_rate"]
+            name = item["name"]
+
+            for v in range(variants_per_clip):
+                # Gain variation (always applied)
+                gain_db  = rng.uniform(-gain_range_db, gain_range_db) if gain_range_db > 0 else 0.0
+                gain_lin = 10.0 ** (gain_db / 20.0)
+                wav_aug  = wav * gain_lin
+
+                # Pitch/stretch via audiomentations
+                if use_am and am_compose is not None:
+                    wav_np = wav_aug[0].numpy()     # [C, L] float32
+                    if wav_np.shape[0] == 1:
+                        wav_np = wav_np[0]          # [L] mono for audiomentations
+                    wav_np = am_compose(samples=wav_np, sample_rate=sr)
+                    if wav_np.ndim == 1:
+                        wav_np = wav_np[np.newaxis, :]  # back to [1, L]
+                    wav_aug = torch.from_numpy(wav_np).unsqueeze(0)  # [1, C, L]
+
+                # Prevent clipping
+                peak = wav_aug.abs().max()
+                if peak > 1.0:
+                    wav_aug = wav_aug / peak
+
+                out.append({
+                    "waveform": wav_aug,
+                    "sample_rate": sr,
+                    "name": f"{name}_aug{v:02d}",
+                    "origin_name": name,
+                })
+
+        print(f"[DatasetAugmenter] {len(dataset)} originals → {len(out)} total clips  "
+              f"gain=±{gain_range_db:.1f}dB"
+              + (f"  pitch=±{pitch_range_semitones:.1f}st" if pitch_range_semitones > 0 else "")
+              + (f"  stretch=±{time_stretch_range:.0%}" if time_stretch_range > 0 else ""),
+              flush=True)
+        return (out,)
@@ -1,14 +1,14 @@
 """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
+using a style loss against target style 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
+against target style 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
@@ -42,35 +42,62 @@ def _load_wav(path):
    return wav, sr


-def _mel_style_loss(mel_gen, ref_mean, ref_gram):
+def _mel_style_loss(mel_gen, ref_mean, ref_gram, gram_weight=0.0):
    """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.
+    ref_mean:    [n_mels]         mean spectrum of reference clips (detached)
+    ref_gram:    [n_mels, n_mels] Gram matrix of reference clips  (detached)
+    gram_weight: weight for Gram matrix component — 0 = mean spectrum only.
+                 Start at 0; enable only if mean-only optimization converges
+                 without noise, then increase slowly (0.01–0.1).
    """
    m = mel_gen.squeeze(0)  # [n_mels, T]

-    # Mean spectrum loss
+    # Mean spectrum loss — captures spectral envelope
    gen_mean = m.mean(dim=-1)  # [n_mels]
    loss_mean = F.l1_loss(gen_mean, ref_mean)

-    # Gram matrix loss (texture, position-invariant)
+    if gram_weight <= 0.0:
+        return loss_mean
+
+    # Gram matrix loss — captures timbral texture (can add noise if too high)
    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
+    return loss_mean + gram_weight * loss_gram
+
+
+def _latent_style_loss(z, ref_mean, ref_gram, gram_weight=0.0):
+    """Style loss computed directly in VAE latent space.
+
+    z:           [T_lat, C_lat]  unnormalized latent at ODE endpoint (with grad)
+    ref_mean:    [C_lat]          mean latent vector of reference clips
+    ref_gram:    [C_lat, C_lat]   Gram matrix of reference latents
+    gram_weight: weight for Gram component — 0 = mean only (recommended start)
+
+    Operating in latent space avoids backprop through the VAE decoder, which
+    is @torch.inference_mode() and produces noisy, unstable gradients.
+    """
+    # Mean latent loss — matches average activation per channel
+    gen_mean  = z.mean(dim=0)                  # [C_lat]
+    loss_mean = F.l1_loss(gen_mean, ref_mean)
+
+    if gram_weight <= 0.0:
+        return loss_mean
+
+    # Gram matrix — inter-channel covariance, position-invariant
+    gram_gen  = (z.T @ z) / z.shape[0]         # [C_lat, C_lat]
+    loss_gram = F.mse_loss(gram_gen, ref_gram)
+
+    return loss_mean + gram_weight * 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.
+    to make the generated audio sound like the target style reference clips.
    No training data or gradient updates to the model — per-video per-run.
    """

@@ -89,7 +116,7 @@ class SelvaDittoOptimizer:
                }),
                "reference_dir": ("STRING", {
                    "default": "",
-                    "tooltip": "Directory with BJ reference audio files (.wav/.flac/.mp3). "
+                    "tooltip": "Directory with target style reference audio files (.wav/.flac/.mp3). "
                               "Reference mel statistics are precomputed from these once.",
                }),
                "n_opt_steps": ("INT", {
@@ -98,9 +125,9 @@ class SelvaDittoOptimizer:
                               "each step requires ~2 DiT forward passes.",
                }),
                "opt_lr": ("FLOAT", {
-                    "default": 0.1, "min": 0.001, "max": 2.0, "step": 0.01,
+                    "default": 0.02, "min": 0.001, "max": 2.0, "step": 0.001,
                    "tooltip": "Adam learning rate for x_0 optimization. "
-                               "0.1 is the DITTO paper default.",
+                               "0.02–0.05 is recommended; 0.1 (paper default) causes oscillation.",
                }),
                "n_ode_steps": ("INT", {
                    "default": 10, "min": 5, "max": 50,
@@ -115,9 +142,22 @@ class SelvaDittoOptimizer:
                               "Must be ≤ n_ode_steps. 5 is a good default.",
                }),
                "style_weight": ("FLOAT", {
+                    "default": 0.1, "min": 0.0, "max": 10.0, "step": 0.05,
+                    "tooltip": "Weight of the target style style loss. High values push harder toward "
+                               "target style style but add noise. Start at 0.1 and increase slowly.",
+                }),
+                "gram_weight": ("FLOAT", {
+                    "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01,
+                    "tooltip": "Weight of the Gram matrix (timbral texture) loss relative to "
+                               "the mean spectrum loss. 0 = mean spectrum only (less noise). "
+                               "0.1 adds texture matching but can introduce white noise.",
+                }),
+                "anchor_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.",
+                    "tooltip": "L2 penalty keeping x0 near its initial N(0,1) noise. "
+                               "Prevents optimization from pushing x0 out of the flow's "
+                               "expected distribution (which causes white noise). "
+                               "Higher = cleaner audio, weaker style. 1.0 is a safe default.",
                }),
                "steps": ("INT", {
                    "default": 25, "min": 1, "max": 200,
@@ -136,19 +176,19 @@ class SelvaDittoOptimizer:

    RETURN_TYPES = ("AUDIO",)
    RETURN_NAMES = ("audio",)
-    OUTPUT_TOOLTIPS = ("DITTO-optimized audio — x_0 steered toward BJ style.",)
+    OUTPUT_TOOLTIPS = ("DITTO-optimized audio — x_0 steered toward target style 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 "
+        "Optimizes the initial noise latent x_0 to match target style 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,
+                 style_weight, gram_weight, anchor_weight, steps, cfg_strength, seed,
                 normalize=True, target_lufs=-27.0):
        import traceback

@@ -177,7 +217,10 @@ class SelvaDittoOptimizer:
        seq_cfg = dataclasses.replace(model["seq_cfg"], duration=duration)
        sample_rate = seq_cfg.sampling_rate

-        # Load and precompute reference mel statistics
+        # Load reference clips and encode to latent space.
+        # Style loss is computed in latent space (after net_generator.unnormalize)
+        # rather than mel space — this avoids backpropagating through the VAE
+        # decoder (which is @torch.inference_mode() and produces noisy gradients).
        ref_dir = Path(reference_dir.strip())
        if not ref_dir.is_absolute():
            ref_dir = Path(folder_paths.models_dir) / ref_dir
@@ -190,39 +233,44 @@ class SelvaDittoOptimizer:
        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)
+        if not hasattr(feature_utils.tod.vae, "encoder"):
+            raise RuntimeError(
+                "[DITTO] VAE encoder not available — model was loaded with need_vae_encoder=False. "
+                "Reload the model with the encoder enabled."
+            )

-        ref_mels = []
+        print(f"[DITTO] Loading {len(ref_files)} reference clips...", flush=True)
+        mel_converter.to(device, torch.float32)  # cuFFT requires float32
+
+        ref_latents = []
        with torch.no_grad():
-            for rf in ref_files[:32]:  # cap at 32 for speed
+            for rf in ref_files:
                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)
+                    wav = wav.squeeze(0).to(device, torch.float32)
+                    mel = mel_converter(wav.unsqueeze(0)).to(dtype)  # [1, n_mels, T_mel]
+                    # encode → sample → VAE latent space (matches unnormalize(x) in loss)
+                    z = feature_utils.tod.encode(mel)               # DiagonalGaussianDistribution
+                    z_sample = z.sample().transpose(1, 2)           # [1, T_lat, C_lat]
+                    ref_latents.append(z_sample.to(dtype).squeeze(0).clone())  # [T_lat, C_lat]
                except Exception as e:
                    print(f"  [DITTO] Skip {rf.name}: {e}", flush=True)

-        if not ref_mels:
+        if not ref_latents:
            raise RuntimeError("[DITTO] No usable reference clips.")

-        # Precompute reference statistics (done once — detached, no grad)
+        # Precompute reference latent 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_means = torch.stack([z.mean(dim=0) for z in ref_latents])
+            ref_mean  = all_means.mean(0)                           # [C_lat]
+            all_grams = [(z.T @ z) / z.shape[0] for z in ref_latents]
+            ref_gram  = torch.stack(all_grams).mean(0)             # [C_lat, C_lat]

-            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  "
+        print(f"[DITTO] Reference latent stats from {len(ref_latents)} clips  "
              f"n_opt={n_opt_steps}  lr={opt_lr}  ode_steps={n_ode_steps}  "
              f"grad_steps={n_grad_steps}", flush=True)

@@ -244,7 +292,7 @@ class SelvaDittoOptimizer:
                    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,
+                    style_weight, gram_weight, anchor_weight, steps, cfg_strength, seed,
                    normalize, target_lufs, pbar,
                )
            except Exception as e:
@@ -270,19 +318,42 @@ def _do_optimize(net_generator, feature_utils, mel_converter,
                 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,
+                 style_weight, gram_weight, anchor_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()
+    # Strip inference flags from ref stats (came from main thread) and cast to
+    # model dtype. ref_mean/ref_gram are float32 (computed via cuFFT mel path);
+    # mel_gen is model dtype (bfloat16). Mixed-dtype loss → float32 gradient →
+    # "Found dtype Float but expected BFloat16" in backward through bfloat16 ops.
+    ref_mean = ref_mean.clone().detach().to(dtype)
+    ref_gram = ref_gram.clone().detach().to(dtype)

    torch.manual_seed(seed)

    clip_f = features["clip_features"].to(device, dtype).clone()
    sync_f = features["sync_features"].to(device, dtype).clone()

+    # Strip inference-mode flags from all model weights and buffers BEFORE any
+    # forward pass. Parameters were loaded in ComfyUI's inference_mode context;
+    # operations on inference tensors produce inference tensors, so conditions
+    # computed from tainted weights would also be tainted. clone() outside
+    # inference_mode produces a normal tensor regardless of the source flag.
+    def _strip_inference(module):
+        for mod in module.modules():
+            for name, buf in list(mod._buffers.items()):
+                if buf is not None:
+                    mod._buffers[name] = buf.clone()
+            for name, param in list(mod._parameters.items()):
+                if param is not None:
+                    mod._parameters[name] = torch.nn.Parameter(
+                        param.data.clone(), requires_grad=False
+                    )
+
+    _strip_inference(net_generator)
+    _strip_inference(feature_utils)
+    _strip_inference(mel_converter)
+
    net_generator.update_seq_lengths(
        latent_seq_len=seq_cfg.latent_seq_len,
        clip_seq_len=clip_f.shape[1],
@@ -300,12 +371,27 @@ def _do_optimize(net_generator, feature_utils, mel_converter,
            bs=1, negative_text_features=neg_text_clip
        )

+    # Clone all tensors inside conditions/empty_conditions to ensure no inference
+    # flags survived from intermediate computations inside preprocess_conditions.
+    def _clone_nested(obj):
+        if isinstance(obj, torch.Tensor):
+            return obj.clone()
+        elif isinstance(obj, dict):
+            return {k: _clone_nested(v) for k, v in obj.items()}
+        elif isinstance(obj, (list, tuple)):
+            return type(obj)(_clone_nested(v) for v in obj)
+        return obj
+
+    conditions       = _clone_nested(conditions)
+    empty_conditions = _clone_nested(empty_conditions)
+
    # 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())
+    x0_init = x0_init.detach()          # anchor — kept fixed, no grad
    optimizer = torch.optim.Adam([x0], lr=opt_lr)

    # n_grad_steps must not exceed n_ode_steps
@@ -364,23 +450,19 @@ def _do_optimize(net_generator, feature_utils, mel_converter,
            )
            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)
+        # ── Style loss in latent space ───────────────────────────────────────
+        # Unnormalize x back to VAE latent space — fully differentiable, no
+        # decode needed. ref_mean/ref_gram are computed from encoded reference
+        # clips in the same space. Avoids backprop through VAE decoder which
+        # is @torch.inference_mode() and produces noisy gradients.
+        x_un = net_generator.unnormalize(x)          # [1, T_lat, C_lat]
+        style_loss = style_weight * _latent_style_loss(x_un.squeeze(0), ref_mean, ref_gram, gram_weight)

-        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)
+        # Anchor regularization — penalize x0 drifting from its initial N(0,1)
+        # value. Flow matching ODE expects x0 ~ N(0,1); large deviations push
+        # the ODE into an out-of-distribution region that decodes as white noise.
+        anchor_loss = anchor_weight * F.mse_loss(x0, x0_init)
+        loss = style_loss + anchor_loss

        optimizer.zero_grad()
        loss.backward()  # gradient flows through Phase 2 + STE back to x0.grad
@@ -390,7 +472,9 @@ def _do_optimize(net_generator, feature_utils, mel_converter,
        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)
+            print(f"[DITTO] {opt_step+1}/{n_opt_steps}  "
+                  f"style={style_loss.item():.4f}  anchor={anchor_loss.item():.4f}  "
+                  f"x0_std={x0.data.std().item():.3f}", flush=True)

    # ── Final generation with optimized x_0 ─────────────────────────────────
    print(f"[DITTO] Optimization done. Final generation ({steps} steps)...", flush=True)
@@ -428,4 +512,4 @@ def _do_optimize(net_generator, feature_utils, mel_converter,
            audio = audio / peak

    print(f"[DITTO] audio: shape={tuple(audio.shape)} sr={sample_rate}", flush=True)
-    return ({"waveform": audio.cpu(), "sample_rate": sample_rate},)
+    return {"waveform": audio.cpu(), "sample_rate": sample_rate}
@@ -316,10 +316,14 @@ class SelvaLoraEvaluator:
                        alpha      = float(meta.get("alpha", float(rank)))
                        target     = list(meta.get("target", ["attn.qkv"]))
                        dropout    = float(meta.get("lora_dropout", 0.0))
+                        use_rslora = meta.get("use_rslora", False)
                        record["meta"] = {"rank": rank, "alpha": alpha, "target": target}

+                        # Always use standard init for loading — PiSSA checkpoints
+                        # include linear.weight (residual) in state_dict
                        n = apply_lora(generator, rank=rank, alpha=alpha,
-                                       target_suffixes=tuple(target), dropout=dropout)
+                                       target_suffixes=tuple(target), dropout=dropout,
+                                       init_mode="standard", use_rslora=use_rslora)
                        if n == 0:
                            raise RuntimeError(
                                f"apply_lora matched 0 layers (target={target})"
@@ -61,16 +61,23 @@ class SelvaLoraLoader:
        rank      = int(meta.get("rank",   16))
        alpha     = float(meta.get("alpha", float(rank)))
        target    = list(meta.get("target", ["attn.qkv"]))
+        init_mode = meta.get("init_mode", "standard")
+        use_rslora = meta.get("use_rslora", False)

        print(f"[SelVA LoRA] Loading adapter: {p.name}", flush=True)
-        print(f"[SelVA LoRA]   rank={rank}  alpha={alpha}  target={target}  strength={strength}",
+        print(f"[SelVA LoRA]   rank={rank}  alpha={alpha}  target={target}  "
+              f"init={init_mode}  rslora={use_rslora}  strength={strength}",
              flush=True)

        # Shallow-copy the model bundle so the original generator is not mutated
        patched = {**model}
        generator = copy.deepcopy(model["generator"])

-        n = apply_lora(generator, rank=rank, alpha=alpha, target_suffixes=tuple(target))
+        # For PiSSA, use standard init (the base weights will be overwritten
+        # by load_state_dict since the checkpoint includes linear.weight)
+        n = apply_lora(generator, rank=rank, alpha=alpha,
+                       target_suffixes=tuple(target),
+                       init_mode="standard", use_rslora=use_rslora)
        if n == 0:
            raise RuntimeError(
                f"[SelVA LoRA] No layers matched target={target}. "
@@ -79,6 +79,10 @@ _PARAM_DEFAULTS = {
    "lora_dropout":        0.0,
    "lora_plus_ratio":     1.0,
    "lr_schedule":         "constant",
+    "init_mode":           "pissa",
+    "use_rslora":          True,
+    "latent_mixup_alpha":  0.0,
+    "latent_noise_sigma":  0.0,
 }

 # Palette for comparison chart: one color per experiment (cycles if > 8)
@@ -388,7 +392,9 @@ class SelvaLoraScheduler:
            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)
+            init_mode   = str(cfg.get("init_mode",            "pissa"))
+            use_rslora  = bool(cfg.get("use_rslora",          True))
+            alpha_val   = alpha if alpha > 0.0 else float(2 * rank)
            target_suffixes = tuple(target.strip().split())

            output_dir = output_root / exp_id
@@ -410,6 +416,7 @@ class SelvaLoraScheduler:
                    "curriculum_switch": curr_switch,
                    "lora_dropout": dropout, "lora_plus_ratio": plus_ratio,
                    "lr_schedule": lr_schedule,
+                    "init_mode": init_mode, "use_rslora": use_rslora,
                },
                "results":     {"status": "running"},
                "adapter_path": None,
@@ -428,7 +435,7 @@ class SelvaLoraScheduler:
                        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,
+                        lr_schedule, init_mode, use_rslora,
                    )

                duration          = time.monotonic() - t_start
@@ -21,7 +21,10 @@ import folder_paths
 from .utils import SELVA_CATEGORY, get_device, soft_empty_cache
 from selva_core.model.utils.features_utils import FeaturesUtils
 from selva_core.model.flow_matching import FlowMatching
-from selva_core.model.lora import apply_lora, get_lora_state_dict, load_lora
+from selva_core.model.lora import (
+    apply_lora, get_lora_state_dict, get_lora_and_base_state_dict, load_lora,
+    spectral_surgery,
+)


 _AUDIO_EXTS = {".wav", ".flac", ".mp3", ".ogg", ".aiff", ".aif"}
@@ -69,6 +72,10 @@ def _load_audio(path: Path, target_sr: int, duration: float) -> torch.Tensor:
        waveform = waveform.mean(0, keepdim=True)
    waveform = waveform.squeeze(0).float()
    if sr != target_sr:
+        if not getattr(_load_audio, "_sr_warned", False):
+            print(f"[LoRA Trainer] WARNING: audio sr={sr} ≠ target {target_sr}, resampling on-the-fly. "
+                  f"Pre-resample with SelVA Dataset Resampler for faster loading.", flush=True)
+            _load_audio._sr_warned = True
        waveform = torchaudio.functional.resample(
            waveform.unsqueeze(0), sr, target_sr).squeeze(0)
    target_len = int(duration * target_sr)
@@ -119,18 +126,22 @@ def _eval_sample(generator, feature_utils_orig, dataset, seq_cfg, device, dtype,

        with torch.no_grad():
            x1_pred   = eval_fm.to_data(velocity_fn, x0)
-            x1_unnorm = generator.unnormalize(x1_pred)
+            # .clone() strips inference-mode flag from x1_pred (computed from
+            # inference-mode weights) so unnormalize's in-place ops don't fail.
+            x1_unnorm = generator.unnormalize(x1_pred.clone())

-            # feature_utils_orig may be on CPU (offload strategy) — move temporarily
-            orig_device = next(feature_utils_orig.parameters()).device
-            if orig_device != device:
-                feature_utils_orig.to(device)
+            # Only move the VAE+vocoder (tod) to GPU — avoids moving the
+            # entire FeaturesUtils (CLIP, T5, Synchformer) which wastes VRAM
+            # and fixes mixed-device state issues where the first parameter
+            # check could miss CPU-stranded vocoder weights.
+            tod = feature_utils_orig.tod
+            tod_orig_device = next(tod.parameters()).device
+            tod.to(device)
            try:
                spec  = feature_utils_orig.decode(x1_unnorm)
                audio = feature_utils_orig.vocode(spec)
            finally:
-                if orig_device != device:
-                    feature_utils_orig.to(orig_device)
+                tod.to(tod_orig_device)

        audio = audio.float().cpu()
        if audio.dim() == 2:
@@ -213,6 +224,77 @@ def _spectral_metrics(wav: torch.Tensor, sr: int) -> dict:
    }


+def _reference_metrics(gen_wav: torch.Tensor, ref_wav: torch.Tensor, sr: int) -> dict:
+    """Compare generated eval sample against the original reference audio.
+
+    gen_wav, ref_wav: [1, L] float32 CPU tensors (mono).
+
+    Returns:
+      log_spectral_distance_db — RMS dB difference across freq bins (lower = better)
+      mel_cepstral_distortion  — scaled L2 in log-mel space (lower = better)
+      per_band_correlation     — mean Pearson r across mel bands (higher = better)
+    """
+    from .selva_audio_preprocessors import _mel_filterbank
+
+    L = min(gen_wav.shape[-1], ref_wav.shape[-1])
+    gen = gen_wav[..., :L].squeeze(0)
+    ref = ref_wav[..., :L].squeeze(0)
+
+    n_fft  = _SPEC_N_FFT
+    hop    = _SPEC_HOP
+    window = torch.hann_window(n_fft)
+
+    gen_stft = torch.stft(gen, n_fft, hop, window=window, return_complex=True)
+    ref_stft = torch.stft(ref, n_fft, hop, window=window, return_complex=True)
+
+    gen_mag = gen_stft.abs()
+    ref_mag = ref_stft.abs()
+
+    T = min(gen_mag.shape[1], ref_mag.shape[1])
+    gen_mag = gen_mag[:, :T]
+    ref_mag = ref_mag[:, :T]
+
+    # Log-spectral distance (dB)
+    gen_db = 20.0 * torch.log10(gen_mag.clamp(min=1e-8))
+    ref_db = 20.0 * torch.log10(ref_mag.clamp(min=1e-8))
+    lsd = float(((gen_db - ref_db) ** 2).mean().sqrt())
+
+    # Mel-scale comparison
+    n_mels = 80
+    fb = _mel_filterbank(sr, n_fft, n_mels, 0, sr // 2)  # [n_mels, n_freqs]
+
+    gen_mel = torch.matmul(fb, gen_mag).clamp(min=1e-8)
+    ref_mel = torch.matmul(fb, ref_mag).clamp(min=1e-8)
+
+    gen_mel_log = torch.log(gen_mel)
+    ref_mel_log = torch.log(ref_mel)
+
+    # Mel cepstral distortion (L2 in log-mel space, standard scaling)
+    mcd = float(
+        (10.0 / np.log(10))
+        * ((gen_mel_log - ref_mel_log) ** 2).mean(0).sqrt().mean()
+    )
+
+    # Per-band Pearson correlation
+    gen_np = gen_mel_log.numpy()
+    ref_np = ref_mel_log.numpy()
+    correlations = []
+    for b in range(n_mels):
+        g, r = gen_np[b], ref_np[b]
+        if g.std() < 1e-8 or r.std() < 1e-8:
+            continue
+        corr = np.corrcoef(g, r)[0, 1]
+        if np.isfinite(corr):
+            correlations.append(corr)
+    mean_corr = float(np.mean(correlations)) if correlations else 0.0
+
+    return {
+        "log_spectral_distance_db": round(lsd, 2),
+        "mel_cepstral_distortion":  round(mcd, 2),
+        "per_band_correlation":     round(mean_corr, 4),
+    }
+
+
 def _save_spectrogram(wav: torch.Tensor, sr: int, path: Path) -> None:
    """Save a log-frequency dB spectrogram PNG for an eval sample.

@@ -482,8 +564,9 @@ class SelvaLoraTrainer:
            },
            "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).",
+                    "default": 0.0, "min": 0.0, "max": 512.0, "step": 0.5,
+                    "tooltip": "LoRA alpha. 0 = use 2×rank (fewer intruder dimensions, "
+                               "arXiv:2410.21228). Set explicitly to override.",
                }),
                "target": ("STRING", {
                    "default": "attn.qkv",
@@ -521,13 +604,27 @@ class SelvaLoraTrainer:
                "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).",
+                               "0=disabled. 0.05–0.1 helps regularize on small datasets (arXiv:2404.09610). "
+                               "Forced to 0 when using PiSSA init.",
                }),
                "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).",
                }),
+                "init_mode": (["standard", "pissa"], {
+                    "default": "pissa",
+                    "tooltip": "LoRA initialization mode. "
+                               "standard: Kaiming-uniform A + zero B (classic LoRA). "
+                               "pissa: A and B from top-r SVD of pretrained weight — starts "
+                               "on-manifold, eliminates intruder dimensions (arXiv:2404.02948). "
+                               "Recommended for audio generation where off-manifold latents cause noise.",
+                }),
+                "use_rslora": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "Rank-stabilized LoRA scaling: alpha/sqrt(rank) instead of alpha/rank. "
+                               "Prevents gradient collapse at high ranks (128+). Recommended (arXiv:2312.03732).",
+                }),
                "lr_schedule": (["constant", "cosine"], {
                    "default": "constant",
                    "tooltip": "LR schedule after warmup. "
@@ -535,6 +632,17 @@ class SelvaLoraTrainer:
                               "cosine: decay from lr to ~0 following a cosine curve — "
                               "prevents oscillation when LR is slightly too high.",
                }),
+                "latent_mixup_alpha": ("FLOAT", {
+                    "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05,
+                    "tooltip": "Beta distribution alpha for latent mixup (MusicLDM, arXiv:2308.01546). "
+                               "0 = disabled. 0.4 recommended. Mixes pairs of training latents "
+                               "to reduce memorization on small datasets.",
+                }),
+                "latent_noise_sigma": ("FLOAT", {
+                    "default": 0.0, "min": 0.0, "max": 0.1, "step": 0.005,
+                    "tooltip": "Additive Gaussian noise on training latents, scaled by x1.std(). "
+                               "0 = disabled. 0.01–0.03 adds mild regularization against overfitting.",
+                }),
            },
        }

@@ -558,7 +666,9 @@ class SelvaLoraTrainer:
              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"):
+              lora_dropout=0.0, lora_plus_ratio=1.0,
+              init_mode="pissa", use_rslora=True, lr_schedule="constant",
+              latent_mixup_alpha=0.0, latent_noise_sigma=0.0):

        torch.manual_seed(seed)
        random.seed(seed)
@@ -591,7 +701,7 @@ class SelvaLoraTrainer:
        output_dir = _out_p
        output_dir.mkdir(parents=True, exist_ok=True)

-        alpha_val      = float(alpha) if alpha > 0.0 else float(rank)
+        alpha_val      = float(alpha) if alpha > 0.0 else float(2 * rank)
        target_suffixes = tuple(target.strip().split())

        dataset = _prepare_dataset(model, data_dir, device)
@@ -609,6 +719,8 @@ class SelvaLoraTrainer:
                grad_accum, save_every, resume_path, seed,
                timestep_mode, logit_normal_sigma, curriculum_switch,
                lora_dropout, lora_plus_ratio, lr_schedule,
+                init_mode, use_rslora,
+                latent_mixup_alpha, latent_noise_sigma,
            )
            return (r["patched_model"], r["adapter_path"], r["loss_curve"])

@@ -620,19 +732,25 @@ class SelvaLoraTrainer:
        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",
+        init_mode="pissa", use_rslora=True,
+        latent_mixup_alpha=0.0, latent_noise_sigma=0.0,
    ):
        # --- 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, dropout=lora_dropout)
+                            target_suffixes=target_suffixes, dropout=lora_dropout,
+                            init_mode=init_mode, use_rslora=use_rslora)
        if n_lora == 0:
            raise RuntimeError(
                f"[LoRA Trainer] No layers matched target={target_suffixes}. "
                "Check the 'target' field."
            )
+        scale_str = f"alpha/√rank={alpha_val/math.sqrt(rank):.2f}" if use_rslora \
+                    else f"alpha/rank={alpha_val/rank:.2f}"
        print(f"[LoRA Trainer] Wrapped {n_lora} layers "
-              f"(rank={rank}, alpha={alpha_val}, dropout={lora_dropout})", flush=True)
+              f"(rank={rank}, alpha={alpha_val}, {scale_str}, "
+              f"init={init_mode}, dropout={lora_dropout})", flush=True)

        for name, p in generator.named_parameters():
            p.requires_grad_("lora_" in name)
@@ -651,7 +769,7 @@ class SelvaLoraTrainer:
        optimizer = torch.optim.AdamW([
            {"params": lora_A_params, "lr": lr},
            {"params": lora_B_params, "lr": lr * lora_plus_ratio},
-        ], weight_decay=1e-2)
+        ], weight_decay=0.0)
        if lora_plus_ratio != 1.0:
            print(f"[LoRA Trainer] LoRA+: lr_A={lr:.2e}  lr_B={lr * lora_plus_ratio:.2e}", flush=True)

@@ -717,6 +835,10 @@ class SelvaLoraTrainer:
            "lora_dropout":        lora_dropout,
            "lora_plus_ratio":     lora_plus_ratio,
            "lr_schedule":         lr_schedule,
+            "init_mode":           init_mode,
+            "use_rslora":          use_rslora,
+            "latent_mixup_alpha":  latent_mixup_alpha,
+            "latent_noise_sigma":  latent_noise_sigma,
        }

        # For curriculum mode: compute the step at which we switch from logit_normal to uniform
@@ -727,10 +849,26 @@ class SelvaLoraTrainer:
              f"(step {start_step + 1} → {steps}, batch_size={batch_size}, "
              f"timestep_mode={timestep_mode})\n", flush=True)

+        # Load reference audio for eval comparison against original
+        ref_wav = None
+        try:
+            npz_files = sorted(data_dir.glob("*.npz"))
+            if npz_files:
+                ref_audio_path = _find_audio(npz_files[0])
+                if ref_audio_path is not None:
+                    ref_wav = _load_audio(ref_audio_path, seq_cfg.sampling_rate,
+                                          seq_cfg.duration).unsqueeze(0)  # [1, L]
+                    print(f"[LoRA Trainer] Reference audio: {ref_audio_path.name}", flush=True)
+        except Exception as e:
+            print(f"[LoRA Trainer] Could not load reference audio: {e}", flush=True)
+
        last_step = start_step
        completed = False
        try:
            for step in range(start_step + 1, steps + 1):
+                if batch_size <= len(dataset):
+                    batch = random.sample(dataset, k=batch_size)
+                else:
                    batch = random.choices(dataset, k=batch_size)
                x1_list, clip_list, sync_list, text_list = zip(*batch)

@@ -741,6 +879,18 @@ class SelvaLoraTrainer:

                x1 = generator.normalize(x1)

+                # Latent mixup (MusicLDM, arXiv:2308.01546)
+                if latent_mixup_alpha > 0 and x1.shape[0] > 1:
+                    lam = torch.distributions.Beta(
+                        latent_mixup_alpha, latent_mixup_alpha
+                    ).sample().to(device)
+                    idx = torch.randperm(x1.shape[0], device=device)
+                    x1 = lam * x1 + (1 - lam) * x1[idx]
+
+                # Latent noise injection
+                if latent_noise_sigma > 0:
+                    x1 = x1 + torch.randn_like(x1) * latent_noise_sigma * x1.std()
+
                if timestep_mode == "logit_normal" or (
                    timestep_mode == "curriculum" and step <= curriculum_switch_step
                ):
@@ -811,8 +961,11 @@ class SelvaLoraTrainer:

                if step % save_every == 0 or step == steps:
                    ckpt_path = output_dir / f"adapter_step{step:05d}.pt"
+                    # PiSSA checkpoints need base weights (residual W_res)
+                    sd = get_lora_and_base_state_dict(generator) if init_mode == "pissa" \
+                         else get_lora_state_dict(generator)
                    torch.save({
-                        "state_dict": get_lora_state_dict(generator),
+                        "state_dict": sd,
                        "optimizer":  optimizer.state_dict(),
                        "scheduler":  scheduler.state_dict(),
                        "step":       step,
@@ -842,6 +995,13 @@ class SelvaLoraTrainer:
                                  f"flatness={metrics['spectral_flatness']:.3f}  "
                                  f"temporal_var={metrics['temporal_variance']:.3f}", flush=True)
                            _save_spectrogram(wav, sr, wav_path)
+                            if ref_wav is not None:
+                                ref_m = _reference_metrics(wav, ref_wav, sr)
+                                spectral_metrics[step].update(ref_m)
+                                print(f"[LoRA Trainer] vs Reference: "
+                                      f"LSD={ref_m['log_spectral_distance_db']:.1f}dB  "
+                                      f"MCD={ref_m['mel_cepstral_distortion']:.2f}  "
+                                      f"corr={ref_m['per_band_correlation']:.3f}", flush=True)
                        except Exception as e:
                            print(f"[LoRA Trainer] Spectral/spectrogram failed: {e}", flush=True)

@@ -850,6 +1010,38 @@ class SelvaLoraTrainer:

            completed = True

+            # ── Post-training Spectral Surgery ────────────────────────────────
+            # Reweight LoRA singular values using gradient sensitivity on the
+            # training set. Suppresses intruder dimensions, amplifies useful ones.
+            # (arXiv:2603.03995). Only run on normal completion.
+            try:
+                print("[LoRA Trainer] Running Spectral Surgery...", flush=True)
+                fm_surgery = FlowMatching(min_sigma=0, inference_mode="euler", num_steps=25)
+
+                def _calibration_fn(model_cal, step_idx):
+                    sample = dataset[step_idx % len(dataset)]
+                    x1_cal, clip_cal, sync_cal, text_cal = sample
+                    x1_b = x1_cal.unsqueeze(0).to(device, dtype) if x1_cal.dim() == 2 \
+                           else x1_cal.to(device, dtype)
+                    x1_b = model_cal.normalize(x1_b.clone())
+                    clip_b = clip_cal.to(device, dtype)
+                    sync_b = sync_cal.to(device, dtype)
+                    text_b = text_cal.to(device, dtype)
+                    t = torch.rand(1, device=device, dtype=dtype)
+                    x0_b = torch.randn_like(x1_b)
+                    xt = fm_surgery.get_conditional_flow(x0_b, x1_b, t)
+                    v_pred = model_cal.forward(xt, clip_b, sync_b, text_b, t)
+                    cal_loss = fm_surgery.loss(v_pred, x0_b, x1_b).mean()
+                    cal_loss.backward()
+
+                n_cal = min(128, len(dataset) * 4)
+                n_surgery = spectral_surgery(generator, _calibration_fn,
+                                             n_calibration=n_cal)
+                print(f"[LoRA Trainer] Spectral Surgery done: {n_surgery} layers processed.",
+                      flush=True)
+            except Exception as e:
+                print(f"[LoRA Trainer] Spectral Surgery failed (non-fatal): {e}", flush=True)
+
        finally:
            # Save adapter and loss curves whether training completed or was cancelled.
            # Skip if we never completed a single step (nothing useful to save).
@@ -868,7 +1060,9 @@ class SelvaLoraTrainer:
                    final_path = output_dir / f"adapter_cancelled_step{last_step:05d}.pt"
                    label = f"Cancelled at step {last_step}"

-                torch.save({"state_dict": get_lora_state_dict(generator), "meta": meta}, final_path)
+                final_sd = get_lora_and_base_state_dict(generator) if init_mode == "pissa" \
+                           else get_lora_state_dict(generator)
+                torch.save({"state_dict": final_sd, "meta": meta}, final_path)
                (output_dir / "meta.json").write_text(json.dumps(meta, indent=2))
                print(f"\n[LoRA Trainer] {label}. Adapter saved to {final_path}", flush=True)

@@ -149,7 +149,7 @@ class SelvaModelLoader:
            enable_conditions=True,
            mode=mode,
            bigvgan_vocoder_ckpt=bigvgan_path,
-            need_vae_encoder=False,
+            need_vae_encoder=True,
        ).to(device, dtype).eval()

        if strategy == "offload_to_cpu":
@@ -93,15 +93,14 @@ def _eval_sample_ti(generator, learned_tokens, n_tokens, inject_mode,
            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)
+            tod = feature_utils_orig.tod
+            tod_orig_dev = next(tod.parameters()).device
+            tod.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)
+                tod.to(tod_orig_dev)

        audio = audio.float().cpu()
        if audio.dim() == 2:
@@ -124,15 +124,14 @@ class SelvaVaeRoundtrip:
                      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)
+                tod = feature_utils.tod
+                tod_orig_device = next(tod.parameters()).device
+                tod.to(device)
                try:
                    spec = feature_utils.decode(latent_unnorm)
                    out  = feature_utils.vocode(spec)
                finally:
-                    if orig_device != device:
-                        feature_utils.to(orig_device)
+                    tod.to(tod_orig_device)

            out = out.float().cpu()
            if out.dim() == 1:
@@ -1,6 +1,17 @@
 """
 LoRA (Low-Rank Adaptation) for SelVA / MMAudio generator.

+Supports two initialization modes:
+  - **standard**: Kaiming-uniform A, zero B (classic LoRA).
+  - **pissa**: A and B from the top-r SVD of the pretrained weight.
+    Starts on-manifold, eliminates intruder dimensions at init
+    (arXiv:2404.02948, NeurIPS 2024 Spotlight).
+
+Supports two scaling modes:
+  - **standard**: alpha / rank
+  - **rslora**: alpha / sqrt(rank) — rank-stabilized scaling that prevents
+    gradient collapse at high ranks (arXiv:2312.03732).
+
 Usage:
    from selva_core.model.lora import apply_lora, get_lora_state_dict, load_lora

@@ -25,14 +36,16 @@ 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) + (dropout(x) @ A.T @ B.T) * (alpha / rank)
+    Output: base(x) + (dropout(x) @ A.T @ B.T) * scale

-    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.
+    Standard init: A is Kaiming uniform, B is zero → adapter starts at zero.
+    PiSSA init: A and B from top-r SVD of pretrained weight → adapter starts
+    at the principal components, base weight stores the residual.
    """

-    def __init__(self, linear: nn.Linear, rank: int, alpha: float, dropout: float = 0.0):
+    def __init__(self, linear: nn.Linear, rank: int, alpha: float,
+                 dropout: float = 0.0, init_mode: str = "standard",
+                 use_rslora: bool = False):
        super().__init__()
        in_f  = linear.in_features
        out_f = linear.out_features
@@ -44,11 +57,35 @@ class LoRALinear(nn.Module):

        ref_dtype  = linear.weight.dtype
        ref_device = linear.weight.device
-        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))
+
+        if use_rslora:
+            self.scale = alpha / math.sqrt(rank)
+        else:
            self.scale = alpha / rank
+
        self.dropout = nn.Dropout(p=dropout) if dropout > 0.0 else nn.Identity()

+        if init_mode == "pissa":
+            # PiSSA: init from top-r SVD of pretrained weight.
+            # SVD in float32 for numerical stability, then cast back.
+            W = linear.weight.data.float()  # [out_f, in_f]
+            U, S, Vt = torch.linalg.svd(W, full_matrices=False)
+
+            sqrt_S = S[:rank].sqrt()
+            # A: [rank, in_f],  B: [out_f, rank]
+            A_init = sqrt_S.unsqueeze(1) * Vt[:rank, :]
+            B_init = U[:, :rank] * sqrt_S.unsqueeze(0)
+
+            # Residual: W_res = W - B_init @ A_init * scale
+            # so that base(x) + LoRA(x) = W_res@x + (B@A)*scale@x = W@x at init
+            linear.weight.data = (W - B_init @ A_init * self.scale).to(ref_dtype)
+
+            self.lora_A = nn.Parameter(A_init.to(dtype=ref_dtype, device=ref_device))
+            self.lora_B = nn.Parameter(B_init.to(dtype=ref_dtype, device=ref_device))
+        else:
+            # Standard LoRA: Kaiming A, zero B → starts at identity
+            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))
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -67,6 +104,8 @@ def apply_lora(
    alpha: float = None,
    target_suffixes: tuple = ("attn.qkv",),
    dropout: float = 0.0,
+    init_mode: str = "standard",
+    use_rslora: bool = False,
 ) -> int:
    """Replace matching nn.Linear layers with LoRALinear in-place.

@@ -80,6 +119,9 @@ def apply_lora(
                         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.
+                         Must be 0 when using PiSSA (principal components shouldn't be dropped).
+        init_mode:       "standard" (Kaiming/zero) or "pissa" (SVD-based).
+        use_rslora:      If True, scale by alpha/sqrt(rank) instead of alpha/rank.

    Returns:
        Number of linear layers wrapped.
@@ -87,6 +129,11 @@ def apply_lora(
    if alpha is None:
        alpha = float(rank)

+    if init_mode == "pissa" and dropout > 0.0:
+        print("[LoRA] Warning: dropout forced to 0 for PiSSA init "
+              "(principal components should not be dropped).")
+        dropout = 0.0
+
    count = 0
    for name, module in list(model.named_modules()):
        if not any(name.endswith(s) for s in target_suffixes):
@@ -98,7 +145,10 @@ def apply_lora(
        parent = model
        for part in parts[:-1]:
            parent = getattr(parent, part)
-        setattr(parent, parts[-1], LoRALinear(module, rank, alpha, dropout=dropout))
+        setattr(parent, parts[-1], LoRALinear(
+            module, rank, alpha, dropout=dropout,
+            init_mode=init_mode, use_rslora=use_rslora,
+        ))
        count += 1

    return count
@@ -109,6 +159,141 @@ def get_lora_state_dict(model: nn.Module) -> dict:
    return {k: v for k, v in model.state_dict().items() if "lora_" in k}


+def get_lora_and_base_state_dict(model: nn.Module) -> dict:
+    """Return state dict with LoRA params AND base linear weights.
+
+    Needed for PiSSA checkpoints where the base weight stores the residual
+    (W - top_r(W)*scale), not the original pretrained weight.
+    """
+    result = {}
+    for name, module in model.named_modules():
+        if isinstance(module, LoRALinear):
+            prefix = name + "."
+            result[prefix + "lora_A"] = module.lora_A.data
+            result[prefix + "lora_B"] = module.lora_B.data
+            result[prefix + "linear.weight"] = module.linear.weight.data
+            if module.linear.bias is not None:
+                result[prefix + "linear.bias"] = module.linear.bias.data
+    return result
+
+
+def spectral_surgery(
+    model: nn.Module,
+    calibration_fn,
+    n_calibration: int = 128,
+    policy: str = "smooth_abs",
+):
+    """Post-training Spectral Surgery: reweight LoRA singular values to suppress
+    intruder dimensions and amplify useful components (arXiv:2603.03995).
+
+    Args:
+        model:          Model with LoRA applied.
+        calibration_fn: Callable that takes (model, step_idx) and runs one forward+backward
+                        pass on a calibration sample. Must call loss.backward().
+        n_calibration:  Number of calibration samples to average gradients over.
+        policy:         Reweighting policy: "smooth_abs" (recommended), "hard" (binary).
+
+    Modifies LoRA A and B in-place. Returns number of layers processed.
+    """
+    model.eval()
+    lora_layers = [(name, mod) for name, mod in model.named_modules()
+                   if isinstance(mod, LoRALinear)]
+
+    if not lora_layers:
+        return 0
+
+    # Accumulate per-layer gradient sensitivity: g_k = u_k^T * (dL/dΔW) * v_k
+    sensitivities = {}
+    for name, mod in lora_layers:
+        sensitivities[name] = None
+
+    for step in range(n_calibration):
+        model.zero_grad()
+        # Enable grad temporarily on LoRA params
+        for _, mod in lora_layers:
+            mod.lora_A.requires_grad_(True)
+            mod.lora_B.requires_grad_(True)
+
+        calibration_fn(model, step)
+
+        for name, mod in lora_layers:
+            A = mod.lora_A.data.float()   # [rank, in_f]
+            B = mod.lora_B.data.float()   # [out_f, rank]
+            # ΔW = B @ A * scale → gradient dL/dΔW ≈ (dL/dB @ A + B^T @ dL/dA) / 2
+            # Per-component sensitivity: project onto SVD directions
+            delta_W = (B @ A * mod.scale).detach()
+            U, S, Vt = torch.linalg.svd(delta_W, full_matrices=False)
+            r = A.shape[0]
+            U_r, S_r, Vt_r = U[:, :r], S[:r], Vt[:r, :]
+
+            # Compute sensitivity from LoRA gradients
+            if mod.lora_A.grad is not None and mod.lora_B.grad is not None:
+                grad_A = mod.lora_A.grad.float()   # [rank, in_f]
+                grad_B = mod.lora_B.grad.float()   # [out_f, rank]
+                # dL/d(ΔW) ≈ grad_B @ A + B^T @ grad_A (chain rule through B@A)
+                grad_dW = grad_B @ A + B.T @ grad_A  # approximate
+                # Per-component: g_k = u_k^T @ grad_dW @ v_k
+                g = torch.einsum("ik,ij,jk->k", U_r, grad_dW, Vt_r.T)  # [r]
+            else:
+                g = torch.zeros(r, device=A.device)
+
+            if sensitivities[name] is None:
+                sensitivities[name] = g
+            else:
+                sensitivities[name] += g
+
+        # Disable grad again
+        for _, mod in lora_layers:
+            mod.lora_A.requires_grad_(False)
+            mod.lora_B.requires_grad_(False)
+
+    # Apply reweighting per layer
+    count = 0
+    for name, mod in lora_layers:
+        g = sensitivities[name] / n_calibration
+        A = mod.lora_A.data.float()
+        B = mod.lora_B.data.float()
+
+        delta_W = B @ A * mod.scale
+        U, S, Vt = torch.linalg.svd(delta_W, full_matrices=False)
+        r = A.shape[0]
+        S_r = S[:r]
+
+        if policy == "hard":
+            # Keep components with positive sensitivity, zero out negative
+            mask = (g > 0).float()
+        else:
+            # smooth_abs: sigmoid-weighted by sensitivity magnitude
+            # Normalize g to [-1, 1] range, apply sigmoid
+            g_norm = g / (g.abs().max() + 1e-8)
+            mask = torch.sigmoid(5.0 * g_norm)  # steep sigmoid
+
+        # L1 norm preservation: scale mask so total nuclear norm is preserved
+        mask = mask * (S_r.sum() / (mask * S_r).sum().clamp(min=1e-8))
+
+        # Reconstruct: ΔW' = U_r @ diag(mask * S_r) @ Vt_r
+        S_new = mask * S_r
+        delta_W_new = U[:, :r] @ torch.diag(S_new) @ Vt[:r, :]
+
+        # Factor back into B' @ A' * scale: use SVD of ΔW'/scale
+        dW_unscaled = delta_W_new / mod.scale
+        U2, S2, Vt2 = torch.linalg.svd(dW_unscaled, full_matrices=False)
+        sqrt_S2 = S2[:r].sqrt()
+        A_new = sqrt_S2.unsqueeze(1) * Vt2[:r, :]
+        B_new = U2[:, :r] * sqrt_S2.unsqueeze(0)
+
+        ref_dtype = mod.lora_A.dtype
+        mod.lora_A.data = A_new.to(ref_dtype)
+        mod.lora_B.data = B_new.to(ref_dtype)
+        count += 1
+
+        kept = (mask > 0.5).sum().item()
+        print(f"[Spectral Surgery] {name}: kept {kept}/{r} components, "
+              f"sensitivity range [{g.min():.3f}, {g.max():.3f}]", flush=True)
+
+    return count
+
+
 def load_lora(model: nn.Module, state_dict: dict) -> None:
    """Load LoRA weights into a model that has already had apply_lora() called.