Add merged checkpoint at step 10000

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Meituan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in 
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,75 @@
+---
+license: other
+license_name: longcat-audiodit-license
+base_model: meituan-longcat/LongCat-AudioDiT-1B
+tags:
+  - audio
+  - text-to-speech
+  - tts
+  - environmental-tts
+  - flow-matching
+  - dit
+library_name: transformers
+pipeline_tag: text-to-speech
+---
+
+# LongCat-AudioDiT Env-TTS — 10000-step fine-tune
+
+Fine-tune of [meituan-longcat/LongCat-AudioDiT-1B](https://huggingface.co/meituan-longcat/LongCat-AudioDiT-1B) for the
+**three-stream env-tts task**: given a reference environment audio, a reference
+speaker audio, and three text streams (env caption / speaker caption / target
+speech text), generate target speech that places the target text in the
+referenced environment with the referenced speaker timbre.
+
+## Differences from the base model
+
+The transformer adds **six learnable boundary tokens** (three latent-space, three
+text-space):
+
+```
+latent sequence : [<boe>  z_env  <bos>  z_spk  <bon>  z_target]
+text sequence   : [<boe_t> env_text_emb <bos_t> spk_text_emb <bon_t> target_text_emb]
+```
+
+`encode_multistream_text(env, spk, target, drop_env_text=…, drop_spk_text=…,
+drop_target_text=…)` is the new entry-point. `AudioDiTModel.forward(...)` also
+accepts a pre-assembled `prompt_latent` (replaces `prompt_audio`) so the inference
+path can feed the boundary-tokenized three-stream prompt directly.
+
+## Training summary
+
+| Field | Value |
+|---|---|
+| Steps | 10000 |
+| Effective batch | 16 × grad_accum 2 × 2 GPU = **64 rows / step** |
+| Learning rate | cosine 5e-5 (warmup 250) |
+| AdamW | β₁=0.9, β₂=0.999, wd=0.01 |
+| EMA | disabled |
+| LoRA | r=32, alpha=32, target = attn + ffn |
+| Full-train | boundary tokens + AdaLN + text_conv + latent_embed + input_embed + output_proj + time_embed |
+| Audio filter | target duration ∈ [3, 45] s |
+| RMS normalize | three-stream independent to **-23 dBFS** (target_rms=0.0708) |
+| Augmentation | noise + RIR on spk_audio (DNS5 64GB) |
+| Data | [ChristianYang/Env-TTS-Clean](https://huggingface.co/datasets/ChristianYang/Env-TTS-Clean) |
+
+## How to load
+
+The model uses **custom code** in this repo, so pass `trust_remote_code=True`:
+
+```python
+from transformers import AutoModel, AutoTokenizer
+
+model = AutoModel.from_pretrained(
+    "meituan-longcat/LongCat-AudioDiT-Env-TTS-1B-10000Step",
+    trust_remote_code=True,
+).cuda().eval()
+
+tokenizer = AutoTokenizer.from_pretrained(model.config.text_encoder_model)
+```
+
+For end-to-end env-tts inference (three-stream prompt + ASR fallback for missing
+env/spk text) see the training repo's `tasks/inference.py`.
+
+## License
+
+Inherits the original [meituan-longcat/LongCat-AudioDiT-1B](https://huggingface.co/meituan-longcat/LongCat-AudioDiT-1B) license.

config.json

@@ -0,0 +1,95 @@
+{
+  "architectures": [
+    "AudioDiTModel"
+  ],
+  "dit_adaln_type": "global",
+  "dit_adaln_use_text_cond": true,
+  "dit_bias": true,
+  "dit_cross_attn": true,
+  "dit_cross_attn_norm": false,
+  "dit_depth": 24,
+  "dit_dim": 1536,
+  "dit_dropout": 0.0,
+  "dit_eps": 1e-06,
+  "dit_ff_mult": 4,
+  "dit_heads": 24,
+  "dit_long_skip": true,
+  "dit_qk_norm": true,
+  "dit_text_conv": true,
+  "dit_text_dim": 768,
+  "dit_use_latent_condition": true,
+  "dtype": "float32",
+  "latent_dim": 64,
+  "latent_hop": 2048,
+  "max_wav_duration": 30,
+  "model_type": "audiodit",
+  "repa_dit_layer": 8,
+  "sampling_rate": 24000,
+  "sigma": 0.0,
+  "text_add_embed": true,
+  "text_encoder_config": {
+    "_name_or_path": "ArthurZ/umt5-base",
+    "architectures": [
+      "UMT5ForConditionalGeneration"
+    ],
+    "classifier_dropout": 0.0,
+    "d_ff": 2048,
+    "d_kv": 64,
+    "d_model": 768,
+    "dense_act_fn": "gelu_new",
+    "dropout_rate": 0.1,
+    "dtype": "float32",
+    "feed_forward_proj": "gated-gelu",
+    "initializer_factor": 1.0,
+    "is_gated_act": true,
+    "layer_norm_epsilon": 1e-06,
+    "model_type": "umt5",
+    "num_decoder_layers": 12,
+    "num_heads": 12,
+    "num_layers": 12,
+    "output_past": true,
+    "relative_attention_max_distance": 128,
+    "relative_attention_num_buckets": 32,
+    "scalable_attention": true,
+    "use_cache": true,
+    "vocab_size": 256384
+  },
+  "text_encoder_model": "google/umt5-base",
+  "text_norm_feat": true,
+  "transformers_version": "4.57.6",
+  "vae_config": {
+    "c_mults": [
+      1,
+      2,
+      4,
+      8,
+      16
+    ],
+    "channels": 128,
+    "downsample_shortcut": "averaging",
+    "downsampling_ratio": 2048,
+    "dtype": "float32",
+    "encoder_latent_dim": 128,
+    "final_tanh": false,
+    "in_channels": 1,
+    "in_shortcut": "duplicating",
+    "latent_dim": 64,
+    "model_type": "audiodit_vae",
+    "out_shortcut": "averaging",
+    "sample_rate": 24000,
+    "scale": 0.71,
+    "strides": [
+      2,
+      4,
+      4,
+      8,
+      8
+    ],
+    "upsample_shortcut": "duplicating",
+    "use_snake": true
+  },
+  "auto_map": {
+    "AutoConfig": "configuration_audiodit.AudioDiTConfig",
+    "AutoModel": "modeling_audiodit.AudioDiTModel"
+  }
+}

configuration_audiodit.py

@@ -0,0 +1,230 @@
+"""AudioDiT model configuration"""
+
+# transformers >= 5.0 renamed PretrainedConfig → PreTrainedConfig; tolerate both.
+try:
+    from transformers import PreTrainedConfig
+except ImportError:
+    from transformers import PretrainedConfig as PreTrainedConfig
+from transformers import logging
+from transformers.models.umt5.configuration_umt5 import UMT5Config
+
+
+logger = logging.get_logger(__name__)
+
+
+class AudioDiTVaeConfig(PreTrainedConfig):
+    r"""
+    Configuration class for the AudioDiT WAV-VAE audio autoencoder.
+
+    Args:
+        in_channels (`int`, *optional*, defaults to 1):
+            Number of input audio channels (mono=1).
+        channels (`int`, *optional*, defaults to 128):
+            Base channel count for encoder/decoder.
+        c_mults (`list[int]`, *optional*, defaults to `[1, 2, 4, 8, 16]`):
+            Channel multipliers for each encoder/decoder stage.
+        strides (`list[int]`, *optional*, defaults to `[2, 4, 4, 8, 8]`):
+            Downsampling strides for each encoder stage.
+        latent_dim (`int`, *optional*, defaults to 64):
+            Dimensionality of the latent space (after VAE bottleneck: encoder outputs 128, split to mean+scale → 64).
+        encoder_latent_dim (`int`, *optional*, defaults to 128):
+            Dimensionality of the encoder output before VAE bottleneck.
+        use_snake (`bool`, *optional*, defaults to `True`):
+            Whether to use Snake activation instead of ELU.
+        downsample_shortcut (`str`, *optional*, defaults to `"averaging"`):
+            Shortcut type for encoder downsampling blocks.
+        upsample_shortcut (`str`, *optional*, defaults to `"duplicating"`):
+            Shortcut type for decoder upsampling blocks.
+        out_shortcut (`str`, *optional*, defaults to `"averaging"`):
+            Shortcut type for encoder output projection.
+        in_shortcut (`str`, *optional*, defaults to `"duplicating"`):
+            Shortcut type for decoder input projection.
+        final_tanh (`bool`, *optional*, defaults to `False`):
+            Whether to apply tanh to the decoder output.
+        downsampling_ratio (`int`, *optional*, defaults to 2048):
+            Total downsampling ratio from audio samples to latent frames.
+        sample_rate (`int`, *optional*, defaults to 24000):
+            Audio sample rate.
+        scale (`float`, *optional*, defaults to 0.71):
+            Scale factor for the latent space.
+    """
+
+    model_type = "audiodit_vae"
+
+    def __init__(
+        self,
+        in_channels: int = 1,
+        channels: int = 128,
+        c_mults: list[int] | None = None,
+        strides: list[int] | None = None,
+        latent_dim: int = 64,
+        encoder_latent_dim: int = 128,
+        use_snake: bool = True,
+        downsample_shortcut: str = "averaging",
+        upsample_shortcut: str = "duplicating",
+        out_shortcut: str = "averaging",
+        in_shortcut: str = "duplicating",
+        final_tanh: bool = False,
+        downsampling_ratio: int = 2048,
+        sample_rate: int = 24000,
+        scale: float = 0.71,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.channels = channels
+        self.c_mults = c_mults if c_mults is not None else [1, 2, 4, 8, 16]
+        self.strides = strides if strides is not None else [2, 4, 4, 8, 8]
+        self.latent_dim = latent_dim
+        self.encoder_latent_dim = encoder_latent_dim
+        self.use_snake = use_snake
+        self.downsample_shortcut = downsample_shortcut
+        self.upsample_shortcut = upsample_shortcut
+        self.out_shortcut = out_shortcut
+        self.in_shortcut = in_shortcut
+        self.final_tanh = final_tanh
+        self.downsampling_ratio = downsampling_ratio
+        self.sample_rate = sample_rate
+        self.scale = scale
+
+
+class AudioDiTConfig(PreTrainedConfig):
+    r"""
+    Configuration class for AudioDiT, a Conditional Flow Matching TTS model based on DiT architecture.
+
+    Args:
+        dit_dim (`int`, *optional*, defaults to 1536):
+            Hidden dimension of the DiT transformer.
+        dit_depth (`int`, *optional*, defaults to 24):
+            Number of transformer layers.
+        dit_heads (`int`, *optional*, defaults to 24):
+            Number of attention heads.
+        dit_ff_mult (`float`, *optional*, defaults to 4.0):
+            Feed-forward network multiplier.
+        dit_text_dim (`int`, *optional*, defaults to 768):
+            Dimension of the text encoder output (UMT5-base).
+        dit_dropout (`float`, *optional*, defaults to 0.0):
+            Dropout rate.
+        dit_bias (`bool`, *optional*, defaults to `True`):
+            Whether to use bias in linear layers.
+        dit_cross_attn (`bool`, *optional*, defaults to `True`):
+            Whether to use cross-attention layers.
+        dit_adaln_type (`str`, *optional*, defaults to `"global"`):
+            Type of adaptive layer norm (`"global"` or `"local"`).
+        dit_adaln_use_text_cond (`bool`, *optional*, defaults to `True`):
+            Whether to condition AdaLN on text embeddings.
+        dit_long_skip (`bool`, *optional*, defaults to `True`):
+            Whether to use long skip connection (input added to output).
+        dit_text_conv (`bool`, *optional*, defaults to `True`):
+            Whether to apply ConvNeXt blocks on text embeddings.
+        dit_qk_norm (`bool`, *optional*, defaults to `True`):
+            Whether to apply RMS normalization to Q and K.
+        dit_cross_attn_norm (`bool`, *optional*, defaults to `False`):
+            Whether to apply layer normalization in cross-attention.
+        dit_eps (`float`, *optional*, defaults to 1e-6):
+            Epsilon for normalization layers.
+        dit_use_latent_condition (`bool`, *optional*, defaults to `True`):
+            Whether to use latent conditioning (for prompt audio).
+        repa_dit_layer (`int`, *optional*, defaults to 8):
+            Layer index for representation alignment.
+        latent_dim (`int`, *optional*, defaults to 64):
+            Dimensionality of the audio latent space.
+        sigma (`float`, *optional*, defaults to 0.0):
+            Noise level for conditional flow matching.
+        sampling_rate (`int`, *optional*, defaults to 24000):
+            Audio sample rate.
+        latent_hop (`int`, *optional*, defaults to 2048):
+            Hop size in audio samples per latent frame.
+        max_wav_duration (`float`, *optional*, defaults to 30.0):
+            Maximum audio duration in seconds.
+        text_encoder_model (`str`, *optional*, defaults to `"google/umt5-base"`):
+            HuggingFace model identifier for the text encoder.
+        text_add_embed (`bool`, *optional*, defaults to `True`):
+            Whether to add the first hidden state to the last hidden state in text encoding.
+        text_norm_feat (`bool`, *optional*, defaults to `True`):
+            Whether to apply layer normalization to text features.
+        vae_config (`AudioDiTVaeConfig` or `dict`, *optional*):
+            Configuration for the WAV-VAE audio autoencoder.
+
+    Example:
+
+    ```python
+    >>> from transformers import AudioDiTConfig, AudioDiTModel
+
+    >>> configuration = AudioDiTConfig()
+    >>> model = AudioDiTModel(configuration)
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "audiodit"
+    sub_configs = {"vae_config": AudioDiTVaeConfig, "text_encoder_config": UMT5Config}
+
+    def __init__(
+        self,
+        dit_dim: int = 1536,
+        dit_depth: int = 24,
+        dit_heads: int = 24,
+        dit_ff_mult: float = 4.0,
+        dit_text_dim: int = 768,
+        dit_dropout: float = 0.0,
+        dit_bias: bool = True,
+        dit_cross_attn: bool = True,
+        dit_adaln_type: str = "global",
+        dit_adaln_use_text_cond: bool = True,
+        dit_long_skip: bool = True,
+        dit_text_conv: bool = True,
+        dit_qk_norm: bool = True,
+        dit_cross_attn_norm: bool = False,
+        dit_eps: float = 1e-6,
+        dit_use_latent_condition: bool = True,
+        repa_dit_layer: int = 8,
+        latent_dim: int = 64,
+        sigma: float = 0.0,
+        sampling_rate: int = 24000,
+        latent_hop: int = 2048,
+        max_wav_duration: float = 30.0,
+        text_encoder_model: str = "google/umt5-base",
+        text_add_embed: bool = True,
+        text_norm_feat: bool = True,
+        vae_config: AudioDiTVaeConfig | dict | None = None,
+        text_encoder_config: UMT5Config | dict | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.dit_dim = dit_dim
+        self.dit_depth = dit_depth
+        self.dit_heads = dit_heads
+        self.dit_ff_mult = dit_ff_mult
+        self.dit_text_dim = dit_text_dim
+        self.dit_dropout = dit_dropout
+        self.dit_bias = dit_bias
+        self.dit_cross_attn = dit_cross_attn
+        self.dit_adaln_type = dit_adaln_type
+        self.dit_adaln_use_text_cond = dit_adaln_use_text_cond
+        self.dit_long_skip = dit_long_skip
+        self.dit_text_conv = dit_text_conv
+        self.dit_qk_norm = dit_qk_norm
+        self.dit_cross_attn_norm = dit_cross_attn_norm
+        self.dit_eps = dit_eps
+        self.dit_use_latent_condition = dit_use_latent_condition
+        self.repa_dit_layer = repa_dit_layer
+        self.latent_dim = latent_dim
+        self.sigma = sigma
+        self.sampling_rate = sampling_rate
+        self.latent_hop = latent_hop
+        self.max_wav_duration = max_wav_duration
+        self.text_encoder_model = text_encoder_model
+        self.text_add_embed = text_add_embed
+        self.text_norm_feat = text_norm_feat
+
+        if isinstance(vae_config, dict):
+            vae_config = AudioDiTVaeConfig(**vae_config)
+        self.vae_config = vae_config if vae_config is not None else AudioDiTVaeConfig()
+
+        if isinstance(text_encoder_config, dict):
+            text_encoder_config = UMT5Config(**text_encoder_config)
+        self.text_encoder_config = text_encoder_config
+
+
+__all__ = ["AudioDiTConfig", "AudioDiTVaeConfig"]

model-00001-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8d814aa8350d765dfb9f12e67db470062a5411ba03f9e1a96af271df8ee5b20
+size 4998324592

model-00002-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6d93027ccaaeec9aab217bccc88da12d7a93ee6d9345ce47c366558097b00b0
+size 681520212

modeling_audiodit.py

@@ -0,0 +1,1416 @@
+"""PyTorch AudioDiT model — Conditional Flow Matching TTS with DiT backbone."""
+
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+# Use the modern parametrizations-based weight_norm. Backward-compatible: it
+# auto-remaps legacy `weight_g`/`weight_v` keys in state_dict to the new
+# `parametrizations.weight.original0`/`original1` names on load (via a
+# state_dict pre-hook registered inside parametrize.register_parametrization).
+# This fixes the silent VAE load failure where transformers 5.x's
+# `Materializing param` path renamed weight_norm modules to parametrizations
+# without remapping, leading to VAE outputs of pure silence.
+from torch.nn.utils.parametrizations import weight_norm
+from torch.nn.utils.rnn import pad_sequence
+
+from transformers import PreTrainedModel, logging
+from transformers.modeling_outputs import ModelOutput
+from .configuration_audiodit import AudioDiTConfig, AudioDiTVaeConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Output dataclass
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class AudioDiTOutput(ModelOutput):
+    """
+    Output of [`AudioDiTModel`].
+
+    Args:
+        waveform (`torch.FloatTensor` of shape `(batch_size, num_samples)`):
+            Generated audio waveform.
+        latent (`torch.FloatTensor` of shape `(batch_size, latent_dim, num_frames)`):
+            Predicted latent representation before VAE decoding.
+    """
+
+    waveform: torch.FloatTensor | None = None
+    latent: torch.FloatTensor | None = None
+
+
+# ---------------------------------------------------------------------------
+# ODE solver (inline Euler — replaces torchdiffeq dependency)
+# ---------------------------------------------------------------------------
+
+
+def odeint_euler(fn, y0, t):
+    """Simple Euler ODE integrator (equivalent to `torchdiffeq.odeint` with `method='euler'`).
+
+    Args:
+        fn: callable(t, y) → dy/dt
+        y0: initial state tensor
+        t: 1-D tensor of time steps (must be monotonically increasing)
+
+    Returns:
+        Tensor of shape `(len(t), *y0.shape)` containing the trajectory.
+    """
+    ys = [y0]
+    y = y0
+    for i in range(len(t) - 1):
+        dt = t[i + 1] - t[i]
+        y = y + fn(t[i], y) * dt
+        ys.append(y)
+    return torch.stack(ys)
+
+
+# ---------------------------------------------------------------------------
+# Utility helpers (from model/utils.py)
+# ---------------------------------------------------------------------------
+
+
+def lens_to_mask(lengths: torch.Tensor, length: int | None = None) -> torch.BoolTensor:
+    if length is None:
+        length = lengths.amax()
+    seq = torch.arange(length, device=lengths.device)
+    return seq[None, :] < lengths[:, None]
+
+
+# ---------------------------------------------------------------------------
+# Low-level modules (from model/modules.py)
+# ---------------------------------------------------------------------------
+
+
+class AudioDiTRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self._norm(x.float()).type_as(x) * self.weight
+
+    def _norm(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class AudioDiTSinusPositionEmbedding(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor, scale: float = 1000.0) -> torch.Tensor:
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        return torch.cat((emb.sin(), emb.cos()), dim=-1)
+
+
+class AudioDiTTimestepEmbedding(nn.Module):
+    def __init__(self, dim: int, freq_embed_dim: int = 256):
+        super().__init__()
+        self.time_embed = AudioDiTSinusPositionEmbedding(freq_embed_dim)
+        self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        time_hidden = self.time_embed(timestep)
+        time_hidden = time_hidden.to(timestep.dtype)
+        return self.time_mlp(time_hidden)
+
+
+class AudioDiTRotaryEmbedding(nn.Module):
+    """Qwen2-style rotary position embedding.
+
+    All state (inv_freq, cos/sin caches) is built lazily on first ``forward``
+    call.  This avoids corruption from ``from_pretrained`` meta-device
+    construction while producing bit-identical results to the original
+    ``Qwen2RotaryEmbedding`` (which creates ``inv_freq`` on CPU then moves
+    the whole model to CUDA with ``.to(device)``).
+    """
+
+    def __init__(self, dim: int, max_position_embeddings: int = 2048, base: float = 100000.0):
+        super().__init__()
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        # Do NOT register any buffers here — they get corrupted by meta-device.
+        # Everything is built lazily in forward().
+        self._cos: torch.Tensor | None = None
+        self._sin: torch.Tensor | None = None
+        self._cached_len: int = 0
+        self._cached_device: torch.device | None = None
+
+    def _build(self, seq_len: int, device: torch.device, dtype: torch.dtype):
+        """Build cos/sin tables entirely on CPU (matching original
+        Qwen2RotaryEmbedding which builds in __init__ on CPU, then the
+        whole model is moved with .to(device)), then move to target."""
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
+        t = torch.arange(seq_len, dtype=torch.int64).type_as(inv_freq)
+        freqs = torch.outer(t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self._cos = emb.cos().to(dtype=dtype, device=device)
+        self._sin = emb.sin().to(dtype=dtype, device=device)
+        self._cached_len = seq_len
+        self._cached_device = device
+
+    def forward(self, x: torch.Tensor, seq_len: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+        if seq_len is None:
+            seq_len = x.shape[1]
+        if self._cos is None or seq_len > self._cached_len or self._cached_device != x.device:
+            self._build(max(seq_len, self.max_position_embeddings), x.device, x.dtype)
+        return (
+            self._cos[:seq_len].to(dtype=x.dtype),
+            self._sin[:seq_len].to(dtype=x.dtype),
+        )
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def _apply_rotary_emb(x: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+    cos, sin = freqs_cis
+    cos = cos[None, None].to(x.device)
+    sin = sin[None, None].to(x.device)
+    return (x.float() * cos + _rotate_half(x).float() * sin).to(x.dtype)
+
+
+# ---------------------------------------------------------------------------
+# GRN + ConvNeXtV2 (for text conv)
+# ---------------------------------------------------------------------------
+
+
+class AudioDiTGRN(nn.Module):
+    """Global Response Normalization."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=1, keepdim=True)
+        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class AudioDiTConvNeXtV2Block(nn.Module):
+    def __init__(self, dim: int, intermediate_dim: int, dilation: int = 1, kernel_size: int = 7, bias: bool = True, eps: float = 1e-6):
+        super().__init__()
+        padding = (dilation * (kernel_size - 1)) // 2
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=padding, groups=dim, dilation=dilation, bias=bias)
+        self.norm = nn.LayerNorm(dim, eps=eps)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim, bias=bias)
+        self.act = nn.SiLU()
+        self.grn = AudioDiTGRN(intermediate_dim)
+        self.pwconv2 = nn.Linear(intermediate_dim, dim, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = x.transpose(1, 2)
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        return residual + x
+
+
+# ---------------------------------------------------------------------------
+# Embedder (shared for input / text / latent)
+# ---------------------------------------------------------------------------
+
+
+class AudioDiTEmbedder(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int):
+        super().__init__()
+        self.proj = nn.Sequential(nn.Linear(in_dim, out_dim), nn.SiLU(), nn.Linear(out_dim, out_dim))
+
+    def forward(self, x: torch.Tensor, mask: torch.BoolTensor | None = None) -> torch.Tensor:
+        if mask is not None:
+            x = x.masked_fill(mask.logical_not().unsqueeze(-1), 0.0)
+        x = self.proj(x)
+        if mask is not None:
+            x = x.masked_fill(mask.logical_not().unsqueeze(-1), 0.0)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# AdaLN modules
+# ---------------------------------------------------------------------------
+
+
+class AudioDiTAdaLNMLP(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, bias: bool = True):
+        super().__init__()
+        self.mlp = nn.Sequential(nn.SiLU(), nn.Linear(in_dim, out_dim, bias=bias))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+
+class AudioDiTAdaLayerNormZeroFinal(nn.Module):
+    def __init__(self, dim: int, bias: bool = True, eps: float = 1e-6):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 2, bias=bias)
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+
+    def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
+        emb = self.linear(self.silu(emb))
+        scale, shift = torch.chunk(emb, 2, dim=-1)
+        x = self.norm(x.float()).type_as(x)
+        if scale.ndim == 2:
+            x = x * (1 + scale)[:, None, :] + shift[:, None, :]
+        else:
+            x = x * (1 + scale) + shift
+        return x
+
+
+# ---------------------------------------------------------------------------
+# Attention
+# ---------------------------------------------------------------------------
+
+
+def _modulate(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    """LayerNorm without affine + modulate."""
+    x = F.layer_norm(x.float(), (x.shape[-1],), eps=eps).type_as(x)
+    if scale.ndim == 2:
+        return x * (1 + scale[:, None]) + shift[:, None]
+    return x * (1 + scale) + shift
+
+
+class AudioDiTSelfAttention(nn.Module):
+    def __init__(self, dim: int, heads: int, dim_head: int, dropout: float = 0.0, bias: bool = True, qk_norm: bool = False, eps: float = 1e-6):
+        super().__init__()
+        self.heads = heads
+        self.inner_dim = dim_head * heads
+        self.to_q = nn.Linear(dim, self.inner_dim, bias=bias)
+        self.to_k = nn.Linear(dim, self.inner_dim, bias=bias)
+        self.to_v = nn.Linear(dim, self.inner_dim, bias=bias)
+        self.qk_norm = qk_norm
+        if qk_norm:
+            self.q_norm = AudioDiTRMSNorm(self.inner_dim, eps=eps)
+            self.k_norm = AudioDiTRMSNorm(self.inner_dim, eps=eps)
+        self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, dim, bias=bias), nn.Dropout(dropout)])
+
+    def forward(self, x: torch.Tensor, mask: torch.BoolTensor | None = None, rope: tuple | None = None) -> torch.Tensor:
+        batch_size = x.shape[0]
+        query = self.to_q(x)
+        key = self.to_k(x)
+        value = self.to_v(x)
+        if self.qk_norm:
+            query = self.q_norm(query)
+            key = self.k_norm(key)
+        head_dim = self.inner_dim // self.heads
+        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        if rope is not None:
+            query = _apply_rotary_emb(query, rope)
+            key = _apply_rotary_emb(key, rope)
+        attn_mask = None
+        if mask is not None:
+            attn_mask = mask.unsqueeze(1).unsqueeze(1).expand(batch_size, self.heads, query.shape[-2], key.shape[-2])
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = x.transpose(1, 2).reshape(batch_size, -1, self.inner_dim).to(query.dtype)
+        x = self.to_out[0](x)
+        x = self.to_out[1](x)
+        return x
+
+
+class AudioDiTCrossAttention(nn.Module):
+    def __init__(self, q_dim: int, kv_dim: int, heads: int, dim_head: int, dropout: float = 0.0, bias: bool = True, qk_norm: bool = False, eps: float = 1e-6):
+        super().__init__()
+        self.heads = heads
+        self.inner_dim = dim_head * heads
+        self.to_q = nn.Linear(q_dim, self.inner_dim, bias=bias)
+        self.to_k = nn.Linear(kv_dim, self.inner_dim, bias=bias)
+        self.to_v = nn.Linear(kv_dim, self.inner_dim, bias=bias)
+        self.qk_norm = qk_norm
+        if qk_norm:
+            self.q_norm = AudioDiTRMSNorm(self.inner_dim, eps=eps)
+            self.k_norm = AudioDiTRMSNorm(self.inner_dim, eps=eps)
+        self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, q_dim, bias=bias), nn.Dropout(dropout)])
+
+    def forward(
+        self, x: torch.Tensor, cond: torch.Tensor, mask: torch.BoolTensor | None = None,
+        cond_mask: torch.BoolTensor | None = None, rope: tuple | None = None, cond_rope: tuple | None = None,
+    ) -> torch.Tensor:
+        batch_size = x.shape[0]
+        query = self.to_q(x)
+        key = self.to_k(cond)
+        value = self.to_v(cond)
+        if self.qk_norm:
+            query = self.q_norm(query)
+            key = self.k_norm(key)
+        head_dim = self.inner_dim // self.heads
+        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        if rope is not None:
+            query = _apply_rotary_emb(query, rope)
+        if cond_rope is not None:
+            key = _apply_rotary_emb(key, cond_rope)
+        attn_mask = None
+        if mask is not None:
+            attn_mask = cond_mask.unsqueeze(1).expand(-1, mask.shape[1], -1).unsqueeze(1)
+            attn_mask = attn_mask.expand(batch_size, self.heads, query.shape[-2], key.shape[-2])
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = x.transpose(1, 2).reshape(batch_size, -1, self.inner_dim).to(query.dtype)
+        x = self.to_out[0](x)
+        x = self.to_out[1](x)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# FeedForward
+# ---------------------------------------------------------------------------
+
+
+class AudioDiTFeedForward(nn.Module):
+    def __init__(self, dim: int, mult: float = 4.0, dropout: float = 0.0, bias: bool = True):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        self.ff = nn.Sequential(
+            nn.Linear(dim, inner_dim, bias=bias),
+            nn.GELU(approximate="tanh"),
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim, bias=bias),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.ff(x)
+
+
+# ---------------------------------------------------------------------------
+# Transformer Block (CrossDiTBlock)
+# ---------------------------------------------------------------------------
+
+
+class AudioDiTBlock(nn.Module):
+    """Single DiT block with self-attention, optional cross-attention, FFN, and AdaLN modulation."""
+
+    def __init__(self, config: AudioDiTConfig):
+        super().__init__()
+        dim = config.dit_dim
+        cond_dim = config.dit_dim  # after text embedding, cond_dim == dim
+        heads = config.dit_heads
+        dim_head = dim // heads
+        bias = config.dit_bias
+        eps = config.dit_eps
+
+        self.adaln_type = config.dit_adaln_type
+        self.adaln_use_text_cond = config.dit_adaln_use_text_cond
+        if config.dit_adaln_type == "local":
+            self.adaln_mlp = AudioDiTAdaLNMLP(dim, dim * 6, bias=True)
+        elif config.dit_adaln_type == "global":
+            self.adaln_scale_shift = nn.Parameter(torch.randn(dim * 6) / dim**0.5)
+
+        self.self_attn = AudioDiTSelfAttention(
+            dim=dim, heads=heads, dim_head=dim_head, dropout=config.dit_dropout,
+            bias=bias, qk_norm=config.dit_qk_norm, eps=eps,
+        )
+
+        self.use_cross_attn = config.dit_cross_attn
+        if config.dit_cross_attn:
+            self.cross_attn = AudioDiTCrossAttention(
+                q_dim=dim, kv_dim=cond_dim, heads=heads, dim_head=dim_head,
+                dropout=config.dit_dropout, bias=bias, qk_norm=config.dit_qk_norm, eps=eps,
+            )
+            self.cross_attn_norm = nn.LayerNorm(dim, elementwise_affine=True, eps=eps) if config.dit_cross_attn_norm else nn.Identity()
+            self.cross_attn_norm_c = nn.LayerNorm(cond_dim, elementwise_affine=True, eps=eps) if config.dit_cross_attn_norm else nn.Identity()
+
+        self.ffn = AudioDiTFeedForward(dim=dim, mult=config.dit_ff_mult, dropout=config.dit_dropout, bias=bias)
+
+    def forward(
+        self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor,
+        mask: torch.BoolTensor | None = None, cond_mask: torch.BoolTensor | None = None,
+        rope: tuple | None = None, cond_rope: tuple | None = None,
+        adaln_global_out: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.adaln_type == "local" and adaln_global_out is None:
+            if self.adaln_use_text_cond:
+                cond_mean = cond.sum(1) / cond_mask.sum(1, keepdim=True)
+                norm_cond = t + cond_mean
+            else:
+                norm_cond = t
+            adaln_out = self.adaln_mlp(norm_cond)
+            gate_sa, scale_sa, shift_sa, gate_ffn, scale_ffn, shift_ffn = torch.chunk(adaln_out, 6, dim=-1)
+        else:
+            from einops import rearrange
+            adaln_out = adaln_global_out + rearrange(self.adaln_scale_shift, "f -> 1 f")
+            gate_sa, scale_sa, shift_sa, gate_ffn, scale_ffn, shift_ffn = torch.chunk(adaln_out, 6, dim=-1)
+
+        # Self-attention
+        norm = _modulate(x, scale_sa, shift_sa)
+        attn_output = self.self_attn(norm, mask=mask, rope=rope)
+        if gate_sa.ndim == 2:
+            gate_sa = gate_sa.unsqueeze(1)
+        x = x + gate_sa * attn_output
+
+        # Cross-attention
+        if self.use_cross_attn:
+            cross_out = self.cross_attn(
+                x=self.cross_attn_norm(x), cond=self.cross_attn_norm_c(cond),
+                mask=mask, cond_mask=cond_mask, rope=rope, cond_rope=cond_rope,
+            )
+            x = x + cross_out
+
+        # FFN
+        norm = _modulate(x, scale_ffn, shift_ffn)
+        ff_output = self.ffn(norm)
+        if gate_ffn.ndim == 2:
+            gate_ffn = gate_ffn.unsqueeze(1)
+        x = x + gate_ffn * ff_output
+        return x
+
+
+# ---------------------------------------------------------------------------
+# AudioDiTTransformer (CrossDiT backbone)
+# ---------------------------------------------------------------------------
+
+
+class AudioDiTTransformer(nn.Module):
+    """The core DiT transformer backbone for AudioDiT."""
+
+    def __init__(self, config: AudioDiTConfig):
+        super().__init__()
+        dim = config.dit_dim
+        latent_dim = config.latent_dim  # 64
+        text_dim = config.dit_text_dim
+        dim_head = dim // config.dit_heads
+
+        self.config = config
+        self.dim = dim
+        self.depth = config.dit_depth
+        self.long_skip = config.dit_long_skip
+        self.adaln_type = config.dit_adaln_type
+        self.adaln_use_text_cond = config.dit_adaln_use_text_cond
+
+        self.time_embed = AudioDiTTimestepEmbedding(dim)
+        self.input_embed = AudioDiTEmbedder(latent_dim, dim)
+        self.text_embed = AudioDiTEmbedder(text_dim, dim)
+        self.rotary_embed = AudioDiTRotaryEmbedding(dim_head, 2048, base=100000.0)
+
+        self.blocks = nn.ModuleList([AudioDiTBlock(config) for _ in range(config.dit_depth)])
+
+        self.norm_out = AudioDiTAdaLayerNormZeroFinal(dim, bias=True, eps=config.dit_eps)
+        self.proj_out = nn.Linear(dim, latent_dim)
+
+        if config.dit_adaln_type == "global":
+            self.adaln_global_mlp = AudioDiTAdaLNMLP(dim, dim * 6, bias=True)
+
+        self.text_conv = config.dit_text_conv
+        if config.dit_text_conv:
+            self.text_conv_layer = nn.Sequential(
+                *[AudioDiTConvNeXtV2Block(dim, dim * 2, bias=config.dit_bias, eps=config.dit_eps) for _ in range(4)]
+            )
+
+        self.use_latent_condition = config.dit_use_latent_condition
+        if config.dit_use_latent_condition:
+            self.latent_embed = AudioDiTEmbedder(latent_dim, dim)
+            self.latent_cond_embedder = AudioDiTEmbedder(dim * 2, dim)
+
+        # Latent-space boundary tokens for env-spk multistream input.
+        # Layout: [<boe>, env_latent, <bos>, spk_latent, <bon>, target_latent].
+        # Each (1, 1, latent_dim), trainable nn.Parameter. Re-initialized in
+        # _initialize_weights() to N(0, 0.02). The trainer / inference code
+        # is responsible for concatenating these into the latent sequence.
+        self.boe_token = nn.Parameter(torch.zeros(1, 1, latent_dim))
+        self.bos_token = nn.Parameter(torch.zeros(1, 1, latent_dim))
+        self.bon_token = nn.Parameter(torch.zeros(1, 1, latent_dim))
+
+        # Text-space boundary tokens (parallel to latent ones but in the UMT5
+        # output space). Used in encode_multistream_text() to build:
+        # [<boe_text>, env_text_emb, <bos_text>, spk_text_emb, <bon_text>, target_text_emb].
+        # Sized to dit_text_dim (UMT5 d_model = 768 for the base config).
+        self.boe_text_token = nn.Parameter(torch.zeros(1, 1, text_dim))
+        self.bos_text_token = nn.Parameter(torch.zeros(1, 1, text_dim))
+        self.bon_text_token = nn.Parameter(torch.zeros(1, 1, text_dim))
+
+        self._initialize_weights()
+
+    def _initialize_weights(self):
+        """Zero-out AdaLN and output projection weights for stable training init."""
+        bias = self.config.dit_bias
+        if self.adaln_type == "local":
+            for block in self.blocks:
+                nn.init.constant_(block.adaln_mlp.mlp[-1].weight, 0)
+                if bias:
+                    nn.init.constant_(block.adaln_mlp.mlp[-1].bias, 0)
+        elif self.adaln_type == "global":
+            nn.init.constant_(self.adaln_global_mlp.mlp[-1].weight, 0)
+            if bias:
+                nn.init.constant_(self.adaln_global_mlp.mlp[-1].bias, 0)
+
+        nn.init.constant_(self.norm_out.linear.weight, 0)
+        nn.init.constant_(self.proj_out.weight, 0)
+        if bias:
+            nn.init.constant_(self.norm_out.linear.bias, 0)
+            nn.init.constant_(self.proj_out.bias, 0)
+
+        for m in self.time_embed.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+        for m in self.text_embed.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+        # Boundary tokens: N(0, 0.02) so they carry non-trivial signal from
+        # step 0. HF from_pretrained's meta-init can leave new params with
+        # garbage (~1e36) past bf16 saturation; this re-init guarantees finite.
+        for tok in (
+            self.boe_token, self.bos_token, self.bon_token,
+            self.boe_text_token, self.bos_text_token, self.bon_text_token,
+        ):
+            nn.init.normal_(tok, mean=0.0, std=0.02)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        text: torch.Tensor,
+        text_len: torch.Tensor,
+        time: torch.Tensor,
+        mask: torch.BoolTensor | None = None,
+        cond_mask: torch.BoolTensor | None = None,
+        return_ith_layer: int | None = None,
+        latent_cond: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor | None]:
+        dtype = next(self.parameters()).dtype
+        x = x.to(dtype)
+        text = text.to(dtype)
+        time = time.to(dtype)
+
+        batch = x.shape[0]
+        text_seq_len = text.shape[1]
+        if time.ndim == 0:
+            time = time.repeat(batch)
+
+        t = self.time_embed(time)
+        text = self.text_embed(text, cond_mask)
+        if self.text_conv:
+            # The text ConvNeXt contains a GRN (ConvNeXtV2) that L2-pools over the
+            # TIME axis. GRN assumes every position is valid (it's an image op);
+            # running it on a zero-padded batch makes the pool length/padding-
+            # dependent → batched output ≠ single-sample. UMT5's own norms are
+            # per-token, so they're already batch-invariant. To match B=1 exactly,
+            # run the conv PER SAMPLE on each sequence's valid tokens (no padding
+            # enters the GRN), then scatter back. Robust to padding side; a no-op
+            # difference vs the old path when B=1 / no padding.
+            conv_out = torch.zeros_like(text)
+            for i in range(text.shape[0]):
+                mi = cond_mask[i]
+                conv_out[i][mi] = self.text_conv_layer(text[i][mi].unsqueeze(0))[0]
+            text = conv_out
+
+        x = self.input_embed(x, mask)
+        if self.use_latent_condition:
+            latent_cond = latent_cond.to(dtype)
+            latent_cond = self.latent_embed(latent_cond, mask)
+            x = self.latent_cond_embedder(torch.cat([x, latent_cond], dim=-1))
+
+        if self.long_skip:
+            x_clone = x.clone()
+
+        seq_len = x.shape[1]
+        rope = self.rotary_embed(x, seq_len)
+        cond_rope = self.rotary_embed(text, text_seq_len)
+
+        if self.adaln_type == "global":
+            if self.adaln_use_text_cond:
+                text_mean = text.sum(1) / text_len.unsqueeze(1).to(text.dtype)
+                norm_cond = t + text_mean
+            else:
+                norm_cond = t
+            adaln_mlp_out = self.adaln_global_mlp(norm_cond)
+        else:
+            adaln_mlp_out = None
+            norm_cond = None
+
+        hidden_state = None
+        for i, block in enumerate(self.blocks):
+            x = block(
+                x=x, t=t, cond=text, mask=mask, cond_mask=cond_mask,
+                rope=rope, cond_rope=cond_rope, adaln_global_out=adaln_mlp_out,
+            )
+            if return_ith_layer == i + 1:
+                hidden_state = x.clone()
+                if self.long_skip:
+                    x = x + x_clone
+
+        if self.long_skip:
+            x = x + x_clone
+
+        x = self.norm_out(x, norm_cond if norm_cond is not None else t)
+        output = self.proj_out(x)
+        return {"last_hidden_state": output, "hidden_state": hidden_state}
+
+
+# ---------------------------------------------------------------------------
+# WAV-VAE components (from wav_vae.py)
+# ---------------------------------------------------------------------------
+
+
+def _snake_beta(x: torch.Tensor, alpha: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:
+    return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2)
+
+
+class AudioDiTSnakeBeta(nn.Module):
+    def __init__(self, in_features: int, alpha_logscale: bool = True):
+        super().__init__()
+        self.alpha_logscale = alpha_logscale
+        self.alpha = nn.Parameter(torch.zeros(in_features))
+        self.beta = nn.Parameter(torch.zeros(in_features))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        return _snake_beta(x, alpha, beta)
+
+
+def _get_vae_activation(activation: str, channels: int | None = None) -> nn.Module:
+    if activation == "elu":
+        return nn.ELU()
+    elif activation == "snake":
+        return AudioDiTSnakeBeta(channels)
+    elif activation == "none":
+        return nn.Identity()
+    raise ValueError(f"Unknown activation {activation}")
+
+
+def _wn_conv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def _wn_conv_transpose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+def _pixel_unshuffle_1d(x: torch.Tensor, factor: int) -> torch.Tensor:
+    b, c, w = x.size()
+    return x.view(b, c, w // factor, factor).permute(0, 1, 3, 2).contiguous().view(b, c * factor, w // factor)
+
+
+def _pixel_shuffle_1d(x: torch.Tensor, factor: int) -> torch.Tensor:
+    b, c, w = x.size()
+    c = c // factor
+    return x.view(b, c, factor, w).permute(0, 1, 3, 2).contiguous().view(b, c, w * factor)
+
+
+class _DownsampleShortcut(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, factor: int):
+        super().__init__()
+        self.factor = factor
+        self.group_size = in_channels * factor // out_channels
+        self.out_channels = out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = _pixel_unshuffle_1d(x, self.factor)
+        b, c, n = x.shape
+        return x.view(b, self.out_channels, self.group_size, n).mean(dim=2)
+
+
+class _UpsampleShortcut(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, factor: int):
+        super().__init__()
+        self.factor = factor
+        self.repeats = out_channels * factor // in_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.repeat_interleave(self.repeats, dim=1)
+        return _pixel_shuffle_1d(x, self.factor)
+
+
+class _VaeResidualUnit(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, dilation: int, kernel_size: int = 7, use_snake: bool = False):
+        super().__init__()
+        padding = (dilation * (kernel_size - 1)) // 2
+        act = "snake" if use_snake else "elu"
+        self.layers = nn.Sequential(
+            _get_vae_activation(act, channels=out_channels),
+            _wn_conv1d(in_channels, out_channels, kernel_size, dilation=dilation, padding=padding),
+            _get_vae_activation(act, channels=out_channels),
+            _wn_conv1d(out_channels, out_channels, kernel_size=1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.layers(x)
+
+
+class _VaeEncoderBlock(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, stride: int, use_snake: bool = False, downsample_shortcut: str = "none"):
+        super().__init__()
+        layers = []
+        for d in [1, 3, 9]:
+            layers.append(_VaeResidualUnit(in_ch, in_ch, dilation=d, use_snake=use_snake))
+        act = "snake" if use_snake else "elu"
+        layers.append(_get_vae_activation(act, channels=in_ch))
+        layers.append(_wn_conv1d(in_ch, out_ch, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)))
+        self.layers = nn.Sequential(*layers)
+        self.res = _DownsampleShortcut(in_ch, out_ch, stride) if downsample_shortcut == "averaging" else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.res is not None:
+            return self.layers(x) + self.res(x)
+        return self.layers(x)
+
+
+class _VaeDecoderBlock(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, stride: int, use_snake: bool = False, upsample_shortcut: str = "none"):
+        super().__init__()
+        act = "snake" if use_snake else "elu"
+        layers = [
+            _get_vae_activation(act, channels=in_ch),
+            _wn_conv_transpose1d(in_ch, out_ch, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)),
+        ]
+        for d in [1, 3, 9]:
+            layers.append(_VaeResidualUnit(out_ch, out_ch, dilation=d, use_snake=use_snake))
+        self.layers = nn.Sequential(*layers)
+        self.res = _UpsampleShortcut(in_ch, out_ch, stride) if upsample_shortcut == "duplicating" else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.res is not None:
+            return self.layers(x) + self.res(x)
+        return self.layers(x)
+
+
+class AudioDiTVaeEncoder(nn.Module):
+    def __init__(self, config: AudioDiTVaeConfig):
+        super().__init__()
+        c_mults = [1] + config.c_mults
+        ch = config.channels
+        layers = [_wn_conv1d(config.in_channels, c_mults[0] * ch, kernel_size=7, padding=3)]
+        for i in range(len(c_mults) - 1):
+            layers.append(_VaeEncoderBlock(c_mults[i] * ch, c_mults[i + 1] * ch, config.strides[i], use_snake=config.use_snake, downsample_shortcut=config.downsample_shortcut))
+        layers.append(_wn_conv1d(c_mults[-1] * ch, config.encoder_latent_dim, kernel_size=3, padding=1))
+        self.layers = nn.Sequential(*layers)
+
+        if config.out_shortcut == "averaging":
+            self.shortcut = _DownsampleShortcut(c_mults[-1] * ch, config.encoder_latent_dim, 1)
+        else:
+            self.shortcut = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.shortcut is None:
+            return self.layers(x)
+        x = self.layers[:-1](x)
+        return self.layers[-1](x) + self.shortcut(x)
+
+
+class AudioDiTVaeDecoder(nn.Module):
+    def __init__(self, config: AudioDiTVaeConfig):
+        super().__init__()
+        c_mults = [1] + config.c_mults
+        ch = config.channels
+
+        if config.in_shortcut == "duplicating":
+            self.shortcut = _UpsampleShortcut(config.latent_dim, c_mults[-1] * ch, 1)
+        else:
+            self.shortcut = None
+
+        layers = [_wn_conv1d(config.latent_dim, c_mults[-1] * ch, kernel_size=7, padding=3)]
+        for i in range(len(c_mults) - 1, 0, -1):
+            layers.append(_VaeDecoderBlock(c_mults[i] * ch, c_mults[i - 1] * ch, config.strides[i - 1], use_snake=config.use_snake, upsample_shortcut=config.upsample_shortcut))
+        act = "snake" if config.use_snake else "elu"
+        layers.append(_get_vae_activation(act, channels=c_mults[0] * ch))
+        layers.append(_wn_conv1d(c_mults[0] * ch, config.in_channels, kernel_size=7, padding=3, bias=False))
+        if config.final_tanh:
+            layers.append(nn.Tanh())
+        else:
+            layers.append(nn.Identity())
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.shortcut is None:
+            return self.layers(x)
+        x_short = self.shortcut(x) + self.layers[0](x)
+        return self.layers[1:](x_short)
+
+
+class AudioDiTVae(nn.Module):
+    """WAV-VAE audio autoencoder with VAE bottleneck and scale factor.
+
+    The original checkpoint runs encode/decode in **float16** (``model_half=True``
+    in ``AutoencoderPretransform``).  We replicate this behaviour so that the
+    outputs are numerically identical to the original codebase.
+    """
+
+    def __init__(self, config: AudioDiTVaeConfig):
+        super().__init__()
+        self.config = config
+        self.encoder = AudioDiTVaeEncoder(config)
+        self.decoder = AudioDiTVaeDecoder(config)
+        self.scale = config.scale
+        self.downsampling_ratio = config.downsampling_ratio
+
+    def to_half(self):
+        """Convert encoder and decoder weights to float16 (matching original behaviour)."""
+        self.encoder.half()
+        self.decoder.half()
+        return self
+
+    def encode(self, audio: torch.Tensor) -> torch.Tensor:
+        """Encode audio to latent space.
+
+        Runs encoder **and** VAE bottleneck in float16 when weights are float16,
+        matching the original ``AutoencoderPretransform(model_half=True)`` +
+        ``AudioAutoencoder.encode`` behaviour where the bottleneck operates on
+        the fp16 encoder output before the final ``.float()`` conversion.
+
+        Args:
+            audio: ``(batch, 1, num_samples)`` raw waveform.
+
+        Returns:
+            Latent tensor ``(batch, latent_dim, num_frames)`` in float32.
+        """
+        is_half = next(self.encoder.parameters()).dtype == torch.float16
+        if is_half:
+            audio = audio.half()
+        latents = self.encoder(audio)
+        # VAE bottleneck runs in the same dtype as encoder output (fp16)
+        # to match original: bottleneck.encode(latents) happens before .float()
+        mean, scale_param = latents.chunk(2, dim=1)
+        stdev = F.softplus(scale_param) + 1e-4
+        latents = torch.randn_like(mean) * stdev + mean
+        # Convert to fp32 after bottleneck, matching original AutoencoderPretransform
+        if is_half:
+            latents = latents.float()
+        return latents / self.scale
+
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        """Decode latents to audio waveform.
+
+        Runs decoder in float16 when weights are float16, matching the original
+        ``AutoencoderPretransform(model_half=True)`` behaviour.
+
+        Args:
+            latents: ``(batch, latent_dim, num_frames)``.
+
+        Returns:
+            Waveform tensor ``(batch, 1, num_samples)`` in float32.
+        """
+        z = latents * self.scale
+        is_half = next(self.decoder.parameters()).dtype == torch.float16
+        if is_half:
+            z = z.half()
+        decoded = self.decoder(z)
+        if is_half:
+            decoded = decoded.float()
+        return decoded
+
+
+# ---------------------------------------------------------------------------
+# Top-level AudioDiTModel
+# ---------------------------------------------------------------------------
+
+
+class AudioDiTPreTrainedModel(PreTrainedModel):
+    config_class = AudioDiTConfig
+    base_model_prefix = "audiodit"
+    supports_gradient_checkpointing = True
+    _supports_sdpa = True
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=0.02)
+        elif isinstance(module, AudioDiTTransformer):
+            # Re-init the boundary tokens after HF from_pretrained — they're
+            # nn.Parameter (not modules) and don't appear in the pretrained
+            # ckpt for the env-tts task, so HF's meta-init path leaves them
+            # with uninitialized memory (~1e32) past bf16 saturation.
+            for tok_name in (
+                "boe_token", "bos_token", "bon_token",
+                "boe_text_token", "bos_text_token", "bon_text_token",
+            ):
+                tok = getattr(module, tok_name, None)
+                if tok is not None:
+                    nn.init.normal_(tok, mean=0.0, std=0.02)
+
+
+class AudioDiTModel(AudioDiTPreTrainedModel):
+    """AudioDiT: Conditional Flow Matching TTS model with DiT backbone, UMT5 text encoder, and WAV-VAE.
+
+    All sub-models (text_encoder, transformer, vae) are constructed from config
+    and their weights are loaded together via ``from_pretrained``.
+
+    Example::
+
+        model = AudioDiTModel.from_pretrained("hf_audiodit_1b")
+        tokenizer = AutoTokenizer.from_pretrained(model.config.text_encoder_model)
+        output = model(text=["Hello world"], tokenizer=tokenizer)
+        waveform = output.waveform  # (B, num_samples)
+    """
+
+    def __init__(self, config: AudioDiTConfig):
+        super().__init__(config)
+        self.config = config
+
+        # Text encoder — constructed from embedded config, weights loaded by from_pretrained
+        from transformers import UMT5EncoderModel, UMT5Config
+
+        if config.text_encoder_config is not None:
+            self.text_encoder = UMT5EncoderModel(config.text_encoder_config)
+        else:
+            te_config = UMT5Config.from_pretrained(config.text_encoder_model)
+            self.text_encoder = UMT5EncoderModel(te_config)
+        self.text_encoder.requires_grad_(False)
+
+        # DiT transformer
+        self.transformer = AudioDiTTransformer(config)
+
+        # WAV-VAE
+        self.vae = AudioDiTVae(config.vae_config)
+        self.vae.requires_grad_(False)
+
+        self.post_init()
+
+    def encode_text(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.LongTensor,
+    ) -> torch.FloatTensor:
+        """Encode tokenized text using the UMT5 text encoder.
+
+        Args:
+            input_ids: Token ids ``(batch, seq_len)``.
+            attention_mask: Attention mask ``(batch, seq_len)``.
+
+        Returns:
+            Text embeddings ``(batch, seq_len, text_dim)`` in float32.
+        """
+        with torch.no_grad():
+            output = self.text_encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+        emb = output.last_hidden_state
+        d_model = self.text_encoder.config.d_model
+
+        if self.config.text_norm_feat:
+            emb = F.layer_norm(emb, (d_model,), eps=1e-6)
+
+        if self.config.text_add_embed:
+            first_hidden = output.hidden_states[0]
+            if self.config.text_norm_feat:
+                first_hidden = F.layer_norm(first_hidden, (d_model,), eps=1e-6)
+            emb = emb + first_hidden
+
+        return emb.float()
+
+    def encode_multistream_text(
+        self,
+        env_input_ids: torch.LongTensor,
+        env_attn: torch.LongTensor,
+        spk_input_ids: torch.LongTensor,
+        spk_attn: torch.LongTensor,
+        target_input_ids: torch.LongTensor,
+        target_attn: torch.LongTensor,
+        drop_env_text: torch.BoolTensor | None = None,
+        drop_spk_text: torch.BoolTensor | None = None,
+        drop_target_text: torch.BoolTensor | None = None,
+    ) -> tuple[torch.FloatTensor, torch.BoolTensor, torch.LongTensor]:
+        """Encode three text streams and assemble with boundary tokens.
+
+        Each segment is independently tokenized + frozen-UMT5 encoded, then
+        concatenated as:
+            [<boe_text>, env_emb, <bos_text>, spk_emb, <bon_text>, tgt_emb]
+        Boundary tokens are always visible in the output mask. If a per-sample
+        drop flag is set, that segment's embedding values are zeroed in place
+        (position + boundary preserved, content zeroed).
+
+        Args:
+            env_input_ids / env_attn:    (B, S_env)
+            spk_input_ids / spk_attn:    (B, S_spk)
+            target_input_ids / target_attn: (B, S_tgt)
+            drop_env_text / drop_spk_text / drop_target_text: (B,) bool or None.
+
+        Returns:
+            text_emb:  (B, 3 + S_env + S_spk + S_tgt, dit_text_dim) float32
+            text_mask: (B, 3 + S_env + S_spk + S_tgt) bool — boundary positions
+                always True; segment positions follow their attention masks
+                (NOT cleared by drop flags, consistent with dit-training CFG
+                null-pass convention which preserves cond_mask).
+            text_len:  (B,) long — sum of text_mask along dim=1.
+        """
+        device = self.device
+
+        # Concat-batch UMT5 encode: pad three streams to a common seq_len then
+        # run ONE encode_text call on shape (3B, S_max). Splits back per-stream
+        # at the end. Saves 2 kernel-launch round-trips per cfm_step.
+        env_ids = env_input_ids.to(device)
+        spk_ids = spk_input_ids.to(device)
+        tgt_ids = target_input_ids.to(device)
+        env_msk = env_attn.to(device)
+        spk_msk = spk_attn.to(device)
+        tgt_msk = target_attn.to(device)
+        S_env, S_spk, S_tgt = env_ids.shape[1], spk_ids.shape[1], tgt_ids.shape[1]
+        S_max = max(S_env, S_spk, S_tgt)
+
+        def _pad(t, s):
+            return F.pad(t, (0, s - t.shape[1])) if t.shape[1] < s else t
+
+        all_ids  = torch.cat([_pad(env_ids, S_max),  _pad(spk_ids, S_max),  _pad(tgt_ids, S_max)],  dim=0)
+        all_msk  = torch.cat([_pad(env_msk, S_max),  _pad(spk_msk, S_max),  _pad(tgt_msk, S_max)],  dim=0)
+        all_emb  = self.encode_text(all_ids, all_msk)   # (3B, S_max, D)
+
+        B = env_ids.shape[0]
+        env_emb = all_emb[0    : B   , :S_env, :]
+        spk_emb = all_emb[B    : 2*B , :S_spk, :]
+        tgt_emb = all_emb[2*B : 3*B , :S_tgt, :]
+        text_dim = env_emb.shape[-1]
+
+        # Apply text-side drop (position-preserving content zero).
+        if drop_env_text is not None:
+            env_emb = env_emb * (~drop_env_text.to(device)).view(B, 1, 1).to(env_emb.dtype)
+        if drop_spk_text is not None:
+            spk_emb = spk_emb * (~drop_spk_text.to(device)).view(B, 1, 1).to(spk_emb.dtype)
+        if drop_target_text is not None:
+            tgt_emb = tgt_emb * (~drop_target_text.to(device)).view(B, 1, 1).to(tgt_emb.dtype)
+
+        # Resolve PEFT-wrapped transformer to access the boundary nn.Parameter.
+        src = getattr(self.transformer, "base_model", None)
+        src = src.model if src is not None else self.transformer
+        boe_t = src.boe_text_token.to(device=device, dtype=env_emb.dtype)
+        bos_t = src.bos_text_token.to(device=device, dtype=env_emb.dtype)
+        bon_t = src.bon_text_token.to(device=device, dtype=env_emb.dtype)
+        # ── Tight per-sample assembly (batch-invariant positions) ─────────
+        # Pack each sample's VALID tokens contiguously [boe|env|bos|spk|bon|tgt]
+        # and end-pad to the batch max. A sample's assembled text — and thus its
+        # cond_rope positions — is therefore INDEPENDENT of other batch members'
+        # stream lengths, so batched == single-sample.
+        #
+        # (The earlier version padded each stream to the batch-max per-stream
+        # length and concatenated WITH that padding inside, interleaving padding
+        # mid-sequence. For any sample shorter than the batch max this shifted the
+        # bos/bon + spk/tgt positions → cond_rope mismatch at B>1 → corrupted
+        # generation that compounds over the ODE. B=1 is unaffected and stays
+        # byte-identical; the single-stream encode_text path is untouched.)
+        boe1, bos1, bon1 = boe_t.reshape(1, text_dim), bos_t.reshape(1, text_dim), bon_t.reshape(1, text_dim)
+        env_m = env_attn.to(device).bool()
+        spk_m = spk_attn.to(device).bool()
+        tgt_m = target_attn.to(device).bool()
+        seqs = [
+            torch.cat([
+                boe1, env_emb[i][env_m[i]],
+                bos1, spk_emb[i][spk_m[i]],
+                bon1, tgt_emb[i][tgt_m[i]],
+            ], dim=0)                                   # (L_i, text_dim), tight
+            for i in range(B)
+        ]
+        L_max = max(s.shape[0] for s in seqs)
+        text_emb = torch.stack(
+            [F.pad(s, (0, 0, 0, L_max - s.shape[0])) for s in seqs], dim=0)   # (B, L_max, text_dim)
+        text_mask = torch.zeros(B, L_max, dtype=torch.bool, device=device)
+        for i, s in enumerate(seqs):
+            text_mask[i, : s.shape[0]] = True
+        text_len = text_mask.sum(dim=1).long()
+        return text_emb.float(), text_mask, text_len
+
+    def encode_prompt_audio(self, prompt_audio: torch.FloatTensor) -> tuple[torch.FloatTensor, int]:
+        """Encode prompt audio to latent space.
+
+        Args:
+            prompt_audio: Waveform tensor ``(batch, 1, num_samples)`` or ``(batch, num_samples)``.
+
+        Returns:
+            Tuple of (prompt_latent ``(batch, num_frames, latent_dim)``, prompt_duration_frames).
+        """
+        full_hop = self.config.latent_hop
+        off = 3
+        wav = prompt_audio.to(self.device)
+        if wav.ndim == 2:
+            wav = wav.unsqueeze(1)
+        if wav.shape[-1] % full_hop != 0:
+            wav = F.pad(wav, (0, full_hop - wav.shape[-1] % full_hop))
+        wav = F.pad(wav, (0, full_hop * off))
+        latent = self.vae.encode(wav)
+        if off != 0:
+            latent = latent[..., :-off]
+        prompt_duration_frames = latent.shape[-1]
+        return latent.permute(0, 2, 1), prompt_duration_frames
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.LongTensor | None = None,
+        text_embedding: torch.FloatTensor | None = None,
+        text_mask: torch.BoolTensor | None = None,
+        prompt_audio: torch.FloatTensor | None = None,
+        prompt_latent: torch.FloatTensor | None = None,
+        prompt_lens: torch.LongTensor | None = None,
+        duration: int | None = None,
+        steps: int = 16,
+        cfg_strength: float = 4.0,
+        guidance_method: str = "cfg",
+        return_dict: bool = True,
+    ) -> AudioDiTOutput | tuple:
+        """Generate audio from text (and optional prompt audio).
+
+        Args:
+            input_ids: Tokenized text ``(batch, seq_len)``. Use with ``attention_mask``.
+            attention_mask: Attention mask ``(batch, seq_len)``.
+            text_embedding: Pre-computed text embeddings ``(batch, seq_len, dim)``. Alternative to input_ids.
+                When supplied alongside ``text_mask`` the model bypasses ``encode_text``
+                entirely — used by the env-tts pipeline which builds a multi-stream
+                text embedding via ``encode_multistream_text``.
+            text_mask: Optional bool mask ``(batch, seq_len)`` for ``text_embedding``.
+                Required when ``text_embedding`` is supplied without ``attention_mask``.
+            prompt_audio: Optional prompt audio ``(batch, 1, num_samples)`` for voice cloning.
+            prompt_latent: Optional pre-assembled prompt latent ``(batch, T_prompt, latent_dim)``,
+                bypassing ``encode_prompt_audio``. Use this for env-tts multi-stream
+                latents already containing latent-space boundary tokens.
+                Mutually exclusive with ``prompt_audio`` — if both are given,
+                ``prompt_latent`` wins.
+            duration: Target duration in latent frames (prompt + gen). If None, uses max_wav_duration.
+            steps: Number of ODE Euler steps (default 16).
+            cfg_strength: Guidance strength for CFG/APG (default 4.0).
+            guidance_method: ``"cfg"`` or ``"apg"`` (default ``"cfg"``).
+            return_dict: Whether to return ``AudioDiTOutput`` or tuple.
+        """
+        device = self.device
+        sr = self.config.sampling_rate
+        full_hop = self.config.latent_hop
+        max_duration_frames = int(self.config.max_wav_duration * sr // full_hop)
+        repa_layer = self.config.repa_dit_layer
+
+        # ── text encoding ─────────────────────────────────────────────
+        if text_embedding is not None:
+            text_condition = text_embedding.to(device, torch.float32)
+            if text_mask is not None:
+                text_condition_len = text_mask.to(device).sum(dim=1).long()
+            elif attention_mask is not None:
+                text_condition_len = attention_mask.sum(dim=1).to(device)
+            else:
+                text_condition_len = torch.full(
+                    (text_condition.shape[0],), text_condition.shape[1], device=device,
+                )
+        else:
+            text_condition = self.encode_text(
+                input_ids.to(device), attention_mask.to(device),
+            )
+            text_condition_len = attention_mask.sum(dim=1).to(device)
+
+        batch = text_condition.shape[0]
+
+        # ── prompt latent / audio encoding ────────────────────────────
+        # Precedence: explicit ``prompt_latent`` > ``prompt_audio`` > empty.
+        # ``prompt_latent`` is used by the env-tts pipeline which builds a
+        # multi-stream latent [boe|z_env|bos|z_spk|bon] externally; the
+        # ``prompt_audio`` path is the single-stream voice-cloning default.
+        has_prompt = prompt_latent is not None or prompt_audio is not None
+        if prompt_latent is not None:
+            prompt_latent = prompt_latent.to(device)
+            prompt_dur = prompt_latent.shape[1]
+        elif prompt_audio is not None:
+            prompt_latent, prompt_dur = self.encode_prompt_audio(prompt_audio)
+        else:
+            prompt_latent = torch.empty(batch, 0, self.config.latent_dim, device=device)
+            prompt_dur = 0
+
+        # ── duration ──────────────────────────────────────────────────
+        # ``duration`` may be a scalar (uniform, the single-sample path) OR a
+        # per-sample 1-D tensor / list of TOTAL frames (prompt+gen). The latter
+        # enables BATCHED generation of variable-length samples — callers pad all
+        # prompts to a common width (so ``prompt_dur`` stays uniform) and pass the
+        # per-sample total lengths here; the transformer ``mask`` + per-sample
+        # ``y0`` already handle ragged gen lengths.
+        if duration is None:
+            duration = max_duration_frames
+        if torch.is_tensor(duration) or isinstance(duration, (list, tuple)):
+            duration_tensor = torch.as_tensor(duration, device=device, dtype=torch.long).clamp(max=max_duration_frames)
+        else:
+            duration_tensor = torch.full((batch,), min(int(duration), max_duration_frames),
+                                         device=device, dtype=torch.long)
+        max_dur = int(duration_tensor.max().item())
+
+        # ── masks & conditioning ──────────────────────────────────────
+        mask = lens_to_mask(duration_tensor, length=max_dur)
+        if text_mask is not None:
+            text_cond_mask = text_mask.to(device).bool()
+        else:
+            text_cond_mask = lens_to_mask(text_condition_len, length=text_condition.shape[1])
+
+        neg_text = torch.zeros_like(text_condition)
+        neg_text_len = text_condition_len
+
+        # ``prompt_lens`` (B,) gives each sample's REAL prompt length so the gen
+        # region starts immediately after that sample's ``bon`` (no padding between
+        # the boundary token and gen). When None, the single-sample uniform path is
+        # used unchanged. ``latent_cond`` is the real prompt at [0:T_p_i] then zeros.
+        latent_len = prompt_dur
+        prompt_mask = None
+        if prompt_lens is not None:
+            prompt_lens = torch.as_tensor(prompt_lens, device=device, dtype=torch.long)
+            prompt_mask = lens_to_mask(prompt_lens, length=max_dur)   # (B, max_dur)
+        if has_prompt:
+            latent_cond = F.pad(prompt_latent, (0, 0, 0, max_dur - prompt_latent.shape[1]))
+            empty_latent_cond = torch.zeros_like(latent_cond)
+        else:
+            latent_cond = torch.zeros(batch, max_dur, self.config.latent_dim, device=device)
+            empty_latent_cond = latent_cond
+
+        # ── APG buffer ────────────────────────────────────────────────
+        if guidance_method == "apg":
+            if prompt_mask is not None:
+                raise NotImplementedError("APG guidance unsupported with batched prompt_lens; use cfg.")
+            apg_buffer = _MomentumBuffer(momentum=-0.3)
+
+        # ── ODE function ──────────────────────────────────────────────
+        def fn(t, x):
+            if prompt_mask is not None:   # per-sample prompt region (in-place reset each step)
+                x[prompt_mask] = (prompt_noise * (1 - t) + latent_cond * t)[prompt_mask]
+            else:
+                x[:, :latent_len] = prompt_noise * (1 - t) + latent_cond[:, :latent_len] * t
+            output = self.transformer(
+                x=x, text=text_condition, text_len=text_condition_len, time=t,
+                mask=mask, cond_mask=text_cond_mask,
+                return_ith_layer=repa_layer, latent_cond=latent_cond,
+            )
+            pred = output["last_hidden_state"]
+
+            if cfg_strength < 1e-5:
+                return pred
+
+            if prompt_mask is not None:
+                x[prompt_mask] = 0
+            else:
+                x[:, :latent_len] = 0
+            null_output = self.transformer(
+                x=x, text=neg_text, text_len=neg_text_len, time=t,
+                mask=mask, cond_mask=text_cond_mask,
+                return_ith_layer=repa_layer, latent_cond=empty_latent_cond,
+            )
+            null_pred = null_output["last_hidden_state"]
+
+            if guidance_method == "cfg":
+                return pred + (pred - null_pred) * cfg_strength
+
+            # APG (single-sample path only)
+            x_s = x[:, latent_len:]
+            pred_s = pred[:, latent_len:]
+            null_s = null_pred[:, latent_len:]
+            pred_sample = x_s + (1 - t) * pred_s
+            null_sample = x_s + (1 - t) * null_s
+            out = _apg_forward(
+                pred_sample, null_sample, cfg_strength, apg_buffer,
+                eta=0.5, norm_threshold=0.0, dims=[-1, -2],
+            )
+            out = (out - x_s) / (1 - t)
+            return F.pad(out, (0, 0, latent_len, 0), value=0.0)
+
+        # ── initial noise ─────────────────────────────────────────────
+        y0 = []
+        for dur in duration_tensor:
+            noise = torch.randn(dur.item(), self.config.latent_dim, device=device)
+            y0.append(noise)
+        y0 = pad_sequence(y0, padding_value=0, batch_first=True)
+
+        # ── ODE solve ─────────────────────────────────────────────────
+        t = torch.linspace(0, 1, steps, device=device)
+        prompt_noise = y0.clone() if prompt_mask is not None else y0[:, :latent_len].clone()
+        trajectory = odeint_euler(fn, y0, t)
+        sampled = trajectory[-1]
+
+        # ── extract gen region + decode ───────────────────────────────
+        if prompt_mask is not None:
+            # Decode each sample's gen latent at its EXACT length (decoding a
+            # zero-PADDED batch latent bleeds VAE-decoder artifacts into the valid
+            # audio), then zero-pad the waveforms; the caller trims to true length.
+            wavs = []
+            for b in range(batch):
+                g = sampled[b, int(prompt_lens[b]):int(duration_tensor[b])]   # (T_gen, D)
+                gl = g.permute(1, 0).unsqueeze(0).float()                     # (1, D, T_gen)
+                wavs.append(self.vae.decode(gl).reshape(-1))                  # (T_gen*hop,)
+            max_w = max((w.shape[0] for w in wavs), default=1)
+            waveform = torch.stack([F.pad(w, (0, max_w - w.shape[0])) for w in wavs], dim=0)
+            pred_latent = None
+        else:
+            pred_latent = sampled
+            if has_prompt:
+                pred_latent = pred_latent[:, prompt_dur:]
+            pred_latent = pred_latent.permute(0, 2, 1).float()
+            waveform = self.vae.decode(pred_latent).squeeze(1)
+
+        if not return_dict:
+            return (waveform, pred_latent)
+        return AudioDiTOutput(waveform=waveform, latent=pred_latent)
+
+
+# ---------------------------------------------------------------------------
+# APG helpers (from model/cfm.py — Adaptive Projected Guidance)
+# ---------------------------------------------------------------------------
+
+
+class _MomentumBuffer:
+    def __init__(self, momentum: float = -0.75):
+        self.momentum = momentum
+        self.running_average = 0
+
+    def update(self, update_value: torch.Tensor):
+        new_average = self.momentum * self.running_average
+        self.running_average = update_value + new_average
+
+
+def _project(v0: torch.Tensor, v1: torch.Tensor, dims=(-1, -2)):
+    dtype = v0.dtype
+    device_type = v0.device.type
+    if device_type == "mps":
+        v0, v1 = v0.cpu(), v1.cpu()
+    v0, v1 = v0.double(), v1.double()
+    v1 = F.normalize(v1, dim=dims)
+    v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
+    v0_orthogonal = v0 - v0_parallel
+    return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(device_type)
+
+
+def _apg_forward(pred_cond, pred_uncond, guidance_scale, momentum_buffer=None, eta=0.0, norm_threshold=2.5, dims=(-1, -2)):
+    diff = pred_cond - pred_uncond
+    if momentum_buffer is not None:
+        momentum_buffer.update(diff)
+        diff = momentum_buffer.running_average
+    if norm_threshold > 0:
+        ones = torch.ones_like(diff)
+        diff_norm = diff.norm(p=2, dim=dims, keepdim=True)
+        scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
+        diff = diff * scale_factor
+    diff_parallel, diff_orthogonal = _project(diff, pred_cond, dims)
+    normalized_update = diff_orthogonal + eta * diff_parallel
+    return pred_cond + guidance_scale * normalized_update
+
+
+__all__ = [
+    "AudioDiTConfig",
+    "AudioDiTVaeConfig",
+    "AudioDiTOutput",
+    "AudioDiTPreTrainedModel",
+    "AudioDiTModel",
+    "AudioDiTTransformer",
+    "AudioDiTVae",
+]