code/scripts/dramabox_finetune_train_multi_aux.py

#!/usr/bin/env python3
"""
DramaBox Voice Cloning LoRA Fine-Tuning with Multi-Auxiliary Losses (Epochs 15-17)

Extends dramabox_finetune_train_clap.py with THREE auxiliary losses:
  1. CLAP Naturalness — CLAP text similarity (positive-negative) + quality MLP
  2. Centroid Real/Fake — cos(emb, real_centroid) - cos(emb, synth_centroid)
  3. Speaker Similarity — WavLM-SV cosine similarity between ref and pred speaker embs

Each auxiliary loss is individually normalized via EMA-based adaptive coefficients
to have approximately the same magnitude as the flow matching loss.

Usage:
    accelerate launch --num_processes=8 scripts/dramabox_finetune_train_multi_aux.py \
        --config configs/finetune_multi_aux.yaml
"""

import os
import sys

# Filter out conda ml-general paths that break native cuDNN libraries
_ld = os.environ.get("LD_LIBRARY_PATH", "")
if _ld:
    _filtered = [p for p in _ld.split(":") if "ml-general" not in p]
    os.environ["LD_LIBRARY_PATH"] = ":".join(_filtered)

# Auto-accept trust_remote_code for HuggingFace models (non-interactive multi-GPU)
os.environ["HF_HUB_TRUST_REMOTE_CODE"] = "1"
os.environ["TRUST_REMOTE_CODE"] = "1"

import argparse
import http.server
import json
import logging
import math
import random
import glob as glob_mod
import shutil
import subprocess as _subprocess
import threading
import time
from collections import defaultdict, deque
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset, Sampler

# ── Path setup ─────────────────────────────────────────────────────────

def setup_paths(dramabox_dir: str):
    """Add DramaBox to sys.path for imports."""
    sys.path.insert(0, os.path.join(dramabox_dir, "ltx2"))
    sys.path.insert(0, os.path.join(dramabox_dir, "src"))


# ── Timestep Samplers (from DramaBox train.py) ────────────────────────

class ShiftedLogitNormalTimestepSampler:
    """Shifted logit-normal distribution, shift depends on sequence length."""

    def __init__(self, std: float = 1.0, eps: float = 1e-3, uniform_prob: float = 0.1):
        self.std = std
        self.eps = eps
        self.uniform_prob = uniform_prob
        self.normal_999_percentile = 3.0902 * std
        self.normal_005_percentile = -2.5758 * std

    def sample(self, batch_size: int, seq_length: int, device: torch.device = None) -> torch.Tensor:
        mu = self._get_shift(seq_length)
        normal = torch.randn(batch_size, device=device) * self.std + mu
        logitnormal = torch.sigmoid(normal)
        p999 = torch.sigmoid(torch.tensor(mu + self.normal_999_percentile, device=device))
        p005 = torch.sigmoid(torch.tensor(mu + self.normal_005_percentile, device=device))
        stretched = (logitnormal - p005) / (p999 - p005)
        stretched = torch.where(stretched >= self.eps, stretched, 2 * self.eps - stretched)
        stretched = stretched.clamp(0, 1)
        uniform = (1 - self.eps) * torch.rand(batch_size, device=device) + self.eps
        prob = torch.rand(batch_size, device=device)
        return torch.where(prob > self.uniform_prob, stretched, uniform)

    @staticmethod
    def _get_shift(seq_length, min_tok=1024, max_tok=4096, min_s=0.95, max_s=2.05):
        m = (max_s - min_s) / (max_tok - min_tok)
        return m * seq_length + (min_s - m * min_tok)


class DistilledTimestepSampler:
    SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0]

    def __init__(self, jitter: float = 0.02):
        self.jitter = jitter

    def sample(self, batch_size: int, seq_length: int = None, device: torch.device = None) -> torch.Tensor:
        n_intervals = len(self.SIGMAS) - 1
        interval_idx = torch.randint(0, n_intervals, (batch_size,), device=device)
        t = torch.rand(batch_size, device=device)
        sigma_high = torch.tensor([self.SIGMAS[i] for i in interval_idx], device=device)
        sigma_low = torch.tensor([self.SIGMAS[i + 1] for i in interval_idx], device=device)
        sigma = sigma_low + t * (sigma_high - sigma_low)
        return sigma.clamp(0.01, 0.99)


# ── Dataset ────────────────────────────────────────────────────────────

class DramaBoxFinetuneDataset(Dataset):
    """Three-mode dataset for DramaBox voice cloning fine-tuning.

    Modes:
        A (voice_clone_fwd): target=part2, ref=part1, text=scene2_expected_text
        B (voice_clone_rev): target=part1, ref=part2, text=scene1_expected_text
        C (unconditional):   target=full,  ref=None,  text=full_prompt
    """

    def __init__(self, preprocessed_dir: str, mode_weights: dict = None,
                 max_ref_tokens: int = 250, overfit_n: int = 0,
                 expand_all_modes: bool = False, index_file: str = None):
        self.data_dir = Path(preprocessed_dir)
        self.latent_dir = self.data_dir / "audio_latents"
        self.cond_dir = self.data_dir / "conditions"

        self.max_ref_tokens = max_ref_tokens
        self.expand_all_modes = expand_all_modes
        self.mode_weights = mode_weights or {
            "voice_clone_fwd": 0.33,
            "voice_clone_rev": 0.33,
            "unconditional": 0.34,
        }

        # Load index — support custom index file for filtered datasets
        index_path = index_file if index_file else str(self.data_dir / "index.json")
        with open(index_path) as f:
            index = json.load(f)

        self.groups = index["groups"]  # prompt_id -> [sample_indices]
        self.group_keys = list(self.groups.keys())
        self.all_samples = index["samples"]  # list of sample metadata

        if overfit_n > 0:
            self.group_keys = self.group_keys[:overfit_n]
            valid_indices = set()
            for k in self.group_keys:
                valid_indices.update(self.groups[k])
            self.all_samples = [s for s in self.all_samples if s["index"] in valid_indices]

        self._build_items()

        logging.info(f"Dataset: {len(self.items)} items from {len(self.group_keys)} groups, "
                     f"{len(self.all_samples)} total samples (index: {os.path.basename(index_path)})")

    def _build_items(self):
        """Build flat item list."""
        self.items = []
        modes = list(self.mode_weights.keys())
        if self.expand_all_modes:
            for sample in self.all_samples:
                source = sample.get("source", "dramabox")
                if source in ("emolia", "augmented", "podcast"):
                    # Emolia/augmented/podcast pairs have direction baked in; use fwd only
                    self.items.append((sample["index"], "voice_clone_fwd"))
                else:
                    for mode in modes:
                        self.items.append((sample["index"], mode))
        else:
            for group_key in self.group_keys:
                for mode in modes:
                    self.items.append((group_key, mode))

    def __len__(self):
        return len(self.items)

    def _load_latent(self, idx: int, part: str) -> torch.Tensor:
        path = self.latent_dir / f"sample_{idx:06d}_{part}.pt"
        lat = torch.load(path, weights_only=True).detach()
        return lat

    def _load_condition(self, idx: int, part: str):
        path = self.cond_dir / f"sample_{idx:06d}_{part}.pt"
        cond = torch.load(path, weights_only=False)
        audio_feats = cond.get("audio_prompt_embeds", cond.get("prompt_embeds")).detach()
        attn_mask = cond.get("prompt_attention_mask").detach()

        # Pad to multiple of 128 for audio_connector
        REG = 128
        L = audio_feats.shape[0]
        target_L = ((L + REG - 1) // REG) * REG
        if target_L != L:
            pad_len = target_L - L
            pad_emb = torch.zeros(pad_len, audio_feats.shape[1], dtype=audio_feats.dtype)
            pad_mask = torch.zeros(pad_len, dtype=attn_mask.dtype)
            audio_feats = torch.cat([pad_emb, audio_feats], dim=0)
            attn_mask = torch.cat([pad_mask, attn_mask], dim=0)

        return audio_feats, attn_mask

    def __getitem__(self, idx):
        item_key, mode = self.items[idx]

        if self.expand_all_modes:
            sample_idx = item_key
        else:
            sample_indices = self.groups[item_key]
            sample_idx = random.choice(sample_indices)

        if mode == "voice_clone_fwd":
            tgt_latent = self._load_latent(sample_idx, "part2")
            ref_latent = self._load_latent(sample_idx, "part1")
            audio_feats, attn_mask = self._load_condition(sample_idx, "part2")
        elif mode == "voice_clone_rev":
            tgt_latent = self._load_latent(sample_idx, "part1")
            ref_latent = self._load_latent(sample_idx, "part2")
            audio_feats, attn_mask = self._load_condition(sample_idx, "part1")
        else:  # unconditional
            tgt_latent = self._load_latent(sample_idx, "full")
            C, F_dim = tgt_latent.shape[0], tgt_latent.shape[2]
            ref_latent = torch.zeros(C, 0, F_dim, dtype=tgt_latent.dtype)
            audio_feats, attn_mask = self._load_condition(sample_idx, "full")

        # Cap reference length
        if ref_latent.shape[1] > self.max_ref_tokens:
            ref_latent = ref_latent[:, :self.max_ref_tokens, :]

        return {
            "tgt_latent": tgt_latent,       # [C=8, T, F=16]
            "ref_latent": ref_latent,        # [C=8, T_ref, F=16] or [C=8, 0, F=16]
            "audio_features": audio_feats,
            "attention_mask": attn_mask,
            "mode": mode,
        }


def collate_fn(batch):
    """Pad variable-length audio to max in batch."""
    max_tgt_T = max(b["tgt_latent"].shape[1] for b in batch)
    max_ref_T = max(b["ref_latent"].shape[1] for b in batch)
    C = batch[0]["tgt_latent"].shape[0]
    F_dim = batch[0]["tgt_latent"].shape[2]

    tgt_list, ref_list, feat_list, mask_list = [], [], [], []
    tgt_lengths, ref_lengths = [], []
    modes = []

    for b in batch:
        tgt = b["tgt_latent"]
        ref = b["ref_latent"]
        tgt_lengths.append(tgt.shape[1])
        ref_lengths.append(ref.shape[1])

        if tgt.shape[1] < max_tgt_T:
            pad = torch.zeros(C, max_tgt_T - tgt.shape[1], F_dim, dtype=tgt.dtype)
            tgt = torch.cat([tgt, pad], dim=1)
        tgt_list.append(tgt)

        if ref.shape[1] < max_ref_T:
            pad = torch.zeros(C, max_ref_T - ref.shape[1], F_dim, dtype=ref.dtype)
            ref = torch.cat([ref, pad], dim=1)
        ref_list.append(ref)

        feat_list.append(b["audio_features"])
        mask_list.append(b["attention_mask"])
        modes.append(b["mode"])

    return {
        "tgt_latent": torch.stack(tgt_list),
        "ref_latent": torch.stack(ref_list),
        "audio_features": torch.stack(feat_list),
        "attention_mask": torch.stack(mask_list),
        "tgt_lengths": torch.tensor(tgt_lengths),
        "ref_lengths": torch.tensor(ref_lengths),
        "modes": modes,
    }


# ── Bucket-Weighted Sampler ───────────────────────────────────────────

class BucketWeightedSampler(Sampler):
    """Weighted sampling from source buckets with round-robin within each bucket.

    Divides dataset items into buckets by source (dramabox, podcast, emolia,
    augmented). Each epoch produces items in proportions matching bucket_weights,
    cycling through each bucket's items in shuffled order.

    One "epoch" = each item in the anchor bucket (podcast) seen once.
    Total training length = anchor_items / anchor_weight.
    """

    def __init__(self, dataset, bucket_weights: dict, anchor_bucket: str = "podcast",
                 seed: int = 42):
        """
        Args:
            dataset: DramaBoxFinetuneDataset with expand_all_modes=True
            bucket_weights: source -> weight, e.g. {"dramabox": 0.5, "podcast": 0.35, ...}
            anchor_bucket: which bucket defines one epoch (each item seen once per epoch)
            seed: base random seed
        """
        super().__init__(dataset)
        self.bucket_weights = bucket_weights
        self.seed = seed
        self._epoch = 0

        # Build source -> item_indices mapping
        sample_idx_to_source = {
            s["index"]: s.get("source", "dramabox") for s in dataset.all_samples
        }
        self.buckets = {}
        for source in bucket_weights:
            self.buckets[source] = []
        for i, (item_key, mode) in enumerate(dataset.items):
            source = sample_idx_to_source.get(item_key, "dramabox")
            if source in self.buckets:
                self.buckets[source].append(i)
            else:
                # Unknown source falls into dramabox bucket
                self.buckets.setdefault("dramabox", []).append(i)

        # Epoch size: anchor bucket items seen exactly once per epoch
        anchor_items = len(self.buckets.get(anchor_bucket, []))
        anchor_weight = bucket_weights.get(anchor_bucket, 0.35)
        self.anchor = anchor_bucket
        self._epoch_size = int(anchor_items / anchor_weight) if anchor_items > 0 else len(dataset)

        # Log bucket info
        info = []
        for source, weight in bucket_weights.items():
            n = len(self.buckets.get(source, []))
            per_epoch = int(self._epoch_size * weight)
            repeats = per_epoch / n if n > 0 else 0
            info.append(f"  {source}: {n} items × {repeats:.1f}× = {per_epoch}/epoch (weight={weight})")
        logging.info(f"BucketWeightedSampler: {self._epoch_size} items/epoch, anchor={self.anchor}")
        for line in info:
            logging.info(line)

    def __len__(self):
        return self._epoch_size

    def set_epoch(self, epoch):
        self._epoch = epoch

    def __iter__(self):
        rng = random.Random(self.seed + self._epoch)
        self._epoch += 1  # auto-increment for next iter()

        indices = []
        for source, weight in self.bucket_weights.items():
            bucket = self.buckets.get(source, [])
            if not bucket or weight <= 0:
                continue
            count = int(self._epoch_size * weight)
            # Repeat bucket items enough times, shuffle, take count
            repeats = (count // len(bucket)) + 1
            pool = []
            for _ in range(repeats):
                shuffled = bucket.copy()
                rng.shuffle(shuffled)
                pool.extend(shuffled)
            indices.extend(pool[:count])

        # Final shuffle to interleave all sources
        rng.shuffle(indices)
        return iter(indices)


# ── Model Building (from DramaBox train.py) ───────────────────────────

def build_audio_only_model(checkpoint_path, device, dtype):
    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
    from ltx_core.loader.registry import DummyRegistry
    from ltx_core.loader.sd_ops import SDOps
    from ltx_core.model.transformer.model import LTXModel, LTXModelType
    from ltx_core.model.transformer.attention import AttentionFunction
    from ltx_core.model.transformer.rope import LTXRopeType

    sd_ops = SDOps("AO").with_matching(
        prefix="model.diffusion_model."
    ).with_replacement("model.diffusion_model.", "")

    class Cfg:
        @classmethod
        def from_config(cls, config):
            from ltx_core.model.model_protocol import ModelConfigurator
            t = config.get("transformer", {})
            cp = None
            if not t.get("caption_proj_before_connector", False):
                from ltx_core.model.transformer.text_projection import create_caption_projection
                with torch.device("meta"):
                    cp = create_caption_projection(t, audio=True)
            return LTXModel(
                model_type=LTXModelType.AudioOnly,
                audio_num_attention_heads=t.get("audio_num_attention_heads", 32),
                audio_attention_head_dim=t.get("audio_attention_head_dim", 64),
                audio_in_channels=t.get("audio_in_channels", 128),
                audio_out_channels=t.get("audio_out_channels", 128),
                num_layers=t.get("num_layers", 48),
                audio_cross_attention_dim=t.get("audio_cross_attention_dim", 2048),
                norm_eps=t.get("norm_eps", 1e-6),
                attention_type=AttentionFunction(t.get("attention_type", "default")),
                positional_embedding_theta=t.get("positional_embedding_theta", 10000.0),
                audio_positional_embedding_max_pos=t.get("audio_positional_embedding_max_pos", [20]),
                timestep_scale_multiplier=t.get("timestep_scale_multiplier", 1000),
                use_middle_indices_grid=t.get("use_middle_indices_grid", True),
                rope_type=LTXRopeType(t.get("rope_type", "interleaved")),
                double_precision_rope=t.get("frequencies_precision", False) == "float64",
                apply_gated_attention=t.get("apply_gated_attention", False),
                audio_caption_projection=cp,
                cross_attention_adaln=t.get("cross_attention_adaln", False),
            )

    builder = Builder(model_path=checkpoint_path, model_class_configurator=Cfg,
                      model_sd_ops=sd_ops, registry=DummyRegistry())
    return builder.build(device=device, dtype=dtype)


def load_audio_connector(checkpoint_path, device, dtype):
    from ltx_trainer.model_loader import load_embeddings_processor
    emb_proc = load_embeddings_processor(checkpoint_path, device=device, dtype=dtype)
    connector = emb_proc.audio_connector
    del emb_proc
    return connector


def apply_lora(model, rank, alpha, dropout=0.0):
    from peft import LoraConfig, get_peft_model
    config = LoraConfig(
        r=rank, lora_alpha=alpha, lora_dropout=dropout, bias="none",
        target_modules=[
            "audio_attn1.to_k", "audio_attn1.to_q",
            "audio_attn1.to_v", "audio_attn1.to_out.0",
            "audio_ff.net.0.proj", "audio_ff.net.2",
        ],
    )
    model = get_peft_model(model, config)
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    logging.info(f"LoRA: {trainable:,} trainable / {total:,} total ({100*trainable/total:.1f}%)")
    return model


@torch.no_grad()
def prepare_audio_context(audio_connector, audio_features, attention_mask, device, dtype):
    from ltx_core.text_encoders.gemma.embeddings_processor import convert_to_additive_mask
    audio_features = audio_features.to(device=device, dtype=dtype)
    attention_mask = attention_mask.to(device=device)
    if audio_features.shape[0] > 1:
        results = []
        for i in range(audio_features.shape[0]):
            feat_i = audio_features[i:i+1]
            mask_i = attention_mask[i:i+1]
            additive = convert_to_additive_mask(mask_i, feat_i.dtype)
            enc_i, _ = audio_connector(feat_i, additive)
            results.append(enc_i)
        return torch.cat(results, dim=0)
    additive_mask = convert_to_additive_mask(attention_mask, audio_features.dtype)
    audio_encoded, _ = audio_connector(audio_features, additive_mask)
    return audio_encoded


def _unwrap_model_safe(model):
    while hasattr(model, "module"):
        model = model.module
    return model


def save_training_state(output_dir, step, epoch, optimizer, scheduler, best_loss,
                        best_step, model, accelerator, tag=""):
    """Save full training state for resumability."""
    suffix = f"_epoch{epoch}" if tag == "" else f"_{tag}"
    state = {
        "step": step,
        "epoch": epoch,
        "best_loss": best_loss,
        "best_step": best_step,
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "rng_python": random.getstate(),
        "rng_torch": torch.random.get_rng_state(),
        "rng_cuda": torch.cuda.get_rng_state(),
    }
    state_path = os.path.join(output_dir, f"training_state{suffix}.pt")
    torch.save(state, state_path)
    # Save LoRA weights alongside
    unwrapped = _unwrap_model_safe(model)
    unwrapped.save_pretrained(output_dir)
    adapter = os.path.join(output_dir, "adapter_model.safetensors")
    lora_path = os.path.join(output_dir, f"lora_epoch{epoch}.safetensors")
    if os.path.exists(adapter):
        shutil.copy(adapter, lora_path)
    logging.info(f"Saved full training state: {state_path}, LoRA: {lora_path}")
    return lora_path


def load_training_state(resume_dir, optimizer, scheduler):
    """Load full training state from checkpoint directory (picks highest step)."""
    state_files = glob_mod.glob(os.path.join(resume_dir, "training_state*.pt"))
    # Filter out broken symlinks
    state_files = [f for f in state_files if os.path.isfile(f)]
    if not state_files:
        raise FileNotFoundError(f"No training state found in {resume_dir}")
    # Pick the checkpoint with the highest step number
    best_path, best_step = None, -1
    for sf in state_files:
        try:
            s = torch.load(sf, map_location="cpu", weights_only=False)
            if s["step"] > best_step:
                best_step = s["step"]
                best_path = sf
            del s
        except Exception as e:
            logging.warning(f"Skipping unreadable state file {sf}: {e}")
    if best_path is None:
        raise FileNotFoundError(f"No valid training state found in {resume_dir}")
    state_path = best_path
    logging.info(f"Resuming from: {state_path} (step {best_step})")
    state = torch.load(state_path, weights_only=False)
    optimizer.load_state_dict(state["optimizer"])
    scheduler.load_state_dict(state["scheduler"])
    random.setstate(state["rng_python"])
    torch.random.set_rng_state(state["rng_torch"])
    torch.cuda.set_rng_state(state["rng_cuda"])
    return state["step"], state["epoch"], state["best_loss"], state["best_step"]


def launch_validation_async(script_dir, output_dir, lora_path, epoch, args):
    """Launch epoch validation as a background subprocess."""
    val_script = os.path.join(script_dir, "run_epoch_validation.py")
    if not os.path.exists(val_script):
        logging.warning(f"Validation script not found: {val_script}")
        return None
    cmd = [
        sys.executable, val_script,
        "--lora", lora_path,
        "--epoch", str(epoch),
        "--output-dir", output_dir,
        "--dramabox-dir", args.dramabox_dir,
        "--checkpoint", args.checkpoint,
        "--full-checkpoint", args.full_checkpoint,
        "--gemma-root", getattr(args, "gemma_root",
            "/home/deployer/.cache/dramabox/models--unsloth--gemma-3-12b-it-bnb-4bit/snapshots/826e729dbaeea4ecb143738eed2bcf3539ebf7bf"),
        "--lora-rank", str(args.lora_rank),
        "--val-samples", str(args.val_samples),
        "--val-refs-dir", args.val_refs_dir,
        "--preprocessed-dir", args.preprocessed_dir,
    ]
    log_path = os.path.join(output_dir, f"val_epoch{epoch}.log")
    log_f = open(log_path, "w")
    proc = _subprocess.Popen(cmd, stdout=log_f, stderr=_subprocess.STDOUT,
                             env={**os.environ, "CUDA_VISIBLE_DEVICES": "0"})
    logging.info(f"Launched validation for epoch {epoch} (PID {proc.pid}), log: {log_path}")
    return proc


# ── Metrics Server ────────────────────────────────────────────────────

class MetricsHandler(http.server.SimpleHTTPRequestHandler):
    """Serves the training monitor HTML and metrics data."""

    metrics_dir = None

    def do_GET(self):
        path = self.path.split("?")[0]

        if path == "/" or path == "/index.html":
            monitor_path = os.path.join(
                os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                "scripts", "dramabox_finetune_monitor.html"
            )
            if os.path.exists(monitor_path):
                self.send_response(200)
                self.send_header("Content-Type", "text/html")
                self.send_header("Access-Control-Allow-Origin", "*")
                self.end_headers()
                with open(monitor_path, "rb") as f:
                    self.wfile.write(f.read())
            else:
                self.send_error(404, "Monitor HTML not found")
        elif path == "/metrics":
            metrics_path = os.path.join(self.metrics_dir, "metrics.jsonl")
            if os.path.exists(metrics_path):
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Access-Control-Allow-Origin", "*")
                self.end_headers()
                with open(metrics_path, "rb") as f:
                    self.wfile.write(f.read())
            else:
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Access-Control-Allow-Origin", "*")
                self.end_headers()
                self.wfile.write(b"")
        elif path == "/status":
            status_path = os.path.join(self.metrics_dir, "status.json")
            if os.path.exists(status_path):
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Access-Control-Allow-Origin", "*")
                self.end_headers()
                with open(status_path, "rb") as f:
                    self.wfile.write(f.read())
            else:
                self.send_error(404)
        elif path.startswith("/val/"):
            rel = path[5:]
            fpath = os.path.join(self.metrics_dir, "val", rel)
            if os.path.exists(fpath):
                self.send_response(200)
                if fpath.endswith(".html"):
                    ct = "text/html"
                elif fpath.endswith(".wav"):
                    ct = "audio/wav"
                elif fpath.endswith(".json"):
                    ct = "application/json"
                else:
                    ct = "application/octet-stream"
                self.send_header("Content-Type", ct)
                self.send_header("Access-Control-Allow-Origin", "*")
                self.end_headers()
                with open(fpath, "rb") as f:
                    self.wfile.write(f.read())
            else:
                self.send_error(404)
        else:
            self.send_error(404)

    def log_message(self, format, *args):
        pass


def start_metrics_server(output_dir: str, port: int = 8765):
    MetricsHandler.metrics_dir = output_dir
    server = http.server.HTTPServer(("0.0.0.0", port), MetricsHandler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    logging.info(f"Training monitor serving on http://0.0.0.0:{port}")
    return server


# ── CLAP + Audio Decoder Setup ────────────────────────────────────────

def setup_clap_models(full_checkpoint, device, dtype):
    """Load AudioDecoder + VoiceCLAP-small for auxiliary losses.

    All models are frozen and in eval mode.
    Returns: (audio_decoder, clap_model, clap_tokenizer)
    """
    from ltx_pipelines.utils.blocks import AudioDecoder

    logging.info("Loading AudioDecoder (VAE decoder + vocoder)...")
    audio_decoder = AudioDecoder(
        checkpoint_path=full_checkpoint,
        dtype=dtype,
        device=device,
        warm=True,
    )

    logging.info("Loading VoiceCLAP-small...")
    from transformers import AutoModel, AutoTokenizer
    clap_model = AutoModel.from_pretrained(
        "laion/voiceclap-small", trust_remote_code=True
    ).eval().to(device)
    clap_tokenizer = AutoTokenizer.from_pretrained("laion/voiceclap-small",
                                                     trust_remote_code=True)

    for p in clap_model.parameters():
        p.requires_grad = False

    logging.info("CLAP + AudioDecoder loaded successfully")
    return audio_decoder, clap_model, clap_tokenizer


def setup_clap_models_large(full_checkpoint, device, dtype, args):
    """Load AudioDecoder + large VoiceCLAP via SentenceTransformer with optional quantization.

    For 7B+ models, INT4 quantization via bitsandbytes keeps VRAM manageable (~4 GB).
    Returns: (audio_decoder, st_model, None) — no separate tokenizer needed.
    """
    from ltx_pipelines.utils.blocks import AudioDecoder

    logging.info("Loading AudioDecoder (VAE decoder + vocoder)...")
    audio_decoder = AudioDecoder(
        checkpoint_path=full_checkpoint,
        dtype=dtype,
        device=device,
        warm=True,
    )

    logging.info(f"Loading large CLAP model: {args.clap_model} (quantize={args.clap_quantize})...")
    from sentence_transformers import SentenceTransformer

    model_kwargs = {"torch_dtype": torch.bfloat16, "trust_remote_code": True}

    if args.clap_quantize in ("int4", "int8"):
        from transformers import BitsAndBytesConfig
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=(args.clap_quantize == "int4"),
            load_in_8bit=(args.clap_quantize == "int8"),
            bnb_4bit_compute_dtype=torch.bfloat16,
        )
        model_kwargs["quantization_config"] = bnb_config

    st_model = SentenceTransformer(
        args.clap_model,
        model_kwargs=model_kwargs,
        trust_remote_code=True,
    )

    # Log embedding dimension
    test_emb = st_model.encode(["test"], convert_to_tensor=True)
    emb_dim = test_emb.shape[-1]
    logging.info(f"Large CLAP loaded: embedding dim = {emb_dim}")

    for p in st_model.parameters():
        p.requires_grad = False

    logging.info("Large CLAP + AudioDecoder loaded successfully")
    return audio_decoder, st_model, None


def encode_text_st(text, st_model):
    """Encode text with SentenceTransformer CLAP, returns normalized embedding [1, D]."""
    with torch.no_grad():
        emb = st_model.encode([text], convert_to_tensor=True)
    emb = F.normalize(emb, p=2, dim=-1)
    return emb  # [1, D]


def encode_audio_st(waveform_np, sr, st_model, rank=0):
    """Encode audio with SentenceTransformer CLAP via temp file, returns normalized [1, D].

    Args:
        waveform_np: numpy array of audio samples (mono)
        sr: sample rate
        st_model: SentenceTransformer model
        rank: GPU rank for unique temp file naming
    """
    import soundfile as sf
    tmp_path = f"/dev/shm/clap_tmp_{rank}.wav"
    try:
        sf.write(tmp_path, waveform_np, sr)
        with torch.no_grad():
            emb = st_model.encode([{"audio": tmp_path}], convert_to_tensor=True)
        emb = F.normalize(emb, p=2, dim=-1)
        return emb  # [1, D]
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)


def encode_clap_text(text, clap_model, clap_tokenizer, device):
    """Encode text with VoiceCLAP-small, returns normalized embedding [1, 768]."""
    enc = clap_tokenizer([text], padding=True, truncation=True, return_tensors="pt").to(device)
    with torch.no_grad():
        emb = clap_model.encode_text(enc.input_ids, enc.attention_mask)
    emb = F.normalize(emb, p=2, dim=-1)
    return emb  # [1, 768]


def encode_clap_waveform_differentiable(waveform, clap_model):
    """Differentiable CLAP audio encoding — bypasses @torch.no_grad() on compute_log_mel.

    VoiceCLAP-small's encode_waveform has @torch.no_grad() on its mel computation,
    which breaks gradient flow. This function replicates the mel computation with
    gradients enabled, then feeds into the audio encoder normally.

    Args:
        waveform: [B, T] or [T] at 16kHz, with requires_grad from upstream
        clap_model: VoiceCLAP-small model (frozen but grad flows through ops)

    Returns:
        Normalized audio embedding [B, 768] with gradient chain intact
    """
    _CHUNK_SAMPLES = 30 * 16000  # 480000

    if waveform.dim() == 1:
        waveform = waveform.unsqueeze(0)
    B, T = waveform.shape
    device = waveform.device

    # Pad to 30s chunks (same as original)
    n_chunks = max(1, (T + _CHUNK_SAMPLES - 1) // _CHUNK_SAMPLES)
    pad = n_chunks * _CHUNK_SAMPLES - T
    if pad > 0:
        waveform = F.pad(waveform, (0, pad))
    chunks = waveform.view(B, n_chunks, _CHUNK_SAMPLES).reshape(B * n_chunks, _CHUNK_SAMPLES)

    # ── Differentiable mel spectrogram (NO @torch.no_grad()) ──
    chunks_f32 = chunks.to(dtype=torch.float32)
    window = torch.hann_window(400, device=device)
    stft = torch.stft(chunks_f32, n_fft=400, hop_length=160, window=window, return_complex=True)
    magnitudes = stft[..., :-1].abs() ** 2

    # Use CLAP's registered mel_filters buffer
    mel_filters = clap_model.mel_filters.to(magnitudes.dtype)
    mel = mel_filters @ magnitudes

    log_spec = torch.clamp(mel, min=1e-10).log10()
    log_spec = torch.maximum(log_spec, log_spec.amax(dim=(-2, -1), keepdim=True) - 8.0)
    log_spec = (log_spec + 4.0) / 4.0

    # ── Audio encoder (already differentiable) ──
    feats = clap_model.audio_encoder(log_spec)   # (B*n_chunks, T', D)
    feats = feats.mean(dim=1)                     # clip-level mean
    feats = clap_model.audio_proj(feats)
    feats = F.normalize(feats, dim=-1)

    if n_chunks == 1:
        return feats
    feats = feats.view(B, n_chunks, -1).mean(dim=1)
    return F.normalize(feats, dim=-1)


# ── Speaker Similarity Model Setup ────────────────────────────────────

def setup_speaker_model(device, dtype):
    """Load WavLM-SV for speaker similarity.

    WavLM-base-plus-sv produces 512-dim speaker embeddings. It is a superior
    alternative to ECAPA-TDNN from the same HuggingFace ecosystem with no extra
    dependencies (uses transformers WavLMForXVector class).
    """
    from transformers import WavLMForXVector, Wav2Vec2FeatureExtractor

    logging.info("Loading WavLM-base-plus-sv for speaker similarity...")
    model = WavLMForXVector.from_pretrained("microsoft/wavlm-base-plus-sv")
    model.eval().to(device=device, dtype=torch.float32)  # keep float32 for stability
    for p in model.parameters():
        p.requires_grad = False

    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("microsoft/wavlm-base-plus-sv")

    logging.info("WavLM-SV loaded successfully (~200MB)")
    return model, feature_extractor


def wavlm_extract_embeddings(model, input_values):
    """Wrapper for WavLM that returns tensor (for grad_checkpoint compatibility)."""
    return model(input_values).embeddings


# ── Centroid + Quality MLP Setup ──────────────────────────────────────

def setup_centroid_and_quality(classifiers_dir, device, dtype):
    """Load centroid embeddings and quality MLP.

    Returns: (real_centroid [1, 768], synth_centroid [1, 768], quality_mlp)
    """
    import torch.nn as nn

    classifiers_dir = Path(classifiers_dir)

    # Centroid embeddings
    logging.info("Loading CLAP centroid embeddings...")
    emb_data = torch.load(classifiers_dir / "clap_embeddings.pt", map_location="cpu",
                          weights_only=False)
    dramabox_embs = emb_data["dramabox_embeddings"]  # [3247, 768]
    emilia_embs = emb_data["emilia_embeddings"]       # [3247, 768]

    # Use 80% for centroids (same train split as classifier training)
    n_train = int(len(dramabox_embs) * 0.8)
    synth_centroid = F.normalize(dramabox_embs[:n_train].float().mean(0, keepdim=True), p=2, dim=-1)
    real_centroid = F.normalize(emilia_embs[:n_train].float().mean(0, keepdim=True), p=2, dim=-1)
    synth_centroid = synth_centroid.to(device=device, dtype=torch.float32)
    real_centroid = real_centroid.to(device=device, dtype=torch.float32)
    logging.info(f"Centroids computed from {n_train} train samples each")

    # Quality MLP
    logging.info("Loading quality classifier MLP...")
    ckpt = torch.load(classifiers_dir / "quality_classifier.pt", map_location="cpu",
                      weights_only=False)

    class BinaryMLP(nn.Module):
        def __init__(self, d_in, h1, h2):
            super().__init__()
            self.net = nn.Sequential(
                nn.Linear(d_in, h1), nn.ReLU(), nn.Dropout(0.3),
                nn.Linear(h1, h2), nn.ReLU(), nn.Dropout(0.3),
                nn.Linear(h2, 1))

        def forward(self, x):
            return self.net(x)

    quality_mlp = BinaryMLP(ckpt["input_dim"], ckpt["hidden1"], ckpt["hidden2"])
    quality_mlp.load_state_dict(ckpt["model_state_dict"])
    quality_mlp.eval().to(device=device, dtype=torch.float32)  # keep float32 for stability
    for p in quality_mlp.parameters():
        p.requires_grad = False
    logging.info("Quality MLP loaded (102K params)")

    return real_centroid, synth_centroid, quality_mlp


# ── Args ──────────────────────────────────────────────────────────────

def parse_args():
    import yaml
    cfg_parser = argparse.ArgumentParser(add_help=False)
    cfg_parser.add_argument("--config", default=None)
    cfg_args, remaining = cfg_parser.parse_known_args()

    yaml_defaults = {}
    if cfg_args.config:
        with open(cfg_args.config) as f:
            yaml_defaults = yaml.safe_load(f) or {}
        yaml_defaults = {k.replace("-", "_"): v for k, v in yaml_defaults.items()}

    def _y(name, fallback):
        return yaml_defaults.get(name, fallback)

    p = argparse.ArgumentParser(parents=[cfg_parser],
                                description="DramaBox LoRA Training with Multi-Auxiliary Losses")
    p.add_argument("--preprocessed-dir", default=_y("preprocessed_dir", "./finetune_data"))
    p.add_argument("--output-dir", default=_y("output_dir", "./finetune_output"))
    p.add_argument("--dramabox-dir", default=_y("dramabox_dir", "/home/deployer/laion/DramaBox"))
    p.add_argument("--checkpoint", default=_y("checkpoint",
                   "/home/deployer/laion/DramaBox/models/ltx-2.3-22b-dev-audio-only-v13-merged.safetensors"))
    p.add_argument("--full-checkpoint", default=_y("full_checkpoint",
                   "/home/deployer/laion/DramaBox/models/ltx-2.3-22b-dev.safetensors"))
    p.add_argument("--gemma-root", default=_y("gemma_root",
                   "/home/deployer/.cache/dramabox/models--unsloth--gemma-3-12b-it-bnb-4bit/snapshots/826e729dbaeea4ecb143738eed2bcf3539ebf7bf"))
    p.add_argument("--base-model", choices=["distilled", "dev"], default=_y("base_model", "dev"))
    p.add_argument("--lora-rank", type=int, default=_y("lora_rank", 128))
    p.add_argument("--lora-alpha", type=int, default=_y("lora_alpha", 128))
    p.add_argument("--lora-dropout", type=float, default=_y("lora_dropout", 0.05))
    p.add_argument("--resume-lora", default=_y("resume_lora", None))
    p.add_argument("--max-ref-tokens", type=int, default=_y("max_ref_tokens", 250))
    p.add_argument("--text-dropout", type=float, default=_y("text_dropout", 0.1))
    p.add_argument("--steps", type=int, default=_y("steps", 15000))
    p.add_argument("--lr", type=float, default=_y("lr", 3e-5))
    p.add_argument("--lr-scheduler", choices=["cosine", "linear", "constant"],
                   default=_y("lr_scheduler", "cosine"))
    p.add_argument("--batch-size", type=int, default=_y("batch_size", 1))
    p.add_argument("--grad-accum", type=int, default=_y("grad_accum", 4))
    p.add_argument("--max-grad-norm", type=float, default=_y("max_grad_norm", 1.0))
    p.add_argument("--save-every", type=int, default=_y("save_every", 500))
    p.add_argument("--log-every", type=int, default=_y("log_every", 25))
    p.add_argument("--seed", type=int, default=_y("seed", 42))
    p.add_argument("--warmup-steps", type=int, default=_y("warmup_steps", 161))
    p.add_argument("--monitor-port", type=int, default=_y("monitor_port", 8765))
    p.add_argument("--test", action="store_true", help="Quick test: 1 GPU, 100 steps")
    p.add_argument("--overfit", type=int, default=0,
                   help="Overfit on N prompt groups (sanity check)")
    p.add_argument("--expand-all-modes", action="store_true",
                   default=bool(_y("expand_all_modes", False)),
                   help="Use all 3 modes per sample (not per group).")
    p.add_argument("--epochs", type=int, default=_y("epochs", 3),
                   help="Train for N epochs (overrides --steps)")
    p.add_argument("--val-samples", type=int, default=_y("val_samples", 10),
                   help="Number of validation samples per epoch")
    p.add_argument("--val-refs-dir", default=_y("val_refs_dir", "/home/deployer/laion/test-refs"),
                   help="Directory with reference audio WAVs for validation")
    p.add_argument("--resume-dir", default=_y("resume_dir", None),
                   help="Resume from full training state (dir with training_state.pt)")

    # Index file
    p.add_argument("--index-file", default=_y("index_file", None),
                   help="Path to custom index.json (e.g. index_filtered.json)")

    # Bucket-weighted sampling
    bw_default = _y("bucket_weights", None)
    p.add_argument("--bucket-weights", default=bw_default,
                   help="Source->weight dict for bucket sampling (YAML dict or JSON string)")
    p.add_argument("--podcast-epochs", type=int, default=_y("podcast_epochs", 0),
                   help="Train until podcast bucket seen N times (overrides --epochs when > 0)")

    # CLAP / multi-auxiliary loss args
    clap_default = _y("clap_loss", True)
    p.add_argument("--clap-loss", action="store_true", default=clap_default,
                   help="Enable CLAP auxiliary losses (default)")
    p.add_argument("--no-clap", action="store_true", default=False,
                   help="Disable all auxiliary losses (for ablation)")
    p.add_argument("--aux-target-ratio", type=float,
                   default=_y("aux_target_ratio", 1.0),
                   help="Target ratio of EACH aux loss to flow matching (default 1.0)")
    p.add_argument("--speaker-sim-ratio", type=float,
                   default=_y("speaker_sim_ratio", 0.0),
                   help="Override target ratio for speaker similarity loss (0 = use aux-target-ratio)")
    p.add_argument("--coeff-cap", type=float,
                   default=_y("coeff_cap", 10.0),
                   help="Max value for adaptive aux coefficients (default 10.0)")
    p.add_argument("--positive-text", default=_y("positive_text",
                   "Realistic, genuine, spontaneous, authentic, sensual, natural voice "
                   "with all imperfections and organic microdistractions a natural situation brings with it"),
                   help="CLAP positive text")
    p.add_argument("--negative-text", default=_y("negative_text",
                   "distorted, unnatural, robotic, distortion"),
                   help="CLAP negative text")
    p.add_argument("--classifiers-dir", default=_y("classifiers_dir", "./classifiers"),
                   help="Directory with quality_classifier.pt and clap_embeddings.pt")

    p.add_argument("--aux-sigma-max", type=float,
                   default=_y("aux_sigma_max", 0.4),
                   help="Only compute aux losses when sigma < this threshold (default 0.4)")
    p.add_argument("--rejection-sampling", action="store_true",
                   default=_y("rejection_sampling", False),
                   help="Enable rejection sampling: only train on above-median reward samples")
    p.add_argument("--rejection-percentile", type=float,
                   default=_y("rejection_percentile", 50.0),
                   help="Percentile threshold for rejection (default 50 = median)")
    p.add_argument("--differentiable-reward", action="store_true",
                   default=_y("differentiable_reward", False),
                   help="Backprop through decoder/CLAP/WavLM for true differentiable rewards (ReFL-style)")
    p.add_argument("--diff-reward-checkpoint", action="store_true",
                   default=_y("diff_reward_checkpoint", False),
                   help="Use gradient checkpointing on decoder/CLAP for VRAM savings")

    # Individual loss toggles (read defaults from YAML)
    p.add_argument("--no-speaker-sim", action="store_true",
                   default=bool(_y("no_speaker_sim", False)),
                   help="Disable speaker similarity loss")
    p.add_argument("--no-centroid", action="store_true",
                   default=bool(_y("no_centroid", False)),
                   help="Disable centroid real/fake loss")
    p.add_argument("--no-quality-mlp", action="store_true",
                   default=bool(_y("no_quality_mlp", False)),
                   help="Disable quality MLP (keep CLAP text similarity only)")

    # Large CLAP model support
    p.add_argument("--clap-model", default=_y("clap_model", "laion/voiceclap-small"),
                   help="CLAP model name (default: laion/voiceclap-small)")
    p.add_argument("--clap-quantize", default=_y("clap_quantize", "none"),
                   choices=["none", "int8", "int4"],
                   help="Quantization for CLAP model: none/int8/int4 (default: none)")
    p.add_argument("--keep-last-n", type=int, default=_y("keep_last_n", 0),
                   help="Rolling checkpoint window (0 = keep all)")
    p.add_argument("--no-save-state", action="store_true",
                   default=bool(_y("no_save_state", False)),
                   help="Skip saving optimizer state .pt files (saves disk)")

    args = p.parse_args(remaining)

    # --no-clap overrides everything
    if args.no_clap:
        args.clap_loss = False

    # Parse bucket_weights if provided as JSON string
    if isinstance(args.bucket_weights, str):
        args.bucket_weights = json.loads(args.bucket_weights)

    return args


# ── Main Training Loop ────────────────────────────────────────────────

def main():
    from accelerate import Accelerator
    from accelerate.utils import set_seed

    args = parse_args()
    setup_paths(args.dramabox_dir)

    from audio_conditioning import AudioConditionByReferenceLatent
    from ltx_core.components.patchifiers import AudioPatchifier
    from ltx_core.guidance.perturbations import BatchedPerturbationConfig
    from ltx_core.tools import AudioLatentTools
    from ltx_core.types import AudioLatentShape, LatentState
    from ltx_pipelines.utils.helpers import modality_from_latent_state

    if args.test:
        args.steps = min(args.steps, 100)
        args.save_every = 50
        args.log_every = 5

    accelerator = Accelerator(
        gradient_accumulation_steps=args.grad_accum,
        mixed_precision="bf16",
    )

    is_main = accelerator.is_main_process
    if is_main:
        logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    else:
        logging.basicConfig(level=logging.WARNING)

    set_seed(args.seed)
    device = accelerator.device
    dtype = torch.bfloat16

    os.makedirs(args.output_dir, exist_ok=True)

    # Start metrics server on main process
    if is_main:
        try:
            start_metrics_server(args.output_dir, args.monitor_port)
        except OSError as e:
            logging.warning(f"Could not start metrics server: {e}")

    # Save training args
    if is_main:
        import yaml
        args_dict = vars(args).copy()
        args_dict["_meta"] = {
            "world_size": accelerator.num_processes,
            "dtype": str(dtype),
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "script": "dramabox_finetune_train_multi_aux.py",
            "pattern": "IC-LoRA 3-mode + 3 auxiliary losses (naturalness, centroid, speaker sim)",
            "aux_enabled": args.clap_loss,
        }
        with open(os.path.join(args.output_dir, "training_args.yaml"), "w") as f:
            yaml.dump(args_dict, f, default_flow_style=False, sort_keys=False)

    # Build model
    if is_main:
        logging.info("Loading audio-only model...")
    model = build_audio_only_model(args.checkpoint, device, dtype)

    if is_main:
        logging.info("Loading audio connector...")
    audio_connector = load_audio_connector(args.full_checkpoint, device, dtype)
    audio_connector.eval()
    for p in audio_connector.parameters():
        p.requires_grad = False

    if is_main:
        logging.info(f"Applying LoRA (rank={args.lora_rank}, alpha={args.lora_alpha})...")
    model = apply_lora(model, args.lora_rank, args.lora_alpha, args.lora_dropout)

    # Resume LoRA
    if args.resume_lora:
        from safetensors.torch import load_file as st_load
        if is_main:
            logging.info(f"Resuming from: {args.resume_lora}")
        lora_sd = st_load(args.resume_lora)
        mapped = {}
        for k, v in lora_sd.items():
            nk = k.replace(".lora_A.weight", ".lora_A.default.weight").replace(
                ".lora_B.weight", ".lora_B.default.weight")
            mapped[nk] = v
        model.load_state_dict(mapped, strict=False)

    model.train()
    model.base_model.model.set_gradient_checkpointing(True)

    # ── Multi-auxiliary loss setup ────────────────────────────────────
    aux_enabled = args.clap_loss
    audio_decoder_clap = None
    clap_model = None
    clap_tokenizer = None
    st_model = None  # SentenceTransformer model for large CLAP
    pos_text_emb = None
    neg_text_emb = None
    real_centroid = None
    synth_centroid = None
    quality_mlp = None
    wavlm_sv = None
    wavlm_fe = None
    use_large_clap = False

    # EMA trackers for adaptive coefficients
    ema_alpha = 0.95
    ema_flow = 0.0
    ema_aux1 = 0.0
    ema_aux2 = 0.0
    ema_aux3 = 0.0
    target_ratio = args.aux_target_ratio
    speaker_sim_ratio = args.speaker_sim_ratio if args.speaker_sim_ratio > 0 else target_ratio
    coeff_cap = args.coeff_cap
    aux_sigma_max = args.aux_sigma_max

    use_centroid = not args.no_centroid
    use_quality_mlp = not args.no_quality_mlp
    use_speaker_sim = not args.no_speaker_sim

    if aux_enabled:
        if is_main:
            logging.info("Setting up multi-auxiliary loss models...")

        # Determine CLAP model type
        use_large_clap = (args.clap_model != "laion/voiceclap-small")

        if use_large_clap:
            # Large CLAP via SentenceTransformer (e.g. 7B with INT4)
            audio_decoder_clap, st_model, _ = setup_clap_models_large(
                args.full_checkpoint, device, dtype, args)
            pos_text_emb = encode_text_st(args.positive_text, st_model)
            neg_text_emb = encode_text_st(args.negative_text, st_model)

            # If centroid or quality_mlp is enabled, also load CLAP-small
            # (these classifiers were trained on 768-dim CLAP-small embeddings)
            if use_centroid or use_quality_mlp:
                if is_main:
                    logging.info("Also loading CLAP-small for centroid/quality_mlp compatibility...")
                from transformers import AutoModel, AutoTokenizer
                clap_model = AutoModel.from_pretrained(
                    "laion/voiceclap-small", trust_remote_code=True
                ).to(device).eval()
                for p in clap_model.parameters():
                    p.requires_grad_(False)
                if is_main:
                    logging.info("CLAP-small loaded alongside CLAP-large for aux classifiers")

            # Warn if differentiable reward requested with large CLAP
            if getattr(args, 'differentiable_reward', False) and is_main:
                logging.warning("Differentiable reward is not supported with large CLAP models "
                                "(7B+ too large for activation storage). "
                                "Falling back to non-differentiable mode.")
                args.differentiable_reward = False
        else:
            # Original VoiceCLAP-small path
            audio_decoder_clap, clap_model, clap_tokenizer = setup_clap_models(
                args.full_checkpoint, device, dtype)
            pos_text_emb = encode_clap_text(args.positive_text, clap_model, clap_tokenizer, device)
            neg_text_emb = encode_clap_text(args.negative_text, clap_model, clap_tokenizer, device)

        # 2. Centroid embeddings + Quality MLP
        if use_centroid or use_quality_mlp:
            real_centroid, synth_centroid, quality_mlp = setup_centroid_and_quality(
                args.classifiers_dir, device, dtype)
            if not use_centroid:
                real_centroid = None
                synth_centroid = None
            if not use_quality_mlp:
                quality_mlp = None

        # 3. WavLM-SV for speaker similarity
        if use_speaker_sim:
            wavlm_sv, wavlm_fe = setup_speaker_model(device, dtype)

        if is_main:
            logging.info(f"Auxiliary losses enabled:")
            logging.info(f"  CLAP model: {args.clap_model} (large={use_large_clap}, quantize={args.clap_quantize})")
            logging.info(f"  Loss 1 (Naturalness): CLAP text={True}, quality_mlp={use_quality_mlp}")
            logging.info(f"  Loss 2 (Centroid): {use_centroid}")
            logging.info(f"  Loss 3 (Speaker Sim): {use_speaker_sim}")
            logging.info(f"  Target ratio (naturalness/centroid): {target_ratio}")
            logging.info(f"  Speaker sim ratio: {speaker_sim_ratio}")
            logging.info(f"  Coefficient cap: {coeff_cap}")
            logging.info(f"  Aux sigma threshold: {aux_sigma_max} (skip aux when sigma >= this)")
            _rej = getattr(args, 'rejection_sampling', False)
            _rej_pct = getattr(args, 'rejection_percentile', 50.0)
            _diff = getattr(args, 'differentiable_reward', False)
            if _diff:
                logging.info(f"  DIFFERENTIABLE REWARD: enabled (ReFL-style backprop through decoder/CLAP/WavLM)")
                logging.info(f"    → Gradients flow: pred_tgt → decoder → waveform → CLAP/WavLM → loss")
                logging.info(f"    → Gradient checkpointing on aux models: {getattr(args, 'diff_reward_checkpoint', False)}")
            elif _rej:
                logging.info(f"  REJECTION SAMPLING: enabled, percentile={_rej_pct}%")
                logging.info(f"    → Only train on top {100-_rej_pct:.0f}% reward samples")
                logging.info(f"    → Flow loss boosted by {100.0/max(100.0-_rej_pct, 1.0):.1f}x to compensate")
            else:
                logging.info(f"  Rejection sampling: disabled (using reward-weighted aux losses)")
            logging.info(f"  Positive text: {args.positive_text[:80]}...")
            logging.info(f"  Negative text: {args.negative_text}")
            if getattr(args, 'keep_last_n', 0) > 0:
                logging.info(f"  Checkpoint management: rolling window of {args.keep_last_n} + keep better older ones")

    # Dataset
    mode_weights = {
        "voice_clone_fwd": 0.33,
        "voice_clone_rev": 0.33,
        "unconditional": 0.34,
    }
    expand_all = getattr(args, "expand_all_modes", False)
    dataset = DramaBoxFinetuneDataset(
        preprocessed_dir=args.preprocessed_dir,
        mode_weights=mode_weights,
        max_ref_tokens=args.max_ref_tokens,
        overfit_n=args.overfit,
        expand_all_modes=expand_all,
        index_file=args.index_file,
    )

    # Build DataLoader (with optional bucket-weighted sampling)
    bucket_sampler = None
    if args.bucket_weights and isinstance(args.bucket_weights, dict):
        bucket_sampler = BucketWeightedSampler(
            dataset, bucket_weights=args.bucket_weights, seed=args.seed,
        )
        dataloader = DataLoader(
            dataset, batch_size=args.batch_size, sampler=bucket_sampler,
            num_workers=2, pin_memory=True, drop_last=True, collate_fn=collate_fn,
        )
    else:
        dataloader = DataLoader(
            dataset, batch_size=args.batch_size, shuffle=True,
            num_workers=2, pin_memory=True, drop_last=True, collate_fn=collate_fn,
        )

    # Compute epochs -> steps
    steps_per_epoch = 0
    if bucket_sampler is not None and args.podcast_epochs > 0:
        # Bucket mode: 1 epoch = 1 pass through sampler (anchor bucket seen once)
        # podcast_epochs = how many times podcast is fully seen
        args.epochs = args.podcast_epochs
        items_per_gpu = math.ceil(len(bucket_sampler) / max(accelerator.num_processes, 1))
        steps_per_epoch = math.ceil(items_per_gpu / args.grad_accum)
        args.steps = steps_per_epoch * args.epochs
        if is_main:
            logging.info(f"Bucket-weighted: {args.epochs} podcast-epochs x {steps_per_epoch} steps/epoch "
                         f"= {args.steps} total optimizer steps")
            logging.info(f"  ({len(bucket_sampler)} items/epoch, {accelerator.num_processes} GPUs, "
                         f"grad_accum={args.grad_accum})")
    elif args.epochs > 0 and not args.overfit:
        items_per_gpu = math.ceil(len(dataset) / max(accelerator.num_processes, 1))
        forward_per_epoch = items_per_gpu  # batch_size=1
        steps_per_epoch = math.ceil(forward_per_epoch / args.grad_accum)
        args.steps = steps_per_epoch * args.epochs
        if is_main:
            logging.info(f"Epoch-based: {args.epochs} epochs x {steps_per_epoch} steps/epoch "
                         f"= {args.steps} total optimizer steps")
            logging.info(f"  ({len(dataset)} items, {accelerator.num_processes} GPUs, "
                         f"grad_accum={args.grad_accum})")

    # Optimizer & Scheduler
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.999), weight_decay=0.01,
    )

    from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR, ConstantLR
    warmup = LinearLR(optimizer, start_factor=0.01, end_factor=1.0, total_iters=args.warmup_steps)
    remaining = args.steps - args.warmup_steps
    if args.lr_scheduler == "cosine":
        hold_steps = max(remaining // 5, 0)
        decay_steps = max(remaining - hold_steps, 1)
        hold_sched = ConstantLR(optimizer, factor=1.0, total_iters=hold_steps)
        decay_sched = CosineAnnealingLR(optimizer, T_max=decay_steps, eta_min=1e-6)
        scheduler = SequentialLR(
            optimizer,
            [warmup, hold_sched, decay_sched],
            milestones=[args.warmup_steps, args.warmup_steps + hold_steps],
        )
    elif args.lr_scheduler == "linear":
        main_sched = LinearLR(optimizer, start_factor=1.0, end_factor=0.01,
                              total_iters=max(remaining, 1))
        scheduler = SequentialLR(optimizer, [warmup, main_sched], milestones=[args.warmup_steps])
    else:
        main_sched = ConstantLR(optimizer, factor=1.0, total_iters=max(remaining, 1))
        scheduler = SequentialLR(optimizer, [warmup, main_sched], milestones=[args.warmup_steps])

    model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

    # Resume from full training state
    resume_step = 0
    resume_epoch = 0
    if args.resume_dir and is_main:
        try:
            resume_step, resume_epoch, best_loss_r, best_step_r = load_training_state(
                args.resume_dir, optimizer, scheduler)
            logging.info(f"Resumed: step={resume_step}, epoch={resume_epoch}, "
                         f"best_loss={best_loss_r:.4f}, best_step={best_step_r}")
        except Exception as e:
            logging.error(f"Failed to resume: {e}")
            resume_step = 0
            resume_epoch = 0

    patchifier = AudioPatchifier(patch_size=1)

    # Timestep sampler
    if args.base_model == "distilled":
        timestep_sampler = DistilledTimestepSampler()
    else:
        timestep_sampler = ShiftedLogitNormalTimestepSampler()

    # Load silence frame for padding
    silence_frame = None
    sf_path = os.path.join(args.dramabox_dir, "assets", "silence_latent_frame.pt")
    if os.path.exists(sf_path):
        silence_frame = torch.load(sf_path, weights_only=True)
        if is_main:
            logging.info(f"Loaded silence latent from {sf_path}")

    # Metrics file
    metrics_path = os.path.join(args.output_dir, "metrics.jsonl")

    if is_main:
        logging.info(f"Training: {args.steps} steps, lr={args.lr}, scheduler={args.lr_scheduler}, "
                     f"batch={args.batch_size}, grad_accum={args.grad_accum}, "
                     f"world_size={accelerator.num_processes}, "
                     f"max_ref_tokens={args.max_ref_tokens}")
        logging.info(f"3-mode IC-LoRA: voice_clone_fwd/rev + unconditional")
        logging.info(f"Multi-auxiliary losses: {'ENABLED' if aux_enabled else 'DISABLED'}")
        if args.overfit:
            logging.info(f"OVERFIT MODE: training on {args.overfit} groups only")

    data_iter = iter(dataloader)
    step = 0
    accum_loss = 0.0
    accum_total_loss = 0.0
    # Per-aux accumulators
    accum_aux1 = 0.0
    accum_aux2 = 0.0
    accum_aux3 = 0.0
    accum_coeff1 = 0.0
    accum_coeff2 = 0.0
    accum_coeff3 = 0.0
    accum_clap_text_reward = 0.0
    accum_quality_prob = 0.0
    accum_naturalness_reward = 0.0
    accum_centroid_score = 0.0
    accum_speaker_sim = 0.0
    aux_log_count = 0
    speaker_sim_count = 0
    # Rejection sampling state
    rejection_enabled = getattr(args, 'rejection_sampling', False)
    rejection_percentile = getattr(args, 'rejection_percentile', 50.0)
    reward_buffer = deque(maxlen=512)  # rolling window for percentile
    rejection_count = 0  # number of rejected micro-batches in log window
    acceptance_count = 0  # number of accepted micro-batches in log window
    # Differentiable reward mode
    differentiable_reward = getattr(args, 'differentiable_reward', False)
    diff_checkpoint = getattr(args, 'diff_reward_checkpoint', False)
    vram_peak_mb = 0.0
    best_loss = float("inf")
    best_step = 0
    t0 = time.time()
    mode_counts = defaultdict(int)
    current_epoch = 0
    last_val_epoch = -1
    val_procs = []
    script_dir = os.path.dirname(os.path.abspath(__file__))
    # Smart checkpoint management
    keep_last_n = getattr(args, 'keep_last_n', 0)
    recent_checkpoints = deque(maxlen=keep_last_n) if keep_last_n > 0 else None
    checkpoint_rewards = {}  # {checkpoint_path: naturalness_reward}

    # Handle resume
    if resume_step > 0:
        step = resume_step
        current_epoch = resume_epoch
        if 'best_loss_r' in dir():
            best_loss = best_loss_r
            best_step = best_step_r
        if is_main:
            logging.info(f"Skipping to step {step}, epoch {current_epoch}")

    total_micro_steps = args.steps * args.grad_accum
    start_micro = step * args.grad_accum

    for micro_step in range(start_micro, total_micro_steps):
        try:
            batch = next(data_iter)
        except StopIteration:
            data_iter = iter(dataloader)
            batch = next(data_iter)

        is_opt_step = (micro_step + 1) % args.grad_accum == 0
        if is_opt_step:
            step += 1

        with accelerator.accumulate(model):
            tgt_latent = batch["tgt_latent"].to(dtype=dtype)  # [B, C, T, F]
            ref_latent = batch["ref_latent"].to(dtype=dtype)
            tgt_lengths = batch["tgt_lengths"].to(device=device)
            B = tgt_latent.shape[0]

            # Track mode distribution
            for m in batch["modes"]:
                mode_counts[m] += 1

            # Random silence padding (0-25 frames)
            max_pad = 25
            pad_frames = random.randint(0, max_pad)
            if pad_frames > 0:
                C, F_dim = tgt_latent.shape[1], tgt_latent.shape[3]
                if silence_frame is not None:
                    sf = silence_frame.to(dtype=dtype, device=device)
                    silence_pad = sf.unsqueeze(0).expand(B, -1, pad_frames, -1)
                else:
                    silence_pad = torch.zeros(B, C, pad_frames, F_dim, dtype=dtype, device=device)
                tgt_latent = torch.cat([silence_pad, tgt_latent], dim=2)

            # Cap reference tokens
            ref_T_frames = min(ref_latent.shape[2], args.max_ref_tokens)
            ref_latent = ref_latent[:, :, :ref_T_frames, :]
            tgt_T_frames = tgt_latent.shape[2]

            # Create target state
            tgt_shape = AudioLatentShape(
                batch=B, channels=tgt_latent.shape[1],
                frames=tgt_T_frames, mel_bins=tgt_latent.shape[3],
            )
            audio_tools = AudioLatentTools(patchifier=patchifier, target_shape=tgt_shape)
            state = audio_tools.create_initial_state(device=device, dtype=dtype, initial_latent=tgt_latent)
            tgt_T = audio_tools.target_shape.token_count()

            # Sample noise + sigma
            total_tokens = tgt_T + ref_T_frames
            sigma = timestep_sampler.sample(B, total_tokens, device=device)
            sigma_exp = sigma.view(-1, 1, 1)

            noise = torch.randn_like(state.latent)
            noisy_tgt = (1 - sigma_exp) * state.latent + sigma_exp * noise

            state = LatentState(
                latent=noisy_tgt,
                denoise_mask=state.denoise_mask,
                positions=state.positions,
                clean_latent=state.clean_latent,
                attention_mask=state.attention_mask,
            )

            # Append reference (skip if ref_T=0 for unconditional mode)
            if ref_T_frames > 0:
                ref_conditioning = AudioConditionByReferenceLatent(
                    latent=ref_latent, strength=1.0,
                )
                state = ref_conditioning.apply_to(latent_state=state, latent_tools=audio_tools)

            # Loss mask
            loss_mask = torch.zeros(B, tgt_T, device=device)
            for b_idx in range(B):
                real_len = min(tgt_lengths[b_idx].item() + pad_frames, tgt_T)
                loss_mask[b_idx, :real_len] = 1.0

            # Text context
            with torch.no_grad():
                audio_context = prepare_audio_context(
                    audio_connector, batch["audio_features"],
                    batch["attention_mask"], device, dtype)
                if args.text_dropout > 0 and random.random() < args.text_dropout:
                    audio_context = torch.zeros_like(audio_context)

            # Build modality
            audio_mod = modality_from_latent_state(
                state=state, context=audio_context, sigma=sigma, enabled=True,
            )

            # Forward pass
            perturbations = BatchedPerturbationConfig.empty(B)
            with torch.autocast(device_type="cuda", dtype=dtype):
                _, velocity_pred = model(video=None, audio=audio_mod, perturbations=perturbations)

            # Loss (IC-LoRA: only on target tokens)
            tgt_patchified = audio_tools.patchifier.patchify(tgt_latent)
            target_velocity = noise - tgt_patchified
            pred_tgt = velocity_pred[:, :tgt_T]
            per_token_mse = (pred_tgt - target_velocity).pow(2).mean(dim=-1)
            loss = per_token_mse.mul(loss_mask).div(loss_mask.mean().clamp(min=1e-6)).mean()

            # ── Multi-Auxiliary Losses ─────────────────────────────────
            aux1_val = 0.0
            aux2_val = 0.0
            aux3_val = 0.0
            coeff1_val = 0.0
            coeff2_val = 0.0
            coeff3_val = 0.0
            clap_text_reward_val = 0.0
            quality_prob_val = 0.0
            naturalness_reward_val = 0.0
            centroid_score_val = 0.0
            speaker_sim_val = None

            # Only compute aux losses at low sigma where x0 prediction is clean enough
            # to produce meaningful decoded audio for reward computation
            sigma_val = sigma.item()
            aux_active = aux_enabled and sigma_val < aux_sigma_max

            if aux_active and differentiable_reward and tgt_T_frames <= 250:
                # ════════════════════════════════════════════════════��══════
                # DIFFERENTIABLE REWARD MODE (ReFL-style)
                # Gradients flow: pred_tgt → decoder → waveform → CLAP/WavLM → loss
                # ═══════════════════════════════════════════════════════════
                import torchaudio
                from torch.utils.checkpoint import checkpoint as grad_checkpoint

                try:
                    torch.cuda.empty_cache()
                    # ── x0 prediction WITH gradient (pred_tgt has grad) ──
                    noisy_tgt_tokens = state.latent[:, :tgt_T].detach()
                    x0_pred_tokens = noisy_tgt_tokens - pred_tgt * sigma_exp

                    x0_latent = patchifier.unpatchify(
                        x0_pred_tokens,
                        AudioLatentShape(
                            batch=B, channels=tgt_latent.shape[1],
                            frames=tgt_T_frames, mel_bins=tgt_latent.shape[3],
                        ),
                    )

                    # ── Decode prediction → waveform (grad flows through decoder) ──
                    if diff_checkpoint:
                        # Gradient checkpointing: recompute decoder activations during backward
                        decoded = grad_checkpoint(
                            audio_decoder_clap, x0_latent[0:1].to(dtype=dtype),
                            use_reentrant=False)
                    else:
                        decoded = audio_decoder_clap(x0_latent[0:1].to(dtype=dtype))

                    pred_wav = decoded.waveform.squeeze(0).float()
                    audio_sr = decoded.sampling_rate
                    pred_mono = pred_wav.mean(0) if pred_wav.ndim > 1 else pred_wav
                    if audio_sr != 16000:
                        pred_mono = torchaudio.functional.resample(pred_mono, audio_sr, 16000)

                    # ── CLAP embedding WITH gradient (differentiable mel computation) ──
                    if diff_checkpoint:
                        audio_emb = grad_checkpoint(
                            encode_clap_waveform_differentiable,
                            pred_mono.unsqueeze(0).to(device),
                            clap_model,
                            use_reentrant=False)
                    else:
                        audio_emb = encode_clap_waveform_differentiable(
                            pred_mono.unsqueeze(0).to(device), clap_model)
                    audio_emb_norm = F.normalize(audio_emb, p=2, dim=-1)  # [1, 768]

                    # ── Differentiable Loss 1: CLAP Naturalness ──
                    clap_pos_sim = (audio_emb_norm @ pos_text_emb.T).squeeze()
                    clap_neg_sim = (audio_emb_norm @ neg_text_emb.T).squeeze()
                    naturalness_loss = -clap_pos_sim + clap_neg_sim

                    # Quality MLP (differentiable — small, no checkpointing needed)
                    quality_logit_t = quality_mlp(audio_emb_norm.float()).squeeze()
                    quality_loss = -quality_logit_t  # maximize P(real)

                    aux1_loss = naturalness_loss + 0.5 * quality_loss

                    # ── Differentiable Loss 2: Centroid Real/Fake ──
                    aux2_loss = torch.tensor(0.0, device=device)
                    if use_centroid and real_centroid is not None:
                        cos_real_t = (audio_emb_norm.float() @ real_centroid.T).squeeze()
                        cos_synth_t = (audio_emb_norm.float() @ synth_centroid.T).squeeze()
                        aux2_loss = -cos_real_t + cos_synth_t

                    # ── Differentiable Loss 3: Speaker Similarity ──
                    aux3_loss = torch.tensor(0.0, device=device)
                    if use_speaker_sim and wavlm_sv is not None and ref_T_frames > 0:
                        # Reference decoding (no grad needed for ref)
                        with torch.no_grad():
                            ref_decoded = audio_decoder_clap(ref_latent[0:1].to(dtype=dtype))
                            ref_wav = ref_decoded.waveform.squeeze(0).float()
                            ref_mono = ref_wav.mean(0) if ref_wav.ndim > 1 else ref_wav
                            if ref_decoded.sampling_rate != 16000:
                                ref_mono = torchaudio.functional.resample(
                                    ref_mono, ref_decoded.sampling_rate, 16000)
                            # WavLM normalization (replaces feature_extractor for differentiability)
                            ref_norm = (ref_mono - ref_mono.mean()) / (ref_mono.std() + 1e-7)
                            ref_spk = wavlm_sv(
                                ref_norm.unsqueeze(0).to(device=device, dtype=torch.float32)
                            ).embeddings
                            ref_spk = F.normalize(ref_spk.float(), p=2, dim=-1)

                        # Prediction speaker embedding WITH grad
                        pred_std = pred_mono.std().detach()  # detach std for stability
                        pred_norm = (pred_mono - pred_mono.mean().detach()) / (pred_std + 1e-7)
                        pred_input = pred_norm.unsqueeze(0).to(device=device, dtype=torch.float32)
                        if diff_checkpoint:
                            pred_spk_emb = grad_checkpoint(
                                wavlm_extract_embeddings,
                                wavlm_sv,
                                pred_input,
                                use_reentrant=False)
                        else:
                            pred_spk_emb = wavlm_sv(pred_input).embeddings
                        pred_spk = F.normalize(pred_spk_emb.float(), p=2, dim=-1)
                        speaker_sim_t = (pred_spk @ ref_spk.T).squeeze()
                        aux3_loss = -speaker_sim_t  # maximize similarity

                    # ── Adaptive coefficients (EMA-based) ──
                    ema_flow = ema_alpha * ema_flow + (1 - ema_alpha) * loss.item()
                    aux1_item = aux1_loss.item()
                    ema_aux1 = ema_alpha * ema_aux1 + (1 - ema_alpha) * abs(aux1_item)
                    coeff1_val = min(target_ratio * ema_flow / max(ema_aux1, 1e-8), coeff_cap)

                    coeff2_val = 0.0
                    if use_centroid and real_centroid is not None:
                        aux2_item = aux2_loss.item()
                        ema_aux2 = ema_alpha * ema_aux2 + (1 - ema_alpha) * abs(aux2_item)
                        coeff2_val = min(target_ratio * ema_flow / max(ema_aux2, 1e-8), coeff_cap)

                    coeff3_val = 0.0
                    if aux3_loss.requires_grad or aux3_loss.item() != 0.0:
                        aux3_item = aux3_loss.item()
                        ema_aux3 = ema_alpha * ema_aux3 + (1 - ema_alpha) * abs(aux3_item)
                        coeff3_val = min(speaker_sim_ratio * ema_flow / max(ema_aux3, 1e-8), coeff_cap)

                    total_loss = loss + coeff1_val * aux1_loss + coeff2_val * aux2_loss + coeff3_val * aux3_loss

                    if torch.isnan(total_loss) or torch.isinf(total_loss):
                        accelerator.backward(loss)
                        total_loss = loss
                    else:
                        accelerator.backward(total_loss)

                    # Log values
                    clap_text_reward_val = (clap_pos_sim - clap_neg_sim).item()
                    quality_prob_val = torch.sigmoid(quality_logit_t.detach()).item()
                    naturalness_reward_val = clap_text_reward_val
                    if use_centroid and real_centroid is not None:
                        centroid_score_val = (cos_real_t - cos_synth_t).item()
                    if aux3_loss.item() != 0.0:
                        speaker_sim_val = -aux3_loss.item()

                    aux1_val = aux1_loss.item()
                    aux2_val = aux2_loss.item()
                    aux3_val = aux3_loss.item()
                    accum_total_loss += total_loss.item()
                    accum_aux1 += aux1_val
                    accum_aux2 += aux2_val
                    accum_aux3 += aux3_val
                    accum_coeff1 += coeff1_val
                    accum_coeff2 += coeff2_val
                    accum_coeff3 += coeff3_val

                    # Track VRAM
                    cur_vram = torch.cuda.max_memory_allocated(device) / 1024**2
                    if cur_vram > vram_peak_mb:
                        vram_peak_mb = cur_vram

                except RuntimeError as e:
                    if "out of memory" in str(e).lower():
                        # OOM: log and fall back to flow-only for this step
                        torch.cuda.empty_cache()
                        if is_main:
                            logging.warning(f"OOM in differentiable reward! VRAM peak: "
                                          f"{torch.cuda.max_memory_allocated(device)/1024**2:.0f}MB. "
                                          f"Falling back to flow-only. Consider --diff-reward-checkpoint")
                        accelerator.backward(loss)
                    else:
                        raise

                # Track reward metrics
                accum_clap_text_reward += clap_text_reward_val
                accum_quality_prob += quality_prob_val
                accum_naturalness_reward += naturalness_reward_val
                accum_centroid_score += centroid_score_val
                aux_log_count += 1
                if speaker_sim_val is not None:
                    accum_speaker_sim += speaker_sim_val
                    speaker_sim_count += 1

            elif aux_active:
                # ══════════════════════���════════════════════════════════════
                # NON-DIFFERENTIABLE MODE (scalar rewards, no grad through aux models)
                # ═══════════════════════════════════════════════════════════
                import torchaudio

                with torch.no_grad():
                    # ── Recover x0 prediction ──
                    noisy_tgt_tokens = state.latent[:, :tgt_T]
                    x0_pred_tokens_ng = noisy_tgt_tokens - pred_tgt.detach() * sigma_exp

                    # Unpatchify to latent shape: [B, 8, T, 16]
                    x0_latent = patchifier.unpatchify(
                        x0_pred_tokens_ng,
                        AudioLatentShape(
                            batch=B, channels=tgt_latent.shape[1],
                            frames=tgt_T_frames, mel_bins=tgt_latent.shape[3],
                        ),
                    )

                    # ── Decode prediction → waveform (first sample only) ──
                    decoded = audio_decoder_clap(x0_latent[0:1].to(dtype=dtype))
                    pred_wav = decoded.waveform.squeeze(0).float()
                    audio_sr = decoded.sampling_rate
                    pred_mono = pred_wav.mean(0) if pred_wav.ndim > 1 else pred_wav
                    if audio_sr != 16000:
                        pred_mono = torchaudio.functional.resample(pred_mono, audio_sr, 16000)

                    # ── CLAP embedding for naturalness ──
                    if use_large_clap and st_model is not None:
                        import numpy as np
                        wav_np = pred_mono.cpu().numpy()
                        audio_emb = encode_audio_st(
                            wav_np, 16000, st_model,
                            rank=accelerator.process_index)
                    else:
                        audio_emb = clap_model.encode_waveform(pred_mono.unsqueeze(0).to(device))
                        audio_emb = F.normalize(audio_emb, p=2, dim=-1)  # [1, 768]

                    # ── CLAP-small embedding for quality_mlp/centroid (768-dim) ──
                    # When using large CLAP, the quality MLP and centroids are trained
                    # on 768-dim CLAP-small embeddings — compute a separate embedding
                    audio_emb_small = None
                    if use_large_clap and clap_model is not None and (use_quality_mlp or use_centroid):
                        audio_emb_small = clap_model.encode_waveform(pred_mono.unsqueeze(0).to(device))
                        audio_emb_small = F.normalize(audio_emb_small, p=2, dim=-1)  # [1, 768]

                    # ── Reward 1: CLAP Naturalness (uses large or small CLAP) ──
                    clap_pos_sim = (audio_emb @ pos_text_emb.T).item()
                    clap_neg_sim = (audio_emb @ neg_text_emb.T).item()
                    clap_text_reward_val = clap_pos_sim - clap_neg_sim  # ~[-0.5, +0.5]

                    if use_quality_mlp and quality_mlp is not None:
                        # Quality MLP uses CLAP-small embeddings (768-dim)
                        emb_for_mlp = audio_emb_small if audio_emb_small is not None else audio_emb
                        quality_logit = quality_mlp(emb_for_mlp.float()).item()
                        quality_prob_val = torch.sigmoid(torch.tensor(quality_logit)).item()
                        naturalness_reward_val = (0.5 * clap_text_reward_val +
                                                  0.5 * (2 * quality_prob_val - 1))
                    else:
                        quality_prob_val = 0.5
                        naturalness_reward_val = clap_text_reward_val

                    # ── Reward 2: Centroid Real/Fake (uses CLAP-small embeddings) ──
                    if use_centroid and real_centroid is not None:
                        emb_for_cent = audio_emb_small if audio_emb_small is not None else audio_emb
                        cos_real = (emb_for_cent @ real_centroid.T).item()
                        cos_synth = (emb_for_cent @ synth_centroid.T).item()
                        centroid_score_val = cos_real - cos_synth  # ~[-0.8, +0.9]

                    # ── Reward 3: Speaker Similarity (only with reference) ──
                    if use_speaker_sim and wavlm_sv is not None and ref_T_frames > 0:
                        # Decode reference → waveform
                        ref_decoded = audio_decoder_clap(ref_latent[0:1].to(dtype=dtype))
                        ref_wav = ref_decoded.waveform.squeeze(0).float()
                        ref_mono = ref_wav.mean(0) if ref_wav.ndim > 1 else ref_wav
                        if ref_decoded.sampling_rate != 16000:
                            ref_mono = torchaudio.functional.resample(
                                ref_mono, ref_decoded.sampling_rate, 16000)

                        # Extract speaker embeddings via WavLM-SV
                        # Process through feature extractor for proper normalization
                        pred_inputs = wavlm_fe(
                            pred_mono.cpu().numpy(), sampling_rate=16000,
                            return_tensors="pt", padding=True)
                        ref_inputs = wavlm_fe(
                            ref_mono.cpu().numpy(), sampling_rate=16000,
                            return_tensors="pt", padding=True)

                        pred_spk = wavlm_sv(
                            pred_inputs.input_values.to(device=device, dtype=torch.float32)
                        ).embeddings
                        ref_spk = wavlm_sv(
                            ref_inputs.input_values.to(device=device, dtype=torch.float32)
                        ).embeddings

                        pred_spk = F.normalize(pred_spk.float(), p=2, dim=-1)
                        ref_spk = F.normalize(ref_spk.float(), p=2, dim=-1)
                        speaker_sim_val = (pred_spk @ ref_spk.T).item()  # [-1, +1]

                # ── Compute composite reward for rejection sampling ──
                composite_reward = naturalness_reward_val
                n_rewards = 1
                if use_centroid and real_centroid is not None:
                    composite_reward += centroid_score_val
                    n_rewards += 1
                if speaker_sim_val is not None:
                    composite_reward += speaker_sim_val
                    n_rewards += 1
                composite_reward /= n_rewards  # normalize to ~[-0.5, +1.0]

                # ── Rejection sampling mode ──
                if rejection_enabled:
                    reward_buffer.append(composite_reward)

                    # Need at least 32 samples before we can reject meaningfully
                    if len(reward_buffer) >= 32:
                        sorted_buf = sorted(reward_buffer)
                        idx = int(len(sorted_buf) * rejection_percentile / 100.0)
                        idx = min(idx, len(sorted_buf) - 1)
                        threshold = sorted_buf[idx]
                    else:
                        threshold = -999.0  # accept everything during warmup

                    if composite_reward >= threshold:
                        # ACCEPTED: train on this sample (flow loss with 2x boost
                        # to compensate for ~50% rejection rate)
                        boost = 100.0 / max(100.0 - rejection_percentile, 1.0)
                        accelerator.backward(loss * boost)
                        acceptance_count += 1
                    else:
                        # REJECTED: zero gradient for this micro-batch
                        accelerator.backward(loss * 0.0)
                        rejection_count += 1

                    total_loss = loss  # for logging purposes
                    accum_total_loss += loss.item()

                else:
                    # ── Original aux loss mode (reward-weighted reconstruction) ──
                    x0_clean = patchifier.patchify(tgt_latent)
                    x0_pred_grad = state.latent[:, :tgt_T].detach() - pred_tgt * sigma_exp
                    x0_recon_loss = ((x0_pred_grad - x0_clean.detach()).pow(2).mean(dim=-1)
                                     * loss_mask).div(loss_mask.mean().clamp(min=1e-6)).mean()

                    # Loss 1: Naturalness
                    w1 = torch.clamp(
                        torch.tensor(0.5 - naturalness_reward_val, device=device),
                        min=0.05, max=2.0)
                    aux1 = w1 * x0_recon_loss

                    # Loss 2: Centroid
                    if use_centroid and real_centroid is not None:
                        w2 = torch.clamp(
                            torch.tensor(0.5 - centroid_score_val, device=device),
                            min=0.05, max=2.0)
                        aux2 = w2 * x0_recon_loss
                    else:
                        aux2 = torch.tensor(0.0, device=device)

                    # Loss 3: Speaker sim
                    if speaker_sim_val is not None:
                        w3 = torch.clamp(
                            torch.tensor(0.5 - speaker_sim_val, device=device),
                            min=0.05, max=2.0)
                        aux3 = w3 * x0_recon_loss
                    else:
                        aux3 = torch.tensor(0.0, device=device)

                    # Adaptive coefficients (EMA-based)
                    ema_flow = ema_alpha * ema_flow + (1 - ema_alpha) * loss.item()
                    ema_aux1 = ema_alpha * ema_aux1 + (1 - ema_alpha) * aux1.item()
                    coeff1_val = min(target_ratio * ema_flow / max(ema_aux1, 1e-8), coeff_cap)

                    coeff2_val = 0.0
                    if use_centroid and real_centroid is not None:
                        ema_aux2 = ema_alpha * ema_aux2 + (1 - ema_alpha) * aux2.item()
                        coeff2_val = min(target_ratio * ema_flow / max(ema_aux2, 1e-8), coeff_cap)

                    coeff3_val = 0.0
                    if speaker_sim_val is not None:
                        ema_aux3 = ema_alpha * ema_aux3 + (1 - ema_alpha) * aux3.item()
                        coeff3_val = min(speaker_sim_ratio * ema_flow / max(ema_aux3, 1e-8), coeff_cap)

                    total_loss = loss + coeff1_val * aux1 + coeff2_val * aux2 + coeff3_val * aux3

                    if torch.isnan(total_loss) or torch.isinf(total_loss):
                        accelerator.backward(loss)
                        total_loss = loss
                    else:
                        accelerator.backward(total_loss)

                    aux1_val = aux1.item() if not math.isnan(aux1.item()) else 0.0
                    aux2_val = aux2.item() if not math.isnan(aux2.item()) else 0.0
                    aux3_val = aux3.item() if not math.isnan(aux3.item()) else 0.0
                    accum_total_loss += total_loss.item() if not math.isnan(total_loss.item()) else loss.item()
                    accum_aux1 += aux1_val
                    accum_aux2 += aux2_val
                    accum_aux3 += aux3_val
                    accum_coeff1 += coeff1_val
                    accum_coeff2 += coeff2_val
                    accum_coeff3 += coeff3_val

                # Track reward metrics for logging (both modes)
                accum_clap_text_reward += clap_text_reward_val
                accum_quality_prob += quality_prob_val
                accum_naturalness_reward += naturalness_reward_val
                accum_centroid_score += centroid_score_val
                aux_log_count += 1
                if speaker_sim_val is not None:
                    accum_speaker_sim += speaker_sim_val
                    speaker_sim_count += 1
            else:
                # aux_enabled=False OR sigma >= aux_sigma_max: flow-only backward
                accelerator.backward(loss)

            if accelerator.sync_gradients and args.max_grad_norm > 0:
                accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)

            optimizer.step()
            optimizer.zero_grad()
            if accelerator.sync_gradients:
                scheduler.step()

        accum_loss += loss.item()

        # Logging
        if is_opt_step and step % args.log_every == 0 and is_main:
            avg_loss = accum_loss / (args.log_every * args.grad_accum)
            lr = optimizer.param_groups[0]["lr"]
            elapsed = time.time() - t0
            sps = step / elapsed if elapsed > 0 else 0
            eta = (args.steps - step) / sps if sps > 0 else 0

            total_modes = sum(mode_counts.values()) or 1
            mode_pcts = {k: f"{100*v/total_modes:.0f}%" for k, v in mode_counts.items()}

            log_msg = (
                f"Step {step}/{args.steps} | loss={avg_loss:.4f} | lr={lr:.2e} | "
                f"tgt_T={tgt_T} ref_T={ref_T_frames} | "
                f"{sps:.2f} steps/s | ETA {eta/60:.0f}min | modes={mode_pcts}"
            )

            # Build metrics dict
            metric = {
                "step": step,
                "flow_loss": round(avg_loss, 6),
                "lr": lr,
                "tgt_tokens": tgt_T,
                "ref_tokens": ref_T_frames,
                "steps_per_sec": round(sps, 3),
                "elapsed_sec": round(elapsed, 1),
                "eta_sec": round(eta, 1),
                "mode_counts": dict(mode_counts),
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            }

            # Add multi-aux metrics
            # aux_log_count = number of micro-batches where sigma < aux_sigma_max
            total_micro_batches = args.log_every * args.grad_accum
            if aux_enabled and aux_log_count > 0:
                avg_total = accum_total_loss / aux_log_count
                avg_aux1 = accum_aux1 / aux_log_count
                avg_aux2 = accum_aux2 / aux_log_count
                avg_aux3 = accum_aux3 / aux_log_count
                avg_coeff1 = accum_coeff1 / aux_log_count
                avg_coeff2 = accum_coeff2 / aux_log_count
                avg_coeff3 = accum_coeff3 / aux_log_count
                avg_clap_text = accum_clap_text_reward / aux_log_count
                avg_quality = accum_quality_prob / aux_log_count
                avg_naturalness = accum_naturalness_reward / aux_log_count
                avg_centroid = accum_centroid_score / aux_log_count
                avg_speaker = (accum_speaker_sim / speaker_sim_count
                               if speaker_sim_count > 0 else None)

                log_msg += (
                    f" | nat={avg_naturalness:.3f} cent={avg_centroid:.3f}"
                )
                if avg_speaker is not None:
                    log_msg += f" spk={avg_speaker:.3f}"
                if rejection_enabled:
                    total_scored = acceptance_count + rejection_count
                    accept_pct = 100 * acceptance_count / max(total_scored, 1)
                    log_msg += (
                        f" | reject={rejection_count}/{total_scored}"
                        f" ({accept_pct:.0f}% accepted)"
                        f" | aux_hit={aux_log_count}/{total_micro_batches}"
                    )
                elif differentiable_reward:
                    log_msg += (
                        f" | c1={avg_coeff1:.2f} c2={avg_coeff2:.2f} c3={avg_coeff3:.2f}"
                        f" | total={avg_total:.4f}"
                        f" | aux_hit={aux_log_count}/{total_micro_batches}"
                        f" | VRAM={vram_peak_mb/1024:.1f}GB"
                    )
                else:
                    log_msg += (
                        f" | c1={avg_coeff1:.2f} c2={avg_coeff2:.2f} c3={avg_coeff3:.2f}"
                        f" | total={avg_total:.4f}"
                        f" | aux_hit={aux_log_count}/{total_micro_batches}"
                    )

                metric.update({
                    "clap_text_reward": round(avg_clap_text, 4),
                    "quality_prob": round(avg_quality, 4),
                    "naturalness_reward": round(avg_naturalness, 4),
                    "centroid_score": round(avg_centroid, 4),
                    "aux1_loss": round(avg_aux1, 6),
                    "aux2_loss": round(avg_aux2, 6),
                    "aux3_loss": round(avg_aux3, 6),
                    "coeff1": round(avg_coeff1, 4),
                    "coeff2": round(avg_coeff2, 4),
                    "coeff3": round(avg_coeff3, 4),
                    "total_loss": round(avg_total, 6),
                })
                if avg_speaker is not None:
                    metric["speaker_sim"] = round(avg_speaker, 4)
                if differentiable_reward and vram_peak_mb > 0:
                    metric["vram_peak_mb"] = round(vram_peak_mb, 0)

                # Reset accumulators
                accum_total_loss = 0.0
                accum_aux1 = 0.0
                accum_aux2 = 0.0
                accum_aux3 = 0.0
                accum_coeff1 = 0.0
                accum_coeff2 = 0.0
                accum_coeff3 = 0.0
                accum_clap_text_reward = 0.0
                accum_quality_prob = 0.0
                accum_naturalness_reward = 0.0
                accum_centroid_score = 0.0
                accum_speaker_sim = 0.0
                aux_log_count = 0
                speaker_sim_count = 0

            logging.info(log_msg)

            with open(metrics_path, "a") as f:
                f.write(json.dumps(metric) + "\n")

            # Update status file
            status = {
                "step": step,
                "total_steps": args.steps,
                "epoch": current_epoch,
                "total_epochs": args.epochs if args.epochs > 0 else 0,
                "steps_per_epoch": steps_per_epoch,
                "flow_loss": round(avg_loss, 6),
                "best_loss": round(best_loss, 6),
                "best_step": best_step,
                "lr": lr,
                "elapsed_sec": round(elapsed, 1),
                "eta_sec": round(eta, 1),
                "steps_per_sec": round(sps, 3),
                "world_size": accelerator.num_processes,
                "mode_counts": dict(mode_counts),
                "aux_enabled": aux_enabled,
                "aux_losses": [l for l, e in [
                    ("naturalness", True), ("quality_mlp", use_quality_mlp),
                    ("centroid", use_centroid), ("speaker_sim", use_speaker_sim),
                ] if e],
            }
            with open(os.path.join(args.output_dir, "status.json"), "w") as f:
                json.dump(status, f, indent=2)

            # Best checkpoint
            if avg_loss < best_loss:
                best_loss = avg_loss
                old_best = os.path.join(args.output_dir, f"best_step_{best_step:05d}.safetensors")
                best_step = step
                new_best = os.path.join(args.output_dir, f"best_step_{best_step:05d}.safetensors")
                unwrapped = _unwrap_model_safe(model)
                unwrapped.save_pretrained(args.output_dir)
                adapter = os.path.join(args.output_dir, "adapter_model.safetensors")
                if os.path.exists(adapter):
                    shutil.copy(adapter, new_best)
                if old_best != new_best and os.path.exists(old_best):
                    os.remove(old_best)
                logging.info(f"  New best: loss={best_loss:.4f} at step {best_step}")

            accum_loss = 0.0
            rejection_count = 0
            acceptance_count = 0

        # Periodic save (with optimizer state for resumability + smart checkpoint management)
        if is_opt_step and step % args.save_every == 0 and is_main:
            save_path = os.path.join(args.output_dir, f"lora_step_{step:05d}.safetensors")
            logging.info(f"Saving: {save_path}")
            unwrapped = _unwrap_model_safe(model)
            unwrapped.save_pretrained(args.output_dir)
            adapter = os.path.join(args.output_dir, "adapter_model.safetensors")
            if os.path.exists(adapter):
                shutil.copy(adapter, save_path)
            if not getattr(args, 'no_save_state', False):
                opt_state = {
                    "step": step,
                    "epoch": current_epoch,
                    "best_loss": best_loss,
                    "best_step": best_step,
                    "optimizer": optimizer.state_dict(),
                    "scheduler": scheduler.state_dict(),
                    "rng_python": random.getstate(),
                    "rng_torch": torch.random.get_rng_state(),
                    "rng_cuda": torch.cuda.get_rng_state(),
                }
                state_path = os.path.join(args.output_dir, f"training_state_step_{step:05d}.pt")
                torch.save(opt_state, state_path)
                latest_path = os.path.join(args.output_dir, "training_state.pt")
                if os.path.islink(latest_path) or os.path.exists(latest_path):
                    os.remove(latest_path)
                os.symlink(os.path.basename(state_path), latest_path)

            # Smart checkpoint management: rolling window + keep better older ones
            if recent_checkpoints is not None:
                # Record reward for this checkpoint (use latest naturalness_reward)
                current_reward = naturalness_reward_val
                checkpoint_rewards[save_path] = current_reward

                # Check if an older checkpoint is being evicted from the deque
                evicted_path = None
                if len(recent_checkpoints) == recent_checkpoints.maxlen:
                    evicted_path = recent_checkpoints[0]  # will be pushed out

                recent_checkpoints.append(save_path)

                if evicted_path and evicted_path in checkpoint_rewards:
                    evicted_reward = checkpoint_rewards[evicted_path]
                    # Compare to min reward in current deque
                    deque_rewards = [checkpoint_rewards.get(p, float('-inf'))
                                     for p in recent_checkpoints]
                    min_deque_reward = min(deque_rewards) if deque_rewards else float('-inf')

                    if evicted_reward < min_deque_reward:
                        # Worse than all current → delete
                        if os.path.exists(evicted_path):
                            os.remove(evicted_path)
                            logging.info(f"  Deleted checkpoint: {os.path.basename(evicted_path)} "
                                         f"(reward={evicted_reward:.4f} < min_deque={min_deque_reward:.4f})")
                        # Delete corresponding training state
                        evicted_step = evicted_path.split("step_")[1].split(".")[0]
                        evicted_state = os.path.join(
                            args.output_dir, f"training_state_step_{evicted_step}.pt")
                        if os.path.exists(evicted_state):
                            os.remove(evicted_state)
                        del checkpoint_rewards[evicted_path]
                    else:
                        logging.info(f"  Keeping older checkpoint: {os.path.basename(evicted_path)} "
                                     f"(reward={evicted_reward:.4f} >= min_deque={min_deque_reward:.4f})")

        # Epoch boundary detection
        if is_opt_step and steps_per_epoch > 0 and is_main:
            new_epoch = step // steps_per_epoch
            if new_epoch > current_epoch:
                current_epoch = new_epoch
                logging.info(f"{'='*60}")
                logging.info(f"EPOCH {current_epoch}/{args.epochs} COMPLETE (step {step})")
                logging.info(f"{'='*60}")
                # Save full state for this epoch
                lora_path = save_training_state(
                    args.output_dir, step, current_epoch, optimizer,
                    scheduler, best_loss, best_step, model, accelerator)
                # NOTE: Validation disabled at epoch boundary to prevent OOM —
                # the validation subprocess uses ~14GB on GPU 0/4, which combined
                # with differentiable reward (~66GB) exceeds 80GB and crashes.
                # Validation runs only after training completes (see below).
                # if args.val_samples > 0:
                #     proc = launch_validation_async(
                #         script_dir, args.output_dir, lora_path,
                #         current_epoch, args)
                #     if proc:
                #         val_procs.append(proc)

    # Final save
    if is_main:
        final_epoch = args.epochs if args.epochs > 0 else 0
        lora_path = save_training_state(
            args.output_dir, step, final_epoch, optimizer,
            scheduler, best_loss, best_step, model, accelerator, tag="final")
        logging.info(f"Training complete! {step} steps in {time.time()-t0:.0f}s")
        logging.info(f"Best loss: {best_loss:.4f} at step {best_step}")
        # Launch final validation
        if args.val_samples > 0 and final_epoch > 0:
            proc = launch_validation_async(
                script_dir, args.output_dir, lora_path,
                final_epoch, args)
            if proc:
                val_procs.append(proc)
        # Wait for all validation processes
        for proc in val_procs:
            proc.wait()
            logging.info(f"Validation PID {proc.pid} finished (rc={proc.returncode})")


if __name__ == "__main__":
    main()