model.py

import math
import os
from dataclasses import dataclass
from typing import Any, Mapping, Optional, Sequence

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio

from data.encodec_utils import (
    decode_codes_to_audio_b1t,
    decode_quantized_latent_to_audio,
    load_target_pca_basis,
    reconstruct_latent_from_pca,
    requantize_latent_to_codes_bct,
    resolve_audio_codec_sample_rate,
    rvq_sum_latents,
    token_ids_to_codebook_embeddings,
)
from data.diffusion_dataset import estimate_target_normalization
from data.seconds_frontend import build_seconds_frontend_from_cfg


# Derived from scripts/analyze_frontend_radii.py on 60k train frames of 4beats_v9.
DEFAULT_FRONTEND_RADII: tuple[int, ...] = (0, 22, 41, 55)
DEFAULT_FRONTEND_PRIMARY_RADIUS = 22
DEFAULT_FRONTEND_VARIANT = "hybrid"
DEFAULT_FRONTEND_EMBED_DIM = 64
DEFAULT_FRONTEND_OUTPUT_KIND = "feat"
DEFAULT_FRONTEND_PADDING_MODE = "reflect"
DEFAULT_FRONTEND_STEP_SECONDS = 0.0
DEFAULT_FRONTEND_CHUNK_SIZE = 0
DEFAULT_FRONTEND_CLASS_LOCAL_DIM = 8
DEFAULT_FRONTEND_CONCAT_MULTISCALE = True
DEFAULT_SAMPLE_X0_CLIP_NORM = 6.0
DEFAULT_AUDIO_WAVE_L1_WEIGHT = 0.0
DEFAULT_AUDIO_MRSTFT_WEIGHT = 0.0
DEFAULT_AUDIO_MRSTFT_RESOLUTIONS: tuple[tuple[int, int], ...] = (
    (512, 128),
    (1024, 256),
    (2048, 512),
)
DEFAULT_INFERENCE_NUM_BEATS = 4
DEFAULT_BEAT_CROSSFADE_MS = 10.0
DEFAULT_TARGET_TOKEN_RATE_HZ = 50.0
DEFAULT_INFERENCE_GUIDANCE_SCALE = 1.0
DEFAULT_POSITIONAL_ENCODING = "seconds"
TIMBRE_NUM_FAMILIES = 8
TIMBRE_MAX_CLASSES = 5
TIMBRE_CLASS_VOCAB_SIZES: tuple[int, ...] = (1, 3, 2, 2, 2, 5, 2, 3)



def masked_mean(x: torch.Tensor, mask: torch.Tensor, dim=None, eps: float = 1e-8):
    mask_f = mask.float()
    if dim is None:
        return (x * mask_f).sum() / mask_f.sum().clamp_min(eps)
    return (x * mask_f).sum(dim=dim) / mask_f.sum(dim=dim).clamp_min(eps)


def apply_seq_mask(x: torch.Tensor, valid_mask_bt: torch.Tensor) -> torch.Tensor:
    return x * valid_mask_bt.unsqueeze(-1).to(x.dtype)


def _normalize_stats_vector(
    value,
    *,
    x_dim: int,
    device: torch.device,
    default_fill: float,
    name: str,
):
    if value is None:
        return torch.full((x_dim,), float(default_fill), dtype=torch.float32, device=device)
    tensor = torch.as_tensor(value, dtype=torch.float32, device=device).view(-1)
    if int(tensor.numel()) != int(x_dim):
        raise ValueError(f"{name} must have {x_dim} values, got {tuple(tensor.shape)}")
    return tensor.contiguous()


def normalize_latent(x, mean, std):
    resolved_mean = _normalize_stats_vector(
        mean,
        x_dim=int(x.shape[-1]),
        device=x.device,
        default_fill=0.0,
        name="target_mean",
    )
    resolved_std = _normalize_stats_vector(
        std,
        x_dim=int(x.shape[-1]),
        device=x.device,
        default_fill=1.0,
        name="target_std",
    ).clamp_min(1.0e-8)
    return (x - resolved_mean.view(1, 1, -1)) / resolved_std.view(1, 1, -1)


def denormalize_latent(x, mean, std):
    resolved_mean = _normalize_stats_vector(
        mean,
        x_dim=int(x.shape[-1]),
        device=x.device,
        default_fill=0.0,
        name="target_mean",
    )
    resolved_std = _normalize_stats_vector(
        std,
        x_dim=int(x.shape[-1]),
        device=x.device,
        default_fill=1.0,
        name="target_std",
    ).clamp_min(1.0e-8)
    return x * resolved_std.view(1, 1, -1) + resolved_mean.view(1, 1, -1)


def timestep_embedding(timesteps: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
    half = dim // 2
    freqs = torch.exp(
        -math.log(max_period) * torch.arange(0, half, dtype=torch.float32, device=timesteps.device) / half
    )
    args = timesteps.float().unsqueeze(1) * freqs.unsqueeze(0)
    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
    if dim % 2:
        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
    return emb


def sinusoidal_positions(length: int, dim: int, device: torch.device) -> torch.Tensor:
    position = torch.arange(length, device=device, dtype=torch.float32).unsqueeze(1)
    half = dim // 2
    div_term = torch.exp(
        torch.arange(0, half, device=device, dtype=torch.float32) * (-math.log(10000.0) / half)
    )
    pe = torch.zeros(length, dim, device=device, dtype=torch.float32)
    pe[:, 0:half] = torch.sin(position * div_term)
    pe[:, half:2 * half] = torch.cos(position * div_term)
    if dim % 2:
        pe[:, -1] = 0
    return pe.unsqueeze(0)


def sinusoidal_time_positions(
    times_sec_bt: torch.Tensor,
    dim: int,
    *,
    rate_hz: float,
) -> torch.Tensor:
    times = torch.as_tensor(times_sec_bt, dtype=torch.float32)
    if int(times.dim()) != 2:
        raise ValueError(f"times_sec_bt must be [B,T], got {tuple(times.shape)}")
    position = times.unsqueeze(-1) * float(max(1.0e-6, float(rate_hz)))
    half = dim // 2
    div_term = torch.exp(
        torch.arange(0, half, device=times.device, dtype=torch.float32) * (-math.log(10000.0) / half)
    ).view(1, 1, -1)
    pe = torch.zeros(
        int(times.shape[0]),
        int(times.shape[1]),
        int(dim),
        device=times.device,
        dtype=torch.float32,
    )
    pe[:, :, 0:half] = torch.sin(position * div_term)
    pe[:, :, half:2 * half] = torch.cos(position * div_term)
    if dim % 2:
        pe[:, :, -1] = 0
    return pe


def build_frontend_cfg_from_batch(
    batch: Mapping[str, Any],
    *,
    variant: str = DEFAULT_FRONTEND_VARIANT,
    embed_dim: int = DEFAULT_FRONTEND_EMBED_DIM,
    output_kind: str = DEFAULT_FRONTEND_OUTPUT_KIND,
    radii: Sequence[int] = DEFAULT_FRONTEND_RADII,
    primary_radius: int = DEFAULT_FRONTEND_PRIMARY_RADIUS,
    padding_mode: str = DEFAULT_FRONTEND_PADDING_MODE,
    step_seconds: float = DEFAULT_FRONTEND_STEP_SECONDS,
    chunk_size: int = DEFAULT_FRONTEND_CHUNK_SIZE,
    class_local_fusion: bool = False,
    class_local_dim: int = DEFAULT_FRONTEND_CLASS_LOCAL_DIM,
) -> dict[str, Any]:
    radii_eff = [int(x) for x in list(radii or ()) if int(x) >= 0]
    if not radii_eff:
        radii_eff = [int(primary_radius)]
    grid = torch.as_tensor(batch["grid"])
    return {
        "input_dim_source": int(grid.shape[1]),
        "class_id_vocab_sizes": [int(x) for x in list(batch.get("class_id_vocab_sizes") or [])],
        "source_feature_names": [str(x) for x in list(batch.get("feature_row_names") or [])],
        "class_names": [str(x) for x in list(batch.get("class_names") or [])],
        "variant": str(variant),
        "embed_dim": int(embed_dim),
        "output_kind": str(output_kind),
        "multiscale_enabled": bool(len(radii_eff) > 1),
        "multiscale_radii": [int(x) for x in list(radii_eff)],
        "primary_radius": int(primary_radius),
        "window_radius": int(primary_radius),
        "padding_mode": str(padding_mode),
        "step_seconds": float(step_seconds),
        "chunk_size": int(chunk_size),
        "class_local_fusion": bool(class_local_fusion),
        "class_local_dim": int(class_local_dim),
    }


def _prepare_batch_tensors(
    batch: Mapping[str, Any],
    device: torch.device,
    *,
    require_target: bool = True,
    require_timing: bool = True,
) -> dict[str, torch.Tensor | None]:
    def _tensor(key: str, dtype: torch.dtype, *, required: bool = True) -> torch.Tensor | None:
        value = batch.get(key)
        if value is None:
            if bool(required):
                raise KeyError(f"batch is missing required key: {key}")
            return None
        return torch.as_tensor(value, device=device, dtype=dtype).contiguous()

    return {
        "grid": _tensor("grid", torch.float32),
        "grid_ids": _tensor("grid_ids", torch.long, required=False),
        "family_onsets_bft": _tensor("family_onsets_bft", torch.bool, required=False),
        "grid_valid_mask": _tensor("grid_valid_mask", torch.bool),
        "grid_times_sec": _tensor("grid_times_sec", torch.float32, required=require_timing),
        "token_times_sec": _tensor("token_times_sec", torch.float32, required=require_timing),
        "beat_boundaries_sec": _tensor("beat_boundaries_sec", torch.float32, required=False),
        "beat_boundaries_valid_mask": _tensor("beat_boundaries_valid_mask", torch.bool, required=False),
        "bpm": _tensor("bpm", torch.float32, required=False),
        "duration_sec": _tensor("duration_sec", torch.float32, required=False),
        "target_btd": _tensor("target_btd", torch.float32, required=require_target),
        "target_sum_btd": _tensor("target_sum_btd", torch.float32, required=False),
        "target_valid_mask_bt": _tensor("target_valid_mask_bt", torch.bool, required=require_target),
        "source_codes_bct": _tensor("source_codes_bct", torch.long, required=False),
        "timbre_bank_latents": _tensor("timbre_bank_latents", torch.float32, required=False),
        "timbre_bank_family_ids": _tensor("timbre_bank_family_ids", torch.long, required=False),
        "timbre_bank_class_ids": _tensor("timbre_bank_class_ids", torch.long, required=False),
        "timbre_bank_velocity": _tensor("timbre_bank_velocity", torch.float32, required=False),
        "timbre_bank_mask": _tensor("timbre_bank_mask", torch.bool, required=False),
        "timbre_dynamic_features": _tensor("timbre_dynamic_features", torch.float32, required=False),
        "timbre_dynamic_mask": _tensor("timbre_dynamic_mask", torch.bool, required=False),
        "timbre_dynamic_counts": _tensor("timbre_dynamic_counts", torch.float32, required=False),
        "timbre_family_default_indices": _tensor("timbre_family_default_indices", torch.long, required=False),
        "timbre_class_token_indices": _tensor("timbre_class_token_indices", torch.long, required=False),
        "reference_timbre_bank_latents": _tensor("reference_timbre_bank_latents", torch.float32, required=False),
        "reference_timbre_bank_family_ids": _tensor("reference_timbre_bank_family_ids", torch.long, required=False),
        "reference_timbre_bank_class_ids": _tensor("reference_timbre_bank_class_ids", torch.long, required=False),
        "reference_timbre_bank_velocity": _tensor("reference_timbre_bank_velocity", torch.float32, required=False),
        "reference_timbre_bank_mask": _tensor("reference_timbre_bank_mask", torch.bool, required=False),
        "reference_timbre_dynamic_features": _tensor("reference_timbre_dynamic_features", torch.float32, required=False),
        "reference_timbre_dynamic_mask": _tensor("reference_timbre_dynamic_mask", torch.bool, required=False),
        "reference_timbre_dynamic_counts": _tensor("reference_timbre_dynamic_counts", torch.float32, required=False),
        "reference_timbre_family_default_indices": _tensor("reference_timbre_family_default_indices", torch.long, required=False),
        "reference_timbre_class_token_indices": _tensor("reference_timbre_class_token_indices", torch.long, required=False),
        "reference_segment_pca144": _tensor("reference_segment_pca144", torch.float32, required=False),
        "x0_prior_btd": _tensor("x0_prior_btd", torch.float32, required=False),
    }


def _slice_prepared_batch(
    prepared: Mapping[str, torch.Tensor | None],
    sample_idx: int,
) -> dict[str, torch.Tensor | None]:
    return {
        key: (
            None
            if value is None
            else value[int(sample_idx) : int(sample_idx) + 1].contiguous()
        )
        for key, value in prepared.items()
    }


def _prepare_geometry_tensors(
    geometry: Mapping[str, Any],
    *,
    device: torch.device,
) -> dict[str, torch.Tensor]:
    return {
        str(key): torch.as_tensor(value, device=device).contiguous()
        for key, value in geometry.items()
    }


def _slice_inference_geometry(
    geometry: Mapping[str, torch.Tensor],
    sample_idx: int,
) -> dict[str, torch.Tensor]:
    sliced: dict[str, torch.Tensor] = {}
    for key, value in geometry.items():
        tensor = torch.as_tensor(value)
        if int(tensor.dim()) > 0:
            if int(tensor.shape[0]) <= int(sample_idx):
                raise IndexError(f"sample_idx={sample_idx} out of range for inference geometry key {key!r}")
            sliced[str(key)] = tensor[int(sample_idx) : int(sample_idx) + 1].contiguous()
        else:
            sliced[str(key)] = tensor.contiguous()
    return sliced


def lengths_to_mask(lengths_b: torch.Tensor, *, max_len: int | None = None) -> torch.Tensor:
    lengths = torch.as_tensor(lengths_b, dtype=torch.long).view(-1)
    if int(lengths.numel()) <= 0:
        resolved_max_len = int(max_len or 0)
        return torch.zeros((0, max(0, resolved_max_len)), dtype=torch.bool, device=lengths.device)
    resolved_max_len = int(max_len) if max_len is not None else int(lengths.max().item())
    if int(resolved_max_len) <= 0:
        return torch.zeros((int(lengths.shape[0]), 0), dtype=torch.bool, device=lengths.device)
    steps = torch.arange(int(resolved_max_len), device=lengths.device, dtype=torch.long).view(1, -1)
    return (steps < lengths.view(-1, 1)).contiguous()


def uniform_frame_times_from_durations(
    frame_counts_b: torch.Tensor,
    duration_sec_b: torch.Tensor,
    *,
    max_num_frames: int | None = None,
) -> torch.Tensor:
    frame_counts = torch.as_tensor(frame_counts_b, dtype=torch.long).view(-1)
    duration_sec = torch.as_tensor(duration_sec_b, dtype=torch.float32, device=frame_counts.device).view(-1)
    if tuple(frame_counts.shape) != tuple(duration_sec.shape):
        raise ValueError(
            f"frame_counts_b and duration_sec_b must match, got {tuple(frame_counts.shape)} / {tuple(duration_sec.shape)}"
        )
    resolved_max_frames = int(max_num_frames) if max_num_frames is not None else int(frame_counts.max().item())
    if int(resolved_max_frames) <= 0:
        return torch.zeros((int(frame_counts.shape[0]), 0), dtype=torch.float32, device=frame_counts.device)
    frame_counts_safe = frame_counts.clamp_min(1).to(dtype=torch.float32).view(-1, 1)
    frame_steps = torch.arange(int(resolved_max_frames), device=frame_counts.device, dtype=torch.float32).view(1, -1)
    centers = ((frame_steps + 0.5) / frame_counts_safe) * duration_sec.view(-1, 1)
    valid_mask_bt = lengths_to_mask(frame_counts, max_len=int(resolved_max_frames))
    return (centers * valid_mask_bt.to(dtype=centers.dtype)).contiguous()


def _metadata_get(codec_metadata: Mapping[str, Any] | Any, key: str, default: Any = None) -> Any:
    if isinstance(codec_metadata, Mapping):
        return codec_metadata.get(key, default)
    return getattr(codec_metadata, key, default)


def _metadata_positive_float(codec_metadata: Mapping[str, Any] | Any, key: str) -> float | None:
    try:
        value = float(_metadata_get(codec_metadata, key))
    except (TypeError, ValueError):
        return None
    if not math.isfinite(value) or value <= 0.0:
        return None
    return float(value)


def _metadata_positive_int(codec_metadata: Mapping[str, Any] | Any, key: str) -> int | None:
    try:
        value = int(_metadata_get(codec_metadata, key))
    except (TypeError, ValueError):
        return None
    if int(value) <= 0:
        return None
    return int(value)


def _legacy_dac_hop_length(codec_metadata: Mapping[str, Any] | Any) -> int | None:
    codec_family = str(_metadata_get(codec_metadata, "codec_family", "") or "").strip().lower()
    codec_model_id = str(_metadata_get(codec_metadata, "codec_model_id", "") or "").strip().lower()
    sample_rate = _metadata_positive_int(codec_metadata, "codec_sample_rate")
    if (
        codec_model_id == "descript/dac_44khz"
        and codec_family in {"", "dac"}
        and int(sample_rate or 0) == 44100
    ):
        return 512
    return None


def resolve_codec_hop_length(codec_metadata: Mapping[str, Any] | Any | None) -> int | None:
    if codec_metadata is None:
        return None
    hop_length = _metadata_positive_int(codec_metadata, "codec_hop_length")
    if hop_length is not None:
        return int(hop_length)
    hop_length = _metadata_positive_int(codec_metadata, "hop_length")
    if hop_length is not None:
        return int(hop_length)
    return _legacy_dac_hop_length(codec_metadata)


def resolve_target_token_rate_hz(
    codec_metadata: Mapping[str, Any] | Any | None,
    *,
    fallback: float = DEFAULT_TARGET_TOKEN_RATE_HZ,
) -> float:
    if codec_metadata is None:
        return float(fallback)
    sample_rate = _metadata_positive_float(codec_metadata, "codec_sample_rate")
    hop_length = resolve_codec_hop_length(codec_metadata)
    if sample_rate is not None and hop_length is not None:
        return float(sample_rate) / float(hop_length)
    rate = _metadata_positive_float(codec_metadata, "codec_frame_rate")
    if rate is None:
        rate = _metadata_positive_float(codec_metadata, "frame_rate")
    if rate is None:
        return float(fallback)
    return float(rate)


def uniform_beat_boundaries_from_durations(
    duration_sec_b: torch.Tensor,
    *,
    num_beats: int,
) -> torch.Tensor:
    num_beats_eff = int(max(1, int(num_beats)))
    duration_sec = torch.as_tensor(duration_sec_b, dtype=torch.float32).view(-1)
    fractions = torch.linspace(
        0.0,
        1.0,
        steps=int(num_beats_eff) + 1,
        device=duration_sec.device,
        dtype=duration_sec.dtype,
    ).view(1, -1)
    return (duration_sec.view(-1, 1) * fractions).contiguous()


def _resolve_duration_from_bpm(
    bpm_b: torch.Tensor | None,
    *,
    num_beats: int,
    fallback_duration_sec_b: torch.Tensor | None = None,
) -> torch.Tensor:
    if bpm_b is None:
        if fallback_duration_sec_b is None:
            raise ValueError("bpm is required when fallback_duration_sec_b is not provided")
        fallback = torch.as_tensor(fallback_duration_sec_b, dtype=torch.float32).view(-1)
        if not bool(torch.all(fallback > 0.0)):
            raise ValueError("fallback_duration_sec_b must be positive")
        return fallback.contiguous()

    bpm = torch.as_tensor(bpm_b, dtype=torch.float32).view(-1)
    duration_sec = torch.full_like(bpm, 0.0)
    valid_bpm_mask = bpm > 1.0e-6
    if bool(valid_bpm_mask.any()):
        duration_sec[valid_bpm_mask] = (float(max(1, int(num_beats))) * 60.0) / bpm[valid_bpm_mask]
    if bool((~valid_bpm_mask).any()):
        if fallback_duration_sec_b is None:
            raise ValueError("bpm must be positive for every example when fallback_duration_sec_b is not provided")
        fallback = torch.as_tensor(
            fallback_duration_sec_b,
            dtype=torch.float32,
            device=bpm.device,
        ).view(-1)
        if tuple(fallback.shape) != tuple(bpm.shape):
            raise ValueError(
                f"fallback_duration_sec_b must match bpm shape, got {tuple(fallback.shape)} / {tuple(bpm.shape)}"
            )
        fallback_invalid = fallback[~valid_bpm_mask]
        if not bool(torch.all(fallback_invalid > 0.0)):
            raise ValueError("fallback_duration_sec_b must be positive for examples with invalid bpm")
        duration_sec[~valid_bpm_mask] = fallback_invalid
    return duration_sec.contiguous()


def resolve_inference_geometry(
    prepared: Mapping[str, torch.Tensor | None],
    *,
    use_bpm_inference_geometry: bool = False,
    inference_num_beats: int = DEFAULT_INFERENCE_NUM_BEATS,
    target_token_rate_hz: float = DEFAULT_TARGET_TOKEN_RATE_HZ,
) -> dict[str, torch.Tensor]:
    grid = prepared.get("grid")
    grid_valid_mask = prepared.get("grid_valid_mask")
    if grid is None or grid_valid_mask is None:
        raise ValueError("prepared batch must include grid and grid_valid_mask")
    if not bool(use_bpm_inference_geometry):
        token_times_sec = prepared.get("token_times_sec")
        target_valid_mask_bt = prepared.get("target_valid_mask_bt")
        beat_boundaries_sec = prepared.get("beat_boundaries_sec")
        beat_boundaries_valid_mask = prepared.get("beat_boundaries_valid_mask")
        grid_times_sec = prepared.get("grid_times_sec")
        duration_sec = prepared.get("duration_sec")
        if (
            token_times_sec is None
            or target_valid_mask_bt is None
            or beat_boundaries_sec is None
            or beat_boundaries_valid_mask is None
            or grid_times_sec is None
            or duration_sec is None
        ):
            raise ValueError(
                "non-derived inference geometry requires grid_times_sec, token_times_sec, "
                "target_valid_mask_bt, beat_boundaries_sec, beat_boundaries_valid_mask, and duration_sec"
            )
        return {
            "grid_times_sec": grid_times_sec.contiguous(),
            "token_times_sec": token_times_sec.contiguous(),
            "target_valid_mask_bt": target_valid_mask_bt.to(dtype=torch.bool).contiguous(),
            "beat_boundaries_sec": beat_boundaries_sec.contiguous(),
            "beat_boundaries_valid_mask": beat_boundaries_valid_mask.to(dtype=torch.bool).contiguous(),
            "duration_sec": duration_sec.contiguous(),
            "target_num_frames_b": target_valid_mask_bt.to(dtype=torch.long).sum(dim=1).contiguous(),
        }

    duration_sec = _resolve_duration_from_bpm(
        prepared.get("bpm"),
        num_beats=int(inference_num_beats),
        fallback_duration_sec_b=prepared.get("duration_sec"),
    )
    grid_num_frames_b = grid_valid_mask.to(dtype=torch.long).sum(dim=1)
    target_num_frames_b = torch.round(duration_sec * float(max(1.0e-6, float(target_token_rate_hz)))).to(dtype=torch.long)
    target_num_frames_b = target_num_frames_b.clamp_min(1)
    max_target_len = int(target_num_frames_b.max().item())
    beat_boundaries_sec = uniform_beat_boundaries_from_durations(
        duration_sec,
        num_beats=int(inference_num_beats),
    )
    return {
        "grid_times_sec": uniform_frame_times_from_durations(
            grid_num_frames_b,
            duration_sec,
            max_num_frames=int(grid_valid_mask.shape[1]),
        ),
        "token_times_sec": uniform_frame_times_from_durations(
            target_num_frames_b,
            duration_sec,
            max_num_frames=int(max_target_len),
        ),
        "target_valid_mask_bt": lengths_to_mask(target_num_frames_b, max_len=int(max_target_len)),
        "beat_boundaries_sec": beat_boundaries_sec,
        "beat_boundaries_valid_mask": torch.ones_like(beat_boundaries_sec, dtype=torch.bool),
        "duration_sec": duration_sec.contiguous(),
        "target_num_frames_b": target_num_frames_b.contiguous(),
    }


def apply_bpm_training_geometry_to_prepared_batch(
    prepared: Mapping[str, torch.Tensor | None],
    *,
    num_beats: int = DEFAULT_INFERENCE_NUM_BEATS,
) -> dict[str, torch.Tensor | None]:
    """Retimes cached training tensors to BPM-derived durations without resizing targets."""
    grid_valid_mask = prepared.get("grid_valid_mask")
    target_valid_mask = prepared.get("target_valid_mask_bt")
    if grid_valid_mask is None or target_valid_mask is None:
        raise ValueError("BPM training geometry requires grid_valid_mask and target_valid_mask_bt")

    duration_sec = _resolve_duration_from_bpm(
        prepared.get("bpm"),
        num_beats=int(num_beats),
        fallback_duration_sec_b=prepared.get("duration_sec"),
    )
    grid_num_frames_b = torch.as_tensor(grid_valid_mask, dtype=torch.bool).to(dtype=torch.long).sum(dim=1)
    target_num_frames_b = torch.as_tensor(target_valid_mask, dtype=torch.bool).to(dtype=torch.long).sum(dim=1)
    beat_boundaries_sec = uniform_beat_boundaries_from_durations(
        duration_sec,
        num_beats=int(num_beats),
    )
    retimed = dict(prepared)
    retimed.update(
        {
            "grid_times_sec": uniform_frame_times_from_durations(
                grid_num_frames_b,
                duration_sec,
                max_num_frames=int(grid_valid_mask.shape[1]),
            ),
            "token_times_sec": uniform_frame_times_from_durations(
                target_num_frames_b,
                duration_sec,
                max_num_frames=int(target_valid_mask.shape[1]),
            ),
            "beat_boundaries_sec": beat_boundaries_sec,
            "beat_boundaries_valid_mask": torch.ones_like(beat_boundaries_sec, dtype=torch.bool),
            "duration_sec": duration_sec.contiguous(),
            "target_valid_mask_bt": torch.as_tensor(target_valid_mask, dtype=torch.bool).contiguous(),
        }
    )
    return retimed


class TimestepMLP(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.SiLU(),
            nn.Linear(d_model * 4, d_model),
        )

    def forward(self, t_emb: torch.Tensor) -> torch.Tensor:
        return self.net(t_emb)


class AdaLNModulation(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.SiLU(),
            nn.Linear(d_model, d_model * 9),
        )

    def forward(self, t_ctx: torch.Tensor):
        return self.net(t_ctx).chunk(9, dim=-1)


def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)


class FeedForward(nn.Module):
    def __init__(self, d_model: int, mlp_ratio: float = 4.0, dropout: float = 0.1):
        super().__init__()
        hidden = int(d_model * mlp_ratio)
        self.net = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden, d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class DiffusionTransformerBlock(nn.Module):
    def __init__(self, d_model: int, num_heads: int, mlp_ratio: float = 4.0, dropout: float = 0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.self_attn = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True,
        )
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True,
        )
        self.mlp = FeedForward(d_model, mlp_ratio=mlp_ratio, dropout=dropout)
        self.mod = AdaLNModulation(d_model)
        self.drop = nn.Dropout(dropout)

    def forward(
        self,
        x: torch.Tensor,
        cond: torch.Tensor,
        t_ctx: torch.Tensor,
        target_valid_mask_bt: torch.Tensor,
        cond_valid_mask_bt: torch.Tensor,
    ) -> torch.Tensor:
        target_pad_bt = ~target_valid_mask_bt
        cond_pad_bt = ~cond_valid_mask_bt

        x = apply_seq_mask(x, target_valid_mask_bt)
        cond = apply_seq_mask(cond, cond_valid_mask_bt)

        (
            shift_sa, scale_sa, gate_sa,
            shift_ca, scale_ca, gate_ca,
            shift_ff, scale_ff, gate_ff,
        ) = self.mod(t_ctx)

        h = self.norm1(x)
        h = modulate(h, shift_sa, scale_sa)
        h, _ = self.self_attn(
            query=h,
            key=h,
            value=h,
            key_padding_mask=target_pad_bt,
            need_weights=False,
        )
        x = x + gate_sa.unsqueeze(1) * self.drop(h)
        x = apply_seq_mask(x, target_valid_mask_bt)

        h = self.norm2(x)
        h = modulate(h, shift_ca, scale_ca)
        h, _ = self.cross_attn(
            query=h,
            key=cond,
            value=cond,
            key_padding_mask=cond_pad_bt,
            need_weights=False,
        )
        x = x + gate_ca.unsqueeze(1) * self.drop(h)
        x = apply_seq_mask(x, target_valid_mask_bt)

        h = self.norm3(x)
        h = modulate(h, shift_ff, scale_ff)
        h = self.mlp(h)
        x = x + gate_ff.unsqueeze(1) * self.drop(h)
        x = apply_seq_mask(x, target_valid_mask_bt)
        return x


class TimbreBankEncoder(nn.Module):
    def __init__(
        self,
        *,
        latent_dim: int,
        d_model: int,
        num_families: int = TIMBRE_NUM_FAMILIES,
        max_classes: int = TIMBRE_MAX_CLASSES,
        velocity_bins: int = 8,
        dropout: float = 0.0,
        bank_mean: Any = None,
        bank_std: Any = None,
    ) -> None:
        super().__init__()
        self.num_families = int(num_families)
        self.max_classes = int(max_classes)
        self.velocity_bins = int(max(1, int(velocity_bins)))
        self.latent_proj = nn.Linear(int(latent_dim), int(d_model))
        self.family_embed = nn.Embedding(int(num_families), int(d_model))
        self.class_embed = nn.Embedding(int(num_families) * int(max_classes), int(d_model))
        self.velocity_embed = nn.Embedding(int(self.velocity_bins), int(d_model))
        self.norm = nn.LayerNorm(int(d_model))
        self.drop = nn.Dropout(float(dropout))
        self.register_buffer(
            "bank_mean",
            _normalize_stats_vector(
                bank_mean,
                x_dim=int(latent_dim),
                device=torch.device("cpu"),
                default_fill=0.0,
                name="bank_mean",
            ),
            persistent=True,
        )
        self.register_buffer(
            "bank_std",
            _normalize_stats_vector(
                bank_std,
                x_dim=int(latent_dim),
                device=torch.device("cpu"),
                default_fill=1.0,
                name="bank_std",
            ).clamp_min(1.0e-6),
            persistent=True,
        )

    def forward(
        self,
        latents_bsd: torch.Tensor,
        family_ids_bs: torch.Tensor,
        class_ids_bs: torch.Tensor,
        velocity_bs: torch.Tensor | None,
        mask_bs: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        latents = torch.as_tensor(latents_bsd, dtype=torch.float32)
        if int(latents.dim()) != 3:
            raise ValueError(f"timbre_bank_latents must be [B,S,D], got {tuple(latents.shape)}")
        latents = (latents - self.bank_mean.to(device=latents.device, dtype=latents.dtype).view(1, 1, -1)) / self.bank_std.to(device=latents.device, dtype=latents.dtype).view(1, 1, -1)
        family_ids = torch.as_tensor(family_ids_bs, dtype=torch.long, device=latents.device).clamp(
            min=0,
            max=int(self.num_families) - 1,
        )
        class_ids = torch.as_tensor(class_ids_bs, dtype=torch.long, device=latents.device).clamp(
            min=0,
            max=int(self.max_classes) - 1,
        )
        if velocity_bs is None:
            velocity = torch.zeros_like(class_ids, dtype=torch.float32, device=latents.device)
        else:
            velocity = torch.as_tensor(velocity_bs, dtype=torch.float32, device=latents.device).clamp(min=0.0, max=1.0)
        if mask_bs is None:
            mask = torch.ones(tuple(class_ids.shape), dtype=torch.bool, device=latents.device)
        else:
            mask = torch.as_tensor(mask_bs, dtype=torch.bool, device=latents.device)
        velocity_bins = torch.clamp(
            torch.floor(velocity * float(max(1, int(self.velocity_bins) - 1))).to(dtype=torch.long),
            min=0,
            max=int(self.velocity_bins) - 1,
        )
        flat_class_ids = (family_ids * int(self.max_classes)) + class_ids
        tokens = (
            self.latent_proj(latents)
            + self.family_embed(family_ids)
            + self.class_embed(flat_class_ids)
            + self.velocity_embed(velocity_bins)
        )
        tokens = self.drop(self.norm(tokens))
        tokens = tokens.masked_fill(~mask[:, :, None], 0.0)
        return tokens.contiguous(), mask.contiguous()


class TimbreDynamicsEncoder(nn.Module):
    def __init__(
        self,
        *,
        feature_dim: int,
        d_model: int,
        num_families: int = TIMBRE_NUM_FAMILIES,
        max_classes: int = TIMBRE_MAX_CLASSES,
        velocity_bins: int = 4,
        dropout: float = 0.0,
        dynamic_mean: Any = None,
        dynamic_std: Any = None,
        dynamic_count_mean: Any = None,
        dynamic_count_std: Any = None,
    ) -> None:
        super().__init__()
        self.feature_dim = int(feature_dim)
        self.num_families = int(num_families)
        self.max_classes = int(max_classes)
        self.velocity_bins = int(max(1, int(velocity_bins)))
        self.input_proj = nn.Linear(int(feature_dim) + 1, int(d_model))
        self.family_embed = nn.Embedding(int(num_families), int(d_model))
        self.class_embed = nn.Embedding(int(num_families) * int(max_classes), int(d_model))
        self.velocity_embed = nn.Embedding(int(self.velocity_bins), int(d_model))
        self.norm = nn.LayerNorm(int(d_model))
        self.drop = nn.Dropout(float(dropout))
        self.register_buffer(
            "dynamic_mean",
            _normalize_stats_vector(
                dynamic_mean,
                x_dim=int(feature_dim),
                device=torch.device("cpu"),
                default_fill=0.0,
                name="dynamic_mean",
            ),
            persistent=True,
        )
        self.register_buffer(
            "dynamic_std",
            _normalize_stats_vector(
                dynamic_std,
                x_dim=int(feature_dim),
                device=torch.device("cpu"),
                default_fill=1.0,
                name="dynamic_std",
            ).clamp_min(1.0e-6),
            persistent=True,
        )
        self.register_buffer(
            "dynamic_count_mean",
            _normalize_stats_vector(
                dynamic_count_mean,
                x_dim=1,
                device=torch.device("cpu"),
                default_fill=0.0,
                name="dynamic_count_mean",
            ),
            persistent=True,
        )
        self.register_buffer(
            "dynamic_count_std",
            _normalize_stats_vector(
                dynamic_count_std,
                x_dim=1,
                device=torch.device("cpu"),
                default_fill=1.0,
                name="dynamic_count_std",
            ).clamp_min(1.0e-6),
            persistent=True,
        )

    def forward(
        self,
        features_bsvd: torch.Tensor,
        family_ids_bs: torch.Tensor,
        class_ids_bs: torch.Tensor,
        counts_bsv: torch.Tensor | None,
        mask_bsv: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        features = torch.as_tensor(features_bsvd, dtype=torch.float32)
        if int(features.dim()) != 4 or int(features.shape[-1]) != int(self.feature_dim):
            raise ValueError(
                f"timbre_dynamic_features must be [B,S,V,{int(self.feature_dim)}], got {tuple(features.shape)}"
            )
        batch_size, slot_count, velocity_bins, _ = tuple(features.shape)
        if int(velocity_bins) != int(self.velocity_bins):
            raise ValueError(f"dynamic velocity bins={int(velocity_bins)}, expected {int(self.velocity_bins)}")
        device = features.device
        family_ids = torch.as_tensor(family_ids_bs, dtype=torch.long, device=device).clamp(
            min=0,
            max=int(self.num_families) - 1,
        )
        class_ids = torch.as_tensor(class_ids_bs, dtype=torch.long, device=device).clamp(
            min=0,
            max=int(self.max_classes) - 1,
        )
        if counts_bsv is None:
            counts = torch.zeros((batch_size, slot_count, velocity_bins), dtype=torch.float32, device=device)
        else:
            counts = torch.as_tensor(counts_bsv, dtype=torch.float32, device=device)
        if mask_bsv is None:
            mask = torch.ones((batch_size, slot_count, velocity_bins), dtype=torch.bool, device=device)
        else:
            mask = torch.as_tensor(mask_bsv, dtype=torch.bool, device=device)
        if tuple(family_ids.shape) != (batch_size, slot_count) or tuple(class_ids.shape) != (batch_size, slot_count):
            raise ValueError("dynamic family/class tensors must be [B,S]")
        if tuple(counts.shape) != (batch_size, slot_count, velocity_bins) or tuple(mask.shape) != (batch_size, slot_count, velocity_bins):
            raise ValueError("dynamic count/mask tensors must be [B,S,V]")

        features = (features - self.dynamic_mean.to(device=device, dtype=features.dtype).view(1, 1, 1, -1)) / self.dynamic_std.to(device=device, dtype=features.dtype).view(1, 1, 1, -1)
        counts = (counts - self.dynamic_count_mean.to(device=device, dtype=features.dtype).view(1, 1, 1)) / self.dynamic_count_std.to(device=device, dtype=features.dtype).view(1, 1, 1)
        inp = torch.cat([features, counts.unsqueeze(-1)], dim=-1)

        family_exp = family_ids[:, :, None].expand(batch_size, slot_count, velocity_bins)
        class_exp = class_ids[:, :, None].expand(batch_size, slot_count, velocity_bins)
        velocity_ids = torch.arange(velocity_bins, dtype=torch.long, device=device).view(1, 1, velocity_bins).expand(batch_size, slot_count, velocity_bins)
        flat_class_ids = (family_exp * int(self.max_classes)) + class_exp
        tokens = (
            self.input_proj(inp)
            + self.family_embed(family_exp)
            + self.class_embed(flat_class_ids)
            + self.velocity_embed(velocity_ids)
        )
        tokens = self.drop(self.norm(tokens))
        return tokens.contiguous(), mask.contiguous()


@dataclass
class DiffusionTransformerConfig:
    x_dim: int = 128
    frontend_cfg: Optional[dict[str, Any]] = None
    concat_multiscale_frontend: bool = DEFAULT_FRONTEND_CONCAT_MULTISCALE
    positional_encoding: str = DEFAULT_POSITIONAL_ENCODING
    positional_rate_hz: float = DEFAULT_TARGET_TOKEN_RATE_HZ
    d_model: int = 256
    num_layers: int = 6
    num_heads: int = 8
    mlp_ratio: float = 4.0
    dropout: float = 0.1
    cond_dropout_prob: float = 0.1
    timbre_conditioning: bool = False
    timbre_bank_dim: int = 0
    timbre_num_families: int = TIMBRE_NUM_FAMILIES
    timbre_max_classes: int = TIMBRE_MAX_CLASSES
    timbre_velocity_bins: int = 8
    timbre_dropout_prob: float = 0.0
    timbre_class_dropout_prob: float = 0.0
    timbre_bank_mean: Optional[Sequence[float]] = None
    timbre_bank_std: Optional[Sequence[float]] = None
    timbre_dynamic_conditioning: bool = False
    timbre_dynamic_dim: int = 0
    timbre_dynamic_velocity_bins: int = 4
    timbre_dynamic_dropout_prob: float = 0.0
    timbre_dynamic_mean: Optional[Sequence[float]] = None
    timbre_dynamic_std: Optional[Sequence[float]] = None
    timbre_dynamic_count_mean: Optional[Sequence[float]] = None
    timbre_dynamic_count_std: Optional[Sequence[float]] = None
    reference_conditioning: bool = False
    reference_source_sampling: str = "random-paired"
    reference_dropout_prob: float = 0.0
    reference_segment_dim: int = 144
    x0_prior_conditioning: bool = False
    x0_prior_dim: int = 72


class ConditionalDiffusionTransformer(nn.Module):
    def __init__(self, cfg: DiffusionTransformerConfig):
        super().__init__()
        self.cfg = cfg
        self.summary_frontend = build_seconds_frontend_from_cfg(cfg.frontend_cfg)
        if self.summary_frontend is None:
            raise ValueError("frontend_cfg is required for seconds-grid conditioning")
        if hasattr(self.summary_frontend, "window_radii"):
            self.frontend_scale_radii = tuple(sorted(int(x) for x in list(getattr(self.summary_frontend, "window_radii"))))
            self.frontend_primary_radius = int(getattr(self.summary_frontend, "primary_radius"))
        else:
            self.frontend_primary_radius = int(getattr(self.summary_frontend, "window_radius", 0))
            self.frontend_scale_radii = (int(self.frontend_primary_radius),)
        self.concat_multiscale_frontend = bool(cfg.concat_multiscale_frontend and hasattr(self.summary_frontend, "forward_multiscale"))
        frontend_output_dim = int(getattr(self.summary_frontend, "output_dim"))
        cond_dim = int(frontend_output_dim) * (int(len(self.frontend_scale_radii)) if bool(self.concat_multiscale_frontend) else 1)
        self.positional_encoding = str(getattr(cfg, "positional_encoding", DEFAULT_POSITIONAL_ENCODING)).strip().lower()
        if self.positional_encoding not in {"index", "seconds"}:
            raise ValueError(f"unsupported positional_encoding={self.positional_encoding!r}")
        self.positional_rate_hz = float(
            max(1.0e-6, float(getattr(cfg, "positional_rate_hz", DEFAULT_TARGET_TOKEN_RATE_HZ)))
        )

        self.x_proj = nn.Linear(cfg.x_dim, cfg.d_model)
        self.cond_proj = nn.Linear(cond_dim, cfg.d_model)
        self.timbre_conditioning = bool(getattr(cfg, "timbre_conditioning", False))
        self.timbre_encoder: TimbreBankEncoder | None = None
        self.timbre_to_cond: nn.Linear | None = None
        if bool(self.timbre_conditioning):
            timbre_bank_dim = int(getattr(cfg, "timbre_bank_dim", 0) or cfg.x_dim)
            self.timbre_encoder = TimbreBankEncoder(
                latent_dim=int(timbre_bank_dim),
                d_model=int(cfg.d_model),
                num_families=int(getattr(cfg, "timbre_num_families", TIMBRE_NUM_FAMILIES)),
                max_classes=int(getattr(cfg, "timbre_max_classes", TIMBRE_MAX_CLASSES)),
                velocity_bins=int(getattr(cfg, "timbre_velocity_bins", 8)),
                dropout=float(getattr(cfg, "dropout", 0.0)),
                bank_mean=getattr(cfg, "timbre_bank_mean", None),
                bank_std=getattr(cfg, "timbre_bank_std", None),
            )
            self.timbre_to_cond = nn.Linear(int(cfg.d_model), int(cond_dim))
        self.timbre_dynamic_conditioning = bool(getattr(cfg, "timbre_dynamic_conditioning", False))
        self.timbre_dynamic_encoder: TimbreDynamicsEncoder | None = None
        self.timbre_dynamic_to_cond: nn.Linear | None = None
        if bool(self.timbre_dynamic_conditioning):
            timbre_dynamic_dim = int(getattr(cfg, "timbre_dynamic_dim", 0) or 32)
            self.timbre_dynamic_encoder = TimbreDynamicsEncoder(
                feature_dim=int(timbre_dynamic_dim),
                d_model=int(cfg.d_model),
                num_families=int(getattr(cfg, "timbre_num_families", TIMBRE_NUM_FAMILIES)),
                max_classes=int(getattr(cfg, "timbre_max_classes", TIMBRE_MAX_CLASSES)),
                velocity_bins=int(getattr(cfg, "timbre_dynamic_velocity_bins", 4)),
                dropout=float(getattr(cfg, "dropout", 0.0)),
                dynamic_mean=getattr(cfg, "timbre_dynamic_mean", None),
                dynamic_std=getattr(cfg, "timbre_dynamic_std", None),
                dynamic_count_mean=getattr(cfg, "timbre_dynamic_count_mean", None),
                dynamic_count_std=getattr(cfg, "timbre_dynamic_count_std", None),
            )
            self.timbre_dynamic_to_cond = nn.Linear(int(cfg.d_model), int(cond_dim))
        self.reference_conditioning = bool(getattr(cfg, "reference_conditioning", False))
        self.reference_timbre_pair_to_cond: nn.Linear | None = None
        self.reference_timbre_to_cond: nn.Linear | None = None
        self.reference_dynamic_pair_to_cond: nn.Linear | None = None
        self.reference_dynamic_to_cond: nn.Linear | None = None
        self.reference_segment_proj: nn.Linear | None = None
        self.reference_segment_norm: nn.LayerNorm | None = None
        self.reference_segment_to_cond: nn.Linear | None = None
        if bool(self.reference_conditioning):
            self.reference_timbre_pair_to_cond = nn.Linear(int(cfg.d_model) * 3, int(cond_dim))
            self.reference_timbre_to_cond = nn.Linear(int(cfg.d_model), int(cond_dim))
            reference_segment_dim = int(getattr(cfg, "reference_segment_dim", 144) or 144)
            self.reference_segment_proj = nn.Linear(reference_segment_dim, int(cfg.d_model))
            self.reference_segment_norm = nn.LayerNorm(int(cfg.d_model))
            self.reference_segment_to_cond = nn.Linear(int(cfg.d_model), int(cond_dim))
            if bool(self.timbre_dynamic_conditioning):
                self.reference_dynamic_pair_to_cond = nn.Linear(int(cfg.d_model) * 3, int(cond_dim))
                self.reference_dynamic_to_cond = nn.Linear(int(cfg.d_model), int(cond_dim))
            for module in (
                self.reference_timbre_pair_to_cond,
                self.reference_timbre_to_cond,
                self.reference_dynamic_pair_to_cond,
                self.reference_dynamic_to_cond,
                self.reference_segment_to_cond,
            ):
                if module is not None:
                    nn.init.zeros_(module.weight)
                    nn.init.zeros_(module.bias)
        self.x0_prior_conditioning = bool(getattr(cfg, "x0_prior_conditioning", False))
        self.x0_prior_proj: nn.Linear | None = None
        self.x0_prior_norm: nn.LayerNorm | None = None
        self.x0_prior_to_cond: nn.Linear | None = None
        if bool(self.x0_prior_conditioning):
            x0_prior_dim = int(getattr(cfg, "x0_prior_dim", 0) or cfg.x_dim)
            self.x0_prior_proj = nn.Linear(int(x0_prior_dim), int(cfg.d_model))
            self.x0_prior_norm = nn.LayerNorm(int(cfg.d_model))
            self.x0_prior_to_cond = nn.Linear(int(cfg.d_model), int(cond_dim))
            nn.init.zeros_(self.x0_prior_to_cond.weight)
            nn.init.zeros_(self.x0_prior_to_cond.bias)
        self.time_mlp = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.d_model),
            nn.SiLU(),
            nn.Linear(cfg.d_model, cfg.d_model),
        )
        self.blocks = nn.ModuleList(
            [
                DiffusionTransformerBlock(
                    d_model=cfg.d_model,
                    num_heads=cfg.num_heads,
                    mlp_ratio=cfg.mlp_ratio,
                    dropout=cfg.dropout,
                )
                for _ in range(cfg.num_layers)
            ]
        )
        self.final_norm = nn.LayerNorm(cfg.d_model)
        self.final_mod = nn.Sequential(
            nn.SiLU(),
            nn.Linear(cfg.d_model, cfg.d_model * 2),
        )
        self.out_proj = nn.Linear(cfg.d_model, cfg.x_dim)

    def _batched_timbre_tensor(
        self,
        value: torch.Tensor | None,
        *,
        batch_size: int,
        device: torch.device,
        dtype: torch.dtype | None = None,
    ) -> torch.Tensor | None:
        if value is None:
            return None
        tensor = torch.as_tensor(value, device=device)
        if dtype is not None:
            tensor = tensor.to(dtype=dtype)
        if int(tensor.dim()) >= 1 and int(tensor.shape[0]) == int(batch_size):
            return tensor.contiguous()
        return tensor.unsqueeze(0).expand(int(batch_size), *tuple(tensor.shape)).contiguous()

    def _encode_timbre_tokens(
        self,
        *,
        timbre_bank_latents: torch.Tensor | None,
        timbre_bank_family_ids: torch.Tensor | None,
        timbre_bank_class_ids: torch.Tensor | None,
        timbre_bank_velocity: torch.Tensor | None,
        timbre_bank_mask: torch.Tensor | None,
        batch_size: int,
        device: torch.device,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        if not bool(self.timbre_conditioning):
            return None, None
        if self.timbre_encoder is None or self.timbre_to_cond is None:
            return None, None
        if timbre_bank_latents is None or timbre_bank_family_ids is None or timbre_bank_class_ids is None:
            return None, None
        latents = self._batched_timbre_tensor(
            timbre_bank_latents,
            batch_size=batch_size,
            device=device,
            dtype=torch.float32,
        )
        family_ids = self._batched_timbre_tensor(
            timbre_bank_family_ids,
            batch_size=batch_size,
            device=device,
            dtype=torch.long,
        )
        class_ids = self._batched_timbre_tensor(
            timbre_bank_class_ids,
            batch_size=batch_size,
            device=device,
            dtype=torch.long,
        )
        velocity = self._batched_timbre_tensor(
            timbre_bank_velocity,
            batch_size=batch_size,
            device=device,
            dtype=torch.float32,
        )
        mask = self._batched_timbre_tensor(
            timbre_bank_mask,
            batch_size=batch_size,
            device=device,
            dtype=torch.bool,
        )
        if latents is None or family_ids is None or class_ids is None:
            return None, None
        return self.timbre_encoder(latents, family_ids, class_ids, velocity, mask)

    def _encode_timbre_dynamic_tokens(
        self,
        *,
        timbre_dynamic_features: torch.Tensor | None,
        timbre_dynamic_mask: torch.Tensor | None,
        timbre_dynamic_counts: torch.Tensor | None,
        timbre_bank_family_ids: torch.Tensor | None,
        timbre_bank_class_ids: torch.Tensor | None,
        batch_size: int,
        device: torch.device,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        if not bool(self.timbre_dynamic_conditioning):
            return None, None
        if self.timbre_dynamic_encoder is None or self.timbre_dynamic_to_cond is None:
            return None, None
        if timbre_dynamic_features is None or timbre_bank_family_ids is None or timbre_bank_class_ids is None:
            return None, None
        features = self._batched_timbre_tensor(
            timbre_dynamic_features,
            batch_size=batch_size,
            device=device,
            dtype=torch.float32,
        )
        family_ids = self._batched_timbre_tensor(
            timbre_bank_family_ids,
            batch_size=batch_size,
            device=device,
            dtype=torch.long,
        )
        class_ids = self._batched_timbre_tensor(
            timbre_bank_class_ids,
            batch_size=batch_size,
            device=device,
            dtype=torch.long,
        )
        mask = self._batched_timbre_tensor(
            timbre_dynamic_mask,
            batch_size=batch_size,
            device=device,
            dtype=torch.bool,
        )
        counts = self._batched_timbre_tensor(
            timbre_dynamic_counts,
            batch_size=batch_size,
            device=device,
            dtype=torch.float32,
        )
        if features is None or family_ids is None or class_ids is None:
            return None, None
        return self.timbre_dynamic_encoder(features, family_ids, class_ids, counts, mask)

    def _encode_reference_segment_token(
        self,
        *,
        reference_segment_pca144: torch.Tensor | None,
        batch_size: int,
        device: torch.device,
    ) -> torch.Tensor | None:
        if not bool(self.reference_conditioning):
            return None
        if self.reference_segment_proj is None or self.reference_segment_norm is None:
            return None
        if reference_segment_pca144 is None:
            return None
        segment = torch.as_tensor(reference_segment_pca144, dtype=torch.float32, device=device)
        if int(segment.dim()) == 1:
            segment = segment.view(1, -1).expand(int(batch_size), -1)
        elif int(segment.dim()) == 2 and int(segment.shape[0]) == int(batch_size):
            segment = segment.contiguous()
        else:
            raise ValueError(f"reference_segment_pca144 must be [D] or [B,D], got {tuple(segment.shape)}")
        expected_dim = int(self.reference_segment_proj.in_features)
        if int(segment.shape[-1]) != expected_dim:
            raise ValueError(f"reference_segment_pca144 must be [B,{expected_dim}], got {tuple(segment.shape)}")
        return self.reference_segment_norm(self.reference_segment_proj(segment)).contiguous()

    def _reference_drop_mask(self, *, batch_size: int, device: torch.device) -> torch.Tensor | None:
        prob = float(getattr(self.cfg, "reference_dropout_prob", 0.0))
        if not self.training or prob <= 0.0:
            return None
        drop = torch.rand(int(batch_size), device=device) < prob
        return drop if bool(drop.any()) else None

    @staticmethod
    def _adapter_has_nonzero_weights(module: nn.Linear | None) -> bool:
        if module is None:
            return False
        with torch.no_grad():
            total = module.weight.detach().abs().sum()
            if module.bias is not None:
                total = total + module.bias.detach().abs().sum()
        return bool(float(total.cpu().item()) > 0.0)

    @staticmethod
    def _propagate_family_onset_metadata(
        *,
        family_ids_t: torch.Tensor,
        onset_t: torch.Tensor,
        activity_t: torch.Tensor,
        fallback_velocity_t: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Carry onset-only class/velocity metadata across the active hit tail."""
        activity = torch.as_tensor(activity_t, dtype=torch.float32)
        device = activity.device
        ids = torch.as_tensor(family_ids_t, dtype=torch.long, device=device)
        onset = torch.as_tensor(onset_t, dtype=torch.float32, device=device).clamp(min=0.0, max=1.0)
        if fallback_velocity_t is None:
            fallback_velocity = activity.clamp(min=0.0, max=1.0)
        else:
            fallback_velocity = torch.as_tensor(
                fallback_velocity_t,
                dtype=torch.float32,
                device=device,
            ).clamp(min=0.0, max=1.0)
        frame_idx = torch.arange(int(activity.shape[0]), device=device, dtype=torch.long)
        active = activity > 0.0
        metadata_source = active & (ids >= 0)
        source_velocity = torch.where(onset > 0.0, onset, fallback_velocity)
        last_metadata_pos = torch.cummax(
            torch.where(metadata_source, frame_idx, torch.full_like(frame_idx, -1)),
            dim=0,
        ).values
        last_inactive_pos = torch.cummax(
            torch.where(~active, frame_idx, torch.full_like(frame_idx, -1)),
            dim=0,
        ).values
        valid = active & (last_metadata_pos >= 0) & (last_metadata_pos > last_inactive_pos)
        safe_pos = last_metadata_pos.clamp_min(0)
        class_ids = torch.where(valid, ids[safe_pos], torch.full_like(ids, -1))
        onset_velocity = torch.where(valid, source_velocity[safe_pos], torch.zeros_like(activity))
        return class_ids, onset_velocity, valid

    def _time_aligned_timbre_tokens(
        self,
        *,
        timbre_tokens_bsd: torch.Tensor,
        timbre_mask_bs: torch.Tensor,
        grid: torch.Tensor,
        grid_ids: Optional[torch.Tensor],
        grid_times_sec: torch.Tensor,
        token_times_sec: torch.Tensor,
        grid_valid_mask_bt: Optional[torch.Tensor],
        target_valid_mask_bt: torch.Tensor,
        timbre_family_default_indices: torch.Tensor | None,
        timbre_class_token_indices: torch.Tensor | None,
    ) -> torch.Tensor:
        batch_size, target_len = int(target_valid_mask_bt.shape[0]), int(target_valid_mask_bt.shape[1])
        family_count = int(min(TIMBRE_NUM_FAMILIES, int(grid.shape[1]) // 3 if int(grid.shape[1]) >= TIMBRE_NUM_FAMILIES * 3 else int(grid.shape[1])))
        if family_count <= 0:
            return timbre_tokens_bsd.new_zeros((batch_size, target_len, int(timbre_tokens_bsd.shape[-1])))
        if grid_ids is None:
            ids = torch.zeros((batch_size, family_count, int(grid.shape[-1])), dtype=torch.long, device=grid.device)
        else:
            ids = torch.as_tensor(grid_ids[:, :family_count, :], dtype=torch.long, device=grid.device)
        if timbre_family_default_indices is None or timbre_class_token_indices is None:
            return timbre_tokens_bsd.new_zeros((batch_size, target_len, int(timbre_tokens_bsd.shape[-1])))
        default_indices = self._batched_timbre_tensor(
            timbre_family_default_indices,
            batch_size=batch_size,
            device=grid.device,
            dtype=torch.long,
        )
        class_indices = self._batched_timbre_tensor(
            timbre_class_token_indices,
            batch_size=batch_size,
            device=grid.device,
            dtype=torch.long,
        )
        if default_indices is None or class_indices is None:
            return timbre_tokens_bsd.new_zeros((batch_size, target_len, int(timbre_tokens_bsd.shape[-1])))
        out = timbre_tokens_bsd.new_zeros((batch_size, target_len, int(timbre_tokens_bsd.shape[-1])))
        denom = timbre_tokens_bsd.new_zeros((batch_size, target_len, 1))
        for batch_idx in range(batch_size):
            grid_valid_len = int(grid.shape[-1])
            if grid_valid_mask_bt is not None:
                grid_valid_len = int(torch.as_tensor(grid_valid_mask_bt[batch_idx], dtype=torch.bool).sum().item())
            grid_valid_len = max(1, min(grid_valid_len, int(grid.shape[-1])))
            distances = (
                token_times_sec[batch_idx, :, None].to(device=grid.device, dtype=torch.float32)
                - grid_times_sec[batch_idx, :grid_valid_len][None, :].to(device=grid.device, dtype=torch.float32)
            ).abs()
            nearest = distances.argmin(dim=1)
            for family_idx in range(family_count):
                if int(grid.shape[1]) >= TIMBRE_NUM_FAMILIES * 3:
                    state = grid[batch_idx, family_idx * 3 + 0, :grid_valid_len].abs()
                    onset = grid[batch_idx, family_idx * 3 + 1, :grid_valid_len].abs()
                    count = (grid[batch_idx, family_idx * 3 + 2, :grid_valid_len] > 0).to(dtype=torch.float32)
                    activity_grid = torch.maximum(torch.maximum(state, onset), count)
                    fallback_velocity_grid = state
                else:
                    activity_grid = grid[batch_idx, family_idx, :grid_valid_len].abs()
                    onset = activity_grid
                    fallback_velocity_grid = activity_grid
                class_id_grid, _onset_velocity_grid, metadata_valid_grid = self._propagate_family_onset_metadata(
                    family_ids_t=ids[batch_idx, family_idx, :grid_valid_len],
                    onset_t=onset,
                    activity_t=activity_grid,
                    fallback_velocity_t=fallback_velocity_grid,
                )
                activity = activity_grid[nearest]
                metadata_valid = metadata_valid_grid[nearest]
                if not bool(((activity > 0.0) & metadata_valid).any()):
                    continue
                class_id_t = class_id_grid[nearest].clamp(min=0, max=int(class_indices.shape[-1]) - 1)
                exact = class_indices[batch_idx, family_idx, class_id_t]
                fallback = default_indices[batch_idx, family_idx].expand_as(exact)
                token_idx = torch.where(exact >= 0, exact, fallback).clamp(min=0, max=int(timbre_tokens_bsd.shape[1]) - 1)
                token_ok = timbre_mask_bs[batch_idx, token_idx].to(dtype=torch.bool)
                active = (activity > 0.0) & metadata_valid & token_ok
                if not bool(active.any()):
                    continue
                gathered = timbre_tokens_bsd[batch_idx, token_idx]
                weight = activity.to(dtype=timbre_tokens_bsd.dtype).view(target_len, 1) * active.to(dtype=timbre_tokens_bsd.dtype).view(target_len, 1)
                out[batch_idx] = out[batch_idx] + (gathered * weight)
                denom[batch_idx] = denom[batch_idx] + weight
        active_mask = denom > 0.0
        out = out / denom.clamp_min(1.0e-8)
        out = out.masked_fill(~active_mask, 0.0)
        out = out.masked_fill(~target_valid_mask_bt[:, :, None].to(device=out.device, dtype=torch.bool), 0.0)
        return out.contiguous()

    def _time_aligned_timbre_dynamic_tokens(
        self,
        *,
        timbre_dynamic_tokens_bsvd: torch.Tensor,
        timbre_dynamic_mask_bsv: torch.Tensor,
        grid: torch.Tensor,
        grid_ids: Optional[torch.Tensor],
        grid_times_sec: torch.Tensor,
        token_times_sec: torch.Tensor,
        grid_valid_mask_bt: Optional[torch.Tensor],
        target_valid_mask_bt: torch.Tensor,
        timbre_family_default_indices: torch.Tensor | None,
        timbre_class_token_indices: torch.Tensor | None,
    ) -> torch.Tensor:
        batch_size, target_len = int(target_valid_mask_bt.shape[0]), int(target_valid_mask_bt.shape[1])
        family_count = int(min(TIMBRE_NUM_FAMILIES, int(grid.shape[1]) // 3 if int(grid.shape[1]) >= TIMBRE_NUM_FAMILIES * 3 else int(grid.shape[1])))
        if family_count <= 0:
            return timbre_dynamic_tokens_bsvd.new_zeros((batch_size, target_len, int(timbre_dynamic_tokens_bsvd.shape[-1])))
        if grid_ids is None:
            ids = torch.zeros((batch_size, family_count, int(grid.shape[-1])), dtype=torch.long, device=grid.device)
        else:
            ids = torch.as_tensor(grid_ids[:, :family_count, :], dtype=torch.long, device=grid.device)
        if timbre_family_default_indices is None or timbre_class_token_indices is None:
            return timbre_dynamic_tokens_bsvd.new_zeros((batch_size, target_len, int(timbre_dynamic_tokens_bsvd.shape[-1])))
        default_indices = self._batched_timbre_tensor(
            timbre_family_default_indices,
            batch_size=batch_size,
            device=grid.device,
            dtype=torch.long,
        )
        class_indices = self._batched_timbre_tensor(
            timbre_class_token_indices,
            batch_size=batch_size,
            device=grid.device,
            dtype=torch.long,
        )
        if default_indices is None or class_indices is None:
            return timbre_dynamic_tokens_bsvd.new_zeros((batch_size, target_len, int(timbre_dynamic_tokens_bsvd.shape[-1])))
        velocity_bins = int(timbre_dynamic_tokens_bsvd.shape[2])
        out = timbre_dynamic_tokens_bsvd.new_zeros((batch_size, target_len, int(timbre_dynamic_tokens_bsvd.shape[-1])))
        denom = timbre_dynamic_tokens_bsvd.new_zeros((batch_size, target_len, 1))
        for batch_idx in range(batch_size):
            grid_valid_len = int(grid.shape[-1])
            if grid_valid_mask_bt is not None:
                grid_valid_len = int(torch.as_tensor(grid_valid_mask_bt[batch_idx], dtype=torch.bool).sum().item())
            grid_valid_len = max(1, min(grid_valid_len, int(grid.shape[-1])))
            distances = (
                token_times_sec[batch_idx, :, None].to(device=grid.device, dtype=torch.float32)
                - grid_times_sec[batch_idx, :grid_valid_len][None, :].to(device=grid.device, dtype=torch.float32)
            ).abs()
            nearest = distances.argmin(dim=1)
            for family_idx in range(family_count):
                if int(grid.shape[1]) >= TIMBRE_NUM_FAMILIES * 3:
                    state = grid[batch_idx, family_idx * 3 + 0, :grid_valid_len].abs()
                    onset = grid[batch_idx, family_idx * 3 + 1, :grid_valid_len].abs()
                    count = (grid[batch_idx, family_idx * 3 + 2, :grid_valid_len] > 0).to(dtype=torch.float32)
                    activity_grid = torch.maximum(torch.maximum(state, onset), count)
                    fallback_velocity_grid = state
                else:
                    activity_grid = grid[batch_idx, family_idx, :grid_valid_len].abs()
                    onset = activity_grid
                    fallback_velocity_grid = activity_grid
                class_id_grid, velocity_grid, metadata_valid_grid = self._propagate_family_onset_metadata(
                    family_ids_t=ids[batch_idx, family_idx, :grid_valid_len],
                    onset_t=onset,
                    activity_t=activity_grid,
                    fallback_velocity_t=fallback_velocity_grid,
                )
                activity = activity_grid[nearest]
                metadata_valid = metadata_valid_grid[nearest]
                if not bool(((activity > 0.0) & metadata_valid).any()):
                    continue
                velocity_t = velocity_grid[nearest]
                dynamic_bin = torch.clamp(
                    torch.floor(velocity_t * float(velocity_bins)).to(dtype=torch.long),
                    min=0,
                    max=int(velocity_bins) - 1,
                )
                class_id_t = class_id_grid[nearest].clamp(min=0, max=int(class_indices.shape[-1]) - 1)
                exact = class_indices[batch_idx, family_idx, class_id_t]
                fallback = default_indices[batch_idx, family_idx].expand_as(exact)
                token_idx = torch.where(exact >= 0, exact, fallback).clamp(min=0, max=int(timbre_dynamic_tokens_bsvd.shape[1]) - 1)
                token_ok = timbre_dynamic_mask_bsv[batch_idx].any(dim=-1)[token_idx].to(dtype=torch.bool)
                active = (activity > 0.0) & metadata_valid & token_ok
                if not bool(active.any()):
                    continue
                gathered = timbre_dynamic_tokens_bsvd[batch_idx, token_idx, dynamic_bin]
                weight = activity.to(dtype=timbre_dynamic_tokens_bsvd.dtype).view(target_len, 1) * active.to(dtype=timbre_dynamic_tokens_bsvd.dtype).view(target_len, 1)
                out[batch_idx] = out[batch_idx] + (gathered * weight)
                denom[batch_idx] = denom[batch_idx] + weight
        active_mask = denom > 0.0
        out = out / denom.clamp_min(1.0e-8)
        out = out.masked_fill(~active_mask, 0.0)
        out = out.masked_fill(~target_valid_mask_bt[:, :, None].to(device=out.device, dtype=torch.bool), 0.0)
        return out.contiguous()

    def encode_conditioning(
        self,
        *,
        grid: torch.Tensor,
        grid_ids: Optional[torch.Tensor],
        grid_times_sec: torch.Tensor,
        token_times_sec: torch.Tensor,
        target_valid_mask_bt: torch.Tensor,
        grid_valid_mask_bt: Optional[torch.Tensor] = None,
        timbre_bank_latents: torch.Tensor | None = None,
        timbre_bank_family_ids: torch.Tensor | None = None,
        timbre_bank_class_ids: torch.Tensor | None = None,
        timbre_bank_velocity: torch.Tensor | None = None,
        timbre_bank_mask: torch.Tensor | None = None,
        timbre_dynamic_features: torch.Tensor | None = None,
        timbre_dynamic_mask: torch.Tensor | None = None,
        timbre_dynamic_counts: torch.Tensor | None = None,
        timbre_family_default_indices: torch.Tensor | None = None,
        timbre_class_token_indices: torch.Tensor | None = None,
        reference_timbre_bank_latents: torch.Tensor | None = None,
        reference_timbre_bank_family_ids: torch.Tensor | None = None,
        reference_timbre_bank_class_ids: torch.Tensor | None = None,
        reference_timbre_bank_velocity: torch.Tensor | None = None,
        reference_timbre_bank_mask: torch.Tensor | None = None,
        reference_timbre_dynamic_features: torch.Tensor | None = None,
        reference_timbre_dynamic_mask: torch.Tensor | None = None,
        reference_timbre_dynamic_counts: torch.Tensor | None = None,
        reference_timbre_family_default_indices: torch.Tensor | None = None,
        reference_timbre_class_token_indices: torch.Tensor | None = None,
        reference_segment_pca144: torch.Tensor | None = None,
        x0_prior_btd: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        frontend_kwargs = {
            "grid_ids_bct": grid_ids,
            "grid_times_sec_bt": grid_times_sec,
            "token_times_sec_bt": token_times_sec,
            "grid_valid_mask_bt": grid_valid_mask_bt,
            "valid_mask_bt": target_valid_mask_bt,
        }
        if bool(self.concat_multiscale_frontend):
            scale_features = {
                int(scale_radius): scale_feat
                for scale_radius, scale_feat in dict(self.summary_frontend.forward_multiscale(grid, **frontend_kwargs)).items()
            }
            cond_btd = torch.cat(
                [scale_features[int(scale_radius)] for scale_radius in list(self.frontend_scale_radii)],
                dim=-1,
            ).contiguous()
        else:
            cond_btd = self.summary_frontend(
                grid,
                **frontend_kwargs,
            )
        cond_valid_mask_bt = target_valid_mask_bt.to(dtype=torch.bool)
        cond_btd = apply_seq_mask(cond_btd, cond_valid_mask_bt)
        target_len = int(target_valid_mask_bt.shape[1])
        batch_size = int(grid.shape[0])
        target_aligned_timbre: torch.Tensor | None = None
        target_aligned_dynamic: torch.Tensor | None = None
        reference_drop = self._reference_drop_mask(batch_size=batch_size, device=grid.device)
        if (
            bool(self.x0_prior_conditioning)
            and x0_prior_btd is not None
            and self.x0_prior_proj is not None
            and self.x0_prior_norm is not None
            and self.x0_prior_to_cond is not None
        ):
            prior = torch.as_tensor(x0_prior_btd, dtype=torch.float32, device=grid.device)
            if int(prior.dim()) != 3:
                raise ValueError(f"x0_prior_btd must be [B,T,D], got {tuple(prior.shape)}")
            if int(prior.shape[0]) != int(batch_size):
                raise ValueError(f"x0_prior_btd batch must be {batch_size}, got {tuple(prior.shape)}")
            expected_dim = int(self.x0_prior_proj.in_features)
            if int(prior.shape[-1]) != expected_dim:
                raise ValueError(f"x0_prior_btd last dim must be {expected_dim}, got {tuple(prior.shape)}")
            if int(prior.shape[1]) != int(target_len):
                prior = F.interpolate(
                    prior.transpose(1, 2),
                    size=int(target_len),
                    mode="linear",
                    align_corners=False,
                ).transpose(1, 2).contiguous()
            prior = apply_seq_mask(prior, target_valid_mask_bt.to(dtype=torch.bool))
            prior_cond = self.x0_prior_to_cond(self.x0_prior_norm(self.x0_prior_proj(prior)))
            cond_btd = cond_btd + prior_cond
        timbre_tokens, timbre_mask = self._encode_timbre_tokens(
            timbre_bank_latents=timbre_bank_latents,
            timbre_bank_family_ids=timbre_bank_family_ids,
            timbre_bank_class_ids=timbre_bank_class_ids,
            timbre_bank_velocity=timbre_bank_velocity,
            timbre_bank_mask=timbre_bank_mask,
            batch_size=batch_size,
            device=grid.device,
        )
        if timbre_tokens is not None and timbre_mask is not None and self.timbre_to_cond is not None:
            if self.training and float(getattr(self.cfg, "timbre_dropout_prob", 0.0)) > 0.0:
                drop = torch.rand(int(grid.shape[0]), device=grid.device) < float(getattr(self.cfg, "timbre_dropout_prob", 0.0))
                if bool(drop.any()):
                    timbre_tokens = timbre_tokens.clone()
                    timbre_tokens[drop] = 0.0
                    timbre_mask = timbre_mask.clone()
                    timbre_mask[drop] = False
            aligned = self._time_aligned_timbre_tokens(
                timbre_tokens_bsd=timbre_tokens,
                timbre_mask_bs=timbre_mask,
                grid=grid,
                grid_ids=grid_ids,
                grid_times_sec=grid_times_sec,
                token_times_sec=token_times_sec,
                grid_valid_mask_bt=grid_valid_mask_bt,
                target_valid_mask_bt=target_valid_mask_bt,
                timbre_family_default_indices=timbre_family_default_indices,
                timbre_class_token_indices=timbre_class_token_indices,
            )
            target_aligned_timbre = aligned
            cond_btd = cond_btd + self.timbre_to_cond(aligned)
            bank_cond = self.timbre_to_cond(timbre_tokens)
            cond_btd = torch.cat([cond_btd, bank_cond], dim=1).contiguous()
            cond_valid_mask_bt = torch.cat([cond_valid_mask_bt, timbre_mask.to(dtype=torch.bool)], dim=1).contiguous()
        reference_tokens, reference_mask = self._encode_timbre_tokens(
            timbre_bank_latents=reference_timbre_bank_latents,
            timbre_bank_family_ids=reference_timbre_bank_family_ids,
            timbre_bank_class_ids=reference_timbre_bank_class_ids,
            timbre_bank_velocity=reference_timbre_bank_velocity,
            timbre_bank_mask=reference_timbre_bank_mask,
            batch_size=batch_size,
            device=grid.device,
        )
        if (
            bool(self.reference_conditioning)
            and reference_tokens is not None
            and reference_mask is not None
            and self.reference_timbre_pair_to_cond is not None
            and self.reference_timbre_to_cond is not None
        ):
            if reference_drop is not None:
                reference_tokens = reference_tokens.clone()
                reference_mask = reference_mask.clone()
                reference_tokens[reference_drop] = 0.0
                reference_mask[reference_drop] = False
            aligned_reference = self._time_aligned_timbre_tokens(
                timbre_tokens_bsd=reference_tokens,
                timbre_mask_bs=reference_mask,
                grid=grid,
                grid_ids=grid_ids,
                grid_times_sec=grid_times_sec,
                token_times_sec=token_times_sec,
                grid_valid_mask_bt=grid_valid_mask_bt,
                target_valid_mask_bt=target_valid_mask_bt,
                timbre_family_default_indices=reference_timbre_family_default_indices,
                timbre_class_token_indices=reference_timbre_class_token_indices,
            )
            target_for_delta = (
                target_aligned_timbre
                if target_aligned_timbre is not None
                else torch.zeros_like(aligned_reference)
            )
            pair = torch.cat(
                [target_for_delta, aligned_reference, target_for_delta - aligned_reference],
                dim=-1,
            )
            cond_btd = torch.cat(
                [cond_btd[:, :target_len] + self.reference_timbre_pair_to_cond(pair), cond_btd[:, target_len:]],
                dim=1,
            ).contiguous()
            if self._adapter_has_nonzero_weights(self.reference_timbre_pair_to_cond):
                reference_cond = self.reference_timbre_to_cond(reference_tokens)
                cond_btd = torch.cat([cond_btd, reference_cond], dim=1).contiguous()
                cond_valid_mask_bt = torch.cat([cond_valid_mask_bt, reference_mask.to(dtype=torch.bool)], dim=1).contiguous()
        reference_segment = self._encode_reference_segment_token(
            reference_segment_pca144=reference_segment_pca144,
            batch_size=batch_size,
            device=grid.device,
        )
        if (
            reference_segment is not None
            and self.reference_segment_to_cond is not None
            and bool(self.reference_conditioning)
        ):
            if reference_drop is not None:
                reference_segment = reference_segment.clone()
                reference_segment[reference_drop] = 0.0
            segment_cond = self.reference_segment_to_cond(reference_segment).unsqueeze(1)
            cond_btd = torch.cat(
                [cond_btd[:, :target_len] + segment_cond, cond_btd[:, target_len:]],
                dim=1,
            ).contiguous()
        dynamic_tokens, dynamic_mask = self._encode_timbre_dynamic_tokens(
            timbre_dynamic_features=timbre_dynamic_features,
            timbre_dynamic_mask=timbre_dynamic_mask,
            timbre_dynamic_counts=timbre_dynamic_counts,
            timbre_bank_family_ids=timbre_bank_family_ids,
            timbre_bank_class_ids=timbre_bank_class_ids,
            batch_size=batch_size,
            device=grid.device,
        )
        if dynamic_tokens is not None and dynamic_mask is not None and self.timbre_dynamic_to_cond is not None:
            if self.training and float(getattr(self.cfg, "timbre_dynamic_dropout_prob", 0.0)) > 0.0:
                drop = torch.rand(int(grid.shape[0]), device=grid.device) < float(getattr(self.cfg, "timbre_dynamic_dropout_prob", 0.0))
                if bool(drop.any()):
                    dynamic_tokens = dynamic_tokens.clone()
                    dynamic_tokens[drop] = 0.0
                    dynamic_mask = dynamic_mask.clone()
                    dynamic_mask[drop] = False
            aligned_dynamic = self._time_aligned_timbre_dynamic_tokens(
                timbre_dynamic_tokens_bsvd=dynamic_tokens,
                timbre_dynamic_mask_bsv=dynamic_mask,
                grid=grid,
                grid_ids=grid_ids,
                grid_times_sec=grid_times_sec,
                token_times_sec=token_times_sec,
                grid_valid_mask_bt=grid_valid_mask_bt,
                target_valid_mask_bt=target_valid_mask_bt,
                timbre_family_default_indices=timbre_family_default_indices,
                timbre_class_token_indices=timbre_class_token_indices,
            )
            target_aligned_dynamic = aligned_dynamic
            dynamic_aligned_cond = self.timbre_dynamic_to_cond(aligned_dynamic)
            cond_btd = torch.cat(
                [cond_btd[:, :target_len] + dynamic_aligned_cond, cond_btd[:, target_len:]],
                dim=1,
            ).contiguous()
            dynamic_flat = dynamic_tokens.reshape(int(dynamic_tokens.shape[0]), -1, int(dynamic_tokens.shape[-1]))
            dynamic_mask_flat = dynamic_mask.reshape(int(dynamic_mask.shape[0]), -1)
            dynamic_cond = self.timbre_dynamic_to_cond(dynamic_flat)
            cond_btd = torch.cat([cond_btd, dynamic_cond], dim=1).contiguous()
            cond_valid_mask_bt = torch.cat([cond_valid_mask_bt, dynamic_mask_flat.to(dtype=torch.bool)], dim=1).contiguous()
        reference_dynamic_tokens, reference_dynamic_mask = self._encode_timbre_dynamic_tokens(
            timbre_dynamic_features=reference_timbre_dynamic_features,
            timbre_dynamic_mask=reference_timbre_dynamic_mask,
            timbre_dynamic_counts=reference_timbre_dynamic_counts,
            timbre_bank_family_ids=reference_timbre_bank_family_ids,
            timbre_bank_class_ids=reference_timbre_bank_class_ids,
            batch_size=batch_size,
            device=grid.device,
        )
        if (
            bool(self.reference_conditioning)
            and reference_dynamic_tokens is not None
            and reference_dynamic_mask is not None
            and self.reference_dynamic_pair_to_cond is not None
            and self.reference_dynamic_to_cond is not None
        ):
            if reference_drop is not None:
                reference_dynamic_tokens = reference_dynamic_tokens.clone()
                reference_dynamic_mask = reference_dynamic_mask.clone()
                reference_dynamic_tokens[reference_drop] = 0.0
                reference_dynamic_mask[reference_drop] = False
            aligned_reference_dynamic = self._time_aligned_timbre_dynamic_tokens(
                timbre_dynamic_tokens_bsvd=reference_dynamic_tokens,
                timbre_dynamic_mask_bsv=reference_dynamic_mask,
                grid=grid,
                grid_ids=grid_ids,
                grid_times_sec=grid_times_sec,
                token_times_sec=token_times_sec,
                grid_valid_mask_bt=grid_valid_mask_bt,
                target_valid_mask_bt=target_valid_mask_bt,
                timbre_family_default_indices=reference_timbre_family_default_indices,
                timbre_class_token_indices=reference_timbre_class_token_indices,
            )
            target_dynamic_for_delta = (
                target_aligned_dynamic
                if target_aligned_dynamic is not None
                else torch.zeros_like(aligned_reference_dynamic)
            )
            dynamic_pair = torch.cat(
                [
                    target_dynamic_for_delta,
                    aligned_reference_dynamic,
                    target_dynamic_for_delta - aligned_reference_dynamic,
                ],
                dim=-1,
            )
            cond_btd = torch.cat(
                [cond_btd[:, :target_len] + self.reference_dynamic_pair_to_cond(dynamic_pair), cond_btd[:, target_len:]],
                dim=1,
            ).contiguous()
            if self._adapter_has_nonzero_weights(self.reference_dynamic_pair_to_cond):
                reference_dynamic_flat = reference_dynamic_tokens.reshape(
                    int(reference_dynamic_tokens.shape[0]),
                    -1,
                    int(reference_dynamic_tokens.shape[-1]),
                )
                reference_dynamic_mask_flat = reference_dynamic_mask.reshape(int(reference_dynamic_mask.shape[0]), -1)
                reference_dynamic_cond = self.reference_dynamic_to_cond(reference_dynamic_flat)
                cond_btd = torch.cat([cond_btd, reference_dynamic_cond], dim=1).contiguous()
                cond_valid_mask_bt = torch.cat(
                    [cond_valid_mask_bt, reference_dynamic_mask_flat.to(dtype=torch.bool)],
                    dim=1,
                ).contiguous()
        return cond_btd.contiguous(), cond_valid_mask_bt.contiguous()

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        *,
        target_valid_mask_bt: torch.Tensor,
        grid: Optional[torch.Tensor] = None,
        grid_ids: Optional[torch.Tensor] = None,
        grid_times_sec: Optional[torch.Tensor] = None,
        token_times_sec: Optional[torch.Tensor] = None,
        grid_valid_mask_bt: Optional[torch.Tensor] = None,
        beat_boundaries_sec: Optional[torch.Tensor] = None,
        beat_boundaries_valid_mask: Optional[torch.Tensor] = None,
        bpm: Optional[torch.Tensor] = None,
        duration_sec: Optional[torch.Tensor] = None,
        cond_btd: Optional[torch.Tensor] = None,
        cond_valid_mask_bt: Optional[torch.Tensor] = None,
        force_uncond: bool = False,
        timbre_bank_latents: torch.Tensor | None = None,
        timbre_bank_family_ids: torch.Tensor | None = None,
        timbre_bank_class_ids: torch.Tensor | None = None,
        timbre_bank_velocity: torch.Tensor | None = None,
        timbre_bank_mask: torch.Tensor | None = None,
        timbre_dynamic_features: torch.Tensor | None = None,
        timbre_dynamic_mask: torch.Tensor | None = None,
        timbre_dynamic_counts: torch.Tensor | None = None,
        timbre_family_default_indices: torch.Tensor | None = None,
        timbre_class_token_indices: torch.Tensor | None = None,
        reference_timbre_bank_latents: torch.Tensor | None = None,
        reference_timbre_bank_family_ids: torch.Tensor | None = None,
        reference_timbre_bank_class_ids: torch.Tensor | None = None,
        reference_timbre_bank_velocity: torch.Tensor | None = None,
        reference_timbre_bank_mask: torch.Tensor | None = None,
        reference_timbre_dynamic_features: torch.Tensor | None = None,
        reference_timbre_dynamic_mask: torch.Tensor | None = None,
        reference_timbre_dynamic_counts: torch.Tensor | None = None,
        reference_timbre_family_default_indices: torch.Tensor | None = None,
        reference_timbre_class_token_indices: torch.Tensor | None = None,
        reference_segment_pca144: torch.Tensor | None = None,
        x0_prior_btd: torch.Tensor | None = None,
    ) -> torch.Tensor:
        del beat_boundaries_sec, beat_boundaries_valid_mask, bpm, duration_sec
        bsz, target_len, _ = x_t.shape
        device = x_t.device

        if cond_btd is None or cond_valid_mask_bt is None:
            missing = [
                name
                for name, value in (
                    ("grid", grid),
                    ("grid_times_sec", grid_times_sec),
                    ("token_times_sec", token_times_sec),
                )
                if value is None
            ]
            if missing:
                raise ValueError(f"missing conditioning inputs: {missing}")
            cond_btd, cond_valid_mask_bt = self.encode_conditioning(
                grid=grid,
                grid_ids=grid_ids,
                grid_times_sec=grid_times_sec,
                token_times_sec=token_times_sec,
                target_valid_mask_bt=target_valid_mask_bt,
                grid_valid_mask_bt=grid_valid_mask_bt,
                timbre_bank_latents=timbre_bank_latents,
                timbre_bank_family_ids=timbre_bank_family_ids,
                timbre_bank_class_ids=timbre_bank_class_ids,
                timbre_bank_velocity=timbre_bank_velocity,
                timbre_bank_mask=timbre_bank_mask,
                timbre_dynamic_features=timbre_dynamic_features,
                timbre_dynamic_mask=timbre_dynamic_mask,
                timbre_dynamic_counts=timbre_dynamic_counts,
                timbre_family_default_indices=timbre_family_default_indices,
                timbre_class_token_indices=timbre_class_token_indices,
                reference_timbre_bank_latents=reference_timbre_bank_latents,
                reference_timbre_bank_family_ids=reference_timbre_bank_family_ids,
                reference_timbre_bank_class_ids=reference_timbre_bank_class_ids,
                reference_timbre_bank_velocity=reference_timbre_bank_velocity,
                reference_timbre_bank_mask=reference_timbre_bank_mask,
                reference_timbre_dynamic_features=reference_timbre_dynamic_features,
                reference_timbre_dynamic_mask=reference_timbre_dynamic_mask,
                reference_timbre_dynamic_counts=reference_timbre_dynamic_counts,
                reference_timbre_family_default_indices=reference_timbre_family_default_indices,
                reference_timbre_class_token_indices=reference_timbre_class_token_indices,
                reference_segment_pca144=reference_segment_pca144,
                x0_prior_btd=x0_prior_btd,
            )

        if int(cond_btd.shape[0]) != int(target_valid_mask_bt.shape[0]) or int(cond_valid_mask_bt.shape[0]) != int(target_valid_mask_bt.shape[0]):
            raise ValueError(
                f"conditioning batch must align with target_valid_mask_bt, got {tuple(cond_btd.shape)} / {tuple(cond_valid_mask_bt.shape)} / {tuple(target_valid_mask_bt.shape)}"
            )

        if self.training and self.cfg.cond_dropout_prob > 0.0:
            drop_mask_b = (torch.rand(bsz, device=device) < self.cfg.cond_dropout_prob)
            if drop_mask_b.any():
                cond_btd = cond_btd.clone()
                cond_btd[drop_mask_b] = 0

        if force_uncond:
            cond_btd = torch.zeros_like(cond_btd)

        if self.positional_encoding == "seconds" and token_times_sec is not None:
            token_pos = sinusoidal_time_positions(
                torch.as_tensor(token_times_sec, dtype=torch.float32, device=device),
                self.cfg.d_model,
                rate_hz=float(self.positional_rate_hz),
            )
            if tuple(token_pos.shape[:2]) != tuple(target_valid_mask_bt.shape):
                raise ValueError(
                    "token_times_sec must align with target_valid_mask_bt for seconds positional encoding, got "
                    f"{tuple(token_pos.shape[:2])} / {tuple(target_valid_mask_bt.shape)}"
                )
            x_pos = token_pos
            if int(cond_btd.shape[1]) == int(target_len):
                c_pos = token_pos
            else:
                extra = cond_btd.new_zeros((bsz, int(cond_btd.shape[1]) - int(target_len), int(self.cfg.d_model)))
                c_pos = torch.cat([token_pos, extra], dim=1).contiguous()
        else:
            x_pos = sinusoidal_positions(target_len, self.cfg.d_model, device)
            c_pos = sinusoidal_positions(int(cond_btd.shape[1]), self.cfg.d_model, device)

        x = self.x_proj(x_t) + x_pos
        c = self.cond_proj(cond_btd) + c_pos
        t_emb = timestep_embedding(t, self.cfg.d_model)
        t_ctx = self.time_mlp(t_emb)

        x = apply_seq_mask(x, target_valid_mask_bt)
        c = apply_seq_mask(c, cond_valid_mask_bt)

        for block in self.blocks:
            x = block(
                x=x,
                cond=c,
                t_ctx=t_ctx,
                target_valid_mask_bt=target_valid_mask_bt,
                cond_valid_mask_bt=cond_valid_mask_bt,
            )

        shift, scale = self.final_mod(t_ctx).chunk(2, dim=-1)
        x = self.final_norm(x)
        x = modulate(x, shift, scale)
        x = self.out_proj(x)
        x = apply_seq_mask(x, target_valid_mask_bt)
        return x


def cosine_beta_schedule(num_steps: int, s: float = 0.008) -> torch.Tensor:
    steps = num_steps + 1
    x = torch.linspace(0, num_steps, steps, dtype=torch.float64)
    alphas_cumprod = torch.cos(((x / num_steps) + s) / (1 + s) * math.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return betas.clamp(1e-5, 0.999).float()


class GaussianDiffusion1D(nn.Module):
    def __init__(self, num_steps: int = 1000):
        super().__init__()
        betas = cosine_beta_schedule(num_steps=num_steps)
        alphas = 1.0 - betas
        alpha_bars = torch.cumprod(alphas, dim=0)

        self.num_steps = num_steps
        self.register_buffer("betas", betas)
        self.register_buffer("alphas", alphas)
        self.register_buffer("alpha_bars", alpha_bars)
        self.register_buffer("sqrt_alpha_bars", torch.sqrt(alpha_bars))
        self.register_buffer("sqrt_one_minus_alpha_bars", torch.sqrt(1.0 - alpha_bars))
        self.register_buffer("sqrt_recip_alphas", torch.sqrt(1.0 / alphas))

        alpha_bars_prev = torch.cat([torch.ones(1, device=betas.device), alpha_bars[:-1]], dim=0)
        posterior_var = betas * (1.0 - alpha_bars_prev) / (1.0 - alpha_bars)
        self.register_buffer("posterior_variance", posterior_var.clamp_min(1e-20))
        posterior_mean_coef1 = betas * torch.sqrt(alpha_bars_prev) / (1.0 - alpha_bars)
        posterior_mean_coef2 = (1.0 - alpha_bars_prev) * torch.sqrt(alphas) / (1.0 - alpha_bars)
        self.register_buffer("posterior_mean_coef1", posterior_mean_coef1)
        self.register_buffer("posterior_mean_coef2", posterior_mean_coef2)

    def q_sample(self, x0: torch.Tensor, t: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
        a = self.sqrt_alpha_bars[t].view(-1, 1, 1)
        b = self.sqrt_one_minus_alpha_bars[t].view(-1, 1, 1)
        return a * x0 + b * noise

    def predict_x0_from_eps(self, x_t: torch.Tensor, t: torch.Tensor, eps: torch.Tensor) -> torch.Tensor:
        a = self.sqrt_alpha_bars[t].view(-1, 1, 1)
        b = self.sqrt_one_minus_alpha_bars[t].view(-1, 1, 1)
        return (x_t - b * eps) / a.clamp_min(1e-8)

    def posterior_mean_from_x0(self, x_t: torch.Tensor, t: torch.Tensor, x0_hat: torch.Tensor) -> torch.Tensor:
        coef1 = self.posterior_mean_coef1[t].view(-1, 1, 1)
        coef2 = self.posterior_mean_coef2[t].view(-1, 1, 1)
        return coef1 * x0_hat + coef2 * x_t


def resolve_valid_audio_num_samples(
    duration_sec_b: torch.Tensor,
    *,
    sample_rate: int,
    max_num_samples: int,
) -> torch.Tensor:
    duration_sec = torch.as_tensor(duration_sec_b, dtype=torch.float32).view(-1)
    valid_num_samples_b = torch.round(duration_sec * float(sample_rate)).to(dtype=torch.long)
    return valid_num_samples_b.clamp(min=1, max=max(1, int(max_num_samples))).contiguous()


def audio_valid_mask(
    valid_num_samples_b: torch.Tensor,
    *,
    max_num_samples: int,
) -> torch.Tensor:
    valid_num_samples = torch.as_tensor(valid_num_samples_b, dtype=torch.long).view(-1)
    return lengths_to_mask(valid_num_samples, max_len=int(max_num_samples))


def masked_audio_l1_per_example(
    pred_audio_bct: torch.Tensor,
    target_audio_bct: torch.Tensor,
    valid_num_samples_b: torch.Tensor,
) -> torch.Tensor:
    pred_audio = torch.as_tensor(pred_audio_bct, dtype=torch.float32)
    target_audio = torch.as_tensor(target_audio_bct, dtype=torch.float32, device=pred_audio.device)
    if tuple(pred_audio.shape) != tuple(target_audio.shape):
        raise ValueError(
            f"pred_audio_bct and target_audio_bct must match, got {tuple(pred_audio.shape)} / {tuple(target_audio.shape)}"
        )
    if int(pred_audio.dim()) != 3:
        raise ValueError(f"expected [B,C,T] audio tensors, got {tuple(pred_audio.shape)}")
    weights_b1t = audio_valid_mask(
        valid_num_samples_b,
        max_num_samples=int(pred_audio.shape[-1]),
    ).to(device=pred_audio.device, dtype=pred_audio.dtype)[:, None, :]
    denom_b = weights_b1t.sum(dim=(1, 2)).clamp_min(1.0)
    return (((pred_audio - target_audio).abs()) * weights_b1t).sum(dim=(1, 2)) / denom_b


def _safe_stft(audio_bt: torch.Tensor, *, n_fft: int, hop: int) -> torch.Tensor:
    audio = torch.as_tensor(audio_bt, dtype=torch.float32)
    if int(audio.dim()) != 2:
        raise ValueError(f"audio_bt must be [B,T], got {tuple(audio.shape)}")
    n_fft_eff = int(max(16, int(n_fft)))
    hop_eff = int(max(1, int(hop)))
    if int(audio.shape[-1]) < int(n_fft_eff):
        audio = F.pad(audio, (0, int(n_fft_eff) - int(audio.shape[-1])))
    window = torch.hann_window(int(n_fft_eff), device=audio.device, dtype=audio.dtype)
    return torch.stft(
        audio,
        n_fft=int(n_fft_eff),
        hop_length=int(hop_eff),
        win_length=int(n_fft_eff),
        window=window,
        center=True,
        return_complex=True,
    )


def mrstft_logmag_l1_per_example(
    pred_audio_bct: torch.Tensor,
    target_audio_bct: torch.Tensor,
    valid_num_samples_b: torch.Tensor,
    *,
    resolutions: Sequence[tuple[int, int]] = DEFAULT_AUDIO_MRSTFT_RESOLUTIONS,
) -> torch.Tensor:
    pred_audio = torch.as_tensor(pred_audio_bct, dtype=torch.float32)
    target_audio = torch.as_tensor(target_audio_bct, dtype=torch.float32, device=pred_audio.device)
    if tuple(pred_audio.shape) != tuple(target_audio.shape):
        raise ValueError(
            f"pred_audio_bct and target_audio_bct must match, got {tuple(pred_audio.shape)} / {tuple(target_audio.shape)}"
        )
    if int(pred_audio.dim()) != 3:
        raise ValueError(f"expected [B,C,T] audio tensors, got {tuple(pred_audio.shape)}")
    pred_audio_bt = pred_audio.mean(dim=1)
    target_audio_bt = target_audio.mean(dim=1)
    valid_mask_bt = audio_valid_mask(
        valid_num_samples_b,
        max_num_samples=int(pred_audio_bt.shape[-1]),
    ).to(device=pred_audio.device)
    pred_audio_bt = pred_audio_bt * valid_mask_bt.to(dtype=pred_audio_bt.dtype)
    target_audio_bt = target_audio_bt * valid_mask_bt.to(dtype=target_audio_bt.dtype)
    total_b = pred_audio_bt.new_zeros((int(pred_audio_bt.shape[0]),), dtype=torch.float32)
    resolutions_eff = tuple((int(n_fft), int(hop)) for n_fft, hop in tuple(resolutions))
    if not resolutions_eff:
        raise ValueError("expected at least one MRSTFT resolution")
    for n_fft, hop in resolutions_eff:
        pred_spec = _safe_stft(pred_audio_bt, n_fft=int(n_fft), hop=int(hop))
        target_spec = _safe_stft(target_audio_bt, n_fft=int(n_fft), hop=int(hop))
        pred_logmag = torch.log1p(pred_spec.abs())
        target_logmag = torch.log1p(target_spec.abs())
        valid_frames_b = 1 + torch.div(valid_mask_bt.sum(dim=1).to(torch.long), int(max(1, hop)), rounding_mode="floor")
        frame_mask_bt = audio_valid_mask(
            valid_frames_b,
            max_num_samples=int(pred_logmag.shape[-1]),
        ).to(device=pred_audio.device, dtype=pred_logmag.dtype)
        weights_bft = frame_mask_bt[:, None, :]
        denom_b = (weights_bft.sum(dim=(1, 2)) * float(pred_logmag.shape[1])).clamp_min(1.0)
        total_b = total_b + (((pred_logmag - target_logmag).abs()) * weights_bft).sum(dim=(1, 2)) / denom_b
    return total_b / float(len(resolutions_eff))


def _onset_boost_for_class_name(
    name: str,
    *,
    kick_snare_boost: float = 3.0,
    hihat_boost: float = 1.0,
) -> float:
    normalized = str(name).strip().lower().replace("-", "_").replace(" ", "_")
    if "kick" in normalized or normalized in {"bd", "bass_drum"}:
        return float(kick_snare_boost)
    if "snare" in normalized or normalized in {"sd"}:
        return float(kick_snare_boost)
    if (
        "hihat" in normalized
        or "hi_hat" in normalized
        or normalized.endswith("_hh")
        or normalized.startswith("hh_")
        or normalized == "hh"
    ):
        return float(hihat_boost)
    return 0.0


def _build_onset_token_weights(
    prepared: Mapping[str, torch.Tensor | None],
    batch: Mapping[str, Any],
    *,
    base_weight: float = 1.0,
    kick_snare_boost: float = 3.0,
    hihat_boost: float = 1.0,
    token_radius: int = 1,
) -> torch.Tensor:
    target_mask = torch.as_tensor(prepared["target_valid_mask_bt"], dtype=torch.bool)
    weights = torch.full_like(target_mask, float(base_weight), dtype=torch.float32)
    weights = weights * target_mask.to(dtype=weights.dtype)

    family_onsets = prepared.get("family_onsets_bft")
    grid_times = prepared.get("grid_times_sec")
    token_times = prepared.get("token_times_sec")
    if family_onsets is None or grid_times is None or token_times is None:
        return weights.contiguous()

    class_names = [str(name) for name in list(batch.get("class_names") or [])]
    if not class_names:
        return weights.contiguous()

    boosts = [
        _onset_boost_for_class_name(
            name,
            kick_snare_boost=float(kick_snare_boost),
            hihat_boost=float(hihat_boost),
        )
        for name in class_names
    ]
    if not any(float(boost) > 0.0 for boost in boosts):
        return weights.contiguous()

    grid_valid = prepared.get("grid_valid_mask")
    radius = max(0, int(token_radius))
    batch_size = int(target_mask.shape[0])
    num_families = int(family_onsets.shape[1])
    token_count = int(target_mask.shape[1])
    inf = torch.tensor(float("inf"), dtype=token_times.dtype, device=token_times.device)

    for batch_idx in range(batch_size):
        target_valid_b = target_mask[batch_idx]
        if not bool(target_valid_b.any()):
            continue
        token_times_b = token_times[batch_idx]
        valid_token_count = int(target_valid_b.shape[0])
        if valid_token_count <= 0:
            continue
        grid_valid_b = (
            grid_valid[batch_idx]
            if grid_valid is not None
            else torch.ones_like(grid_times[batch_idx], dtype=torch.bool)
        )
        for family_idx in range(min(len(boosts), num_families)):
            boost = float(boosts[family_idx])
            if boost <= 0.0:
                continue
            onset_mask = family_onsets[batch_idx, family_idx] & grid_valid_b
            if not bool(onset_mask.any()):
                continue
            onset_times = grid_times[batch_idx][onset_mask]
            for onset_time in onset_times:
                distances = (token_times_b - onset_time).abs().masked_fill(~target_valid_b, inf)
                center_idx = int(distances.argmin().item())
                start = max(0, center_idx - radius)
                stop = min(token_count, center_idx + radius + 1)
                if stop <= start:
                    continue
                weights[batch_idx, start:stop] = weights[batch_idx, start:stop] + (
                    boost * target_valid_b[start:stop].to(dtype=weights.dtype)
                )
    return weights.contiguous()


def _masked_token_mean(
    per_token_bt: torch.Tensor,
    valid_mask_bt: torch.Tensor,
    token_weights_bt: torch.Tensor | None = None,
) -> torch.Tensor:
    per_token = torch.as_tensor(per_token_bt, dtype=torch.float32)
    mask = torch.as_tensor(valid_mask_bt, dtype=torch.bool, device=per_token.device)
    weights = mask.to(dtype=per_token.dtype)
    if token_weights_bt is not None:
        weights = weights * torch.as_tensor(token_weights_bt, dtype=per_token.dtype, device=per_token.device)
    return (per_token * weights).sum() / weights.sum().clamp_min(1.0e-8)


def _resolve_timbre_projection(
    timbre_projection: torch.Tensor | None,
    *,
    x_dim: int,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor | None:
    if timbre_projection is None:
        return None
    projection = torch.as_tensor(timbre_projection, dtype=dtype, device=device).detach()
    if int(projection.dim()) == 1:
        projection = projection.view(1, -1)
    if int(projection.dim()) != 2:
        raise ValueError(f"timbre_projection must be [K,{x_dim}], got {tuple(projection.shape)}")
    if int(projection.shape[1]) != int(x_dim):
        raise ValueError(f"timbre_projection must have {x_dim} columns, got {tuple(projection.shape)}")
    if int(projection.shape[0]) <= 0:
        raise ValueError("timbre_projection must contain at least one row")
    if not bool(torch.isfinite(projection).all()):
        raise ValueError("timbre_projection contains non-finite values")
    return projection.contiguous()


def _resolve_codebook_embeddings(
    quant_codebook_embed_ckd: torch.Tensor | None,
    *,
    x_dim: int,
    device: torch.device,
) -> torch.Tensor:
    if quant_codebook_embed_ckd is None:
        raise ValueError("quant_codebook_embed_ckd is required when codebook auxiliary losses are enabled")
    codebook_embed = torch.as_tensor(
        quant_codebook_embed_ckd,
        dtype=torch.float32,
        device=device,
    ).detach()
    if int(codebook_embed.dim()) != 3 or int(codebook_embed.shape[-1]) != int(x_dim):
        raise ValueError(
            "quant_codebook_embed_ckd must have shape [C,K,D] with "
            f"D={int(x_dim)}, got {tuple(codebook_embed.shape)}"
        )
    if int(codebook_embed.shape[0]) <= 0 or int(codebook_embed.shape[1]) <= 0:
        raise ValueError(f"quant_codebook_embed_ckd must be non-empty, got {tuple(codebook_embed.shape)}")
    return codebook_embed.contiguous()


def _resolve_rvq_target_codes_bct(
    *,
    prepared: Mapping[str, torch.Tensor | None],
    encodec_model: Any | None,
    target_codec_raw: torch.Tensor,
    target_mask: torch.Tensor,
    device: torch.device,
    target_pca_basis: Mapping[str, Any] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    source_codes = prepared.get("source_codes_bct")
    if source_codes is not None:
        target_codes = torch.as_tensor(source_codes, dtype=torch.long, device=device)
        if int(target_codes.dim()) != 3:
            raise ValueError(f"source_codes_bct must be [B,C,T], got {tuple(target_codes.shape)}")
    else:
        if encodec_model is None:
            raise ValueError("encodec_model is required for RVQ CE when source_codes_bct is absent")
        target_codes = requantize_latent_to_codes_bct(
            encodec_model,
            apply_seq_mask(target_codec_raw, target_mask),
            device=device,
            target_pca_basis=target_pca_basis,
        )

    compared_frames = int(min(int(target_codes.shape[-1]), int(target_mask.shape[-1]), int(target_codec_raw.shape[1])))
    target_codes = target_codes[:, :, :compared_frames].contiguous()
    valid_mask = target_mask[:, :compared_frames].to(dtype=torch.bool).contiguous()
    valid_mask = valid_mask & target_codes.ge(0).all(dim=1)
    target_codes = target_codes.clamp_min(0).contiguous()
    return target_codes, valid_mask


def _resolve_target_pca_basis(
    target_pca_basis: Mapping[str, Any] | None,
    *,
    device: torch.device,
    dtype: torch.dtype,
) -> dict[str, Any] | None:
    if target_pca_basis is None:
        return None
    return load_target_pca_basis(target_pca_basis, device=device, dtype=dtype)


def _rvq_cross_entropy_loss(
    pred_latent_btd: torch.Tensor,
    target_codes_bct: torch.Tensor,
    codebook_embed_ckd: torch.Tensor,
    valid_mask_bt: torch.Tensor,
    token_weights_bt: torch.Tensor | None = None,
) -> torch.Tensor:
    if int(pred_latent_btd.dim()) != 3:
        raise ValueError(f"pred_latent_btd must be [B,T,D], got {tuple(pred_latent_btd.shape)}")
    if int(target_codes_bct.dim()) != 3:
        raise ValueError(f"target_codes_bct must be [B,C,T], got {tuple(target_codes_bct.shape)}")
    if int(codebook_embed_ckd.dim()) != 3:
        raise ValueError(f"codebook_embed_ckd must be [C,K,D], got {tuple(codebook_embed_ckd.shape)}")
    batch_size, num_frames, x_dim = [int(x) for x in list(pred_latent_btd.shape)]
    if int(target_codes_bct.shape[0]) != int(batch_size):
        raise ValueError(
            f"target_codes_bct batch must match pred_latent_btd, got {tuple(target_codes_bct.shape)} "
            f"vs {tuple(pred_latent_btd.shape)}"
        )
    if int(codebook_embed_ckd.shape[0]) != int(target_codes_bct.shape[1]):
        raise ValueError(
            f"codebook count mismatch: codes={tuple(target_codes_bct.shape)} "
            f"embeddings={tuple(codebook_embed_ckd.shape)}"
        )
    if int(codebook_embed_ckd.shape[-1]) != int(x_dim):
        raise ValueError(
            f"codebook embedding dimension must be {int(x_dim)}, got {tuple(codebook_embed_ckd.shape)}"
        )

    compared_frames = int(min(int(num_frames), int(target_codes_bct.shape[-1]), int(valid_mask_bt.shape[-1])))
    pred_latent = pred_latent_btd[:, :compared_frames, :].to(dtype=torch.float32)
    target_codes = target_codes_bct[:, :, :compared_frames].to(device=pred_latent.device, dtype=torch.long)
    valid_mask = valid_mask_bt[:, :compared_frames].to(device=pred_latent.device, dtype=torch.bool)
    valid_mask = valid_mask & target_codes.ge(0).all(dim=1)
    target_codes = target_codes.clamp_min(0).contiguous()
    if not bool(valid_mask.any()):
        return pred_latent.sum() * 0.0

    token_weights = None
    if token_weights_bt is not None:
        token_weights = torch.as_tensor(
            token_weights_bt[:, :compared_frames],
            dtype=pred_latent.dtype,
            device=pred_latent.device,
        )
        token_weights = token_weights * valid_mask.to(dtype=token_weights.dtype)

    target_codebook_latents = token_ids_to_codebook_embeddings(
        target_codes,
        codebook_embed_ckd.to(device=pred_latent.device, dtype=torch.float32),
        valid_bt=valid_mask,
    ).detach()

    valid_flat = valid_mask.reshape(-1)
    weights_flat = None if token_weights is None else token_weights.reshape(-1)
    prev_sum = torch.zeros_like(pred_latent)
    total = pred_latent.new_zeros(())
    denom = pred_latent.new_zeros(())
    for codebook_idx in range(int(target_codes.shape[1])):
        residual_flat = (pred_latent - prev_sum).reshape(-1, int(x_dim))[valid_flat]
        labels_flat = target_codes[:, int(codebook_idx), :].reshape(-1)[valid_flat]
        embed_kd = codebook_embed_ckd[int(codebook_idx)].to(device=pred_latent.device, dtype=torch.float32)
        dist_sq = (
            residual_flat.square().sum(dim=-1, keepdim=True)
            + embed_kd.square().sum(dim=-1).view(1, -1)
            - (2.0 * residual_flat.matmul(embed_kd.transpose(0, 1)))
        )
        logits = -torch.sqrt(dist_sq.clamp_min(1.0e-12))
        ce = F.cross_entropy(logits, labels_flat, reduction="none")
        if weights_flat is None:
            total = total + ce.sum()
            denom = denom + torch.as_tensor(float(int(ce.numel())), dtype=denom.dtype, device=denom.device)
        else:
            weights_valid = weights_flat[valid_flat]
            total = total + (ce * weights_valid).sum()
            denom = denom + weights_valid.sum()
        prev_sum = prev_sum + target_codebook_latents[:, int(codebook_idx), :, :]
    return total / denom.clamp_min(1.0e-8)


def diffusion_train_step(
    model: ConditionalDiffusionTransformer,
    diffusion: GaussianDiffusion1D,
    batch: Mapping[str, Any],
    device: torch.device,
    *,
    target_mean=None,
    target_std=None,
    encodec_model: Any | None = None,
    audio_sample_rate: int | None = None,
    audio_wave_l1_weight: float = DEFAULT_AUDIO_WAVE_L1_WEIGHT,
    audio_mrstft_weight: float = DEFAULT_AUDIO_MRSTFT_WEIGHT,
    audio_mrstft_resolutions: Sequence[tuple[int, int]] = DEFAULT_AUDIO_MRSTFT_RESOLUTIONS,
    x0_clip_norm: float | None = DEFAULT_SAMPLE_X0_CLIP_NORM,
    timbre_projection: torch.Tensor | None = None,
    x0_mse_weight: float = 0.0,
    timbre_proj_mse_weight: float = 0.0,
    quant_embed_mse_weight: float = 0.0,
    rvq_ce_weight: float = 0.0,
    quant_codebook_embed_ckd: torch.Tensor | None = None,
    onset_loss_weighting: bool = False,
    onset_token_radius: int = 1,
    target_pca_basis: Mapping[str, Any] | None = None,
    use_bpm_training_geometry: bool = False,
    bpm_geometry_num_beats: int = DEFAULT_INFERENCE_NUM_BEATS,
):
    prepared = _prepare_batch_tensors(batch, device)
    if bool(use_bpm_training_geometry):
        prepared = apply_bpm_training_geometry_to_prepared_batch(
            prepared,
            num_beats=int(bpm_geometry_num_beats),
        )
    target_raw = prepared["target_btd"]
    target = normalize_latent(target_raw, target_mean, target_std)
    target_mask = prepared["target_valid_mask_bt"]
    resolved_target_pca_basis = _resolve_target_pca_basis(
        target_pca_basis,
        device=device,
        dtype=target_raw.dtype,
    )
    target_codec_raw: torch.Tensor | None = None

    def _target_codec_latent() -> torch.Tensor:
        nonlocal target_codec_raw
        if target_codec_raw is not None:
            return target_codec_raw
        target_sum = prepared.get("target_sum_btd")
        if target_sum is not None:
            target_codec_raw = apply_seq_mask(target_sum, target_mask)
        else:
            target_codec_raw = apply_seq_mask(
                reconstruct_latent_from_pca(target_raw, resolved_target_pca_basis),
                target_mask,
            )
        return target_codec_raw

    noise = torch.randn_like(target)
    noise = apply_seq_mask(noise, target_mask)
    x0_prior = prepared.get("x0_prior_btd")
    if x0_prior is not None:
        x0_prior = normalize_latent(x0_prior, target_mean, target_std)
        x0_prior = apply_seq_mask(x0_prior, target_mask)

    batch_size = int(target.shape[0])
    t = torch.randint(0, diffusion.num_steps, (batch_size,), device=device)

    x_t = diffusion.q_sample(target, t, noise)
    x_t = apply_seq_mask(x_t, target_mask)

    pred_eps = model(
        x_t=x_t,
        t=t,
        target_valid_mask_bt=target_mask,
        grid=prepared["grid"],
        grid_ids=prepared["grid_ids"],
        grid_times_sec=prepared["grid_times_sec"],
        token_times_sec=prepared["token_times_sec"],
        grid_valid_mask_bt=prepared["grid_valid_mask"],
        beat_boundaries_sec=prepared["beat_boundaries_sec"],
        beat_boundaries_valid_mask=prepared["beat_boundaries_valid_mask"],
        bpm=prepared["bpm"],
        duration_sec=prepared["duration_sec"],
        timbre_bank_latents=prepared.get("timbre_bank_latents"),
        timbre_bank_family_ids=prepared.get("timbre_bank_family_ids"),
        timbre_bank_class_ids=prepared.get("timbre_bank_class_ids"),
        timbre_bank_velocity=prepared.get("timbre_bank_velocity"),
        timbre_bank_mask=prepared.get("timbre_bank_mask"),
        timbre_dynamic_features=prepared.get("timbre_dynamic_features"),
        timbre_dynamic_mask=prepared.get("timbre_dynamic_mask"),
        timbre_dynamic_counts=prepared.get("timbre_dynamic_counts"),
        timbre_family_default_indices=prepared.get("timbre_family_default_indices"),
        timbre_class_token_indices=prepared.get("timbre_class_token_indices"),
        reference_timbre_bank_latents=prepared.get("reference_timbre_bank_latents"),
        reference_timbre_bank_family_ids=prepared.get("reference_timbre_bank_family_ids"),
        reference_timbre_bank_class_ids=prepared.get("reference_timbre_bank_class_ids"),
        reference_timbre_bank_velocity=prepared.get("reference_timbre_bank_velocity"),
        reference_timbre_bank_mask=prepared.get("reference_timbre_bank_mask"),
        reference_timbre_dynamic_features=prepared.get("reference_timbre_dynamic_features"),
        reference_timbre_dynamic_mask=prepared.get("reference_timbre_dynamic_mask"),
        reference_timbre_dynamic_counts=prepared.get("reference_timbre_dynamic_counts"),
        reference_timbre_family_default_indices=prepared.get("reference_timbre_family_default_indices"),
        reference_timbre_class_token_indices=prepared.get("reference_timbre_class_token_indices"),
        reference_segment_pca144=prepared.get("reference_segment_pca144"),
        x0_prior_btd=x0_prior,
    )

    loss_per_bt = ((pred_eps - noise) ** 2).mean(dim=-1)
    diffusion_loss = loss_per_bt[target_mask].mean()
    x0_hat = diffusion.predict_x0_from_eps(x_t, t, pred_eps)
    if x0_clip_norm is not None:
        x0_hat = x0_hat.clamp(min=-float(x0_clip_norm), max=float(x0_clip_norm))
    x0_hat = apply_seq_mask(x0_hat, target_mask)

    loss = diffusion_loss
    x0_loss = x0_hat.new_zeros(())
    timbre_proj_mse = x0_hat.new_zeros(())
    quant_embed_mse = x0_hat.new_zeros(())
    rvq_ce = x0_hat.new_zeros(())
    onset_weighted_x0 = x0_hat.new_zeros(())
    per_tok_x0 = ((x0_hat - target) ** 2).mean(dim=-1)
    use_x0_loss = float(x0_mse_weight) > 0.0
    use_timbre_proj_loss = float(timbre_proj_mse_weight) > 0.0
    use_quant_embed_loss = float(quant_embed_mse_weight) > 0.0
    use_rvq_ce_loss = float(rvq_ce_weight) > 0.0
    token_weights = None
    if bool(onset_loss_weighting) and (
        use_x0_loss or use_timbre_proj_loss or use_quant_embed_loss or use_rvq_ce_loss
    ):
        token_weights = _build_onset_token_weights(
            prepared,
            batch,
            token_radius=int(onset_token_radius),
        ).to(device=x0_hat.device, dtype=x0_hat.dtype)
        onset_weighted_x0 = _masked_token_mean(per_tok_x0, target_mask, token_weights)

    if use_x0_loss:
        x0_loss = _masked_token_mean(per_tok_x0, target_mask, None)
        x0_objective = onset_weighted_x0 if token_weights is not None else x0_loss
        loss = loss + (float(x0_mse_weight) * x0_objective)

    if use_timbre_proj_loss:
        projection = _resolve_timbre_projection(
            timbre_projection,
            x_dim=int(x0_hat.shape[-1]),
            device=x0_hat.device,
            dtype=x0_hat.dtype,
        )
        if projection is None:
            raise ValueError("timbre_projection is required when timbre_proj_mse_weight > 0")
        projected_error = torch.matmul(x0_hat - target, projection.transpose(0, 1))
        per_tok_proj = projected_error.square().sum(dim=-1)
        timbre_proj_mse = _masked_token_mean(per_tok_proj, target_mask, token_weights)
        loss = loss + (float(timbre_proj_mse_weight) * timbre_proj_mse)

    pred_latent_raw: torch.Tensor | None = None
    pred_codec_latent_raw: torch.Tensor | None = None
    codebook_embed: torch.Tensor | None = None
    if use_quant_embed_loss or use_rvq_ce_loss:
        codebook_embed = _resolve_codebook_embeddings(
            quant_codebook_embed_ckd,
            x_dim=int(_target_codec_latent().shape[-1]),
            device=device,
        )

    if use_quant_embed_loss:
        if encodec_model is None:
            raise ValueError("encodec_model is required when quant_embed_mse_weight > 0")
        if codebook_embed is None:
            raise AssertionError("codebook embeddings should have been resolved")
        pred_latent_raw = denormalize_latent(x0_hat, target_mean, target_std)
        pred_latent_raw = apply_seq_mask(pred_latent_raw, target_mask)
        pred_codec_latent_raw = reconstruct_latent_from_pca(
            pred_latent_raw,
            resolved_target_pca_basis,
        )
        pred_codec_latent_raw = apply_seq_mask(pred_codec_latent_raw, target_mask)
        with torch.no_grad():
            target_requant_codes = requantize_latent_to_codes_bct(
                encodec_model,
                _target_codec_latent(),
                device=device,
            )
            compared_frames = int(min(int(target_requant_codes.shape[-1]), int(target_mask.shape[-1])))
            target_requant_codes = target_requant_codes[:, :, :compared_frames]
            quant_valid_mask = target_mask[:, :compared_frames]
            target_codebook_latents = token_ids_to_codebook_embeddings(
                target_requant_codes,
                codebook_embed,
                valid_bt=quant_valid_mask,
            )
            target_requant_sum = rvq_sum_latents(
                target_codebook_latents,
                valid_bt=quant_valid_mask,
            )
        pred_quant_aligned = pred_codec_latent_raw[:, :compared_frames, :]
        per_tok_quant = ((pred_quant_aligned - target_requant_sum) ** 2).mean(dim=-1)
        token_weights_quant = None if token_weights is None else token_weights[:, :compared_frames]
        quant_embed_mse = _masked_token_mean(per_tok_quant, quant_valid_mask, token_weights_quant)
        loss = loss + (float(quant_embed_mse_weight) * quant_embed_mse)

    if use_rvq_ce_loss:
        if codebook_embed is None:
            raise AssertionError("codebook embeddings should have been resolved")
        if pred_latent_raw is None:
            pred_latent_raw = denormalize_latent(x0_hat, target_mean, target_std)
            pred_latent_raw = apply_seq_mask(pred_latent_raw, target_mask)
        if pred_codec_latent_raw is None:
            pred_codec_latent_raw = reconstruct_latent_from_pca(
                pred_latent_raw,
                resolved_target_pca_basis,
            )
            pred_codec_latent_raw = apply_seq_mask(pred_codec_latent_raw, target_mask)
        target_rvq_codes, rvq_valid_mask = _resolve_rvq_target_codes_bct(
            prepared=prepared,
            encodec_model=encodec_model,
            target_codec_raw=_target_codec_latent(),
            target_mask=target_mask,
            device=device,
            target_pca_basis=resolved_target_pca_basis,
        )
        token_weights_rvq = None if token_weights is None else token_weights[:, : int(rvq_valid_mask.shape[-1])]
        rvq_ce = _rvq_cross_entropy_loss(
            pred_codec_latent_raw[:, : int(rvq_valid_mask.shape[-1]), :],
            target_rvq_codes,
            codebook_embed,
            rvq_valid_mask,
            token_weights_rvq,
        )
        loss = loss + (float(rvq_ce_weight) * rvq_ce)

    audio_wave_l1 = x0_hat.new_zeros(())
    audio_mrstft = x0_hat.new_zeros(())
    if encodec_model is not None and (
        float(audio_wave_l1_weight) > 0.0 or float(audio_mrstft_weight) > 0.0
    ):
        if pred_latent_raw is None:
            pred_latent_raw = denormalize_latent(x0_hat, target_mean, target_std)
            pred_latent_raw = apply_seq_mask(pred_latent_raw, target_mask)
        if pred_codec_latent_raw is None:
            pred_codec_latent_raw = reconstruct_latent_from_pca(
                pred_latent_raw,
                resolved_target_pca_basis,
            )
            pred_codec_latent_raw = apply_seq_mask(pred_codec_latent_raw, target_mask)
        # EnCodec's decoder contains recurrent layers. When the frozen model stays in eval mode,
        # CuDNN RNN backward can fail on CUDA, so route these loss decodes through the non-CuDNN path.
        with torch.backends.cudnn.flags(enabled=False):
            pred_audio_bct = decode_latent_to_audio(
                pred_latent_raw,
                encodec_model,
                target_pca_basis=resolved_target_pca_basis,
            )
            with torch.no_grad():
                target_audio_bct = decode_latent_to_audio(
                    _target_codec_latent(),
                    encodec_model,
                )
        max_num_samples = int(min(pred_audio_bct.shape[-1], target_audio_bct.shape[-1]))
        valid_num_samples_b = resolve_valid_audio_num_samples(
            prepared["duration_sec"],
            sample_rate=int(audio_sample_rate or resolve_encodec_sample_rate(encodec_model)),
            max_num_samples=int(max_num_samples),
        )
        pred_audio_eff = pred_audio_bct[..., : int(max_num_samples)]
        target_audio_eff = target_audio_bct[..., : int(max_num_samples)]
        audio_wave_l1 = masked_audio_l1_per_example(
            pred_audio_eff,
            target_audio_eff,
            valid_num_samples_b,
        ).mean()
        audio_mrstft = mrstft_logmag_l1_per_example(
            pred_audio_eff,
            target_audio_eff,
            valid_num_samples_b,
            resolutions=audio_mrstft_resolutions,
        ).mean()
        loss = loss + (float(audio_wave_l1_weight) * audio_wave_l1) + (float(audio_mrstft_weight) * audio_mrstft)

    with torch.no_grad():
        per_tok = ((x0_hat - target) ** 2).mean(dim=-1)
        per_ex = []
        for idx in range(per_tok.shape[0]):
            per_ex.append(per_tok[idx][target_mask[idx]].mean())
        x0_mse_median = torch.stack(per_ex).median()

    return {
        "loss": loss,
        "diffusion_loss": diffusion_loss,
        "audio_wave_l1": audio_wave_l1,
        "audio_mrstft": audio_mrstft,
        "x0_loss": x0_loss,
        "timbre_proj_mse": timbre_proj_mse,
        "quant_embed_mse": quant_embed_mse,
        "rvq_ce": rvq_ce,
        "onset_weighted_x0": onset_weighted_x0,
        "x0_mse_median": x0_mse_median,
        "t": t,
    }


@torch.no_grad()
def sample_ddpm(
    model: ConditionalDiffusionTransformer,
    diffusion: GaussianDiffusion1D,
    batch: Mapping[str, Any],
    device: torch.device,
    guidance_scale: float = 1.0,
    x0_clip_norm: float | None = DEFAULT_SAMPLE_X0_CLIP_NORM,
    sample_idx: int | None = None,
    start_noise: torch.Tensor | None = None,
    step_noises: Mapping[int, torch.Tensor] | None = None,
    sample_seed: int | None = None,
    use_bpm_inference_geometry: bool = False,
    inference_num_beats: int = DEFAULT_INFERENCE_NUM_BEATS,
    target_token_rate_hz: float = DEFAULT_TARGET_TOKEN_RATE_HZ,
    inference_geometry: Mapping[str, Any] | None = None,
):
    prepared = _prepare_batch_tensors(
        batch,
        device,
        require_target=not bool(use_bpm_inference_geometry),
        require_timing=not bool(use_bpm_inference_geometry),
    )
    if sample_idx is not None:
        grid = prepared["grid"]
        if grid is None:
            raise ValueError("prepared batch is missing grid")
        batch_size = int(grid.shape[0])
        if not (0 <= int(sample_idx) < int(batch_size)):
            raise IndexError(f"sample_idx={sample_idx} out of range for batch size={int(batch_size)}")
        prepared = _slice_prepared_batch(prepared, int(sample_idx))
    if inference_geometry is None:
        geometry = resolve_inference_geometry(
            prepared,
            use_bpm_inference_geometry=bool(use_bpm_inference_geometry),
            inference_num_beats=int(inference_num_beats),
            target_token_rate_hz=float(target_token_rate_hz),
        )
    else:
        geometry = _prepare_geometry_tensors(inference_geometry, device=device)
        if sample_idx is not None:
            geometry = _slice_inference_geometry(geometry, int(sample_idx))
    target_mask = geometry["target_valid_mask_bt"]
    grid = prepared["grid"]
    grid_ids = prepared["grid_ids"]
    grid_valid_mask = prepared["grid_valid_mask"]
    if grid is None or grid_valid_mask is None:
        raise ValueError("prepared batch is missing grid or grid_valid_mask")
    cond_btd, cond_valid_mask_bt = model.encode_conditioning(
        grid=grid,
        grid_ids=grid_ids,
        grid_times_sec=geometry["grid_times_sec"],
        token_times_sec=geometry["token_times_sec"],
        target_valid_mask_bt=target_mask,
        grid_valid_mask_bt=grid_valid_mask,
        timbre_bank_latents=prepared.get("timbre_bank_latents"),
        timbre_bank_family_ids=prepared.get("timbre_bank_family_ids"),
        timbre_bank_class_ids=prepared.get("timbre_bank_class_ids"),
        timbre_bank_velocity=prepared.get("timbre_bank_velocity"),
        timbre_bank_mask=prepared.get("timbre_bank_mask"),
        timbre_dynamic_features=prepared.get("timbre_dynamic_features"),
        timbre_dynamic_mask=prepared.get("timbre_dynamic_mask"),
        timbre_dynamic_counts=prepared.get("timbre_dynamic_counts"),
        timbre_family_default_indices=prepared.get("timbre_family_default_indices"),
        timbre_class_token_indices=prepared.get("timbre_class_token_indices"),
        reference_timbre_bank_latents=prepared.get("reference_timbre_bank_latents"),
        reference_timbre_bank_family_ids=prepared.get("reference_timbre_bank_family_ids"),
        reference_timbre_bank_class_ids=prepared.get("reference_timbre_bank_class_ids"),
        reference_timbre_bank_velocity=prepared.get("reference_timbre_bank_velocity"),
        reference_timbre_bank_mask=prepared.get("reference_timbre_bank_mask"),
        reference_timbre_dynamic_features=prepared.get("reference_timbre_dynamic_features"),
        reference_timbre_dynamic_mask=prepared.get("reference_timbre_dynamic_mask"),
        reference_timbre_dynamic_counts=prepared.get("reference_timbre_dynamic_counts"),
        reference_timbre_family_default_indices=prepared.get("reference_timbre_family_default_indices"),
        reference_timbre_class_token_indices=prepared.get("reference_timbre_class_token_indices"),
        reference_segment_pca144=prepared.get("reference_segment_pca144"),
    )

    batch_size = int(target_mask.shape[0])
    target_len = int(target_mask.shape[1])
    latent_dim = int(model.cfg.x_dim)
    noise_generator = None
    if sample_seed is not None:
        noise_generator = torch.Generator(device=device)
        noise_generator.manual_seed(int(sample_seed))
    if start_noise is None:
        x = torch.randn(
            batch_size,
            target_len,
            latent_dim,
            device=device,
            generator=noise_generator,
        )
    else:
        x = torch.as_tensor(start_noise, dtype=torch.float32, device=device).clone()
        expected_shape = (batch_size, target_len, latent_dim)
        if tuple(x.shape) != expected_shape:
            raise ValueError(f"start_noise must have shape {expected_shape}, got {tuple(x.shape)}")
    x = apply_seq_mask(x, target_mask)

    for step in reversed(range(diffusion.num_steps)):
        t = torch.full((batch_size,), step, device=device, dtype=torch.long)
        if float(guidance_scale) == 1.0:
            eps = model(
                x_t=x,
                t=t,
                target_valid_mask_bt=target_mask,
                token_times_sec=geometry["token_times_sec"],
                cond_btd=cond_btd,
                cond_valid_mask_bt=cond_valid_mask_bt,
                force_uncond=False,
            )
        else:
            eps_cond = model(
                x_t=x,
                t=t,
                target_valid_mask_bt=target_mask,
                token_times_sec=geometry["token_times_sec"],
                cond_btd=cond_btd,
                cond_valid_mask_bt=cond_valid_mask_bt,
                force_uncond=False,
            )
            eps_uncond = model(
                x_t=x,
                t=t,
                target_valid_mask_bt=target_mask,
                token_times_sec=geometry["token_times_sec"],
                cond_btd=cond_btd,
                cond_valid_mask_bt=cond_valid_mask_bt,
                force_uncond=True,
            )
            eps = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
        x0_hat = diffusion.predict_x0_from_eps(x, t, eps)
        if x0_clip_norm is not None:
            x0_hat = x0_hat.clamp(min=-float(x0_clip_norm), max=float(x0_clip_norm))
        x0_hat = apply_seq_mask(x0_hat, target_mask)
        mean = diffusion.posterior_mean_from_x0(x, t, x0_hat)

        if step > 0:
            if step_noises is None or int(step) not in step_noises:
                z = torch.randn(
                    tuple(x.shape),
                    dtype=x.dtype,
                    device=x.device,
                    generator=noise_generator,
                )
            else:
                z = torch.as_tensor(step_noises[int(step)], dtype=torch.float32, device=device).clone()
                if tuple(z.shape) != tuple(x.shape):
                    raise ValueError(f"step_noises[{step}] must have shape {tuple(x.shape)}, got {tuple(z.shape)}")
            var = diffusion.posterior_variance[t].view(-1, 1, 1)
            x = mean + torch.sqrt(var) * z
        else:
            x = mean
        x = apply_seq_mask(x, target_mask)

    return x


def _plot_matrix(
    ax: Any,
    matrix_td: torch.Tensor,
    *,
    title: str,
    token_times_sec_t: torch.Tensor | None = None,
    vabs: float | None = None,
    ylabel: str = "latent dim",
    cmap: str = "coolwarm",
    vmin: float | None = None,
    vmax: float | None = None,
    transpose: bool = True,
) -> None:
    matrix = torch.as_tensor(matrix_td, dtype=torch.float32).detach().cpu()
    image_data = matrix.T.numpy() if bool(transpose) else matrix.numpy()
    extent = None
    if token_times_sec_t is not None:
        times = torch.as_tensor(token_times_sec_t, dtype=torch.float32).detach().cpu().view(-1).numpy()
        if int(times.shape[0]) == int(image_data.shape[1]) and int(times.shape[0]) > 0:
            lo = float(times[0])
            hi = float(times[-1]) if int(times.shape[0]) > 1 else float(times[0] + 1.0)
            extent = (lo, hi, -0.5, float(image_data.shape[0]) - 0.5)
    image = ax.imshow(
        image_data,
        aspect="auto",
        origin="lower",
        interpolation="nearest",
        cmap=str(cmap),
        vmin=float(vmin) if vmin is not None else (-float(vabs) if vabs is not None else None),
        vmax=float(vmax) if vmax is not None else (float(vabs) if vabs is not None else None),
        extent=extent,
    )
    ax.set_title(title)
    ax.set_xlabel("time (seconds)" if extent is not None else "time")
    ax.set_ylabel(str(ylabel))
    plt.colorbar(image, ax=ax, fraction=0.02, pad=0.01)


def save_prediction_diagnostic_plot(
    out_path: str | os.PathLike[str],
    *,
    input_features_td: torch.Tensor | None,
    pred_latent_td: torch.Tensor,
    target_latent_td: torch.Tensor,
    target_codes_ct: torch.Tensor | None = None,
    pred_codes_ct: torch.Tensor | None = None,
    pred_target_token_acc: float | None = None,
    token_times_sec_t: torch.Tensor | None = None,
) -> None:
    input_features = None if input_features_td is None else torch.as_tensor(input_features_td, dtype=torch.float32)
    pred = torch.as_tensor(pred_latent_td, dtype=torch.float32)
    target = torch.as_tensor(target_latent_td, dtype=torch.float32)
    target_codes = None if target_codes_ct is None else torch.as_tensor(target_codes_ct, dtype=torch.float32)
    pred_codes = None if pred_codes_ct is None else torch.as_tensor(pred_codes_ct, dtype=torch.float32)
    diff = pred - target

    target_vabs = max(float(target.abs().max().item()) if int(target.numel()) > 0 else 0.0, 1.0e-6)
    pred_vabs = max(float(pred.abs().max().item()) if int(pred.numel()) > 0 else 0.0, 1.0e-6)
    diff_vabs = max(float(diff.abs().max().item()) if int(diff.numel()) > 0 else 0.0, 1.0e-6)
    has_codes = target_codes is not None and pred_codes is not None
    num_latent_rows = 4 if input_features is not None else 3
    total_rows = num_latent_rows + (1 if has_codes else 0)
    fig = plt.figure(figsize=(18, 3.5 * total_rows), constrained_layout=True)
    grid = fig.add_gridspec(total_rows, 2)

    row_idx = 0
    if input_features is not None:
        input_vabs = max(float(input_features.abs().max().item()), 1.0e-6)
        _plot_matrix(
            fig.add_subplot(grid[row_idx, :]),
            input_features,
            title=f"Conditioning input [T,{int(input_features.shape[-1])}]",
            token_times_sec_t=token_times_sec_t,
            vabs=input_vabs,
            ylabel="feature dim",
        )
        row_idx += 1

    _plot_matrix(
        fig.add_subplot(grid[row_idx, :]),
        target,
        title=f"Target latent [T,128] | absmax={target_vabs:.2f}",
        token_times_sec_t=token_times_sec_t,
        vabs=target_vabs,
    )
    _plot_matrix(
        fig.add_subplot(grid[row_idx + 1, :]),
        pred,
        title=f"Predicted latent [T,128] | absmax={pred_vabs:.2f}",
        token_times_sec_t=token_times_sec_t,
        vabs=pred_vabs,
    )
    _plot_matrix(
        fig.add_subplot(grid[row_idx + 2, :]),
        diff,
        title=f"Prediction error [T,128] | absmax={diff_vabs:.2f}",
        token_times_sec_t=token_times_sec_t,
        vabs=diff_vabs,
    )
    row_idx += 3

    if has_codes:
        target_ax = fig.add_subplot(grid[row_idx, 0])
        pred_ax = fig.add_subplot(grid[row_idx, 1])
        codebook_max = max(
            int(target_codes.max().item()) if int(target_codes.numel()) > 0 else 0,
            int(pred_codes.max().item()) if int(pred_codes.numel()) > 0 else 0,
        )
        _plot_matrix(
            target_ax,
            target_codes,
            title="Target requantized codes [C,T]",
            token_times_sec_t=token_times_sec_t,
            ylabel="codebook",
            cmap="tab20",
            vmin=-1.0,
            vmax=float(max(1, codebook_max)),
            transpose=False,
        )
        token_acc_text = "n/a" if pred_target_token_acc is None else f"{float(pred_target_token_acc):.4f}"
        _plot_matrix(
            pred_ax,
            pred_codes,
            title=f"Predicted requantized codes [C,T] | target_acc={token_acc_text}",
            token_times_sec_t=token_times_sec_t,
            ylabel="codebook",
            cmap="tab20",
            vmin=-1.0,
            vmax=float(max(1, codebook_max)),
            transpose=False,
        )

    out_path_str = os.fspath(out_path)
    out_dir = os.path.dirname(out_path_str)
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)
    fig.savefig(out_path_str, dpi=160)
    plt.close(fig)


@torch.no_grad()
def save_eval_plot_multi_t(
    model,
    diffusion,
    batch,
    device,
    epoch,
    out_dir="eval_plots",
    sample_idx=0,
    t_values=(100, 300, 600),
    fixed_noises=None,
    target_mean=None,
    target_std=None,
    x0_clip_norm: float | None = DEFAULT_SAMPLE_X0_CLIP_NORM,
    use_bpm_training_geometry: bool = False,
    bpm_geometry_num_beats: int = DEFAULT_INFERENCE_NUM_BEATS,
):
    os.makedirs(out_dir, exist_ok=True)
    model.eval()

    prepared = _prepare_batch_tensors(batch, device)
    if bool(use_bpm_training_geometry):
        prepared = apply_bpm_training_geometry_to_prepared_batch(
            prepared,
            num_beats=int(bpm_geometry_num_beats),
        )
    batch_size = int(prepared["target_btd"].shape[0])
    if not (0 <= int(sample_idx) < int(batch_size)):
        raise IndexError(f"sample_idx={sample_idx} out of range for batch size={int(batch_size)}")

    single = _slice_prepared_batch(prepared, sample_idx)
    target_i = single["target_btd"]
    target_mask_i = single["target_valid_mask_bt"]
    target_norm_i = normalize_latent(target_i, target_mean, target_std)
    cond_i, cond_mask_i = model.encode_conditioning(
        grid=single["grid"],
        grid_ids=single["grid_ids"],
        grid_times_sec=single["grid_times_sec"],
        token_times_sec=single["token_times_sec"],
        target_valid_mask_bt=target_mask_i,
        grid_valid_mask_bt=single["grid_valid_mask"],
        timbre_bank_latents=single.get("timbre_bank_latents"),
        timbre_bank_family_ids=single.get("timbre_bank_family_ids"),
        timbre_bank_class_ids=single.get("timbre_bank_class_ids"),
        timbre_bank_velocity=single.get("timbre_bank_velocity"),
        timbre_bank_mask=single.get("timbre_bank_mask"),
        timbre_dynamic_features=single.get("timbre_dynamic_features"),
        timbre_dynamic_mask=single.get("timbre_dynamic_mask"),
        timbre_dynamic_counts=single.get("timbre_dynamic_counts"),
        timbre_family_default_indices=single.get("timbre_family_default_indices"),
        timbre_class_token_indices=single.get("timbre_class_token_indices"),
        reference_timbre_bank_latents=single.get("reference_timbre_bank_latents"),
        reference_timbre_bank_family_ids=single.get("reference_timbre_bank_family_ids"),
        reference_timbre_bank_class_ids=single.get("reference_timbre_bank_class_ids"),
        reference_timbre_bank_velocity=single.get("reference_timbre_bank_velocity"),
        reference_timbre_bank_mask=single.get("reference_timbre_bank_mask"),
        reference_timbre_dynamic_features=single.get("reference_timbre_dynamic_features"),
        reference_timbre_dynamic_mask=single.get("reference_timbre_dynamic_mask"),
        reference_timbre_dynamic_counts=single.get("reference_timbre_dynamic_counts"),
        reference_timbre_family_default_indices=single.get("reference_timbre_family_default_indices"),
        reference_timbre_class_token_indices=single.get("reference_timbre_class_token_indices"),
        reference_segment_pca144=single.get("reference_segment_pca144"),
    )

    valid_len = int(target_mask_i[0].sum().item())
    target_np = target_i[0, :valid_len].detach().cpu().T.numpy()

    nrows = 1 + len(t_values)
    fig, axes = plt.subplots(nrows, 1, figsize=(14, 3 * nrows), squeeze=False)
    ax = axes[0, 0]
    im = ax.imshow(target_np, aspect="auto", origin="lower")
    ax.set_title(f"Epoch {epoch} - Target")
    ax.set_ylabel("latent dim")
    plt.colorbar(im, ax=ax, fraction=0.02, pad=0.01)

    for row, t_eval in enumerate(t_values, start=1):
        if not (0 <= int(t_eval) < int(diffusion.num_steps)):
            raise ValueError(
                f"t_eval={t_eval} out of range for diffusion.num_steps={int(diffusion.num_steps)}; "
                f"expected values in [0, {int(diffusion.num_steps) - 1}]"
            )
        if fixed_noises is not None and t_eval in fixed_noises:
            noise = fixed_noises[t_eval].to(device)
        else:
            noise = torch.randn_like(target_i)
        if noise.shape != target_i.shape:
            raise ValueError(f"Noise shape {noise.shape} != target shape {target_i.shape}")

        noise = noise * target_mask_i.unsqueeze(-1)
        t = torch.full((1,), t_eval, device=device, dtype=torch.long)
        x_t = diffusion.q_sample(target_norm_i, t, noise)
        x_t = x_t * target_mask_i.unsqueeze(-1)

        pred_eps = model(
            x_t=x_t,
            t=t,
            target_valid_mask_bt=target_mask_i,
            token_times_sec=single["token_times_sec"],
            cond_btd=cond_i,
            cond_valid_mask_bt=cond_mask_i,
        )
        x0_hat_norm = diffusion.predict_x0_from_eps(x_t, t, pred_eps)
        if x0_clip_norm is not None:
            x0_hat_norm = x0_hat_norm.clamp(min=-float(x0_clip_norm), max=float(x0_clip_norm))
        x0_hat_norm = x0_hat_norm * target_mask_i.unsqueeze(-1)
        x0_hat = denormalize_latent(x0_hat_norm, target_mean, target_std)
        x0_hat = x0_hat * target_mask_i.unsqueeze(-1)

        pred_np = x0_hat[0, :valid_len].detach().cpu().T.numpy()
        mse_val = ((x0_hat - target_i)[0, :valid_len].pow(2).mean()).item()

        ax = axes[row, 0]
        im = ax.imshow(pred_np, aspect="auto", origin="lower")
        ax.set_title(f"Predicted x0_hat at t={t_eval}  |  mse={mse_val:.4f}")
        ax.set_ylabel("latent dim")
        if row == nrows - 1:
            ax.set_xlabel("time frame")
        plt.colorbar(im, ax=ax, fraction=0.02, pad=0.01)

    plt.tight_layout()
    save_path = os.path.join(out_dir, f"epoch_{epoch:03d}_multi_t.png")
    plt.savefig(save_path, dpi=150)
    plt.close(fig)
    return save_path


def decode_latent_to_audio(
    pred_latent_btd,
    encodec_model,
    *,
    target_pca_basis: Mapping[str, Any] | None = None,
):
    latent = reconstruct_latent_from_pca(
        torch.as_tensor(pred_latent_btd, dtype=torch.float32),
        target_pca_basis,
    )
    return decode_quantized_latent_to_audio(encodec_model, latent)


def resolve_encodec_sample_rate(encodec_model, default: int = 32000) -> int:
    return resolve_audio_codec_sample_rate(encodec_model, default=default)


def stitch_audio_segments_with_crossfade(
    audio_segments_ct: Sequence[torch.Tensor],
    *,
    crossfade_num_samples: int,
) -> torch.Tensor:
    if not audio_segments_ct:
        raise ValueError("expected at least one audio segment")
    segments = [torch.as_tensor(segment, dtype=torch.float32).contiguous() for segment in list(audio_segments_ct)]
    out = segments[0]
    if int(out.dim()) != 2:
        raise ValueError(f"expected audio segments with shape [C,T], got {tuple(out.shape)}")
    crossfade = int(max(0, int(crossfade_num_samples)))
    for segment in list(segments[1:]):
        if tuple(segment.shape[:-1]) != tuple(out.shape[:-1]):
            raise ValueError(
                f"audio segment channel dimensions must match, got {tuple(out.shape)} / {tuple(segment.shape)}"
            )
        overlap = int(min(crossfade, int(out.shape[-1]), int(segment.shape[-1])))
        if int(overlap) <= 0:
            out = torch.cat((out, segment), dim=-1).contiguous()
            continue
        fade_t = torch.linspace(0.0, 1.0, steps=int(overlap), device=out.device, dtype=out.dtype)
        fade_out = torch.cos(0.5 * math.pi * fade_t).view(1, -1)
        fade_in = torch.sin(0.5 * math.pi * fade_t).view(1, -1)
        blended = (out[..., -int(overlap) :] * fade_out) + (segment[..., : int(overlap)] * fade_in)
        out = torch.cat((out[..., :-int(overlap)], blended, segment[..., int(overlap) :]), dim=-1).contiguous()
    return out.contiguous()


def apply_beat_crossfade(
    audio_ct: torch.Tensor,
    beat_boundaries_sec_t: torch.Tensor,
    *,
    sample_rate: int,
    beat_crossfade_ms: float = DEFAULT_BEAT_CROSSFADE_MS,
) -> torch.Tensor:
    audio = torch.as_tensor(audio_ct, dtype=torch.float32).contiguous()
    if int(audio.dim()) != 2:
        raise ValueError(f"audio_ct must be [C,T], got {tuple(audio.shape)}")
    crossfade_num_samples = int(round((max(0.0, float(beat_crossfade_ms)) / 1000.0) * float(sample_rate)))
    if int(crossfade_num_samples) <= 0:
        return audio.contiguous()

    boundaries_sec = torch.as_tensor(
        beat_boundaries_sec_t,
        dtype=torch.float32,
        device=audio.device,
    ).view(-1)
    if int(boundaries_sec.numel()) < 2:
        return audio.contiguous()
    total_num_samples = int(audio.shape[-1])
    boundaries = torch.round(boundaries_sec * float(sample_rate)).to(dtype=torch.long)
    boundaries = boundaries.clamp(min=0, max=max(0, total_num_samples))
    boundaries[0] = 0
    boundaries[-1] = int(total_num_samples)
    if int(boundaries.numel()) == 2:
        return audio.contiguous()

    left_context = int(crossfade_num_samples // 2)
    right_context = int(crossfade_num_samples - left_context)
    segments: list[torch.Tensor] = []
    last_beat_idx = int(boundaries.numel()) - 2
    for beat_idx in range(int(boundaries.numel()) - 1):
        nominal_lo = int(boundaries[int(beat_idx)].item())
        nominal_hi = int(boundaries[int(beat_idx) + 1].item())
        seg_lo = int(nominal_lo if int(beat_idx) == 0 else max(0, nominal_lo - left_context))
        seg_hi = int(nominal_hi if int(beat_idx) == int(last_beat_idx) else min(total_num_samples, nominal_hi + right_context))
        if int(seg_hi) <= int(seg_lo):
            continue
        segments.append(audio[..., int(seg_lo) : int(seg_hi)].contiguous())
    if not segments:
        return audio.contiguous()

    smoothed = stitch_audio_segments_with_crossfade(
        segments,
        crossfade_num_samples=int(crossfade_num_samples),
    )
    if int(smoothed.shape[-1]) > int(total_num_samples):
        smoothed = smoothed[..., : int(total_num_samples)]
    elif int(smoothed.shape[-1]) < int(total_num_samples):
        smoothed = F.pad(smoothed, (0, int(total_num_samples) - int(smoothed.shape[-1])))
    return smoothed.contiguous()


def _code_accuracy_stats(pred_codes_bct: torch.Tensor, ref_codes_bct: torch.Tensor) -> dict[str, Any]:
    pred_codes = torch.as_tensor(pred_codes_bct, dtype=torch.long)
    ref_codes = torch.as_tensor(ref_codes_bct, dtype=torch.long)
    if int(pred_codes.dim()) != 3 or int(ref_codes.dim()) != 3:
        raise ValueError(f"expected [B,C,T] code tensors, got {tuple(pred_codes.shape)} / {tuple(ref_codes.shape)}")
    if tuple(pred_codes.shape[:2]) != tuple(ref_codes.shape[:2]):
        raise ValueError(
            f"pred_codes and ref_codes must match on batch/codebook dims, got {tuple(pred_codes.shape)} / {tuple(ref_codes.shape)}"
        )
    compared_num_frames = int(min(int(pred_codes.shape[-1]), int(ref_codes.shape[-1])))
    if int(compared_num_frames) <= 0:
        raise ValueError(
            f"pred_codes and ref_codes must have at least one overlapping frame, got {tuple(pred_codes.shape)} / {tuple(ref_codes.shape)}"
        )
    pred_codes = pred_codes[:, :, : int(compared_num_frames)]
    ref_codes = ref_codes[:, :, : int(compared_num_frames)]
    equal = pred_codes.eq(ref_codes)
    return {
        "token_acc": float(equal.float().mean().item()),
        "per_codebook_token_acc": [
            float(equal[:, int(codebook_idx), :].float().mean().item())
            for codebook_idx in range(int(equal.shape[1]))
        ],
        "exact_match": bool(equal.all().item()),
        "compared_num_frames": int(compared_num_frames),
        "pred_num_frames": int(pred_codes_bct.shape[-1]),
        "ref_num_frames": int(ref_codes_bct.shape[-1]),
        "shape_match": bool(tuple(pred_codes_bct.shape) == tuple(ref_codes_bct.shape)),
    }


@torch.no_grad()
def save_inference_wav(
    model,
    diffusion,
    encodec_model,
    batch,
    device,
    epoch,
    target_mean,
    target_std,
    sample_rate=None,
    out_dir="best_samples",
    sample_idx=0,
    guidance_scale=1.0,
    start_noise=None,
    step_noises: Mapping[int, torch.Tensor] | None = None,
    x0_clip_norm: float | None = DEFAULT_SAMPLE_X0_CLIP_NORM,
    use_bpm_inference_geometry: bool = True,
    inference_num_beats: int = DEFAULT_INFERENCE_NUM_BEATS,
    target_token_rate_hz: float = DEFAULT_TARGET_TOKEN_RATE_HZ,
    beat_crossfade_ms: float = DEFAULT_BEAT_CROSSFADE_MS,
    target_pca_basis: Mapping[str, Any] | None = None,
):
    os.makedirs(out_dir, exist_ok=True)

    pred_latent = sample_ddpm(
        model=model,
        diffusion=diffusion,
        batch=batch,
        device=device,
        sample_idx=sample_idx,
        guidance_scale=guidance_scale,
        start_noise=start_noise,
        step_noises=step_noises,
        x0_clip_norm=x0_clip_norm,
        use_bpm_inference_geometry=bool(use_bpm_inference_geometry),
        inference_num_beats=int(inference_num_beats),
        target_token_rate_hz=float(target_token_rate_hz),
    )
    prepared = _prepare_batch_tensors(
        batch,
        device,
        require_target=False,
        require_timing=not bool(use_bpm_inference_geometry),
    )
    single = _slice_prepared_batch(prepared, int(sample_idx))
    geometry = resolve_inference_geometry(
        single,
        use_bpm_inference_geometry=bool(use_bpm_inference_geometry),
        inference_num_beats=int(inference_num_beats),
        target_token_rate_hz=float(target_token_rate_hz),
    )
    target_mask_i = geometry["target_valid_mask_bt"]
    pred_latent = denormalize_latent(pred_latent, target_mean, target_std)
    pred_latent = pred_latent * target_mask_i.unsqueeze(-1)
    resolved_target_pca_basis = _resolve_target_pca_basis(
        target_pca_basis,
        device=device,
        dtype=pred_latent.dtype,
    )

    audio = decode_latent_to_audio(
        pred_latent,
        encodec_model,
        target_pca_basis=resolved_target_pca_basis,
    )
    if audio.dim() == 3:
        wav = audio[0]
    elif audio.dim() == 2:
        wav = audio[0].unsqueeze(0)
    else:
        raise ValueError(f"Unexpected audio shape: {audio.shape}")

    write_sample_rate = int(sample_rate) if sample_rate is not None else resolve_encodec_sample_rate(encodec_model)
    if float(beat_crossfade_ms) > 0.0:
        wav = apply_beat_crossfade(
            wav,
            geometry["beat_boundaries_sec"][0],
            sample_rate=int(write_sample_rate),
            beat_crossfade_ms=float(beat_crossfade_ms),
        )
    wav = wav.detach().cpu()
    peak = wav.abs().max().clamp_min(1e-8)
    wav = 0.95 * wav / peak

    save_path = os.path.join(out_dir, f"best_epoch_{epoch:03d}.wav")
    torchaudio.save(save_path, wav, sample_rate=write_sample_rate)

    if "source_codes_bct" in batch and "target_btd" in batch:
        target_mask_single = geometry["target_valid_mask_bt"]
        cond_i, _ = model.encode_conditioning(
            grid=single["grid"],
            grid_ids=single["grid_ids"],
            grid_times_sec=geometry["grid_times_sec"],
            token_times_sec=geometry["token_times_sec"],
            target_valid_mask_bt=target_mask_single,
            grid_valid_mask_bt=single["grid_valid_mask"],
            timbre_bank_latents=single.get("timbre_bank_latents"),
            timbre_bank_family_ids=single.get("timbre_bank_family_ids"),
            timbre_bank_class_ids=single.get("timbre_bank_class_ids"),
            timbre_bank_velocity=single.get("timbre_bank_velocity"),
            timbre_bank_mask=single.get("timbre_bank_mask"),
            timbre_dynamic_features=single.get("timbre_dynamic_features"),
            timbre_dynamic_mask=single.get("timbre_dynamic_mask"),
            timbre_dynamic_counts=single.get("timbre_dynamic_counts"),
            timbre_family_default_indices=single.get("timbre_family_default_indices"),
            timbre_class_token_indices=single.get("timbre_class_token_indices"),
            reference_timbre_bank_latents=single.get("reference_timbre_bank_latents"),
            reference_timbre_bank_family_ids=single.get("reference_timbre_bank_family_ids"),
            reference_timbre_bank_class_ids=single.get("reference_timbre_bank_class_ids"),
            reference_timbre_bank_velocity=single.get("reference_timbre_bank_velocity"),
            reference_timbre_bank_mask=single.get("reference_timbre_bank_mask"),
            reference_timbre_dynamic_features=single.get("reference_timbre_dynamic_features"),
            reference_timbre_dynamic_mask=single.get("reference_timbre_dynamic_mask"),
            reference_timbre_dynamic_counts=single.get("reference_timbre_dynamic_counts"),
            reference_timbre_family_default_indices=single.get("reference_timbre_family_default_indices"),
            reference_timbre_class_token_indices=single.get("reference_timbre_class_token_indices"),
            reference_segment_pca144=single.get("reference_segment_pca144"),
        )

        source_codes = torch.as_tensor(batch["source_codes_bct"][int(sample_idx) : int(sample_idx) + 1], dtype=torch.long, device=device)
        target_plot_ref = torch.as_tensor(batch["target_btd"][int(sample_idx) : int(sample_idx) + 1], dtype=torch.float32, device=device)
        target_ref = torch.as_tensor(
            batch.get("target_sum_btd", batch["target_btd"])[int(sample_idx) : int(sample_idx) + 1],
            dtype=torch.float32,
            device=device,
        )
        source_audio = decode_codes_to_audio_b1t(encodec_model, source_codes, device=device)
        target_direct_audio = decode_latent_to_audio(target_ref, encodec_model)
        target_requant_codes = requantize_latent_to_codes_bct(encodec_model, target_ref, device=device)
        pred_requant_codes = requantize_latent_to_codes_bct(
            encodec_model,
            pred_latent,
            device=device,
            target_pca_basis=resolved_target_pca_basis,
        )
        target_requant_audio = decode_codes_to_audio_b1t(encodec_model, target_requant_codes, device=device)
        pred_requant_audio = decode_codes_to_audio_b1t(encodec_model, pred_requant_codes, device=device)
        valid_len = int(
            min(
                int(target_mask_single[0].sum().item()),
                int(target_plot_ref.shape[1]),
                int(target_requant_codes.shape[-1]),
                int(pred_requant_codes.shape[-1]),
            )
        )

        save_prediction_diagnostic_plot(
            os.path.join(out_dir, f"best_epoch_{epoch:03d}_pred_latent.png"),
            input_features_td=cond_i[0, :valid_len],
            pred_latent_td=pred_latent[0, :valid_len],
            target_latent_td=target_plot_ref[0, :valid_len],
            target_codes_ct=target_requant_codes[0, :, :valid_len],
            pred_codes_ct=pred_requant_codes[0, :, :valid_len],
            pred_target_token_acc=float(
                pred_requant_codes[:, :, :valid_len].eq(target_requant_codes[:, :, :valid_len]).float().mean().item()
            ),
            token_times_sec_t=geometry["token_times_sec"][0, :valid_len],
        )

        for suffix, audio_tensor in (
            ("source_codes", source_audio),
            ("target_direct", target_direct_audio),
            ("target_requant", target_requant_audio),
            ("pred_requant", pred_requant_audio),
        ):
            wav_i = audio_tensor[0].detach().cpu()
            peak_i = wav_i.abs().max().clamp_min(1e-8)
            wav_i = 0.95 * wav_i / peak_i
            torchaudio.save(
                os.path.join(out_dir, f"best_epoch_{epoch:03d}_{suffix}.wav"),
                wav_i,
                sample_rate=write_sample_rate,
            )

        debug_payload = {
            "sample_idx": int(sample_idx),
            "write_sample_rate": int(write_sample_rate),
            "use_bpm_inference_geometry": bool(use_bpm_inference_geometry),
            "inference_num_beats": int(inference_num_beats),
            "target_token_rate_hz": float(target_token_rate_hz),
            "beat_crossfade_ms": float(beat_crossfade_ms),
            "pred_requant_vs_source": _code_accuracy_stats(pred_requant_codes, source_codes),
            "pred_requant_vs_target_requant": _code_accuracy_stats(pred_requant_codes, target_requant_codes),
            "target_requant_vs_source": _code_accuracy_stats(target_requant_codes, source_codes),
            "pred_direct_vs_pred_requant_audio_l1": float(
                (
                    decode_latent_to_audio(
                        pred_latent,
                        encodec_model,
                        target_pca_basis=resolved_target_pca_basis,
                    )
                    - pred_requant_audio
                ).abs().mean().item()
            ),
            "target_direct_vs_source_audio_l1": float((target_direct_audio - source_audio).abs().mean().item()),
            "target_requant_vs_source_audio_l1": float((target_requant_audio - source_audio).abs().mean().item()),
            "pred_requant_codes_shape": list(pred_requant_codes.shape),
            "target_requant_codes_shape": list(target_requant_codes.shape),
            "source_codes_shape": list(source_codes.shape),
        }
        import json

        with open(os.path.join(out_dir, f"best_epoch_{epoch:03d}_decode_debug.json"), "w", encoding="utf-8") as handle:
            json.dump(debug_payload, handle, indent=2, sort_keys=True)

    return save_path


def load_or_compute_target_normalization(cache_root: str, train_loader, *, device: torch.device, x_dim: int):
    stats_path = os.path.join(cache_root, "target_stats.pt")
    if os.path.exists(stats_path):
        payload = torch.load(stats_path, map_location="cpu", weights_only=False)
        mean = torch.as_tensor(payload["target_mean"], dtype=torch.float32, device=device).view(-1)
        std = torch.as_tensor(payload["target_std"], dtype=torch.float32, device=device).view(-1).clamp_min(1.0e-6)
        if int(mean.numel()) != int(x_dim) or int(std.numel()) != int(x_dim):
            raise RuntimeError(
                f"cached target stats under {stats_path} do not match x_dim={x_dim}: "
                f"mean={tuple(mean.shape)} std={tuple(std.shape)}"
            )
        print(f"loaded target stats from {stats_path}")
        return mean.contiguous(), std.contiguous()

    print("estimating target normalization from train split")
    mean, std = estimate_target_normalization(train_loader, device=device)
    torch.save(
        {
            "target_mean": mean.detach().cpu(),
            "target_std": std.detach().cpu(),
            "x_dim": int(x_dim),
        },
        stats_path,
    )
    print(f"saved target stats to {stats_path}")
    return mean.contiguous(), std.contiguous()