train_engine.py

"""
Standalone ACE-Step CPU LoRA Training Engine.

Ported from Side-Step (koda-dernet/Side-Step) into a single self-contained
module. No external Side-Step dependency required.

Exports:
    preprocess_audio()       - 2-pass sequential preprocessing
    train_lora_generator()   - Generator-based LoRA training loop
    cancel_training()        - Set the cancel flag
    get_trained_loras()      - List saved adapters
"""

from __future__ import annotations

import gc
import json
import logging
import math
import os
import random
import re
import shutil
import sys
import tempfile
import time
import types
import unicodedata
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, Generator, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import (
    CosineAnnealingLR,
    LinearLR,
    SequentialLR,
)
from torch.utils.data import DataLoader, Dataset

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configurable caps (edit these at the top of the file)
# ---------------------------------------------------------------------------

MAX_AUDIO_DURATION = 240.0    # seconds, cap per audio file
MAX_TRAINING_TIME = 28800     # 8 hours hard timeout
TARGET_SR = 48000
AUDIO_EXTENSIONS = frozenset({".wav", ".mp3", ".flac", ".ogg", ".opus", ".m4a", ".aac"})

# bfloat16 deadlocks on CPU (known PyTorch bug) -- force float32
CPU_DTYPE = torch.float32

import threading
_training_cancel = threading.Event()


def cancel_training() -> None:
    _training_cancel.set()


# ============================================================================
# CONFIGS
# ============================================================================

@dataclass
class LoRAConfig:
    r: int = 64
    alpha: int = 128
    dropout: float = 0.1
    target_modules: List[str] = field(default_factory=lambda: [
        "q_proj", "k_proj", "v_proj", "o_proj",
    ])
    bias: str = "none"
    attention_type: str = "both"
    target_mlp: bool = True


# ============================================================================
# TIMESTEP SAMPLING & CFG DROPOUT
# ============================================================================

def sample_timesteps(
    batch_size: int,
    device: torch.device,
    dtype: torch.dtype,
    timestep_mu: float = -0.4,
    timestep_sigma: float = 1.0,
) -> Tuple[torch.Tensor, torch.Tensor]:
    t = torch.sigmoid(
        torch.randn((batch_size,), device=device, dtype=dtype) * timestep_sigma + timestep_mu
    )
    r = torch.sigmoid(
        torch.randn((batch_size,), device=device, dtype=dtype) * timestep_sigma + timestep_mu
    )
    t, r = torch.maximum(t, r), torch.minimum(t, r)
    # use_meanflow=False forces r=t (ACE-Step convention)
    return t, t


def apply_cfg_dropout(
    encoder_hidden_states: torch.Tensor,
    null_condition_emb: torch.Tensor,
    cfg_ratio: float = 0.15,
) -> torch.Tensor:
    bsz = encoder_hidden_states.shape[0]
    device = encoder_hidden_states.device
    dtype = encoder_hidden_states.dtype
    mask = torch.where(
        torch.rand(size=(bsz,), device=device, dtype=dtype) < cfg_ratio,
        torch.zeros(size=(bsz,), device=device, dtype=dtype),
        torch.ones(size=(bsz,), device=device, dtype=dtype),
    ).view(-1, 1, 1)
    return torch.where(
        mask > 0,
        encoder_hidden_states,
        null_condition_emb.expand_as(encoder_hidden_states),
    )


# ============================================================================
# OPTIMIZER (Adafactor preferred for CPU -- 1.5 bytes/param)
# ============================================================================

def build_optimizer(
    params, lr: float = 1e-4, weight_decay: float = 0.01,
) -> torch.optim.Optimizer:
    try:
        from transformers.optimization import Adafactor
        logger.info("Using Adafactor optimizer (minimal state memory)")
        return Adafactor(
            params, lr=lr, weight_decay=weight_decay,
            scale_parameter=False, relative_step=False,
        )
    except ImportError:
        logger.warning("transformers not installed, falling back to AdamW")
        return AdamW(params, lr=lr, weight_decay=weight_decay)


def build_scheduler(
    optimizer, total_steps: int, warmup_steps: int, lr: float,
):
    _max_warmup = max(1, total_steps // 10)
    if warmup_steps > _max_warmup:
        warmup_steps = _max_warmup

    warmup = LinearLR(optimizer, start_factor=0.1, end_factor=1.0, total_iters=warmup_steps)
    remaining = max(1, total_steps - warmup_steps)
    main = CosineAnnealingLR(optimizer, T_max=remaining, eta_min=lr * 0.01)
    return SequentialLR(optimizer, [warmup, main], milestones=[warmup_steps])


# ============================================================================
# DATASET
# ============================================================================

def _collate_batch(batch: List[Dict]) -> Dict[str, torch.Tensor]:
    max_t = max(s["target_latents"].shape[0] for s in batch)
    max_e = max(s["encoder_hidden_states"].shape[0] for s in batch)

    def pad(t, max_len, dim=0):
        diff = max_len - t.shape[dim]
        if diff <= 0:
            return t
        shape = list(t.shape)
        shape[dim] = diff
        return torch.cat([t, t.new_zeros(*shape)], dim=dim)

    return {
        "target_latents": torch.stack([pad(s["target_latents"], max_t) for s in batch]),
        "attention_mask": torch.stack([pad(s["attention_mask"], max_t) for s in batch]),
        "encoder_hidden_states": torch.stack([pad(s["encoder_hidden_states"], max_e) for s in batch]),
        "encoder_attention_mask": torch.stack([pad(s["encoder_attention_mask"], max_e) for s in batch]),
        "context_latents": torch.stack([pad(s["context_latents"], max_t) for s in batch]),
    }


class TensorDataset(Dataset):
    _REQUIRED = frozenset([
        "target_latents", "attention_mask", "encoder_hidden_states",
        "encoder_attention_mask", "context_latents",
    ])

    def __init__(self, tensor_dir: str):
        self.paths: List[str] = []
        for f in sorted(os.listdir(tensor_dir)):
            if f.endswith(".pt") and not f.endswith(".tmp.pt") and f != "manifest.json":
                self.paths.append(str(Path(tensor_dir) / f))

    def __len__(self) -> int:
        return len(self.paths)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        data = torch.load(self.paths[idx], map_location="cpu", weights_only=True)
        missing = self._REQUIRED - data.keys()
        if missing:
            raise KeyError(f"Missing keys {sorted(missing)} in {self.paths[idx]}")
        for k in ("target_latents", "encoder_hidden_states", "context_latents"):
            t = data[k]
            if torch.isnan(t).any() or torch.isinf(t).any():
                t.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0)
        return {k: data[k] for k in self._REQUIRED}


# ============================================================================
# GRADIENT CHECKPOINTING
# ============================================================================

def _find_decoder_layers(decoder: nn.Module) -> Optional[nn.ModuleList]:
    for attr in ("layers", "blocks", "transformer_blocks"):
        c = getattr(decoder, attr, None)
        if isinstance(c, nn.ModuleList) and len(c) > 0:
            return c
    for child in decoder.children():
        for attr in ("layers", "blocks", "transformer_blocks"):
            c = getattr(child, attr, None)
            if isinstance(c, nn.ModuleList) and len(c) > 0:
                return c
    return None


def enable_gradient_checkpointing(decoder: nn.Module) -> bool:
    """Enable gradient checkpointing on the decoder to save memory."""
    enabled = False

    # Walk wrapper chain
    stack = [decoder]
    visited = set()
    while stack:
        mod = stack.pop()
        if not isinstance(mod, nn.Module):
            continue
        mid = id(mod)
        if mid in visited:
            continue
        visited.add(mid)

        if hasattr(mod, "gradient_checkpointing_enable"):
            try:
                mod.gradient_checkpointing_enable()
                enabled = True
            except Exception:
                pass
        elif hasattr(mod, "gradient_checkpointing"):
            try:
                mod.gradient_checkpointing = True
                enabled = True
            except Exception:
                pass

        if hasattr(mod, "enable_input_require_grads"):
            try:
                mod.enable_input_require_grads()
            except Exception:
                pass

        cfg = getattr(mod, "config", None)
        if cfg is not None and hasattr(cfg, "use_cache"):
            try:
                cfg.use_cache = False
            except Exception:
                pass

        for a in ("_forward_module", "_orig_mod", "base_model", "model", "module"):
            child = getattr(mod, a, None)
            if isinstance(child, nn.Module):
                stack.append(child)

    return enabled


def force_disable_cache(decoder: nn.Module) -> None:
    stack = [decoder]
    visited = set()
    while stack:
        mod = stack.pop()
        if not isinstance(mod, nn.Module):
            continue
        mid = id(mod)
        if mid in visited:
            continue
        visited.add(mid)
        cfg = getattr(mod, "config", None)
        if cfg is not None and hasattr(cfg, "use_cache"):
            try:
                cfg.use_cache = False
            except Exception:
                pass
        for a in ("_forward_module", "_orig_mod", "base_model", "model", "module"):
            child = getattr(mod, a, None)
            if isinstance(child, nn.Module):
                stack.append(child)


# ============================================================================
# LORA INJECTION (PEFT only -- no DoRA/LoKR/LoHA/OFT)
# ============================================================================

def _unwrap_decoder(model):
    decoder = model.decoder if hasattr(model, "decoder") else model
    while hasattr(decoder, "_forward_module"):
        decoder = decoder._forward_module
    if hasattr(decoder, "base_model"):
        bm = decoder.base_model
        decoder = bm.model if hasattr(bm, "model") else bm
    if hasattr(decoder, "model") and isinstance(decoder.model, nn.Module):
        decoder = decoder.model
    return decoder


def inject_lora(model, lora_cfg: LoRAConfig) -> Tuple[Any, Dict[str, Any]]:
    from peft import get_peft_model, LoraConfig as PeftLoraConfig, TaskType

    decoder = _unwrap_decoder(model)
    model.decoder = decoder

    # Guard enable_input_require_grads for DiT (no get_input_embeddings)
    if hasattr(decoder, "enable_input_require_grads"):
        orig = decoder.enable_input_require_grads

        def _safe(self):
            try:
                return orig()
            except NotImplementedError:
                return None

        decoder.enable_input_require_grads = types.MethodType(_safe, decoder)

    if hasattr(decoder, "is_gradient_checkpointing"):
        try:
            decoder.is_gradient_checkpointing = False
        except Exception:
            pass

    peft_cfg = PeftLoraConfig(
        r=lora_cfg.r,
        lora_alpha=lora_cfg.alpha,
        lora_dropout=lora_cfg.dropout,
        target_modules=lora_cfg.target_modules,
        bias=lora_cfg.bias,
        task_type=TaskType.FEATURE_EXTRACTION,
    )

    model.decoder = get_peft_model(decoder, peft_cfg)

    for name, param in model.named_parameters():
        if "lora_" not in name:
            param.requires_grad = False

    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)

    return model, {
        "total_params": total,
        "trainable_params": trainable,
        "trainable_ratio": trainable / total if total > 0 else 0,
    }


def save_lora_adapter(model, output_dir: str) -> None:
    os.makedirs(output_dir, exist_ok=True)
    decoder = model.decoder if hasattr(model, "decoder") else model
    while hasattr(decoder, "_forward_module"):
        decoder = decoder._forward_module

    if hasattr(decoder, "save_pretrained"):
        decoder.save_pretrained(output_dir)
        # Scrub base_model path for portability
        cfg_path = os.path.join(output_dir, "adapter_config.json")
        if os.path.isfile(cfg_path):
            try:
                with open(cfg_path, "r") as f:
                    cfg = json.load(f)
                if cfg.get("base_model_name_or_path"):
                    cfg["base_model_name_or_path"] = ""
                    with open(cfg_path, "w") as f:
                        json.dump(cfg, f, indent=2)
            except Exception:
                pass
        logger.info("LoRA adapter saved to %s", output_dir)
    else:
        # Fallback: manual extraction
        state = {}
        for name, param in decoder.named_parameters():
            if "lora_" in name:
                state[name] = param.data.clone()
        if state:
            try:
                from safetensors.torch import save_file
                save_file(state, str(Path(output_dir) / "adapter_model.safetensors"))
            except ImportError:
                torch.save(state, str(Path(output_dir) / "lora_weights.pt"))
        logger.info("LoRA adapter saved (fallback) to %s", output_dir)


# ============================================================================
# MODEL LOADING (FA2 -> SDPA -> eager fallback)
# ============================================================================

_VARIANT_DIR = {
    "turbo": "acestep-v15-turbo",
    "xl-turbo": "acestep-v15-xl-turbo",
    "base": "acestep-v15-base",
    "xl-base": "acestep-v15-xl-base",
    "sft": "acestep-v15-sft",
    "xl-sft": "acestep-v15-xl-sft",
}


def _resolve_model_dir(checkpoint_dir: str, variant: str) -> Path:
    base = Path(checkpoint_dir).resolve()
    subdir = _VARIANT_DIR.get(variant)
    if subdir:
        p = (Path(checkpoint_dir) / subdir).resolve()
        if p.is_dir():
            return p
    p = (Path(checkpoint_dir) / variant).resolve()
    if p.is_dir():
        return p
    raise FileNotFoundError(
        f"Model directory not found: tried {_VARIANT_DIR.get(variant, variant)!r} "
        f"and {variant!r} under {checkpoint_dir}"
    )


def _ensure_acestep_imports():
    """Register stub modules so AutoModel can load ACE-Step checkpoints."""
    for name in (
        "acestep", "acestep.models", "acestep.models.common",
        "acestep.models.xl_base", "acestep.models.xl_turbo", "acestep.models.xl_sft",
    ):
        if name not in sys.modules:
            stub = types.ModuleType(name)
            stub.__path__ = []
            sys.modules[name] = stub

    # Try to load real modules from adjacent ACE-Step checkout
    for name in (
        "acestep.models.common.configuration_acestep_v15",
        "acestep.models.common.apg_guidance",
    ):
        if name not in sys.modules:
            sys.modules[name] = types.ModuleType(name)


def _attn_candidates(device: str) -> List[str]:
    """FA2 -> SDPA -> eager, filtered by availability."""
    candidates = []
    if device.startswith("cuda"):
        try:
            import flash_attn  # noqa: F401
            dev_idx = int(device.split(":")[1]) if ":" in device else 0
            props = torch.cuda.get_device_properties(dev_idx)
            if props.major >= 8:
                candidates.append("flash_attention_2")
        except (ImportError, Exception):
            pass
    candidates.extend(["sdpa", "eager"])
    return candidates


def load_model_for_training(
    checkpoint_dir: str, variant: str = "base", device: str = "cpu",
) -> Any:
    from transformers import AutoModel

    model_dir = _resolve_model_dir(checkpoint_dir, variant)
    # CPU always uses float32
    dtype = CPU_DTYPE if device == "cpu" else torch.bfloat16

    _ensure_acestep_imports()

    candidates = _attn_candidates(device)
    model = None
    last_err = None

    for idx, attn in enumerate(candidates):
        try:
            load_kwargs = dict(
                trust_remote_code=True,
                attn_implementation=attn,
                torch_dtype=dtype,
                low_cpu_mem_usage=False,
            )
            if device != "cpu":
                load_kwargs["device_map"] = {"": device}
            model = AutoModel.from_pretrained(str(model_dir), **load_kwargs)
            logger.info("Model loaded with attn_implementation=%s", attn)
            break
        except Exception as exc:
            err_text = str(exc)
            if "packages that were not found" in err_text or "No module named" in err_text:
                raise RuntimeError(
                    f"Model files in {model_dir} require a missing Python package.\n"
                    f"  Original error: {err_text}"
                ) from exc
            last_err = exc
            logger.warning("attn backend '%s' failed: %s", attn, exc)

    if model is None:
        raise RuntimeError(f"Failed to load model from {model_dir}: {last_err}") from last_err

    for param in model.parameters():
        param.requires_grad = False
    model.eval()
    return model


def load_vae(checkpoint_dir: str, device: str = "cpu"):
    from diffusers.models import AutoencoderOobleck

    vae_path = Path(checkpoint_dir) / "vae"
    if not vae_path.is_dir():
        raise FileNotFoundError(f"VAE directory not found: {vae_path}")

    dtype = CPU_DTYPE if device == "cpu" else torch.bfloat16
    vae = AutoencoderOobleck.from_pretrained(str(vae_path), torch_dtype=dtype)
    vae = vae.to(device=device)
    vae.eval()
    return vae


def load_text_encoder(checkpoint_dir: str, device: str = "cpu"):
    from transformers import AutoModel, AutoTokenizer

    text_path = Path(checkpoint_dir) / "Qwen3-Embedding-0.6B"
    if not text_path.is_dir():
        raise FileNotFoundError(f"Text encoder not found: {text_path}")

    dtype = CPU_DTYPE if device == "cpu" else torch.bfloat16
    tokenizer = AutoTokenizer.from_pretrained(str(text_path))
    encoder = AutoModel.from_pretrained(str(text_path), torch_dtype=dtype)
    encoder = encoder.to(device=device)
    encoder.eval()
    return tokenizer, encoder


def load_silence_latent(
    checkpoint_dir: str, device: str = "cpu", variant: str = "base",
) -> torch.Tensor:
    ckpt = Path(checkpoint_dir)
    dtype = CPU_DTYPE if device == "cpu" else torch.bfloat16

    candidates = [ckpt / "silence_latent.pt"]
    subdir = _VARIANT_DIR.get(variant)
    if subdir:
        candidates.append(ckpt / subdir / "silence_latent.pt")
    for sd in _VARIANT_DIR.values():
        candidates.append(ckpt / sd / "silence_latent.pt")

    for c in candidates:
        if c.is_file():
            sl = torch.load(str(c), weights_only=True).transpose(1, 2)
            return sl.to(device=device, dtype=dtype)

    raise FileNotFoundError(f"silence_latent.pt not found under {ckpt}")


def unload_models(*models) -> None:
    for obj in models:
        if obj is None:
            continue
        if hasattr(obj, "to"):
            try:
                obj.to("cpu")
            except Exception:
                pass
        del obj
    gc.collect()


# ============================================================================
# AUDIO LOADING
# ============================================================================

def load_audio_stereo(
    audio_path: str, target_sr: int, max_duration: float,
) -> Tuple[torch.Tensor, int]:
    import numpy as np

    try:
        import soundfile as sf
        data, sr = sf.read(audio_path, dtype="float32", always_2d=True)
        audio_np = np.ascontiguousarray(data.T)
        sr = int(sr)
        if sr != target_sr:
            import librosa
            audio_np = librosa.resample(audio_np, orig_sr=sr, target_sr=target_sr, axis=1)
            sr = target_sr
        audio = torch.from_numpy(np.ascontiguousarray(audio_np))
    except Exception:
        import torchaudio
        audio, sr = torchaudio.load(audio_path)
        sr = int(sr)
        if sr != target_sr:
            audio = torchaudio.transforms.Resample(sr, target_sr)(audio)
            sr = target_sr

    if audio.shape[0] == 1:
        audio = audio.repeat(2, 1)
    elif audio.shape[0] > 2:
        audio = audio[:2, :]

    max_samples = int(max_duration * target_sr)
    if audio.shape[1] > max_samples:
        audio = audio[:, :max_samples]

    return audio, sr


# ============================================================================
# TEXT / LYRICS ENCODING
# ============================================================================

def encode_text(text_encoder, tokenizer, text_prompt: str, device, dtype):
    inputs = tokenizer(
        text_prompt, padding="max_length", max_length=256,
        truncation=True, return_tensors="pt",
    )
    ids = inputs.input_ids.to(device)
    mask = inputs.attention_mask.to(device).to(dtype)

    enc_dev = next(text_encoder.parameters()).device
    if ids.device != enc_dev:
        ids = ids.to(enc_dev)
        mask = mask.to(enc_dev)

    with torch.no_grad():
        hs = text_encoder(ids).last_hidden_state.to(dtype)
    return hs, mask


def encode_lyrics(text_encoder, tokenizer, lyrics: str, device, dtype):
    inputs = tokenizer(
        lyrics, padding="max_length", max_length=512,
        truncation=True, return_tensors="pt",
    )
    ids = inputs.input_ids.to(device)
    mask = inputs.attention_mask.to(device).to(dtype)

    enc_dev = next(text_encoder.parameters()).device
    if ids.device != enc_dev:
        ids = ids.to(enc_dev)
        mask = mask.to(enc_dev)

    with torch.no_grad():
        hs = text_encoder.embed_tokens(ids).to(dtype)
    return hs, mask


# ============================================================================
# VAE TILED ENCODING
# ============================================================================

def tiled_vae_encode(
    vae, audio: torch.Tensor, dtype: torch.dtype,
    chunk_size: Optional[int] = None, overlap: int = 96000,
) -> torch.Tensor:
    vae_device = next(vae.parameters()).device
    vae_dtype = vae.dtype

    if chunk_size is None:
        chunk_size = TARGET_SR * 30

    B, C, S = audio.shape

    if S <= chunk_size:
        vae_input = audio.to(vae_device, dtype=vae_dtype)
        with torch.inference_mode():
            latents = vae.encode(vae_input).latent_dist.sample()
        return latents.transpose(1, 2).to(dtype)

    stride = chunk_size - 2 * overlap
    if stride <= 0:
        raise ValueError(f"chunk_size ({chunk_size}) must be > 2 * overlap ({overlap})")

    num_steps = math.ceil(S / stride)
    ds_factor = None
    write_pos = 0
    final = None

    for i in range(num_steps):
        core_start = i * stride
        core_end = min(core_start + stride, S)
        win_start = max(0, core_start - overlap)
        win_end = min(S, core_end + overlap)

        chunk = audio[:, :, win_start:win_end].to(vae_device, dtype=vae_dtype)
        with torch.inference_mode():
            lat = vae.encode(chunk).latent_dist.sample()

        if ds_factor is None:
            ds_factor = chunk.shape[-1] / lat.shape[-1]
            total_len = int(round(S / ds_factor))
            final = torch.zeros(B, lat.shape[1], total_len, dtype=lat.dtype, device="cpu")

        trim_start = int(round((core_start - win_start) / ds_factor))
        trim_end = int(round((win_end - core_end) / ds_factor))
        end_idx = lat.shape[-1] - trim_end if trim_end > 0 else lat.shape[-1]
        core = lat[:, :, trim_start:end_idx]
        core_len = core.shape[-1]
        final[:, :, write_pos:write_pos + core_len] = core.cpu()
        write_pos += core_len
        del chunk, lat, core

    final = final[:, :, :write_pos]
    return final.transpose(1, 2).to(dtype)


# ============================================================================
# ENCODER / CONTEXT HELPERS
# ============================================================================

def run_encoder(
    model, text_hs, text_mask, lyric_hs, lyric_mask, device, dtype,
):
    refer = torch.zeros(1, 1, 64, device=device, dtype=dtype)
    order_mask = torch.zeros(1, device=device, dtype=torch.long)

    with torch.no_grad():
        enc_hs, enc_mask = model.encoder(
            text_hidden_states=text_hs,
            text_attention_mask=text_mask,
            lyric_hidden_states=lyric_hs,
            lyric_attention_mask=lyric_mask,
            refer_audio_acoustic_hidden_states_packed=refer,
            refer_audio_order_mask=order_mask,
        )
    return enc_hs, enc_mask


def build_context_latents(silence_latent, latent_length: int, device, dtype):
    src = silence_latent[:, :latent_length, :].to(dtype)
    if src.shape[0] < 1:
        src = src.expand(1, -1, -1)
    if src.shape[1] < latent_length:
        pad_len = latent_length - src.shape[1]
        src = torch.cat([src, silence_latent[:, :pad_len, :].expand(1, -1, -1).to(dtype)], dim=1)
    elif src.shape[1] > latent_length:
        src = src[:, :latent_length, :]
    masks = torch.ones(1, latent_length, 64, device=device, dtype=dtype)
    return torch.cat([src, masks], dim=-1)


# ============================================================================
# AUDIO DISCOVERY
# ============================================================================

def _discover_audio_files(audio_dir: str) -> List[Path]:
    files = []
    for root, _, names in os.walk(audio_dir):
        for name in sorted(names):
            if Path(name).suffix.lower() in AUDIO_EXTENSIONS:
                files.append(Path(root) / name)
    return files


def _detect_max_duration(files: List[Path]) -> float:
    """Return the longest audio file duration (capped at MAX_AUDIO_DURATION)."""
    max_dur = 0.0
    try:
        import soundfile as sf
        for f in files[:50]:
            try:
                info = sf.info(str(f))
                max_dur = max(max_dur, info.duration)
            except Exception:
                pass
    except ImportError:
        pass
    return min(max_dur if max_dur > 0 else MAX_AUDIO_DURATION, MAX_AUDIO_DURATION)


# ============================================================================
# AUDIO ANALYSIS (ported from Side-Step -- faf / mid / sas modes)
# ============================================================================
#
# faf  ("Fast As F*ck")       ~2-3 s/file   - single-method, no Demucs
# mid                         ~12 s/file     - 3-method ensemble, Demucs stems
# sas  ("Smart/Slow As Sh*t") ~30 s/file     - deep multi-technique + chunked
#
# Demucs on CPU is SLOW (~2-5 min/file).  mid/sas are designed for GPU
# stem separation but will still work on CPU -- just much slower.
# ============================================================================

_ANALYSIS_MODES = ("faf", "mid", "sas")
_SAS_NUM_CHUNKS = 5
_SAS_CHUNK_SECONDS = 15  # seconds per analysis window

# Key profile families for multi-profile voting (mid / sas)
_KEY_PROFILES = {
    "krumhansl": {
        "major": [6.35, 2.23, 3.48, 2.33, 4.38, 4.09,
                  2.52, 5.19, 2.39, 3.66, 2.29, 2.88],
        "minor": [6.33, 2.68, 3.52, 5.38, 2.60, 3.53,
                  2.54, 4.75, 3.98, 2.69, 3.34, 3.17],
    },
    "temperley": {
        "major": [5.0, 2.0, 3.5, 2.0, 4.5, 4.0,
                  2.0, 4.5, 2.0, 3.5, 1.5, 4.0],
        "minor": [5.0, 2.0, 3.5, 4.5, 2.0, 3.5,
                  2.0, 4.5, 3.5, 2.0, 1.5, 4.0],
    },
    "albrecht": {
        "major": [0.238, 0.006, 0.111, 0.006, 0.137, 0.094,
                  0.016, 0.214, 0.009, 0.080, 0.008, 0.081],
        "minor": [0.220, 0.006, 0.104, 0.123, 0.019, 0.103,
                  0.012, 0.214, 0.062, 0.022, 0.061, 0.052],
    },
}

_PITCH_CLASSES = ["C", "C#", "D", "D#", "E", "F",
                  "F#", "G", "G#", "A", "A#", "B"]

# Filename pattern: "Artist - Title"
_FILENAME_RE = re.compile(r"^(.+?)\s*[-–—]\s*(.+)$")


# ---- Demucs stem separation (mid / sas) --------------------------------

def separate_stems(
    audio_path: Path,
    tmp_dir: Path,
    device: str = "cpu",
) -> Tuple[Path, Path]:
    """Run Demucs HTDemucs and return (drums_path, harmonics_path).

    Harmonics = bass + other stems summed.  Vocals are discarded.

    WARNING: On CPU this takes ~2-5 minutes per file.
    """
    import torchaudio
    from demucs.pretrained import get_model
    from demucs.apply import apply_model

    torch_device = torch.device(device)

    logger.info("Loading Demucs HTDemucs model on %s", device)
    if device == "cpu":
        logger.warning(
            "Demucs on CPU is slow (~2-5 min per file). "
            "Consider using 'faf' mode or running on a GPU machine."
        )
    model = get_model("htdemucs")
    model.to(torch_device)
    model.eval()

    wav, sr = torchaudio.load(str(audio_path))

    # Resample to model's expected rate (44100 Hz) if needed
    if sr != model.samplerate:
        wav = torchaudio.functional.resample(wav, sr, model.samplerate)
        sr = model.samplerate

    # HTDemucs requires stereo input
    if wav.shape[0] == 1:
        wav = wav.repeat(2, 1)

    wav = wav.unsqueeze(0).to(torch_device)

    logger.info("Separating stems for %s", audio_path.name)
    with torch.no_grad():
        sources = apply_model(model, wav, device=torch_device)

    source_map = {name: i for i, name in enumerate(model.sources)}
    drums = sources[0, source_map["drums"]].cpu()
    bass = sources[0, source_map["bass"]].cpu()
    other = sources[0, source_map["other"]].cpu()
    harmonics = bass + other

    drums_path = tmp_dir / "drums.wav"
    harmonics_path = tmp_dir / "harmonics.wav"
    torchaudio.save(str(drums_path), drums, sr)
    torchaudio.save(str(harmonics_path), harmonics, sr)

    del model, sources, wav, drums, bass, other, harmonics
    gc.collect()

    logger.info("Stems written: %s, %s", drums_path, harmonics_path)
    return drums_path, harmonics_path


# ---- Chunk selection (sas mode) ----------------------------------------

def _select_chunks(
    y,  # np.ndarray
    sr: int,
    n_chunks: int = _SAS_NUM_CHUNKS,
    chunk_sec: float = _SAS_CHUNK_SECONDS,
    min_gap_sec: float = 10.0,
    use_onset: bool = True,
) -> list:
    """Select the most informative audio chunks for sas analysis.

    Energy-gated + spread: rank windows by onset density (or RMS),
    discard below-median, then greedily pick chunks maximally spread apart.
    """
    import librosa
    import numpy as np

    chunk_samples = int(chunk_sec * sr)
    hop_samples = chunk_samples // 2
    if len(y) < chunk_samples:
        return [y]

    candidates = []
    for start in range(0, len(y) - chunk_samples + 1, hop_samples):
        window = y[start : start + chunk_samples]
        if use_onset:
            onset_env = librosa.onset.onset_strength(y=window, sr=sr)
            score = float(np.mean(onset_env))
        else:
            score = float(np.sqrt(np.mean(window ** 2)))
        candidates.append((start, score))

    if not candidates:
        return [y]

    scores = np.array([s for _, s in candidates])
    median_score = float(np.median(scores))
    gated = [(start, score) for start, score in candidates if score >= median_score]
    if not gated:
        gated = candidates

    gated.sort(key=lambda x: x[1], reverse=True)

    min_gap_samples = int(min_gap_sec * sr)
    selected_starts = []

    for start, score in gated:
        centre = start + chunk_samples // 2
        too_close = any(
            abs(centre - (s + chunk_samples // 2)) < min_gap_samples
            for s in selected_starts
        )
        if not too_close:
            selected_starts.append(start)
            if len(selected_starts) >= n_chunks:
                break

    if len(selected_starts) < n_chunks:
        for start, score in gated:
            if start not in selected_starts:
                selected_starts.append(start)
                if len(selected_starts) >= n_chunks:
                    break

    selected_starts.sort()
    return [y[s : s + chunk_samples] for s in selected_starts]


# ---- BPM helpers --------------------------------------------------------

def _octave_correct_bpm(bpm: float, lo: float = 70.0, hi: float = 180.0) -> float:
    """Fold BPM into the musical sweet-spot range [lo, hi]."""
    if bpm <= 0:
        return bpm
    candidate = bpm
    while candidate > hi:
        candidate /= 2.0
    while candidate < lo:
        candidate *= 2.0
    if candidate < lo or candidate > hi:
        return bpm
    return candidate


def _bpm_core_ensemble(y, sr) -> list:
    """Run the 3-method BPM ensemble on a single audio buffer (mid/sas).

    Returns a list of octave-corrected BPM estimates.
    """
    import librosa
    import numpy as np

    estimates = []

    # Method A: beat_track
    try:
        tempo_a, _ = librosa.beat.beat_track(y=y, sr=sr)
        val_a = float(np.atleast_1d(tempo_a)[0])
        if val_a > 0:
            estimates.append(_octave_correct_bpm(val_a))
    except Exception:
        pass

    # Method B: tempogram peak
    try:
        onset_env = librosa.onset.onset_strength(y=y, sr=sr)
        tempogram = librosa.feature.tempogram(onset_envelope=onset_env, sr=sr)
        avg_tempogram = np.mean(tempogram, axis=1)
        bpm_axis = librosa.tempo_frequencies(tempogram.shape[0], sr=sr)
        valid = (bpm_axis >= 30) & (bpm_axis <= 300)
        if np.any(valid):
            masked = avg_tempogram.copy()
            masked[~valid] = 0
            peak_idx = np.argmax(masked)
            val_b = float(bpm_axis[peak_idx])
            if val_b > 0:
                estimates.append(_octave_correct_bpm(val_b))
    except Exception:
        pass

    # Method C: onset autocorrelation
    try:
        onset_env = librosa.onset.onset_strength(y=y, sr=sr)
        ac = librosa.autocorrelate(onset_env, max_size=len(onset_env))
        hop = 512
        min_lag = int(60.0 * sr / (300.0 * hop))
        max_lag = int(60.0 * sr / (30.0 * hop))
        max_lag = min(max_lag, len(ac) - 1)
        if min_lag < max_lag and max_lag > 0:
            segment = ac[min_lag:max_lag + 1]
            peak_offset = np.argmax(segment)
            peak_lag = min_lag + peak_offset
            if peak_lag > 0:
                val_c = 60.0 * sr / (peak_lag * hop)
                if val_c > 0:
                    estimates.append(_octave_correct_bpm(val_c))
    except Exception:
        pass

    return estimates


def _bpm_consensus(estimates: list) -> Tuple[Optional[int], str]:
    """Find consensus BPM from a list of estimates + assign confidence."""
    import numpy as np

    if not estimates:
        return None, "low"

    estimates_arr = np.array(estimates)
    best_cluster = []
    for ref in estimates_arr:
        cluster = [e for e in estimates_arr
                   if abs(e - ref) / max(ref, 1) < 0.08]
        if len(cluster) > len(best_cluster):
            best_cluster = cluster

    consensus = float(np.median(best_cluster)) if best_cluster else estimates[0]
    bpm = int(round(consensus))
    if bpm <= 0:
        return None, "low"

    n_agree = len(best_cluster)
    n_total = len(estimates)
    if n_total >= 6:
        # sas thresholds (many data points)
        if n_agree / n_total >= 0.7:
            confidence = "high"
        elif n_agree / n_total >= 0.4:
            confidence = "medium"
        else:
            confidence = "low"
    else:
        # mid thresholds
        if n_agree >= 3:
            confidence = "high"
        elif n_agree >= 2:
            confidence = "medium"
        else:
            confidence = "low"

    return bpm, confidence


# ---- Unified BPM detection ---------------------------------------------

def _detect_bpm(y, sr, mode: str = "faf") -> Tuple[Optional[int], str]:
    """Detect BPM with quality controlled by mode.

    faf: Single beat_track + octave correction.
    mid: 3-method ensemble (beat_track + tempogram + onset-AC).
    sas: mid ensemble + PLP + multi-hop + chunked analysis.

    Returns (bpm, confidence).
    """
    import librosa
    import numpy as np

    try:
        # faf: single method
        if mode == "faf":
            try:
                tempo, _ = librosa.beat.beat_track(y=y, sr=sr)
                val = float(np.atleast_1d(tempo)[0])
                if val > 0:
                    bpm = int(round(_octave_correct_bpm(val)))
                    logger.info("BPM faf: %d (raw: %.1f)", bpm, val)
                    return bpm, "low"
            except Exception:
                pass
            return None, "low"

        # mid: 3-method ensemble
        estimates = _bpm_core_ensemble(y, sr)

        # sas: additional techniques
        ibi_cv = 0.5
        if mode == "sas":
            # PLP (Predominant Local Pulse)
            try:
                onset_env = librosa.onset.onset_strength(y=y, sr=sr)
                pulse = librosa.beat.plp(onset_envelope=onset_env, sr=sr)
                plp_ac = librosa.autocorrelate(pulse, max_size=len(pulse))
                hop = 512
                min_lag = int(60.0 * sr / (300.0 * hop))
                max_lag = int(60.0 * sr / (30.0 * hop))
                max_lag = min(max_lag, len(plp_ac) - 1)
                if min_lag < max_lag and max_lag > 0:
                    seg = plp_ac[min_lag:max_lag + 1]
                    peak_lag = min_lag + np.argmax(seg)
                    if peak_lag > 0:
                        plp_bpm = 60.0 * sr / (peak_lag * hop)
                        if plp_bpm > 0:
                            estimates.append(_octave_correct_bpm(plp_bpm))
            except Exception:
                pass

            # Multi-hop beat_track (256, 1024)
            for extra_hop in (256, 1024):
                try:
                    tempo_h, _ = librosa.beat.beat_track(y=y, sr=sr, hop_length=extra_hop)
                    val_h = float(np.atleast_1d(tempo_h)[0])
                    if val_h > 0:
                        estimates.append(_octave_correct_bpm(val_h))
                except Exception:
                    pass

            # Chunked ensemble
            chunks = _select_chunks(y, sr, n_chunks=_SAS_NUM_CHUNKS, use_onset=True)
            for chunk in chunks:
                chunk_estimates = _bpm_core_ensemble(chunk, sr)
                estimates.extend(chunk_estimates)

            # IBI stability
            try:
                _, beat_frames = librosa.beat.beat_track(y=y, sr=sr)
                if beat_frames is not None and len(beat_frames) > 4:
                    beat_times = librosa.frames_to_time(beat_frames, sr=sr)
                    ibis = np.diff(beat_times)
                    ibi_cv = float(np.std(ibis) / (np.mean(ibis) + 1e-10))
                else:
                    ibi_cv = 0.5
            except Exception:
                ibi_cv = 0.5

        bpm, confidence = _bpm_consensus(estimates)

        # sas: IBI stability can upgrade medium->high or downgrade
        if mode == "sas" and bpm is not None:
            if ibi_cv < 0.10 and confidence == "medium":
                confidence = "high"
            elif ibi_cv > 0.30 and confidence == "high":
                confidence = "medium"

        logger.info(
            "BPM [%s]: %s (estimates=%s, conf=%s)",
            mode, bpm,
            [round(e, 1) for e in estimates[:10]],
            confidence,
        )
        return bpm, confidence

    except Exception as exc:
        logger.warning("BPM detection failed: %s", exc)
        return None, "low"


# ---- Key detection helpers ----------------------------------------------

def _best_key_for_profile(chroma_avg, major_profile, minor_profile):
    """Find the best key match for a single profile family.

    Returns (key_label, correlation).
    """
    import numpy as np

    major_norm = np.array(major_profile, dtype=float)
    major_norm = major_norm / major_norm.sum()
    minor_norm = np.array(minor_profile, dtype=float)
    minor_norm = minor_norm / minor_norm.sum()

    best_corr = -2.0
    best_key = "C major"

    for shift in range(12):
        rotated = np.roll(chroma_avg, -shift)
        corr_maj = float(np.corrcoef(rotated, major_norm)[0, 1])
        if corr_maj > best_corr:
            best_corr = corr_maj
            best_key = f"{_PITCH_CLASSES[shift]} major"
        corr_min = float(np.corrcoef(rotated, minor_norm)[0, 1])
        if corr_min > best_corr:
            best_corr = corr_min
            best_key = f"{_PITCH_CLASSES[shift]} minor"

    return best_key, best_corr


def _key_votes_from_chroma(chroma_avg, profiles=None) -> list:
    """Vote on key from a single chroma vector using specified profiles.

    Returns list of (key_label, correlation) -- one per profile family.
    """
    if profiles is None:
        profiles = _KEY_PROFILES
    results = []
    for name, pf in profiles.items():
        key_label, corr = _best_key_for_profile(
            chroma_avg, pf["major"], pf["minor"],
        )
        results.append((key_label, corr))
    return results


def _energy_weighted_chroma(chroma, y_harmonic):
    """Compute an energy-weighted average chroma vector.

    Returns normalized chroma_avg or None if zero energy.
    """
    import librosa
    import numpy as np

    rms = librosa.feature.rms(y=y_harmonic, frame_length=2048, hop_length=512)
    rms_vec = rms[0]
    min_len = min(chroma.shape[1], len(rms_vec))
    chroma = chroma[:, :min_len]
    rms_vec = rms_vec[:min_len]

    weights = rms_vec / (rms_vec.sum() + 1e-10)
    chroma_avg = (chroma * weights[None, :]).sum(axis=1)

    s = chroma_avg.sum()
    if s == 0:
        return None
    return chroma_avg / s


# ---- Unified key detection ----------------------------------------------

def _detect_key(y, sr, mode: str = "faf") -> Tuple[Optional[str], str]:
    """Detect musical key with quality controlled by mode.

    faf: Single Krumhansl profile on chroma_cens.
    mid: 3-profile x energy-weighted chroma_cens x 8s segment voting.
    sas: mid + multi-chroma fusion + tonnetz + tuning correction +
         ending resolution + chunked voting.

    Returns (key, confidence).
    """
    import librosa
    import numpy as np
    from collections import Counter

    try:
        # Harmonic enhancement
        margin = 4.0 if mode != "faf" else 2.0
        y_harmonic = librosa.effects.harmonic(y, margin=margin)

        # sas: tuning correction
        tuning = 0.0
        if mode == "sas":
            try:
                tuning = float(librosa.estimate_tuning(y=y_harmonic, sr=sr))
            except Exception:
                tuning = 0.0

        # faf: single chroma, single profile
        if mode == "faf":
            chroma = librosa.feature.chroma_cens(y=y_harmonic, sr=sr)
            chroma_avg = _energy_weighted_chroma(chroma, y_harmonic)
            if chroma_avg is None:
                return None, "low"
            kr = _KEY_PROFILES["krumhansl"]
            key_label, corr = _best_key_for_profile(
                chroma_avg, kr["major"], kr["minor"],
            )
            logger.info("Key faf: %s (corr=%.3f)", key_label, corr)
            return key_label, "low"

        # mid / sas: multi-profile voting
        all_votes = []
        all_weights = []

        if mode == "sas":
            chroma_types = {
                "cens": lambda: librosa.feature.chroma_cens(
                    y=y_harmonic, sr=sr, tuning=tuning,
                ),
                "cqt": lambda: librosa.feature.chroma_cqt(
                    y=y_harmonic, sr=sr, tuning=tuning,
                ),
                "stft": lambda: librosa.feature.chroma_stft(
                    y=y_harmonic, sr=sr, tuning=tuning,
                ),
            }
        else:
            chroma_types = {
                "cens": lambda: librosa.feature.chroma_cens(
                    y=y_harmonic, sr=sr,
                ),
            }

        for chroma_name, chroma_fn in chroma_types.items():
            try:
                chroma = chroma_fn()
            except Exception:
                continue

            chroma_avg = _energy_weighted_chroma(chroma, y_harmonic)
            if chroma_avg is None:
                continue

            # Global multi-profile vote
            for key_label, corr in _key_votes_from_chroma(chroma_avg):
                all_votes.append(key_label)
                all_weights.append(1.0)

            # Segment-based voting
            rms = librosa.feature.rms(
                y=y_harmonic, frame_length=2048, hop_length=512,
            )
            rms_vec = rms[0]
            min_len = min(chroma.shape[1], len(rms_vec))
            chroma_s = chroma[:, :min_len]
            rms_s = rms_vec[:min_len]

            seg_frames = int(8.0 * sr / 512)
            n_segments = max(1, chroma_s.shape[1] // seg_frames)

            for seg_i in range(n_segments):
                start = seg_i * seg_frames
                end = min(start + seg_frames, chroma_s.shape[1])
                seg_chroma = chroma_s[:, start:end]
                seg_w = rms_s[start:end]

                w_sum = seg_w.sum()
                if w_sum < 1e-10:
                    continue

                seg_w_norm = seg_w / w_sum
                seg_avg = (seg_chroma * seg_w_norm[None, :]).sum(axis=1)
                s = seg_avg.sum()
                if s < 1e-10:
                    continue
                seg_avg = seg_avg / s

                for key_label, _ in _key_votes_from_chroma(seg_avg):
                    all_votes.append(key_label)
                    all_weights.append(1.0)

        # sas-only extras
        if mode == "sas":
            # Tonnetz -- weighted vote for major/minor disambiguation
            try:
                tonnetz = librosa.feature.tonnetz(y=y_harmonic, sr=sr)
                tonnetz_avg = np.mean(tonnetz, axis=1)
                major_energy = float(np.sum(tonnetz_avg[4:6] ** 2))
                minor_energy = float(np.sum(tonnetz_avg[2:4] ** 2))
                tonnetz_ratio = major_energy / (minor_energy + 1e-10)

                if all_votes:
                    temp_counts = Counter(all_votes)
                    leader = temp_counts.most_common(1)[0][0]
                    leader_is_major = "major" in leader
                    tonnetz_says_major = tonnetz_ratio > 1.0

                    if leader_is_major == tonnetz_says_major:
                        all_votes.extend([leader] * 3)
                        all_weights.extend([1.5] * 3)
                    else:
                        alt_mode = "minor" if leader_is_major else "major"
                        chroma_cens = librosa.feature.chroma_cens(
                            y=y_harmonic, sr=sr, tuning=tuning,
                        )
                        ca = _energy_weighted_chroma(chroma_cens, y_harmonic)
                        if ca is not None:
                            for name, pf in _KEY_PROFILES.items():
                                prof = np.array(pf[alt_mode], dtype=float)
                                prof_norm = prof / prof.sum()
                                best_corr = -2.0
                                best_k = ""
                                for shift in range(12):
                                    rotated = np.roll(ca, -shift)
                                    c = float(np.corrcoef(rotated, prof_norm)[0, 1])
                                    if c > best_corr:
                                        best_corr = c
                                        best_k = f"{_PITCH_CLASSES[shift]} {alt_mode}"
                                if best_k:
                                    all_votes.append(best_k)
                                    all_weights.append(1.0)
            except Exception:
                pass

            # Ending resolution -- last ~5 s weighted extra
            try:
                end_samples = min(int(5.0 * sr), len(y_harmonic))
                y_end = y_harmonic[-end_samples:]
                chroma_end = librosa.feature.chroma_cens(
                    y=y_end, sr=sr, tuning=tuning,
                )
                end_avg = np.mean(chroma_end, axis=1)
                s = end_avg.sum()
                if s > 1e-10:
                    end_avg = end_avg / s
                    for key_label, _ in _key_votes_from_chroma(end_avg):
                        all_votes.append(key_label)
                        all_weights.append(2.0)
            except Exception:
                pass

            # Chunked voting
            chunks = _select_chunks(
                y_harmonic, sr, n_chunks=_SAS_NUM_CHUNKS, use_onset=False,
            )
            for chunk in chunks:
                try:
                    ch_chroma = librosa.feature.chroma_cens(
                        y=chunk, sr=sr, tuning=tuning,
                    )
                    ch_avg = _energy_weighted_chroma(ch_chroma, chunk)
                    if ch_avg is not None:
                        for key_label, _ in _key_votes_from_chroma(ch_avg):
                            all_votes.append(key_label)
                            all_weights.append(1.0)
                except Exception:
                    pass

        # Final weighted majority vote
        if not all_votes:
            return None, "low"

        weighted_counts = {}
        for vote, w in zip(all_votes, all_weights):
            weighted_counts[vote] = weighted_counts.get(vote, 0.0) + w

        best_key = max(weighted_counts, key=weighted_counts.get)
        total_weight = sum(all_weights)
        best_weight = weighted_counts[best_key]
        share = best_weight / total_weight

        if share >= 0.55:
            confidence = "high"
        elif share >= 0.35:
            confidence = "medium"
        else:
            confidence = "low"

        logger.info(
            "Key [%s]: %s (share=%.0f%%, votes=%d, conf=%s)",
            mode, best_key, share * 100, len(all_votes), confidence,
        )
        return best_key, confidence

    except Exception as exc:
        logger.warning("Key detection failed: %s", exc)
        return None, "low"


# ---- Time-signature helpers ---------------------------------------------

def _timesig_core_scores(y, sr) -> dict:
    """Compute 3-signal time-signature scores on a single buffer (mid/sas).

    Returns dict mapping signature labels to raw scores.
    """
    import librosa
    import numpy as np

    scores = {}

    tempo, beat_frames = librosa.beat.beat_track(y=y, sr=sr)
    if beat_frames is None or len(beat_frames) < 8:
        return scores

    onset_env = librosa.onset.onset_strength(y=y, sr=sr)
    beat_strengths = onset_env[beat_frames[beat_frames < len(onset_env)]]
    if len(beat_strengths) < 8:
        return scores

    # Signal 1: Accent pattern analysis
    for label, grouping in [("3/4", 3), ("4/4", 4), ("6/8", 6)]:
        if len(beat_strengths) < grouping * 2:
            scores[label] = 0.0
            continue
        usable = len(beat_strengths) - (len(beat_strengths) % grouping)
        grouped = beat_strengths[:usable].reshape(-1, grouping)
        downbeat_mean = float(np.mean(grouped[:, 0]))
        offbeat_mean = float(np.mean(grouped[:, 1:]))
        contrast = downbeat_mean / offbeat_mean if offbeat_mean > 0 else 1.0
        scores[label] = contrast

    # Signal 2: Autocorrelation at meter periods
    hop = 512
    beat_times = librosa.frames_to_time(beat_frames, sr=sr)
    intervals = np.diff(beat_times)
    if len(intervals) > 0:
        median_interval = float(np.median(intervals))
        beat_period = int(round(median_interval * sr / hop))
        if beat_period > 0:
            ac = librosa.autocorrelate(onset_env, max_size=len(onset_env))
            for label, mult in [("3/4", 3), ("4/4", 4), ("6/8", 6)]:
                period = beat_period * mult
                if period < len(ac):
                    lo = max(0, period - 2)
                    hi = min(len(ac), period + 3)
                    ac_score = float(np.mean(ac[lo:hi]))
                    if ac[0] > 0:
                        ac_score /= float(ac[0])
                    scores[label] = scores.get(label, 0.0) + ac_score

    # Signal 3: Beat-strength variance ratio
    for label, grouping in [("3/4", 3), ("4/4", 4)]:
        usable = len(beat_strengths) - (len(beat_strengths) % grouping)
        if usable >= grouping * 2:
            grouped = beat_strengths[:usable].reshape(-1, grouping)
            row_vars = np.var(grouped, axis=1)
            scores[label] = scores.get(label, 0.0) + float(np.mean(row_vars))

    return scores


# ---- Unified time-signature detection -----------------------------------

def _detect_time_sig(y, sr, mode: str = "faf") -> Tuple[Optional[str], str]:
    """Estimate time signature with quality controlled by mode.

    faf: Hardcoded "4/4" (correct ~80%+ of the time).
    mid: Beat-sync accent + AC + variance + 4/4 prior.
    sas: mid signals + PLP periodicity + multi-band onset +
         tempogram harmonic ratios + chunked voting.

    Returns (signature, confidence).
    """
    if mode == "faf":
        return "4/4", "low"

    import librosa
    import numpy as np

    try:
        # mid: core 3-signal scoring
        scores = _timesig_core_scores(y, sr)

        # sas: additional techniques
        if mode == "sas":
            onset_env = librosa.onset.onset_strength(y=y, sr=sr)

            # PLP periodicity
            try:
                pulse = librosa.beat.plp(onset_envelope=onset_env, sr=sr)
                plp_ac = librosa.autocorrelate(pulse, max_size=len(pulse))
                tempo_est, _ = librosa.beat.beat_track(y=y, sr=sr)
                tempo_val = float(np.atleast_1d(tempo_est)[0])
                if tempo_val > 0:
                    hop = 512
                    bp = int(round(60.0 / tempo_val * sr / hop))
                    if bp > 0:
                        for label, mult in [("3/4", 3), ("4/4", 4), ("6/8", 6)]:
                            lag = bp * mult
                            if lag < len(plp_ac):
                                lo = max(0, lag - 2)
                                hi = min(len(plp_ac), lag + 3)
                                s = float(np.mean(plp_ac[lo:hi]))
                                if plp_ac[0] > 0:
                                    s /= float(plp_ac[0])
                                scores[label] = scores.get(label, 0.0) + s
            except Exception:
                pass

            # Multi-band onset analysis (low/mid/high)
            try:
                S = np.abs(librosa.stft(y))
                n_bins = S.shape[0]
                third = n_bins // 3
                bands = {
                    "low": S[:third, :],
                    "mid_band": S[third:2*third, :],
                    "high": S[2*third:, :],
                }
                for band_name, band_S in bands.items():
                    band_onset = librosa.onset.onset_strength(S=band_S, sr=sr)
                    band_ac = librosa.autocorrelate(
                        band_onset, max_size=len(band_onset),
                    )
                    tempo_val2 = float(np.atleast_1d(tempo_est)[0])
                    if tempo_val2 > 0:
                        hop = 512
                        bp2 = int(round(60.0 / tempo_val2 * sr / hop))
                        if bp2 > 0 and band_ac[0] > 0:
                            for label, mult in [("3/4", 3), ("4/4", 4)]:
                                lag = bp2 * mult
                                if lag < len(band_ac):
                                    lo = max(0, lag - 2)
                                    hi = min(len(band_ac), lag + 3)
                                    s = float(np.mean(band_ac[lo:hi]))
                                    s /= float(band_ac[0])
                                    w = 1.5 if band_name == "low" else 1.0
                                    scores[label] = scores.get(label, 0.0) + s * w
            except Exception:
                pass

            # Tempogram harmonic ratios
            try:
                tempogram = librosa.feature.tempogram(
                    onset_envelope=onset_env, sr=sr,
                )
                avg_tg = np.mean(tempogram, axis=1)
                bpm_axis = librosa.tempo_frequencies(tempogram.shape[0], sr=sr)
                if tempo_val > 0:
                    for mult_label, t_mult in [("duple", 2.0), ("triple", 3.0)]:
                        target_bpm = tempo_val * t_mult
                        if target_bpm < 300:
                            idx = np.argmin(np.abs(bpm_axis - target_bpm))
                            energy = float(avg_tg[idx])
                            base_idx = np.argmin(np.abs(bpm_axis - tempo_val))
                            base_energy = float(avg_tg[base_idx]) + 1e-10
                            ratio = energy / base_energy
                            if t_mult == 2.0:
                                scores["4/4"] = scores.get("4/4", 0.0) + ratio
                            else:
                                scores["3/4"] = scores.get("3/4", 0.0) + ratio
            except Exception:
                pass

            # Chunked voting
            chunks = _select_chunks(y, sr, n_chunks=_SAS_NUM_CHUNKS, use_onset=True)
            chunk_votes = []
            for chunk in chunks:
                cs = _timesig_core_scores(chunk, sr)
                if cs:
                    cs["4/4"] = cs.get("4/4", 0.0) * 1.15
                    best_c = max(cs, key=cs.get)
                    chunk_votes.append(best_c)
            for vote in chunk_votes:
                scores[vote] = scores.get(vote, 0.0) + 1.0

        # Bayesian prior: bias toward 4/4
        scores["4/4"] = scores.get("4/4", 0.0) * 1.15

        if not scores:
            return "4/4", "low"

        best = max(scores, key=scores.get)

        # Confidence: margin between top 2
        sorted_scores = sorted(scores.values(), reverse=True)
        if len(sorted_scores) >= 2 and sorted_scores[1] > 0:
            margin = sorted_scores[0] / sorted_scores[1]
        else:
            margin = 1.0

        if margin > 1.4:
            confidence = "high"
        elif margin > 1.15:
            confidence = "medium"
        else:
            confidence = "low"

        logger.info(
            "TimeSig [%s]: %s (scores=%s, margin=%.2f, conf=%s)",
            mode, best,
            {k: round(v, 3) for k, v in scores.items()},
            margin, confidence,
        )
        return best, confidence

    except Exception as exc:
        logger.warning("Time signature detection failed: %s", exc)
        return "4/4", "low"


def _sanitize_tag(value: str) -> str:
    """Normalize a tag value: NFKC normalize, strip invisible chars."""
    value = unicodedata.normalize("NFKC", value)
    value = (
        value
        .replace("", "").replace("￾", "")
        .replace("​", "").replace("‌", "")
        .replace("‍", "").replace("‎", "")
        .replace("‏", "").replace("‪", "")
        .replace("‬", "")
    )
    value = "".join(
        c for c in value
        if c in ("\n", "\r", "\t", " ") or unicodedata.category(c)[0] != "C"
    )
    return value.strip()


def _extract_metadata_from_tags(audio_path: Path) -> tuple:
    """Extract (title, artist) from audio tags via mutagen, fallback to filename."""
    title, artist = None, None
    try:
        import mutagen
        mf = mutagen.File(str(audio_path))
        if mf is not None and mf.tags is not None:
            # ID3 (MP3, AIFF)
            for key in ("TIT2",):
                val = mf.tags.get(key)
                if val:
                    title = _sanitize_tag(str(val))
                    break
            for key in ("TPE1", "TPE2"):
                val = mf.tags.get(key)
                if val:
                    artist = _sanitize_tag(str(val))
                    break
            # Vorbis (FLAC, OGG) and MP4 atoms
            if title is None:
                for key in ("title", "\xa9nam"):
                    vals = mf.tags.get(key)
                    if vals:
                        raw = str(vals[0]) if isinstance(vals, list) else str(vals)
                        title = _sanitize_tag(raw)
                        break
            if artist is None:
                for key in ("artist", "\xa9ART", "albumartist", "aART"):
                    vals = mf.tags.get(key)
                    if vals:
                        raw = str(vals[0]) if isinstance(vals, list) else str(vals)
                        artist = _sanitize_tag(raw)
                        break
    except Exception:
        pass

    # Fallback to filename parsing
    if not title:
        stem = audio_path.stem
        match = _FILENAME_RE.match(stem)
        if match:
            artist = artist or match.group(1).strip()
            title = match.group(2).strip()
        else:
            title = stem.strip()

    return title or audio_path.stem, artist or ""


def analyze_and_caption(
    audio_path: str,
    mode: str = "faf",
    device: str = "cpu",
) -> Dict[str, Any]:
    """Analyze an audio file and build a training caption.

    Supports three quality modes:
        faf  - CPU, ~2-3s/file.  Single-method detection on raw mix.
        mid  - ~12s/file.  Demucs stems + 3-method ensemble.
        sas  - ~30s/file.  Deep multi-technique + chunked analysis.

    For mid/sas, Demucs separates drums and harmonics stems first.
    On CPU, Demucs adds ~2-5 minutes per file.

    Args:
        audio_path: Path to the audio file.
        mode: Analysis mode ("faf", "mid", or "sas").
        device: Torch device for Demucs ("cpu").

    Returns:
        Dict with keys: caption, bpm, key, signature, lyrics, title, artist,
        confidence (dict of per-field confidence levels).
    """
    import librosa
    import numpy as np

    audio_path = Path(audio_path)

    if mode not in _ANALYSIS_MODES:
        logger.warning("Unknown analysis mode '%s', falling back to 'faf'", mode)
        mode = "faf"

    # Load audio once, reuse for all detectors
    try:
        y, sr = librosa.load(str(audio_path), sr=None, mono=True)
        # Trim silence + peak normalize
        y_trimmed, _ = librosa.effects.trim(y, top_db=30)
        if len(y_trimmed) >= sr:
            y = y_trimmed
        peak = np.max(np.abs(y))
        if peak > 0:
            y = y / peak
    except Exception as exc:
        logger.warning("Could not load audio for analysis: %s: %s", audio_path.name, exc)
        title, artist = _extract_metadata_from_tags(audio_path)
        return {
            "caption": f"A track by {artist}" if artist else f"A track titled {title}",
            "bpm": None, "key": None, "signature": "4/4",
            "lyrics": "[Instrumental]", "title": title, "artist": artist,
            "confidence": {},
        }

    confidence = {}
    tmp_dir = None

    try:
        if mode in ("mid", "sas"):
            # Demucs stem separation -- run BPM/timesig on drums,
            # key detection on harmonics
            tmp_dir = Path(tempfile.mkdtemp(prefix="ace_analysis_"))
            try:
                drums_path, harmonics_path = separate_stems(
                    audio_path, tmp_dir, device=device,
                )
                # Load separated stems for analysis
                y_drums, sr_drums = librosa.load(
                    str(drums_path), sr=None, mono=True,
                )
                y_harmonics, sr_harmonics = librosa.load(
                    str(harmonics_path), sr=None, mono=True,
                )
                # Preprocess stems
                y_drums_trimmed, _ = librosa.effects.trim(y_drums, top_db=30)
                if len(y_drums_trimmed) >= sr_drums:
                    y_drums = y_drums_trimmed
                peak_d = np.max(np.abs(y_drums))
                if peak_d > 0:
                    y_drums = y_drums / peak_d

                y_harm_trimmed, _ = librosa.effects.trim(y_harmonics, top_db=30)
                if len(y_harm_trimmed) >= sr_harmonics:
                    y_harmonics = y_harm_trimmed
                peak_h = np.max(np.abs(y_harmonics))
                if peak_h > 0:
                    y_harmonics = y_harmonics / peak_h

                # BPM + time sig on drums stem
                bpm, bpm_conf = _detect_bpm(y_drums, sr_drums, mode)
                signature, sig_conf = _detect_time_sig(y_drums, sr_drums, mode)
                # Key on harmonics stem
                key, key_conf = _detect_key(y_harmonics, sr_harmonics, mode)

                confidence = {"bpm": bpm_conf, "key": key_conf, "signature": sig_conf}

            except Exception as exc:
                logger.warning(
                    "Demucs separation failed for %s: %s -- "
                    "falling back to analysis on raw mix",
                    audio_path.name, exc,
                )
                # Fallback: run detectors on raw mix
                bpm, bpm_conf = _detect_bpm(y, sr, mode)
                key, key_conf = _detect_key(y, sr, mode)
                signature, sig_conf = _detect_time_sig(y, sr, mode)
                confidence = {"bpm": bpm_conf, "key": key_conf, "signature": sig_conf}
        else:
            # faf: all detectors on raw mix
            bpm, bpm_conf = _detect_bpm(y, sr, mode)
            key, key_conf = _detect_key(y, sr, mode)
            signature, sig_conf = _detect_time_sig(y, sr, mode)
            confidence = {"bpm": bpm_conf, "key": key_conf, "signature": sig_conf}

    finally:
        if tmp_dir is not None:
            try:
                shutil.rmtree(tmp_dir)
            except OSError as exc:
                logger.debug("Could not clean temp dir %s: %s", tmp_dir, exc)

    title, artist = _extract_metadata_from_tags(audio_path)

    # Build caption string for ACE-Step training
    parts = ["A"]
    if artist:
        parts.append(f"track by {artist}")
    else:
        parts.append("track")
    if bpm:
        parts.append(f"at {bpm} BPM")
    if key:
        parts.append(f"in {key}")
    parts.append(f"{signature} time")
    caption = " ".join(parts)

    lyrics = "[Instrumental]"

    result = {
        "caption": caption,
        "bpm": bpm,
        "key": key,
        "signature": signature,
        "lyrics": lyrics,
        "title": title,
        "artist": artist,
        "confidence": confidence,
    }

    logger.info("Auto-caption [%s] for %s: %s", mode, audio_path.name, caption)
    return result


def _write_caption_sidecar(audio_path: Path, analysis: Dict[str, Any]) -> Path:
    """Write analysis results as a .json sidecar next to the audio file."""
    sidecar_path = audio_path.with_suffix(".json")
    with open(sidecar_path, "w", encoding="utf-8") as f:
        json.dump(analysis, f, indent=2, ensure_ascii=False)
    logger.info("Wrote caption sidecar: %s", sidecar_path)
    return sidecar_path


def _read_caption_sidecar(audio_path: Path) -> Optional[Dict[str, Any]]:
    """Read an existing .json caption sidecar if it exists."""
    sidecar_path = audio_path.with_suffix(".json")
    if not sidecar_path.is_file():
        return None
    try:
        with open(sidecar_path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


# ============================================================================
# PREPROCESSING (2-pass sequential)
# ============================================================================

def preprocess_audio(
    audio_dir: str,
    output_dir: str,
    checkpoint_dir: str,
    device: str = "cpu",
    variant: str = "base",
    max_duration: float = 0,
    progress_callback: Optional[Callable] = None,
    cancel_check: Optional[Callable] = None,
) -> Dict[str, Any]:
    """2-pass sequential preprocessing.

    Pass 1: Load VAE + text encoder, encode audio + text, save intermediates.
    Pass 2: Load DIT model, run encoder, build context, save final .pt files.
    """
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)

    # Clean orphaned staging files
    for orphan in out.glob("*.__writing__"):
        try:
            orphan.unlink()
        except OSError:
            pass

    audio_files = _discover_audio_files(audio_dir)
    if not audio_files:
        return {"processed": 0, "failed": 0, "total": 0, "output_dir": str(out)}

    total = len(audio_files)

    if max_duration <= 0:
        max_duration = _detect_max_duration(audio_files)

    dtype = CPU_DTYPE if device == "cpu" else torch.bfloat16

    # ---- Pass 1: VAE + Text Encoder ----
    logger.info("Pass 1/2: Loading VAE + Text Encoder...")
    vae = load_vae(checkpoint_dir, device)
    tokenizer, text_enc = load_text_encoder(checkpoint_dir, device)
    silence_lat = load_silence_latent(checkpoint_dir, device, variant=variant)

    intermediates: List[Path] = []
    p1_failed = 0

    try:
        for i, af in enumerate(audio_files):
            if cancel_check and cancel_check():
                break

            stem = af.stem
            final_pt = out / f"{stem}.pt"
            if final_pt.exists():
                continue

            try:
                audio, _ = load_audio_stereo(str(af), TARGET_SR, max_duration)
                audio = audio.unsqueeze(0).to(device=device, dtype=vae.dtype)

                with torch.no_grad():
                    target_latents = tiled_vae_encode(vae, audio, dtype)
                del audio

                if torch.isnan(target_latents).any() or torch.isinf(target_latents).any():
                    p1_failed += 1
                    del target_latents
                    continue

                lat_len = target_latents.shape[1]
                att_mask = torch.ones(1, lat_len, device=device, dtype=dtype)

                # Auto-caption: read existing sidecar or analyze
                sidecar = _read_caption_sidecar(af)
                if sidecar and sidecar.get("caption"):
                    caption = sidecar["caption"]
                    lyrics = sidecar.get("lyrics", "[Instrumental]")
                    logger.info("[Caption] %s: using existing sidecar", af.name)
                else:
                    # Auto-select analysis mode based on dataset size
                    if total <= 20:
                        analysis_mode = "sas"
                    elif total <= 100:
                        analysis_mode = "mid"
                    else:
                        analysis_mode = "faf"

                    # Log mode selection with reasoning (first file only)
                    if i == 0:
                        _MODE_DESC = {
                            "faf": "fast, ~3s/file",
                            "mid": "balanced, ~12s/file",
                            "sas": "best quality, ~30s/file on GPU, slower on CPU",
                        }
                        logger.info(
                            "[Analysis] Mode auto-selected: '%s' (%s) "
                            "for %d files (<=20: sas, 21-100: mid, 100+: faf)",
                            analysis_mode, _MODE_DESC[analysis_mode], total,
                        )
                        if analysis_mode in ("mid", "sas") and device == "cpu":
                            logger.warning(
                                "[Analysis] Mode '%s' uses Demucs stem separation "
                                "which is SLOW on CPU (~2-5 min/file). "
                                "Total estimated time: ~%d-%d min for %d files. "
                                "Use 'faf' mode or a GPU machine for faster processing.",
                                analysis_mode,
                                total * 2, total * 5, total,
                            )

                    try:
                        logger.info("[Caption] %s: analyzing (mode=%s)...", af.name, analysis_mode)
                        analysis = analyze_and_caption(
                            str(af), mode=analysis_mode, device=device,
                        )
                        caption = analysis["caption"]
                        lyrics = analysis.get("lyrics", "[Instrumental]")
                        _write_caption_sidecar(af, analysis)
                        logger.info("[Caption] %s: %s", af.name, caption)
                    except Exception as exc:
                        logger.warning("[Caption] %s: analysis failed (%s), using filename", af.name, exc)
                        caption = af.stem
                        lyrics = "[Instrumental]"
                text_prompt = caption

                with torch.no_grad():
                    text_hs, text_mask = encode_text(text_enc, tokenizer, text_prompt, device, dtype)
                    lyric_hs, lyric_mask = encode_lyrics(text_enc, tokenizer, lyrics, device, dtype)

                has_bad = any(
                    torch.isnan(t).any() or torch.isinf(t).any()
                    for t in [text_hs, lyric_hs]
                )
                if has_bad:
                    p1_failed += 1
                    del target_latents, att_mask, text_hs, text_mask, lyric_hs, lyric_mask
                    continue

                tmp_path = out / f"{stem}.tmp.pt"
                torch.save({
                    "target_latents": target_latents.squeeze(0).cpu(),
                    "attention_mask": att_mask.squeeze(0).cpu(),
                    "text_hidden_states": text_hs.cpu(),
                    "text_attention_mask": text_mask.cpu(),
                    "lyric_hidden_states": lyric_hs.cpu(),
                    "lyric_attention_mask": lyric_mask.cpu(),
                    "silence_latent": silence_lat.cpu(),
                    "latent_length": lat_len,
                    "metadata": {
                        "audio_path": str(af),
                        "filename": af.name,
                        "caption": caption,
                        "lyrics": lyrics,
                    },
                }, tmp_path)

                del target_latents, att_mask, text_hs, text_mask, lyric_hs, lyric_mask
                intermediates.append(tmp_path)

                if progress_callback:
                    progress_callback(i + 1, total, f"[Pass 1] {af.name}")

            except Exception as exc:
                p1_failed += 1
                logger.error("Pass 1 FAIL %s: %s", af.name, exc)
    finally:
        logger.info("Unloading VAE + Text Encoder...")
        unload_models(vae, text_enc, tokenizer, silence_lat)

    # ---- Pass 2: DIT Encoder ----
    if not intermediates:
        return {"processed": 0, "failed": p1_failed, "total": total, "output_dir": str(out)}

    logger.info("Pass 2/2: Loading DIT model (variant=%s)...", variant)
    model = load_model_for_training(checkpoint_dir, variant, device)

    processed = 0
    p2_failed = 0
    p2_total = len(intermediates)

    try:
        for i, tmp_path in enumerate(intermediates):
            if cancel_check and cancel_check():
                break

            try:
                data = torch.load(str(tmp_path), weights_only=True)
                m_device = next(model.parameters()).device
                m_dtype = next(model.parameters()).dtype

                text_hs = data["text_hidden_states"].to(m_device, dtype=m_dtype)
                text_mask = data["text_attention_mask"].to(m_device, dtype=m_dtype)
                lyric_hs = data["lyric_hidden_states"].to(m_device, dtype=m_dtype)
                lyric_mask = data["lyric_attention_mask"].to(m_device, dtype=m_dtype)
                silence_lat = data["silence_latent"].to(m_device, dtype=m_dtype)
                lat_len = data["latent_length"]

                enc_hs, enc_mask = run_encoder(
                    model, text_hs, text_mask, lyric_hs, lyric_mask,
                    str(m_device), m_dtype,
                )
                del text_hs, text_mask, lyric_hs, lyric_mask

                if silence_lat.dim() == 2:
                    silence_lat = silence_lat.unsqueeze(0)
                ctx = build_context_latents(silence_lat, lat_len, str(m_device), m_dtype)
                del silence_lat

                has_bad = any(
                    torch.isnan(t).any() or torch.isinf(t).any()
                    for t in [enc_hs, ctx]
                )
                if has_bad:
                    p2_failed += 1
                    del enc_hs, enc_mask, ctx, data
                    continue

                base_name = tmp_path.name.replace(".tmp.pt", ".pt")
                final_path = out / base_name
                staging_path = out / (base_name + ".__writing__")

                torch.save({
                    "target_latents": data["target_latents"],
                    "attention_mask": data["attention_mask"],
                    "encoder_hidden_states": enc_hs.squeeze(0).cpu(),
                    "encoder_attention_mask": enc_mask.squeeze(0).cpu(),
                    "context_latents": ctx.squeeze(0).cpu(),
                    "metadata": data.get("metadata", {}),
                }, staging_path)
                os.replace(staging_path, final_path)

                del enc_hs, enc_mask, ctx, data
                tmp_path.unlink(missing_ok=True)
                processed += 1

                if progress_callback:
                    progress_callback(i + 1, p2_total, f"[Pass 2] {tmp_path.stem}")

            except Exception as exc:
                p2_failed += 1
                logger.error("Pass 2 FAIL %s: %s", tmp_path.stem, exc)
    finally:
        logger.info("Unloading DIT model...")
        unload_models(model)

    failed = p1_failed + p2_failed
    return {"processed": processed, "failed": failed, "total": total, "output_dir": str(out)}


# ============================================================================
# TRAINING LOOP (generator for Gradio compatibility)
# ============================================================================

def train_lora_generator(
    dataset_dir: str,
    output_dir: str,
    checkpoint_dir: str,
    epochs: int = 1000,
    lr: float = 3e-4,
    rank: int = 64,
    alpha: int = 128,
    dropout: float = 0.1,
    batch_size: int = 1,
    gradient_accumulation_steps: int = 4,
    warmup_steps: int = 100,
    weight_decay: float = 0.01,
    max_grad_norm: float = 1.0,
    save_every_n_epochs: int = 0,
    seed: int = 42,
    variant: str = "base",
    device: str = "cpu",
    cfg_ratio: float = 0.15,
    timestep_mu: float = -0.4,
    timestep_sigma: float = 1.0,
    target_modules: Optional[List[str]] = None,
    log_every: int = 10,
    resume_from: Optional[str] = None,
) -> Generator[str, None, None]:
    """Run LoRA training, yielding progress strings each epoch.

    This is a generator for Gradio live-update compatibility.
    Call cancel_training() to stop after the current epoch.
    """
    _training_cancel.clear()
    train_start = time.time()

    if target_modules is None:
        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"]

    ds_path = Path(dataset_dir)
    if not ds_path.is_dir():
        yield f"[FAIL] Dataset directory not found: {ds_path}"
        return

    out_path = Path(output_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    yield "[INFO] Loading model..."

    try:
        model = load_model_for_training(checkpoint_dir, variant, device)
    except Exception as exc:
        yield f"[FAIL] Model load failed: {exc}"
        return

    # float32 on CPU (bfloat16 deadlocks)
    dtype = CPU_DTYPE if device == "cpu" else torch.bfloat16
    model = model.to(dtype=dtype)

    yield "[INFO] Injecting LoRA..."

    lora_cfg = LoRAConfig(
        r=rank, alpha=alpha, dropout=dropout,
        target_modules=target_modules, bias="none",
    )

    try:
        model, info = inject_lora(model, lora_cfg)
    except Exception as exc:
        yield f"[FAIL] LoRA injection failed: {exc}"
        unload_models(model)
        return

    yield f"[OK] LoRA injected: {info['trainable_params']:,} trainable params"

    # Gradient checkpointing + cache disable
    force_disable_cache(model.decoder)
    ckpt_ok = enable_gradient_checkpointing(model.decoder)
    force_input_grads = ckpt_ok
    if ckpt_ok:
        yield "[INFO] Gradient checkpointing enabled"

    # Dataset
    dataset = TensorDataset(dataset_dir)
    if len(dataset) == 0:
        yield "[FAIL] No valid .pt files found in dataset directory"
        unload_models(model)
        return

    yield f"[OK] Loaded {len(dataset)} preprocessed samples"

    loader = DataLoader(
        dataset, batch_size=batch_size, shuffle=True,
        num_workers=0, collate_fn=_collate_batch, drop_last=False,
    )

    # Optimizer & scheduler
    torch.manual_seed(seed)
    random.seed(seed)

    trainable_params = [p for p in model.parameters() if p.requires_grad]
    if not trainable_params:
        yield "[FAIL] No trainable parameters found"
        unload_models(model)
        return

    optimizer = build_optimizer(trainable_params, lr=lr, weight_decay=weight_decay)
    steps_per_epoch = max(1, math.ceil(len(loader) / gradient_accumulation_steps))
    total_steps = steps_per_epoch * epochs
    scheduler = build_scheduler(optimizer, total_steps, warmup_steps, lr)

    yield f"[INFO] Training {sum(p.numel() for p in trainable_params):,} params for {epochs} epochs"
    yield f"[INFO] Steps/epoch: {steps_per_epoch}, total: {total_steps}"

    # Null condition embedding for CFG dropout
    null_cond = getattr(model, "null_condition_emb", None)

    # Resume checkpoint
    start_epoch = 0
    global_step = 0

    if resume_from and Path(resume_from).exists():
        try:
            yield f"[INFO] Resuming from {resume_from}"
            ckpt_dir = Path(resume_from)
            if ckpt_dir.is_file():
                ckpt_dir = ckpt_dir.parent

            # Load adapter weights
            aw = ckpt_dir / "adapter_model.safetensors"
            if aw.exists():
                from safetensors.torch import load_file
                state = load_file(str(aw))
                decoder = model.decoder
                while hasattr(decoder, "_forward_module"):
                    decoder = decoder._forward_module
                decoder.load_state_dict(state, strict=False)

            # Load training state
            ts = ckpt_dir / "training_state.pt"
            if ts.exists():
                tstate = torch.load(str(ts), map_location=device, weights_only=True)
                start_epoch = tstate.get("epoch", 0)
                global_step = tstate.get("global_step", 0)
                if "optimizer_state_dict" in tstate:
                    try:
                        optimizer.load_state_dict(tstate["optimizer_state_dict"])
                    except Exception:
                        pass
                if "scheduler_state_dict" in tstate:
                    try:
                        scheduler.load_state_dict(tstate["scheduler_state_dict"])
                    except Exception:
                        pass

            yield f"[OK] Resumed from epoch {start_epoch}, step {global_step}"
        except Exception as exc:
            yield f"[WARN] Checkpoint load failed: {exc}, starting fresh"
            start_epoch = 0
            global_step = 0

    # Training loop
    model.decoder.train()
    acc_step = 0
    acc_loss = 0.0
    optimizer.zero_grad(set_to_none=True)

    best_loss = float("inf")
    best_epoch = 0
    consecutive_nan = 0
    MAX_NAN = 10

    for epoch in range(start_epoch, epochs):
        # Cancel check
        if _training_cancel.is_set():
            _training_cancel.clear()
            early_path = str(out_path / "early_exit")
            model.decoder.eval()
            save_lora_adapter(model, early_path)
            model.decoder.train()
            yield f"[OK] Cancelled at epoch {epoch + 1}, saved to {early_path}"
            yield "[DONE]"
            unload_models(model)
            return

        # Timeout check
        elapsed = time.time() - train_start
        if elapsed > MAX_TRAINING_TIME:
            early_path = str(out_path / "timeout_exit")
            model.decoder.eval()
            save_lora_adapter(model, early_path)
            yield f"[WARN] Training timed out after {int(elapsed)}s, saved to {early_path}"
            yield "[DONE]"
            unload_models(model)
            return

        epoch_loss = 0.0
        num_updates = 0
        epoch_start = time.time()

        for batch in loader:
            # Forward
            nb = device != "cpu"
            tgt = batch["target_latents"].to(device, dtype=dtype, non_blocking=nb)
            att = batch["attention_mask"].to(device, dtype=dtype, non_blocking=nb)
            enc_hs = batch["encoder_hidden_states"].to(device, dtype=dtype, non_blocking=nb)
            enc_mask = batch["encoder_attention_mask"].to(device, dtype=dtype, non_blocking=nb)
            ctx = batch["context_latents"].to(device, dtype=dtype, non_blocking=nb)

            bsz = tgt.shape[0]

            # CFG dropout
            if null_cond is not None and cfg_ratio > 0:
                enc_hs = apply_cfg_dropout(enc_hs, null_cond, cfg_ratio)

            # Timestep sampling
            t, _r = sample_timesteps(bsz, torch.device(device), dtype, timestep_mu, timestep_sigma)

            # Flow matching noise
            x1 = torch.randn_like(tgt)
            x0 = tgt
            t_ = t.unsqueeze(-1).unsqueeze(-1)
            xt = t_ * x1 + (1.0 - t_) * x0

            if force_input_grads:
                xt = xt.requires_grad_(True)

            # Decoder forward
            dec_out = model.decoder(
                hidden_states=xt,
                timestep=t,
                timestep_r=t,
                attention_mask=att,
                encoder_hidden_states=enc_hs,
                encoder_attention_mask=enc_mask,
                context_latents=ctx,
            )

            flow = x1 - x0
            loss = F.mse_loss(dec_out[0], flow)
            loss = loss.float()  # fp32 for stable backward

            # NaN guard
            if torch.isnan(loss) or torch.isinf(loss):
                consecutive_nan += 1
                del loss, tgt, att, enc_hs, enc_mask, ctx, xt, dec_out, flow
                if consecutive_nan >= MAX_NAN:
                    yield f"[FAIL] {consecutive_nan} consecutive NaN losses, halting"
                    unload_models(model)
                    return
                if acc_step > 0:
                    optimizer.zero_grad(set_to_none=True)
                    acc_loss = 0.0
                    acc_step = 0
                continue
            consecutive_nan = 0

            loss = loss / gradient_accumulation_steps
            loss.backward()
            acc_loss += loss.item()
            del loss, tgt, att, enc_hs, enc_mask, ctx, xt, dec_out, flow
            acc_step += 1

            if acc_step >= gradient_accumulation_steps:
                torch.nn.utils.clip_grad_norm_(trainable_params, max_grad_norm)
                optimizer.step()
                scheduler.step()
                global_step += 1

                avg_loss = acc_loss * gradient_accumulation_steps / acc_step

                if global_step % log_every == 0:
                    current_lr = scheduler.get_last_lr()[0]
                    yield (
                        f"Epoch {epoch + 1}/{epochs}, "
                        f"Step {global_step}, "
                        f"Loss: {avg_loss:.4f}, "
                        f"LR: {current_lr:.2e}"
                    )

                optimizer.zero_grad(set_to_none=True)
                epoch_loss += avg_loss
                num_updates += 1
                acc_loss = 0.0
                acc_step = 0

        # Flush remainder
        if acc_step > 0:
            torch.nn.utils.clip_grad_norm_(trainable_params, max_grad_norm)
            optimizer.step()
            scheduler.step()
            global_step += 1
            avg_loss = acc_loss * gradient_accumulation_steps / acc_step
            optimizer.zero_grad(set_to_none=True)
            epoch_loss += avg_loss
            num_updates += 1
            acc_loss = 0.0
            acc_step = 0

        epoch_time = time.time() - epoch_start
        avg_epoch_loss = epoch_loss / max(num_updates, 1)

        is_best = avg_epoch_loss < best_loss - 0.001
        if is_best:
            best_loss = avg_epoch_loss
            best_epoch = epoch + 1

        best_str = f" (best: {best_loss:.4f} @ ep{best_epoch})" if best_epoch > 0 else ""
        yield (
            f"[OK] Epoch {epoch + 1}/{epochs} in {epoch_time:.1f}s, "
            f"Loss: {avg_epoch_loss:.4f}{best_str}"
        )

        # Save best
        if is_best and epoch + 1 >= 10:
            best_path = str(out_path / "best")
            model.decoder.eval()
            save_lora_adapter(model, best_path)
            model.decoder.train()
            yield f"[OK] Best model saved (epoch {epoch + 1}, loss: {best_loss:.4f})"

        # Periodic checkpoint (0 = disabled, only save on cancel/finish)
        if save_every_n_epochs > 0 and (epoch + 1) % save_every_n_epochs == 0:
            ckpt_path = str(out_path / "checkpoints" / f"epoch_{epoch + 1}")
            model.decoder.eval()
            save_lora_adapter(model, ckpt_path)

            tstate = {
                "epoch": epoch + 1,
                "global_step": global_step,
                "optimizer_state_dict": optimizer.state_dict(),
                "scheduler_state_dict": scheduler.state_dict(),
            }
            os.makedirs(ckpt_path, exist_ok=True)
            torch.save(tstate, str(Path(ckpt_path) / "training_state.pt"))
            model.decoder.train()
            yield f"[OK] Checkpoint saved at epoch {epoch + 1}"

    # Sanity check
    if global_step == 0:
        yield "[FAIL] Training completed 0 steps -- no batches processed"
        unload_models(model)
        return

    # Final save (directly to output_dir, not a subdirectory)
    model.decoder.eval()
    save_lora_adapter(model, str(out_path))

    final_loss = avg_epoch_loss if num_updates > 0 else 0.0
    best_note = ""
    if best_epoch > 0 and Path(out_path / "best").exists():
        best_note = f"\n  Best: {out_path / 'best'} (epoch {best_epoch}, loss: {best_loss:.4f})"
    yield (
        f"[OK] Training complete! LoRA saved to {out_path}{best_note}\n"
        f"  Adapter ready for inference."
    )
    yield "[DONE]"
    unload_models(model)


# ============================================================================
# ADAPTER LISTING
# ============================================================================

def get_trained_loras(adapter_dir: str) -> List[str]:
    """List all saved LoRA adapter directories under adapter_dir."""
    result = []
    base = Path(adapter_dir)
    if not base.is_dir():
        return result

    for root, dirs, files in os.walk(str(base)):
        for f in files:
            if f in ("adapter_config.json", "adapter_model.safetensors", "lora_weights.pt"):
                result.append(root)
                break

    return sorted(set(result))