Upload folder using huggingface_hub

README.md

@@ -1,13 +1,18 @@
 ---
-title: Irodori TTS 500M V3 Demo
-emoji: 🐨
-colorFrom: red
-colorTo: yellow
+title: Irodori-TTS-500M-v3 Demo
+emoji: 🎤
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 6.14.0
+sdk_version: 6.6.0
 python_version: '3.12'
 app_file: app.py
-pinned: false
+pinned: true
+license: mit
+short_description: TTS demo for Irodori-TTS-500M-v3
+models:
+  - Aratako/Irodori-TTS-500M-v3
+  - Aratako/Semantic-DACVAE-Japanese-32dim
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,436 @@
+import io
+import os
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+from huggingface_hub import hf_hub_download
+
+from irodori_tts.inference_runtime import (
+    InferenceRuntime,
+    RuntimeKey,
+    SamplingRequest,
+)
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+MODEL_REPO = os.environ.get("MODEL_REPO", "Aratako/Irodori-TTS-500M-v3")
+CODEC_REPO = "Aratako/Semantic-DACVAE-Japanese-32dim"
+MAX_GRADIO_CANDIDATES = int(os.environ.get("MAX_GRADIO_CANDIDATES", "32"))
+GRADIO_AUDIO_COLS_PER_ROW = 8
+
+# Global state
+_runtime: InferenceRuntime | None = None
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _parse_optional_float(raw: str | None, label: str) -> float | None:
+    if raw is None:
+        return None
+    text = str(raw).strip()
+    if text == "" or text.lower() == "none":
+        return None
+    try:
+        return float(text)
+    except ValueError as exc:
+        raise ValueError(f"{label} must be a float or blank.") from exc
+
+
+def _parse_optional_int(raw: str | None, label: str) -> int | None:
+    if raw is None:
+        return None
+    text = str(raw).strip()
+    if text == "" or text.lower() == "none":
+        return None
+    try:
+        return int(text)
+    except ValueError as exc:
+        raise ValueError(f"{label} must be an int or blank.") from exc
+
+
+def _format_timings(stage_timings: list[tuple[str, float]], total_to_decode: float) -> str:
+    lines = [
+        "[timing] ---- request ----",
+        *[f"[timing] {name}: {sec * 1000.0:.1f} ms" for name, sec in stage_timings],
+        f"[timing] total_to_decode: {total_to_decode:.3f} s",
+    ]
+    return "\n".join(lines)
+
+
+# ---------------------------------------------------------------------------
+# Model Loading
+# ---------------------------------------------------------------------------
+
+
+def load_models():
+    global _runtime
+
+    if _runtime is not None:
+        return
+
+    print(f"[Info] Downloading checkpoint from {MODEL_REPO}...")
+    checkpoint_path = hf_hub_download(repo_id=MODEL_REPO, filename="model.safetensors")
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    precision = "bf16" if device == "cuda" else "fp32"
+
+    key = RuntimeKey(
+        checkpoint=checkpoint_path,
+        model_device=device,
+        codec_repo=CODEC_REPO,
+        model_precision=precision,
+        codec_device=device,
+        codec_precision=precision,
+    )
+
+    print("[Info] Building runtime...")
+    _runtime = InferenceRuntime.from_key(key)
+    print("[Info] All models loaded successfully.")
+
+
+# Load models at startup
+load_models()
+
+
+# ---------------------------------------------------------------------------
+# GPU-decorated Inference
+# ---------------------------------------------------------------------------
+
+
+@spaces.GPU(duration=120)
+def run_inference_gpu(
+    text: str,
+    uploaded_audio: str | None,
+    num_steps: int,
+    num_candidates: int,
+    seed_raw: str,
+    seconds_raw: str,
+    duration_scale: float,
+    cfg_guidance_mode: str,
+    cfg_scale_text: float,
+    cfg_scale_speaker: float,
+    cfg_scale_raw: str,
+    cfg_min_t: float,
+    cfg_max_t: float,
+    context_kv_cache: bool,
+    truncation_factor_raw: str,
+    rescale_k_raw: str,
+    rescale_sigma_raw: str,
+    speaker_kv_scale_raw: str,
+    speaker_kv_min_t_raw: str,
+    speaker_kv_max_layers_raw: str,
+) -> tuple[list[tuple[int, np.ndarray]], str]:
+    load_models()
+
+    log_buffer = io.StringIO()
+
+    def stdout_log(msg: str) -> None:
+        print(msg, flush=True)
+        log_buffer.write(msg + "\n")
+
+    if not str(text).strip():
+        raise gr.Error("Please enter text to synthesize.")
+
+    cfg_scale = _parse_optional_float(cfg_scale_raw, "cfg_scale")
+    truncation_factor = _parse_optional_float(truncation_factor_raw, "truncation_factor")
+    rescale_k = _parse_optional_float(rescale_k_raw, "rescale_k")
+    rescale_sigma = _parse_optional_float(rescale_sigma_raw, "rescale_sigma")
+    speaker_kv_scale = _parse_optional_float(speaker_kv_scale_raw, "speaker_kv_scale")
+    speaker_kv_min_t = _parse_optional_float(speaker_kv_min_t_raw, "speaker_kv_min_t")
+    speaker_kv_max_layers = _parse_optional_int(speaker_kv_max_layers_raw, "speaker_kv_max_layers")
+    seed = _parse_optional_int(seed_raw, "seed")
+    manual_seconds = _parse_optional_float(seconds_raw, "seconds")
+    requested_candidates = int(num_candidates)
+    if requested_candidates <= 0:
+        raise gr.Error("num_candidates must be >= 1.")
+    if requested_candidates > MAX_GRADIO_CANDIDATES:
+        raise gr.Error(f"num_candidates must be <= {MAX_GRADIO_CANDIDATES}.")
+
+    ref_wav: str | None = None
+    no_ref = True
+    if uploaded_audio is not None and str(uploaded_audio).strip() != "":
+        ref_wav = str(uploaded_audio)
+        no_ref = False
+
+    stdout_log(
+        (
+            "[Info] request: mode={} seconds={} duration_scale={} "
+            "steps={} seed={} no_ref={} candidates={}"
+        ).format(
+            cfg_guidance_mode,
+            "auto" if manual_seconds is None else manual_seconds,
+            float(duration_scale),
+            int(num_steps),
+            "random" if seed is None else seed,
+            no_ref,
+            requested_candidates,
+        )
+    )
+
+    result = _runtime.synthesize(
+        SamplingRequest(
+            text=str(text),
+            ref_wav=ref_wav,
+            ref_latent=None,
+            no_ref=bool(no_ref),
+            ref_normalize_db=-16.0,
+            ref_ensure_max=True,
+            num_candidates=requested_candidates,
+            decode_mode="sequential",
+            seconds=manual_seconds,
+            duration_scale=float(duration_scale),
+            max_ref_seconds=30.0,
+            max_text_len=None,
+            num_steps=int(num_steps),
+            seed=None if seed is None else int(seed),
+            cfg_guidance_mode=str(cfg_guidance_mode),
+            cfg_scale_text=float(cfg_scale_text),
+            cfg_scale_speaker=float(cfg_scale_speaker),
+            cfg_scale=cfg_scale,
+            cfg_min_t=float(cfg_min_t),
+            cfg_max_t=float(cfg_max_t),
+            truncation_factor=truncation_factor,
+            rescale_k=rescale_k,
+            rescale_sigma=rescale_sigma,
+            context_kv_cache=bool(context_kv_cache),
+            speaker_kv_scale=speaker_kv_scale,
+            speaker_kv_min_t=speaker_kv_min_t,
+            speaker_kv_max_layers=speaker_kv_max_layers,
+            trim_tail=True,
+        ),
+        log_fn=stdout_log,
+    )
+
+    sample_rate = result.sample_rate
+    audio_results: list[tuple[int, np.ndarray]] = []
+    for audio in result.audios:
+        waveform = audio.squeeze(0).float().numpy()
+        audio_results.append((sample_rate, waveform))
+    stdout_log(f"[Info] seed_used: {result.used_seed}")
+    stdout_log(f"[Info] candidates: {len(result.audios)}")
+    for message in result.messages:
+        stdout_log(message)
+    stdout_log(_format_timings(result.stage_timings, result.total_to_decode))
+    return audio_results, log_buffer.getvalue()
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+
+def build_demo():
+    MODEL_LINK = f"https://huggingface.co/{MODEL_REPO}"
+    GITHUB_REPO = "https://github.com/Aratako/Irodori-TTS"
+
+    title = "# Irodori-TTS-500M-v3 Demo"
+    description = f"""\
+[Model]({MODEL_LINK}) | [GitHub]({GITHUB_REPO})
+
+Flow-matching based Japanese TTS model (500M parameters). \
+Generates speech from text using rectified flow over DACVAE latents.
+
+- **Reference audio**: Optional. Upload to condition the speaker voice. \
+Leave blank for unconditional generation.
+- **Duration**: By default, v3 predicts the output duration automatically. \
+Use Duration Scale for small adjustments or Seconds for exact manual control.
+"""
+
+    with gr.Blocks() as demo:
+        gr.Markdown(title)
+        gr.Markdown(description)
+
+        text = gr.Textbox(label="Text", lines=4)
+        uploaded_audio = gr.Audio(
+            label="Reference Audio Upload (optional, blank = no-reference mode)",
+            type="filepath",
+        )
+
+        with gr.Accordion("Sampling", open=True):
+            with gr.Row():
+                num_steps = gr.Slider(
+                    label="Num Steps",
+                    minimum=1,
+                    maximum=120,
+                    value=40,
+                    step=1,
+                )
+                num_candidates = gr.Slider(
+                    label="Num Candidates",
+                    minimum=1,
+                    maximum=MAX_GRADIO_CANDIDATES,
+                    value=1,
+                    step=1,
+                )
+                seed_raw = gr.Textbox(
+                    label="Seed (blank=random)",
+                    value="",
+                )
+                seconds_raw = gr.Textbox(
+                    label="Seconds (blank=auto)",
+                    value="",
+                )
+                duration_scale = gr.Slider(
+                    label="Duration Scale",
+                    minimum=0.5,
+                    maximum=1.5,
+                    value=1.0,
+                    step=0.01,
+                )
+
+            with gr.Row():
+                cfg_guidance_mode = gr.Dropdown(
+                    label="CFG Guidance Mode",
+                    choices=["independent", "joint", "alternating"],
+                    value="independent",
+                )
+                cfg_scale_text = gr.Slider(
+                    label="CFG Scale Text",
+                    minimum=0.0,
+                    maximum=10.0,
+                    value=3.0,
+                    step=0.1,
+                )
+                cfg_scale_speaker = gr.Slider(
+                    label="CFG Scale Speaker",
+                    minimum=0.0,
+                    maximum=10.0,
+                    value=5.0,
+                    step=0.1,
+                )
+
+        with gr.Accordion("Advanced (Optional)", open=False):
+            cfg_scale_raw = gr.Textbox(label="CFG Scale Override (optional)", value="")
+            with gr.Row():
+                cfg_min_t = gr.Number(label="CFG Min t", value=0.5)
+                cfg_max_t = gr.Number(label="CFG Max t", value=1.0)
+                context_kv_cache = gr.Checkbox(label="Context KV Cache", value=True)
+            with gr.Row():
+                truncation_factor_raw = gr.Textbox(label="Truncation Factor (optional)", value="")
+                rescale_k_raw = gr.Textbox(label="Rescale k (optional)", value="")
+                rescale_sigma_raw = gr.Textbox(label="Rescale sigma (optional)", value="")
+            with gr.Row():
+                speaker_kv_scale_raw = gr.Textbox(label="Speaker KV Scale (optional)", value="")
+                speaker_kv_min_t_raw = gr.Textbox(label="Speaker KV Min t (optional)", value="0.9")
+                speaker_kv_max_layers_raw = gr.Textbox(
+                    label="Speaker KV Max Layers (optional)", value=""
+                )
+
+        generate_btn = gr.Button("Generate", variant="primary")
+
+        out_audios: list[gr.Audio] = []
+        num_rows = (
+            MAX_GRADIO_CANDIDATES + GRADIO_AUDIO_COLS_PER_ROW - 1
+        ) // GRADIO_AUDIO_COLS_PER_ROW
+        with gr.Column():
+            for row_idx in range(num_rows):
+                with gr.Row():
+                    for col_idx in range(GRADIO_AUDIO_COLS_PER_ROW):
+                        i = row_idx * GRADIO_AUDIO_COLS_PER_ROW + col_idx
+                        if i >= MAX_GRADIO_CANDIDATES:
+                            break
+                        out_audios.append(
+                            gr.Audio(
+                                label=f"Generated Audio {i + 1}",
+                                type="numpy",
+                                visible=(i == 0),
+                            )
+                        )
+        out_log = gr.Textbox(label="Run Log", lines=6)
+
+        def gradio_inference(
+            text,
+            uploaded_audio,
+            num_steps,
+            num_candidates,
+            seed_raw,
+            seconds_raw,
+            duration_scale,
+            cfg_guidance_mode,
+            cfg_scale_text,
+            cfg_scale_speaker,
+            cfg_scale_raw,
+            cfg_min_t,
+            cfg_max_t,
+            context_kv_cache,
+            truncation_factor_raw,
+            rescale_k_raw,
+            rescale_sigma_raw,
+            speaker_kv_scale_raw,
+            speaker_kv_min_t_raw,
+            speaker_kv_max_layers_raw,
+        ):
+            try:
+                audio_results, log_text = run_inference_gpu(
+                    text=text,
+                    uploaded_audio=uploaded_audio,
+                    num_steps=num_steps,
+                    num_candidates=num_candidates,
+                    seed_raw=seed_raw,
+                    seconds_raw=seconds_raw,
+                    duration_scale=duration_scale,
+                    cfg_guidance_mode=cfg_guidance_mode,
+                    cfg_scale_text=cfg_scale_text,
+                    cfg_scale_speaker=cfg_scale_speaker,
+                    cfg_scale_raw=cfg_scale_raw,
+                    cfg_min_t=cfg_min_t,
+                    cfg_max_t=cfg_max_t,
+                    context_kv_cache=context_kv_cache,
+                    truncation_factor_raw=truncation_factor_raw,
+                    rescale_k_raw=rescale_k_raw,
+                    rescale_sigma_raw=rescale_sigma_raw,
+                    speaker_kv_scale_raw=speaker_kv_scale_raw,
+                    speaker_kv_min_t_raw=speaker_kv_min_t_raw,
+                    speaker_kv_max_layers_raw=speaker_kv_max_layers_raw,
+                )
+                audio_updates: list[object] = []
+                for i in range(MAX_GRADIO_CANDIDATES):
+                    if i < len(audio_results):
+                        audio_updates.append(gr.update(value=audio_results[i], visible=True))
+                    else:
+                        audio_updates.append(gr.update(value=None, visible=False))
+                return (*audio_updates, log_text)
+            except Exception as e:
+                raise gr.Error(str(e)) from e
+
+        generate_btn.click(
+            fn=gradio_inference,
+            inputs=[
+                text,
+                uploaded_audio,
+                num_steps,
+                num_candidates,
+                seed_raw,
+                seconds_raw,
+                duration_scale,
+                cfg_guidance_mode,
+                cfg_scale_text,
+                cfg_scale_speaker,
+                cfg_scale_raw,
+                cfg_min_t,
+                cfg_max_t,
+                context_kv_cache,
+                truncation_factor_raw,
+                rescale_k_raw,
+                rescale_sigma_raw,
+                speaker_kv_scale_raw,
+                speaker_kv_min_t_raw,
+                speaker_kv_max_layers_raw,
+            ],
+            outputs=[*out_audios, out_log],
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = build_demo()
+    demo.queue(default_concurrency_limit=1)
+    demo.launch()

irodori_tts/__init__.py

@@ -0,0 +1,16 @@
+"""Irodori-TTS package: text-conditioned RF diffusion over DACVAE latents."""
+
+from .config import ModelConfig, SamplingConfig, TrainConfig
+from .lora import LORA_TARGET_PRESETS
+from .model import TextToLatentRFDiT
+from .tokenizer import ByteTokenizer, PretrainedTextTokenizer
+
+__all__ = [
+    "ByteTokenizer",
+    "LORA_TARGET_PRESETS",
+    "ModelConfig",
+    "PretrainedTextTokenizer",
+    "SamplingConfig",
+    "TextToLatentRFDiT",
+    "TrainConfig",
+]

irodori_tts/codec.py

@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torchaudio
+from huggingface_hub import hf_hub_download
+
+_CODEC_DEFAULT = object()
+
+
+def patchify_latent(latent: torch.Tensor, patch_size: int) -> torch.Tensor:
+    """
+    Convert latent from (B, T, D) -> (B, T//patch, D*patch).
+    Extra tail tokens are dropped.
+    """
+    if patch_size <= 1:
+        return latent
+    bsz, seq_len, dim = latent.shape
+    usable = (seq_len // patch_size) * patch_size
+    latent = latent[:, :usable]
+    latent = latent.reshape(bsz, usable // patch_size, dim * patch_size)
+    return latent
+
+
+def unpatchify_latent(patched: torch.Tensor, patch_size: int, latent_dim: int) -> torch.Tensor:
+    """
+    Convert latent from (B, T_p, D*patch) -> (B, T_p*patch, D).
+    """
+    if patch_size <= 1:
+        return patched
+    return patched.reshape(patched.shape[0], patched.shape[1] * patch_size, latent_dim)
+
+
+@dataclass
+class DACVAECodec:
+    model: torch.nn.Module
+    sample_rate: int
+    latent_dim: int
+    device: torch.device
+    dtype: torch.dtype
+    deterministic_encode: bool
+    deterministic_decode: bool
+    normalize_db: float | None
+
+    @classmethod
+    def load(
+        cls,
+        repo_id: str = "Aratako/Semantic-DACVAE-Japanese-32dim",
+        device: str = "cuda",
+        dtype: torch.dtype | None = None,
+        deterministic_encode: bool = True,
+        deterministic_decode: bool = True,
+        normalize_db: float | None = -16.0,
+    ) -> DACVAECodec:
+        # Prefer installed package; fallback to local clone at ../dacvae.
+        try:
+            from dacvae import DACVAE
+        except ImportError:
+            local_repo = Path(__file__).resolve().parents[2] / "dacvae"
+            if local_repo.exists():
+                sys.path.insert(0, str(local_repo))
+            from dacvae import DACVAE
+
+        location = str(repo_id).strip()
+        if location.startswith("hf://"):
+            location = location[len("hf://") :]
+        if not Path(location).exists() and "/" in location and not location.endswith(".pth"):
+            try:
+                location = hf_hub_download(repo_id=location, filename="weights.pth")
+                print(f"[codec] dacvae: hf://{repo_id} -> {location}", flush=True)
+            except Exception:
+                # Let DACVAE.load surface a clearer error if this is not a valid path/repo.
+                pass
+
+        model = DACVAE.load(location).eval().to(device)
+        if dtype is not None:
+            model = model.to(dtype=dtype)
+
+        decoder = getattr(model, "decoder", None)
+        if decoder is not None and hasattr(decoder, "alpha"):
+            decoder.alpha = 0.0
+            if hasattr(decoder, "wm_model"):
+                # Irodori checkpoints were trained without the DACVAE watermark branch.
+                # Keep decode output mono while skipping that encode/decode path.
+                def _watermark_passthrough(
+                    x: torch.Tensor,
+                    message: torch.Tensor | None = None,
+                    _decoder=decoder,
+                ) -> torch.Tensor:
+                    del message
+                    return _decoder.wm_model.encoder_block.forward_no_conv(x)
+
+                decoder.watermark = _watermark_passthrough
+
+        if deterministic_decode:
+            cls._configure_deterministic_decode(model=model, device=device)
+
+        model_dtype = next(model.parameters()).dtype
+        # Infer latent dimension by encoding a tiny random signal.
+        dummy = torch.zeros(1, 1, 2048, device=device, dtype=model_dtype)
+        with torch.inference_mode():
+            z = model.encode(dummy)  # (B, D, T)
+        return cls(
+            model=model,
+            sample_rate=int(model.sample_rate),
+            latent_dim=int(z.shape[1]),
+            device=torch.device(device),
+            dtype=model_dtype,
+            deterministic_encode=bool(deterministic_encode),
+            deterministic_decode=bool(deterministic_decode),
+            normalize_db=None if normalize_db is None else float(normalize_db),
+        )
+
+    @staticmethod
+    def _configure_deterministic_decode(model: torch.nn.Module, device: str | torch.device) -> None:
+        decoder = getattr(model, "decoder", None)
+        wm_model = getattr(decoder, "wm_model", None)
+        msg_processor = getattr(wm_model, "msg_processor", None)
+        if msg_processor is None:
+            return
+        nbits = int(msg_processor.nbits)
+        message_device = torch.device(device)
+
+        def _fixed_message(batch_size: int) -> torch.Tensor:
+            return torch.zeros((batch_size, nbits), dtype=torch.float32, device=message_device)
+
+        wm_model.random_message = _fixed_message
+
+    @staticmethod
+    def _normalize_loudness(
+        wav: torch.Tensor, sample_rate: int, target_db: float | None
+    ) -> torch.Tensor:
+        if target_db is None:
+            return wav
+        wav_device = wav.device
+        wav = wav.to(dtype=torch.float32)
+        if wav.ndim == 2:
+            if wav.shape[0] == 1:
+                wav = wav[0]
+            elif wav.shape[1] == 1:
+                wav = wav[:, 0]
+            else:
+                wav = wav.mean(dim=0)
+        if wav.ndim != 1:
+            raise ValueError(
+                "normalize_loudness expects a mono waveform with shape (T,) "
+                f"or singleton-channel (1, T)/(T, 1), got {tuple(wav.shape)}"
+            )
+
+        try:
+            from audiotools import AudioSignal
+        except Exception as exc:
+            raise RuntimeError(
+                "audiotools is required when normalize_db is set. "
+                "Install audiotools or disable normalize_db."
+            ) from exc
+
+        signal = AudioSignal(wav.unsqueeze(0).unsqueeze(0), int(sample_rate))
+        signal.normalize(float(target_db))
+        signal.ensure_max_of_audio()
+        normalized = signal.audio_data
+        if not isinstance(normalized, torch.Tensor):
+            normalized = torch.as_tensor(normalized)
+        normalized = normalized.to(dtype=torch.float32, device=wav_device)
+        normalized = normalized.squeeze()
+        if normalized.ndim != 1:
+            raise RuntimeError(
+                "audiotools normalization returned an unexpected waveform shape "
+                f"{tuple(normalized.shape)}"
+            )
+        return normalized
+
+    @torch.inference_mode()
+    def encode_waveform(
+        self,
+        waveform: torch.Tensor,
+        sample_rate: int,
+        *,
+        normalize_db: float | None | object = _CODEC_DEFAULT,
+        ensure_max: bool | None = None,
+    ) -> torch.Tensor:
+        """
+        Input:
+          waveform: (B, C, T) or (C, T)
+          normalize_db: Optional target loudness (LUFS-like dB) applied before encode
+          ensure_max: If True and normalize_db is None, scale down only when abs peak exceeds 1.0
+        Output:
+          latent: (B, T_latent, D_latent)
+        """
+        if waveform.ndim == 2:
+            waveform = waveform.unsqueeze(0)
+        if waveform.ndim != 3:
+            raise ValueError(f"Expected waveform ndim=3, got shape={tuple(waveform.shape)}")
+
+        if waveform.shape[1] != 1:
+            waveform = waveform.mean(dim=1, keepdim=True)
+        if sample_rate != self.sample_rate:
+            waveform = torchaudio.functional.resample(waveform, sample_rate, self.sample_rate)
+
+        if normalize_db is _CODEC_DEFAULT:
+            effective_normalize_db = self.normalize_db
+        elif normalize_db is None:
+            effective_normalize_db = None
+        else:
+            effective_normalize_db = float(normalize_db)
+        # audiotools normalization already applies ensure_max_of_audio(), so codec-side
+        # peak scaling is only needed when normalization is disabled.
+        effective_ensure_max = (
+            effective_normalize_db is None and bool(ensure_max) if ensure_max is not None else False
+        )
+
+        waveform = waveform.to(dtype=torch.float32)
+        if effective_normalize_db is not None or effective_ensure_max:
+            # Keep behavior deterministic per utterance by normalizing each waveform independently.
+            processed: list[torch.Tensor] = []
+            for wav in waveform.squeeze(1):
+                if effective_normalize_db is not None:
+                    wav = self._normalize_loudness(
+                        wav, sample_rate=self.sample_rate, target_db=effective_normalize_db
+                    )
+                wav = wav.squeeze()
+                if wav.ndim != 1:
+                    raise RuntimeError(
+                        "Expected mono per-item waveform after preprocessing, "
+                        f"got shape={tuple(wav.shape)}"
+                    )
+                if effective_ensure_max:
+                    peak = wav.abs().max()
+                    if torch.isfinite(peak) and peak > 1.0:
+                        wav = wav * (1.0 / float(peak))
+                processed.append(wav)
+            waveform = torch.stack(processed, dim=0).unsqueeze(1)
+
+        waveform = waveform.to(self.device, dtype=self.dtype)
+        if self.deterministic_encode:
+            required_paths_present = (
+                hasattr(self.model, "encoder")
+                and hasattr(self.model, "_pad")
+                and hasattr(self.model, "quantizer")
+                and hasattr(self.model.quantizer, "in_proj")
+            )
+            if not required_paths_present:
+                raise RuntimeError(
+                    "deterministic_encode=True requires encoder/_pad/quantizer.in_proj on DACVAE model."
+                )
+            z = self.model.encoder(self.model._pad(waveform))
+            mean, _scale = self.model.quantizer.in_proj(z).chunk(2, dim=1)
+            encoded = mean
+        else:
+            encoded = self.model.encode(waveform)  # (B, D, T)
+        return encoded.transpose(1, 2).contiguous()  # (B, T, D)
+
+    @torch.inference_mode()
+    def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        """
+        Input:
+          latent: (B, T, D)
+        Output:
+          audio: (B, 1, samples)
+        """
+        if latent.ndim != 3:
+            raise ValueError(f"Expected latent ndim=3, got shape={tuple(latent.shape)}")
+        z = latent.transpose(1, 2).contiguous().to(self.device, dtype=self.dtype)  # (B, D, T)
+        return self.model.decode(z)
+
+    def encode_file(self, path: str | Path) -> torch.Tensor:
+        try:
+            wav, sr = torchaudio.load(str(path))
+        except RuntimeError:
+            import soundfile as sf
+
+            data, sr = sf.read(str(path), dtype="float32")
+            wav = torch.from_numpy(data)
+            if wav.ndim == 1:
+                wav = wav.unsqueeze(0)
+            else:
+                wav = wav.T
+        wav = wav.unsqueeze(0)  # (1, C, T)
+        return self.encode_waveform(wav, sr).cpu()

irodori_tts/config.py

@@ -0,0 +1,257 @@
+import json
+from dataclasses import asdict, dataclass, fields
+from pathlib import Path
+from typing import Any, TypeVar
+
+
+@dataclass
+class ModelConfig:
+    latent_dim: int = 128
+    latent_patch_size: int = 1
+    model_dim: int = 2048
+    num_layers: int = 24
+    num_heads: int = 16
+    mlp_ratio: float = 2.875
+    text_mlp_ratio: float | None = 2.6
+    speaker_mlp_ratio: float | None = 2.6
+    dropout: float = 0.0
+    text_vocab_size: int = 102400
+    text_tokenizer_repo: str = "sbintuitions/sarashina2.2-0.5b"
+    text_add_bos: bool = True
+    text_dim: int = 1280
+    text_layers: int = 14
+    text_heads: int = 10
+    use_caption_condition: bool = False
+    caption_vocab_size: int | None = None
+    caption_tokenizer_repo: str | None = None
+    caption_add_bos: bool | None = None
+    caption_dim: int | None = None
+    caption_layers: int | None = None
+    caption_heads: int | None = None
+    caption_mlp_ratio: float | None = None
+    speaker_dim: int = 1280
+    speaker_layers: int = 14
+    speaker_heads: int = 10
+    speaker_patch_size: int = 1
+    timestep_embed_dim: int = 512
+    adaln_rank: int = 256
+    norm_eps: float = 1e-5
+    use_duration_predictor: bool = False
+    duration_aux_dim: int = 14
+    duration_hidden_dim: int = 1024
+    duration_layers: int = 3
+    duration_dropout: float = 0.1
+    duration_attention_heads: int = 8
+    duration_architecture: str = "token_sum_adarn_zero_no_aux"
+    duration_token_init_frames: float = 9.0
+    duration_speaker_fusion: str = "adarn_zero"
+
+    @property
+    def patched_latent_dim(self) -> int:
+        return self.latent_dim * self.latent_patch_size
+
+    @property
+    def speaker_patched_latent_dim(self) -> int:
+        return self.patched_latent_dim * self.speaker_patch_size
+
+    @property
+    def use_speaker_condition(self) -> bool:
+        # Voice-design checkpoints are caption-driven and intentionally omit
+        # reference-speaker conditioning to avoid the easier shortcut.
+        return not bool(self.use_caption_condition)
+
+    @property
+    def text_mlp_ratio_resolved(self) -> float:
+        if self.text_mlp_ratio is None:
+            return self.mlp_ratio
+        return float(self.text_mlp_ratio)
+
+    @property
+    def caption_vocab_size_resolved(self) -> int:
+        if self.caption_vocab_size is None:
+            return int(self.text_vocab_size)
+        return int(self.caption_vocab_size)
+
+    @property
+    def caption_tokenizer_repo_resolved(self) -> str:
+        if self.caption_tokenizer_repo is None:
+            return self.text_tokenizer_repo
+        return str(self.caption_tokenizer_repo)
+
+    @property
+    def caption_add_bos_resolved(self) -> bool:
+        if self.caption_add_bos is None:
+            return bool(self.text_add_bos)
+        return bool(self.caption_add_bos)
+
+    @property
+    def caption_dim_resolved(self) -> int:
+        if self.caption_dim is None:
+            return int(self.text_dim)
+        return int(self.caption_dim)
+
+    @property
+    def caption_layers_resolved(self) -> int:
+        if self.caption_layers is None:
+            return int(self.text_layers)
+        return int(self.caption_layers)
+
+    @property
+    def caption_heads_resolved(self) -> int:
+        if self.caption_heads is None:
+            return int(self.text_heads)
+        return int(self.caption_heads)
+
+    @property
+    def caption_mlp_ratio_resolved(self) -> float:
+        if self.caption_mlp_ratio is None:
+            return self.text_mlp_ratio_resolved
+        return float(self.caption_mlp_ratio)
+
+    @property
+    def speaker_mlp_ratio_resolved(self) -> float:
+        if self.speaker_mlp_ratio is None:
+            return self.mlp_ratio
+        return float(self.speaker_mlp_ratio)
+
+
+@dataclass
+class TrainConfig:
+    manifest_path: str = ""
+    output_dir: str = "outputs"
+    batch_size: int = 8
+    num_workers: int = 2
+    dataloader_persistent_workers: bool = False
+    dataloader_prefetch_factor: int = 2
+    allow_tf32: bool = False
+    compile_model: bool = False
+    train_mode: str = "rf"
+    learning_rate: float = 1e-4
+    weight_decay: float = 0.01
+    optimizer: str = "muon"
+    adam_beta1: float = 0.9
+    adam_beta2: float = 0.999
+    adam_eps: float = 1e-8
+    muon_momentum: float = 0.95
+    muon_adjust_lr_fn: str = "match_rms_adamw"
+    lr_scheduler: str = "none"
+    warmup_steps: int = 0
+    caption_warmup: bool = False
+    caption_warmup_steps: int = 0
+    stable_steps: int = 0
+    min_lr_scale: float = 0.1
+    max_steps: int = 200000
+    log_every: int = 100
+    save_every: int = 1000
+    checkpoint_best_n: int = 0
+    valid_ratio: float = 0.0
+    valid_every: int = 0
+    progress: bool = True
+    progress_all_ranks: bool = False
+    precision: str = "bf16"
+    grad_clip_norm: float = 1.0
+    gradient_accumulation_steps: int = 1
+    max_text_len: int = 256
+    max_caption_len: int | None = None
+    text_condition_dropout: float = 0.1
+    caption_condition_dropout: float = 0.1
+    speaker_condition_dropout: float = 0.1
+    max_latent_steps: int = 750
+    fixed_target_latent_steps: int | None = 750
+    fixed_target_full_mask: bool = True
+    rf_loss_mode: str = "echo"
+    duration_loss_weight: float = 0.1
+    duration_speaker_dropout: float = 0.1
+    duration_huber_delta: float = 0.1
+    timestep_logit_mean: float = 0.0
+    timestep_logit_std: float = 1.0
+    timestep_stratified: bool = True
+    timestep_min: float = 0.001
+    timestep_max: float = 0.999
+    wandb_enabled: bool = False
+    wandb_project: str = "Irodori-TTS"
+    wandb_entity: str | None = None
+    wandb_run_name: str | None = None
+    wandb_mode: str = "online"
+    ddp_find_unused_parameters: bool = False
+    lora_enabled: bool = False
+    lora_r: int = 16
+    lora_alpha: int = 32
+    lora_dropout: float = 0.0
+    lora_bias: str = "none"
+    lora_target_modules: str = "diffusion_attn"
+    lora_modules_to_save: str | None = "auto"
+    seed: int = 0
+
+
+@dataclass
+class SamplingConfig:
+    num_steps: int = 40
+    cfg_scale_text: float = 3.0
+    cfg_scale_caption: float = 3.0
+    cfg_scale_speaker: float = 5.0
+    cfg_guidance_mode: str = "independent"
+    cfg_scale: float | None = None
+    cfg_min_t: float = 0.5
+    cfg_max_t: float = 1.0
+    truncation_factor: float | None = None
+    rescale_k: float | None = None
+    rescale_sigma: float | None = None
+    context_kv_cache: bool = True
+    speaker_kv_scale: float | None = None
+    speaker_kv_min_t: float | None = 0.9
+    speaker_kv_max_layers: int | None = None
+    # Deprecated: inference length is derived from --seconds and codec hop_length.
+    sequence_length: int | None = None
+    seed: int = 0
+
+
+def save_json(path: str | Path, payload: dict) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+
+
+def dump_configs(path: str | Path, model_cfg: ModelConfig, train_cfg: TrainConfig) -> None:
+    save_json(path, {"model": asdict(model_cfg), "train": asdict(train_cfg)})
+
+
+T = TypeVar("T")
+
+
+def load_experiment_yaml(path: str | Path) -> dict[str, Any]:
+    """
+    Load experiment config YAML. Returns {} for an empty document.
+    """
+    try:
+        import yaml
+    except ImportError as exc:
+        raise RuntimeError(
+            "PyYAML is required for --config support. Install with `pip install pyyaml`."
+        ) from exc
+
+    payload = yaml.safe_load(Path(path).read_text(encoding="utf-8"))
+    if payload is None:
+        return {}
+    if not isinstance(payload, dict):
+        raise ValueError(f"Config root must be a mapping: {path}")
+    return payload
+
+
+def merge_dataclass_overrides(base: T, overrides: dict[str, Any] | None, section: str) -> T:
+    """
+    Merge mapping overrides into a dataclass instance with key validation.
+    """
+    if overrides is None:
+        return base
+    if not isinstance(overrides, dict):
+        raise ValueError(f"Config section '{section}' must be a mapping.")
+
+    allowed = {f.name for f in fields(base)}
+    unknown = sorted(set(overrides) - allowed)
+    if unknown:
+        raise ValueError(f"Unknown keys in '{section}' config: {unknown}")
+
+    merged = asdict(base)
+    merged.update(overrides)
+    return type(base)(**merged)

irodori_tts/duration.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import math
+import re
+from collections.abc import Iterable, Sequence
+
+import torch
+
+ALLOWED_ANNOTATION_EMOJIS: tuple[str, ...] = (
+    "⏩",
+    "⏱️",
+    "⏸️",
+    "🌬️",
+    "🍭",
+    "🎛️",
+    "🎭",
+    "🎵",
+    "🐢",
+    "🐱",
+    "👂",
+    "👃",
+    "👅",
+    "👌",
+    "👏",
+    "💋",
+    "💥",
+    "💦",
+    "💪",
+    "📄",
+    "📞",
+    "📢",
+    "📣",
+    "😆",
+    "😊",
+    "😌",
+    "😎",
+    "😏",
+    "😒",
+    "😖",
+    "😟",
+    "😠",
+    "😪",
+    "😭",
+    "😮",
+    "😮‍💨",
+    "😰",
+    "😱",
+    "😲",
+    "😴",
+    "🙄",
+    "🙏",
+    "🤐",
+    "🤔",
+    "🤢",
+    "🤧",
+    "🤭",
+    "🥤",
+    "🥱",
+    "🥴",
+    "🥵",
+    "🥹",
+    "🥺",
+    "🫣",
+    "🫶",
+    "📖",
+)
+
+_ALLOWED_ANNOTATION_EMOJI_PATTERN = re.compile(
+    "|".join(sorted((re.escape(x) for x in ALLOWED_ANNOTATION_EMOJIS), key=len, reverse=True))
+)
+
+
+def _log1p_cap(count: int, cap: int) -> float:
+    return math.log1p(float(min(max(int(count), 0), int(cap)))) / math.log1p(float(cap))
+
+
+def _log1p_cap_float(value: float, cap: float) -> float:
+    value = min(max(float(value), 0.0), float(cap))
+    return math.log1p(value) / math.log1p(float(cap))
+
+
+def _is_kana(ch: str) -> bool:
+    code = ord(ch)
+    return (0x3040 <= code <= 0x309F) or (0x30A0 <= code <= 0x30FF)
+
+
+def _is_kanji(ch: str) -> bool:
+    code = ord(ch)
+    return (
+        (0x3400 <= code <= 0x4DBF)
+        or (0x4E00 <= code <= 0x9FFF)
+        or (0xF900 <= code <= 0xFAFF)
+        or (0x20000 <= code <= 0x2FA1F)
+    )
+
+
+def _is_alnum(ch: str) -> bool:
+    return ch.isascii() and ch.isalnum()
+
+
+def count_annotation_emojis(text: str) -> int:
+    return len(_ALLOWED_ANNOTATION_EMOJI_PATTERN.findall(text))
+
+
+def build_duration_features(
+    texts: Sequence[str] | Iterable[str],
+    *,
+    token_counts: Sequence[int] | torch.Tensor,
+    max_text_len: int,
+    has_speaker: Sequence[bool] | torch.Tensor,
+) -> torch.Tensor:
+    text_list = [str(x) for x in texts]
+    if isinstance(token_counts, torch.Tensor):
+        token_count_list = [int(x) for x in token_counts.detach().cpu().tolist()]
+    else:
+        token_count_list = [int(x) for x in token_counts]
+    if isinstance(has_speaker, torch.Tensor):
+        has_speaker_list = [bool(x) for x in has_speaker.detach().cpu().tolist()]
+    else:
+        has_speaker_list = [bool(x) for x in has_speaker]
+
+    if len(text_list) != len(token_count_list) or len(text_list) != len(has_speaker_list):
+        raise ValueError(
+            "Duration feature inputs must have matching lengths: "
+            f"texts={len(text_list)} token_counts={len(token_count_list)} "
+            f"has_speaker={len(has_speaker_list)}"
+        )
+    if max_text_len <= 0:
+        raise ValueError(f"max_text_len must be > 0, got {max_text_len}")
+
+    rows: list[list[float]] = []
+    for text, token_count, speaker_available in zip(
+        text_list, token_count_list, has_speaker_list, strict=True
+    ):
+        char_count = max(len(text), 1)
+        kana_count = sum(1 for ch in text if _is_kana(ch))
+        kanji_count = sum(1 for ch in text if _is_kanji(ch))
+        alnum_count = sum(1 for ch in text if _is_alnum(ch))
+        emoji_count = count_annotation_emojis(text)
+
+        period_count = text.count("。") + text.count(".")
+        comma_count = text.count("、") + text.count(",")
+        long_vowel_count = text.count("ー")
+        ellipsis_count = text.count("…")
+        exclamation_count = text.count("！") + text.count("!")
+        question_count = text.count("？") + text.count("?")
+
+        rows.append(
+            [
+                min(max(float(token_count), 0.0), float(max_text_len)) / float(max_text_len),
+                _log1p_cap_float(float(char_count), 512.0),
+                float(token_count) / float(char_count),
+                _log1p_cap(period_count, 8),
+                _log1p_cap(comma_count, 16),
+                _log1p_cap(long_vowel_count, 8),
+                _log1p_cap(ellipsis_count, 8),
+                _log1p_cap(exclamation_count, 8),
+                _log1p_cap(question_count, 8),
+                _log1p_cap(emoji_count, 8),
+                float(kana_count) / float(char_count),
+                float(kanji_count) / float(char_count),
+                float(alnum_count) / float(char_count),
+                1.0 if speaker_available else 0.0,
+            ]
+        )
+
+    return torch.tensor(rows, dtype=torch.float32)
+
+
+def set_duration_has_speaker_feature(
+    features: torch.Tensor,
+    has_speaker: torch.Tensor,
+) -> torch.Tensor:
+    if features.ndim != 2:
+        raise ValueError(f"Expected duration features shape (B, D), got {tuple(features.shape)}")
+    if features.shape[1] <= 0:
+        raise ValueError(
+            f"duration features must have at least one column, got {features.shape[1]}"
+        )
+    out = features.clone()
+    out[:, -1] = has_speaker.to(device=features.device, dtype=features.dtype)
+    return out

irodori_tts/inference_runtime.py

@@ -0,0 +1,1063 @@
+from __future__ import annotations
+
+import gc
+import json
+import math
+import secrets
+import threading
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torchaudio
+from safetensors import safe_open
+from safetensors.torch import load_file as load_safetensors_file
+
+from .codec import DACVAECodec, patchify_latent, unpatchify_latent
+from .config import ModelConfig
+from .duration import build_duration_features
+from .lora import checkpoint_state_uses_lora
+from .model import TextToLatentRFDiT
+from .rf import sample_euler_rf_cfg
+from .text_normalization import normalize_text
+from .tokenizer import PretrainedTextTokenizer
+from .watermark import SilentCipherWatermarker
+
+
+def _is_mps_available() -> bool:
+    backends = getattr(torch, "backends", None)
+    if backends is None or not hasattr(backends, "mps"):
+        return False
+    return bool(torch.backends.mps.is_available())
+
+
+def resolve_runtime_device(device: str | torch.device) -> torch.device:
+    resolved = torch.device(device)
+    if resolved.type == "cpu":
+        return resolved
+    if resolved.type == "cuda":
+        if not torch.cuda.is_available():
+            raise ValueError("CUDA device requested but torch.cuda.is_available() is False.")
+        return resolved
+    if resolved.type == "mps":
+        if resolved.index is not None:
+            raise ValueError("MPS device index is not supported. Use 'mps'.")
+        if not _is_mps_available():
+            raise ValueError("MPS device requested but torch.backends.mps.is_available() is False.")
+        return torch.device("mps")
+    raise ValueError(f"Unsupported inference device={resolved!s}. Expected one of: cpu, cuda, mps.")
+
+
+def list_available_runtime_devices() -> list[str]:
+    devices: list[str] = []
+    if torch.cuda.is_available():
+        devices.append("cuda")
+    if _is_mps_available():
+        devices.append("mps")
+    devices.append("cpu")
+    return devices
+
+
+def default_runtime_device() -> str:
+    return list_available_runtime_devices()[0]
+
+
+def list_available_runtime_precisions(device: str | torch.device) -> list[str]:
+    resolved = resolve_runtime_device(device)
+    if resolved.type == "cuda":
+        return ["fp32", "bf16"]
+    return ["fp32"]
+
+
+def _sync_device(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+    elif device.type == "mps":
+        mps = getattr(torch, "mps", None)
+        if mps is not None and hasattr(mps, "synchronize"):
+            mps.synchronize()
+
+
+def _sync_devices(*devices: torch.device) -> None:
+    seen: set[tuple[str, int | None]] = set()
+    for device in devices:
+        key = (device.type, device.index)
+        if key in seen:
+            continue
+        _sync_device(device)
+        seen.add(key)
+
+
+def _measure_start(device: torch.device, *extra_devices: torch.device) -> float:
+    _sync_devices(device, *extra_devices)
+    return time.perf_counter()
+
+
+def _measure_end(device: torch.device, t0: float, *extra_devices: torch.device) -> float:
+    _sync_devices(device, *extra_devices)
+    return time.perf_counter() - t0
+
+
+def _coerce_latent_shape(latent: torch.Tensor, latent_dim: int) -> torch.Tensor:
+    if latent.ndim == 3 and latent.shape[0] == 1:
+        latent = latent[0]
+    if latent.ndim != 2:
+        raise ValueError(f"Unsupported latent shape: {tuple(latent.shape)}")
+    if latent.shape[1] == latent_dim:
+        return latent
+    if latent.shape[0] == latent_dim:
+        return latent.transpose(0, 1).contiguous()
+    raise ValueError(
+        f"Could not infer latent layout for shape={tuple(latent.shape)} and latent_dim={latent_dim}"
+    )
+
+
+def find_flattening_point(
+    latent: torch.Tensor,
+    target_value: float = 0.0,
+    window_size: int = 20,
+    std_threshold: float = 0.05,
+    mean_threshold: float = 0.1,
+) -> int:
+    """
+    Echo-style heuristic: find first index where a trailing window becomes near-flat and near-zero.
+
+    Args:
+      latent: (T, D) latent sequence.
+    Returns:
+      Flattening index in [0, T].
+    """
+    if latent.ndim != 2:
+        raise ValueError(f"Expected latent shape (T, D), got {tuple(latent.shape)}")
+    total_steps = int(latent.shape[0])
+    if total_steps <= 0 or window_size <= 0:
+        return total_steps
+
+    pad = torch.zeros(
+        (window_size, latent.shape[1]),
+        device=latent.device,
+        dtype=latent.dtype,
+    )
+    padded = torch.cat([latent, pad], dim=0)
+    for i in range(padded.shape[0] - window_size):
+        window = padded[i : i + window_size]
+        window_std = window.std(unbiased=False)
+        window_mean = window.mean()
+        if window_std < std_threshold and torch.abs(window_mean - target_value) < mean_threshold:
+            return int(i)
+    return total_steps
+
+
+@dataclass(frozen=True)
+class RuntimeKey:
+    checkpoint: str
+    model_device: str
+    codec_repo: str = "Aratako/Semantic-DACVAE-Japanese-32dim"
+    model_precision: str = "fp32"
+    codec_device: str = "cpu"
+    codec_precision: str = "fp32"
+    codec_deterministic_encode: bool = True
+    codec_deterministic_decode: bool = True
+    compile_model: bool = False
+    compile_dynamic: bool = False
+
+
+@dataclass
+class SamplingRequest:
+    text: str
+    caption: str | None = None
+    ref_wav: str | None = None
+    ref_latent: str | None = None
+    no_ref: bool = False
+    ref_normalize_db: float | None = -16.0
+    ref_ensure_max: bool = True
+    num_candidates: int = 1
+    decode_mode: str = "sequential"
+    seconds: float | None = None
+    duration_scale: float = 1.0
+    min_seconds: float = 0.5
+    max_seconds: float = 30.0
+    max_ref_seconds: float | None = 30.0
+    max_text_len: int | None = None
+    max_caption_len: int | None = None
+    num_steps: int = 40
+    cfg_scale_text: float = 3.0
+    cfg_scale_caption: float = 3.0
+    cfg_scale_speaker: float = 5.0
+    cfg_guidance_mode: str = "independent"
+    cfg_scale: float | None = None
+    cfg_min_t: float = 0.5
+    cfg_max_t: float = 1.0
+    truncation_factor: float | None = None
+    rescale_k: float | None = None
+    rescale_sigma: float | None = None
+    context_kv_cache: bool = True
+    speaker_kv_scale: float | None = None
+    speaker_kv_min_t: float | None = None
+    speaker_kv_max_layers: int | None = None
+    seed: int | None = None
+    trim_tail: bool = True
+    tail_window_size: int = 20
+    tail_std_threshold: float = 0.05
+    tail_mean_threshold: float = 0.1
+
+
+@dataclass
+class SamplingResult:
+    audio: torch.Tensor
+    audios: list[torch.Tensor]
+    sample_rate: int
+    stage_timings: list[tuple[str, float]]
+    total_to_decode: float
+    used_seed: int
+    messages: list[str]
+
+
+def _maybe_compile_inference_model(
+    model: TextToLatentRFDiT,
+    *,
+    enabled: bool,
+    dynamic: bool,
+) -> TextToLatentRFDiT:
+    if not enabled:
+        return model
+    if not hasattr(torch, "compile"):
+        raise RuntimeError("compile_model=True requires torch.compile (PyTorch 2+).")
+    compile_kwargs = {"dynamic": bool(dynamic)}
+    model.encode_conditions = torch.compile(model.encode_conditions, **compile_kwargs)
+    model.build_context_kv_cache = torch.compile(model.build_context_kv_cache, **compile_kwargs)
+    model.forward_with_encoded_conditions = torch.compile(
+        model.forward_with_encoded_conditions,
+        **compile_kwargs,
+    )
+    return model
+
+
+def resolve_runtime_dtype(*, precision: str, device: torch.device) -> torch.dtype:
+    mode = str(precision).strip().lower()
+    if mode == "fp32":
+        return torch.float32
+    if mode == "bf16":
+        if device.type != "cuda":
+            raise ValueError("precision='bf16' currently requires CUDA device.")
+        return torch.bfloat16
+    raise ValueError(f"Unsupported precision={precision!r}. Expected one of: fp32, bf16.")
+
+
+def resolve_cfg_scales(
+    *,
+    cfg_guidance_mode: str,
+    cfg_scale_text: float,
+    cfg_scale_caption: float,
+    cfg_scale_speaker: float,
+    cfg_scale: float | None,
+    use_caption_condition: bool = True,
+    use_speaker_condition: bool = True,
+) -> tuple[float, float, float, list[str]]:
+    """Normalize/validate CFG scales for guidance mode."""
+    messages: list[str] = []
+    text_val = float(cfg_scale_text)
+    caption_val = float(cfg_scale_caption)
+    speaker_val = float(cfg_scale_speaker)
+
+    if cfg_scale is not None:
+        text_val = float(cfg_scale)
+        caption_val = float(cfg_scale)
+        speaker_val = float(cfg_scale)
+    if not use_speaker_condition:
+        if speaker_val > 0.0:
+            messages.append(
+                "info: speaker conditioning is disabled for this checkpoint; ignoring cfg_scale_speaker."
+            )
+        speaker_val = 0.0
+
+    mode = str(cfg_guidance_mode).strip().lower()
+    enabled_vals = [value for value in (text_val, speaker_val) if value > 0.0]
+    if use_caption_condition and caption_val > 0.0:
+        enabled_vals.append(caption_val)
+    if mode == "joint" and enabled_vals and (max(enabled_vals) - min(enabled_vals) > 1e-6):
+        raise ValueError(
+            "cfg_guidance_mode='joint' requires equal enabled cfg_scale_text/cfg_scale_caption/cfg_scale_speaker, "
+            "or set cfg_scale."
+        )
+
+    return text_val, caption_val, speaker_val, messages
+
+
+def _load_torch_checkpoint_payload(path: Path) -> dict:
+    payload = torch.load(path, map_location="cpu", weights_only=True)
+    if not isinstance(payload, dict):
+        raise ValueError(f"Unsupported checkpoint payload type: {type(payload)!r}")
+    return payload
+
+
+_CONFIG_META_KEY = "config_json"
+_INFERENCE_CONFIG_KEYS = {"max_text_len", "max_caption_len", "fixed_target_latent_steps"}
+
+
+def _load_checkpoint_from_pt(path: Path) -> tuple[dict[str, torch.Tensor], dict, dict | None]:
+    ckpt = _load_torch_checkpoint_payload(path)
+    model_state = ckpt.get("model")
+    model_cfg = ckpt.get("model_config")
+    train_cfg = ckpt.get("train_config")
+
+    if not isinstance(model_state, dict):
+        raise ValueError(f"Checkpoint missing model weights dictionary: {path}")
+    if not isinstance(model_cfg, dict):
+        raise ValueError(f"Checkpoint missing model_config dictionary: {path}")
+    if train_cfg is not None and not isinstance(train_cfg, dict):
+        raise ValueError(f"Checkpoint train_config must be a dictionary when present: {path}")
+
+    if checkpoint_state_uses_lora(model_state):
+        raise ValueError(
+            f"LoRA checkpoints must be loaded from adapter directories or merged safetensors: {path}"
+        )
+    return model_state, model_cfg, _extract_inference_train_config(train_cfg)
+
+
+def _parse_json_mapping(
+    raw: str | None,
+    *,
+    field: str,
+    path: Path,
+    required: bool = False,
+) -> dict | None:
+    if raw is None:
+        if required:
+            raise ValueError(f"Missing required metadata field '{field}' in checkpoint: {path}")
+        return None
+    try:
+        payload = json.loads(raw)
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"Invalid JSON in '{field}' metadata for checkpoint: {path}") from exc
+    if not isinstance(payload, dict):
+        raise ValueError(f"Metadata field '{field}' must decode to an object: {path}")
+    return payload
+
+
+def _extract_inference_train_config(raw: dict | None) -> dict | None:
+    if raw is None:
+        return None
+
+    inference_cfg: dict[str, int] = {}
+    for key in _INFERENCE_CONFIG_KEYS:
+        value = raw.get(key)
+        if value is None:
+            continue
+        if not isinstance(value, int):
+            raise ValueError(f"Inference config key '{key}' must be int, got {type(value)!r}.")
+        inference_cfg[key] = int(value)
+
+    return inference_cfg or None
+
+
+def _split_flat_checkpoint_config(path: Path, flat_config: dict) -> tuple[dict, dict | None]:
+    model_cfg: dict[str, object] = {}
+    inference_cfg: dict[str, int] = {}
+    for key, value in flat_config.items():
+        if key in _INFERENCE_CONFIG_KEYS:
+            if not isinstance(value, int):
+                raise ValueError(
+                    f"Inference config key '{key}' must be int in checkpoint metadata: {path}"
+                )
+            inference_cfg[key] = int(value)
+            continue
+        model_cfg[key] = value
+    return model_cfg, (inference_cfg or None)
+
+
+def _load_checkpoint_from_safetensors(
+    path: Path,
+) -> tuple[dict[str, torch.Tensor], dict, dict | None]:
+    model_state = load_safetensors_file(str(path), device="cpu")
+    if not isinstance(model_state, dict) or not model_state:
+        raise ValueError(f"Safetensors checkpoint has no model weights: {path}")
+
+    with safe_open(str(path), framework="pt", device="cpu") as handle:
+        metadata = handle.metadata() or {}
+
+    flat_config = _parse_json_mapping(
+        metadata.get(_CONFIG_META_KEY),
+        field=_CONFIG_META_KEY,
+        path=path,
+        required=True,
+    )
+    model_cfg, inference_cfg = _split_flat_checkpoint_config(path=path, flat_config=flat_config)
+    return model_state, model_cfg, inference_cfg
+
+
+def _load_checkpoint_for_inference(path: Path) -> tuple[dict[str, torch.Tensor], dict, dict | None]:
+    if path.suffix.lower() == ".safetensors":
+        return _load_checkpoint_from_safetensors(path)
+    return _load_checkpoint_from_pt(path)
+
+
+class InferenceRuntime:
+    def __init__(
+        self,
+        *,
+        key: RuntimeKey,
+        model_cfg: ModelConfig,
+        train_cfg: dict | None,
+        model: TextToLatentRFDiT,
+        tokenizer: PretrainedTextTokenizer,
+        caption_tokenizer: PretrainedTextTokenizer | None,
+        codec: DACVAECodec,
+        default_text_max_len: int,
+        default_caption_max_len: int,
+    ) -> None:
+        self.key = key
+        self.model_device = resolve_runtime_device(key.model_device)
+        self.codec_device = resolve_runtime_device(key.codec_device)
+        self.model_cfg = model_cfg
+        self.train_cfg = train_cfg
+        self.model = model
+        self.tokenizer = tokenizer
+        self.caption_tokenizer = caption_tokenizer
+        self.codec = codec
+        self.default_text_max_len = default_text_max_len
+        self.default_caption_max_len = default_caption_max_len
+        self.watermarker = SilentCipherWatermarker(device=str(self.codec_device))
+        self._infer_lock = threading.Lock()
+
+    @classmethod
+    def from_key(cls, key: RuntimeKey) -> InferenceRuntime:
+        model_device = resolve_runtime_device(key.model_device)
+        codec_device = resolve_runtime_device(key.codec_device)
+        model_dtype = resolve_runtime_dtype(
+            precision=key.model_precision,
+            device=model_device,
+        )
+        codec_dtype = resolve_runtime_dtype(
+            precision=key.codec_precision,
+            device=codec_device,
+        )
+
+        model_state, model_cfg_dict, train_cfg = _load_checkpoint_for_inference(
+            Path(key.checkpoint)
+        )
+        model_cfg = ModelConfig(**model_cfg_dict)
+
+        model = TextToLatentRFDiT(model_cfg).to(model_device)
+        model.load_state_dict(model_state)
+        model = model.to(dtype=model_dtype)
+        model.eval()
+        model = _maybe_compile_inference_model(
+            model,
+            enabled=bool(key.compile_model),
+            dynamic=bool(key.compile_dynamic),
+        )
+
+        tokenizer = PretrainedTextTokenizer.from_pretrained(
+            repo_id=model_cfg.text_tokenizer_repo,
+            add_bos=bool(model_cfg.text_add_bos),
+            local_files_only=False,
+        )
+        if tokenizer.vocab_size != model_cfg.text_vocab_size:
+            raise ValueError(
+                f"text_vocab_size mismatch: checkpoint text_vocab_size={model_cfg.text_vocab_size} but tokenizer "
+                f"({model_cfg.text_tokenizer_repo}) vocab_size={tokenizer.vocab_size}."
+            )
+        caption_tokenizer = None
+        if model_cfg.use_caption_condition:
+            caption_tokenizer = PretrainedTextTokenizer.from_pretrained(
+                repo_id=model_cfg.caption_tokenizer_repo_resolved,
+                add_bos=model_cfg.caption_add_bos_resolved,
+                local_files_only=False,
+            )
+            if caption_tokenizer.vocab_size != model_cfg.caption_vocab_size_resolved:
+                raise ValueError(
+                    f"caption_vocab_size mismatch: checkpoint caption_vocab_size={model_cfg.caption_vocab_size_resolved} but tokenizer ({model_cfg.caption_tokenizer_repo_resolved}) "
+                    f"vocab_size={caption_tokenizer.vocab_size}."
+                )
+
+        default_text_max_len = 256
+        default_caption_max_len = default_text_max_len
+        if isinstance(train_cfg, dict):
+            ckpt_text_max_len = train_cfg.get("max_text_len")
+            if isinstance(ckpt_text_max_len, int) and ckpt_text_max_len > 0:
+                default_text_max_len = int(ckpt_text_max_len)
+            ckpt_caption_max_len = train_cfg.get("max_caption_len")
+            if isinstance(ckpt_caption_max_len, int) and ckpt_caption_max_len > 0:
+                default_caption_max_len = int(ckpt_caption_max_len)
+            else:
+                default_caption_max_len = default_text_max_len
+
+        codec = DACVAECodec.load(
+            repo_id=key.codec_repo,
+            device=str(codec_device),
+            dtype=codec_dtype,
+            deterministic_encode=bool(key.codec_deterministic_encode),
+            deterministic_decode=bool(key.codec_deterministic_decode),
+        )
+        if model_cfg.latent_dim != codec.latent_dim:
+            raise ValueError(
+                f"Latent dimension mismatch: checkpoint latent_dim={model_cfg.latent_dim} but codec latent_dim={codec.latent_dim}. "
+                "Use a compatible codec/checkpoint pair."
+            )
+
+        return cls(
+            key=key,
+            model_cfg=model_cfg,
+            train_cfg=train_cfg if isinstance(train_cfg, dict) else None,
+            model=model,
+            tokenizer=tokenizer,
+            caption_tokenizer=caption_tokenizer,
+            codec=codec,
+            default_text_max_len=default_text_max_len,
+            default_caption_max_len=default_caption_max_len,
+        )
+
+    def _load_reference_latent(
+        self,
+        *,
+        req: SamplingRequest,
+        batch_size: int,
+        messages: list[str],
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        runtime_dtype = next(self.model.parameters()).dtype
+        if not self.model_cfg.use_speaker_condition:
+            if req.ref_wav is not None or req.ref_latent is not None:
+                messages.append(
+                    "info: speaker conditioning is disabled for this checkpoint; ignoring reference input."
+                )
+            return None, None
+        if req.no_ref:
+            ref_len = max(1, int(self.model_cfg.speaker_patch_size))
+            ref_latent_patched = torch.zeros(
+                (
+                    batch_size,
+                    ref_len,
+                    self.model_cfg.latent_dim * self.model_cfg.latent_patch_size,
+                ),
+                device=self.model_device,
+                dtype=runtime_dtype,
+            )
+            ref_mask = torch.zeros(
+                (batch_size, ref_len), dtype=torch.bool, device=self.model_device
+            )
+            return ref_latent_patched, ref_mask
+
+        if req.ref_wav is None and req.ref_latent is None:
+            raise ValueError("Specify either ref_wav/ref_latent, or set no_ref=True.")
+
+        max_ref_latent_steps = None
+        if req.max_ref_seconds is not None and req.max_ref_seconds > 0:
+            max_ref_latent_steps = max(
+                1,
+                math.ceil(
+                    float(req.max_ref_seconds)
+                    * float(self.codec.sample_rate)
+                    / float(int(self.codec.model.hop_length))
+                ),
+            )
+
+        if req.ref_latent is not None:
+            latent_raw = torch.load(req.ref_latent, map_location="cpu", weights_only=True)
+            ref_latent = _coerce_latent_shape(
+                latent_raw, latent_dim=self.model_cfg.latent_dim
+            ).unsqueeze(0)
+            ref_latent = ref_latent.to(dtype=runtime_dtype)
+        else:
+            wav, sr = _load_audio(req.ref_wav)
+            if req.max_ref_seconds is not None and req.max_ref_seconds > 0:
+                max_ref_samples = max(1, int(float(req.max_ref_seconds) * float(sr)))
+                if wav.shape[1] > max_ref_samples:
+                    messages.append(
+                        f"warning: reference audio exceeds max_ref_seconds ({req.max_ref_seconds}s). "
+                        f"Trimming from {float(wav.shape[1]) / float(sr):.2f}s to {float(max_ref_samples) / float(sr):.2f}s."
+                    )
+                    wav = wav[:, :max_ref_samples]
+            if req.ref_normalize_db is not None:
+                messages.append(
+                    f"info: reference loudness normalize enabled (target_db={float(req.ref_normalize_db):.2f}, includes peak safety scaling)."
+                )
+            elif req.ref_ensure_max:
+                messages.append("info: reference peak safety scaling enabled (ensure_max=True).")
+            ref_latent = self.codec.encode_waveform(
+                wav.unsqueeze(0),
+                sample_rate=int(sr),
+                normalize_db=req.ref_normalize_db,
+                ensure_max=bool(req.ref_ensure_max),
+            ).cpu()
+
+        if max_ref_latent_steps is not None and ref_latent.shape[1] > max_ref_latent_steps:
+            messages.append(
+                f"warning: reference latent steps ({ref_latent.shape[1]}) exceed max_ref_seconds bound ({max_ref_latent_steps} steps). "
+                "Trimming reference latent."
+            )
+            ref_latent = ref_latent[:, :max_ref_latent_steps]
+
+        ref_latent_patched = patchify_latent(ref_latent, self.model_cfg.latent_patch_size).to(
+            self.model_device
+        )
+        if ref_latent_patched.shape[1] == 0:
+            raise ValueError(
+                "Reference latent length became zero after patchify. Use longer reference audio."
+            )
+        if batch_size > 1:
+            ref_latent_patched = ref_latent_patched.repeat(batch_size, 1, 1)
+        ref_mask = torch.ones(
+            (batch_size, ref_latent_patched.shape[1]), dtype=torch.bool, device=self.model_device
+        )
+        return ref_latent_patched, ref_mask
+
+    def synthesize(
+        self,
+        req: SamplingRequest,
+        *,
+        log_fn: Callable[[str], None] | None = None,
+    ) -> SamplingResult:
+        def _log(msg: str) -> None:
+            if log_fn is not None:
+                log_fn(msg)
+
+        messages: list[str] = []
+        _log(
+            (
+                "[runtime] start synthesize "
+                "model_device={} model_precision={} codec_device={} codec_precision={} "
+                "silentcipher_watermark={} mode={} seconds={} steps={} seed={} candidates={} decode_mode={}"
+            ).format(
+                self.key.model_device,
+                self.key.model_precision,
+                self.key.codec_device,
+                self.key.codec_precision,
+                self.watermarker.ready,
+                req.cfg_guidance_mode,
+                req.seconds,
+                req.num_steps,
+                "random" if req.seed is None else int(req.seed),
+                req.num_candidates,
+                req.decode_mode,
+            )
+        )
+
+        manual_seconds = None if req.seconds is None else float(req.seconds)
+        if manual_seconds is not None and manual_seconds <= 0:
+            raise ValueError(f"seconds must be > 0 when provided, got {req.seconds}")
+        duration_scale = float(req.duration_scale)
+        if duration_scale <= 0:
+            raise ValueError(f"duration_scale must be > 0, got {duration_scale}")
+        min_seconds = float(req.min_seconds)
+        max_seconds = float(req.max_seconds)
+        if min_seconds <= 0:
+            raise ValueError(f"min_seconds must be > 0, got {min_seconds}")
+        if max_seconds < min_seconds:
+            raise ValueError(
+                f"max_seconds must be >= min_seconds, got min={min_seconds} max={max_seconds}"
+            )
+        num_candidates = int(req.num_candidates)
+        if num_candidates <= 0:
+            raise ValueError(f"num_candidates must be > 0, got {num_candidates}")
+        decode_mode = str(req.decode_mode).strip().lower()
+        if decode_mode not in {"sequential", "batch"}:
+            raise ValueError(
+                f"Unsupported decode_mode={req.decode_mode!r}. Expected one of: sequential, batch."
+            )
+
+        raw_text = str(req.text)
+        normalized_text = normalize_text(raw_text).strip()
+        if normalized_text == "":
+            raise ValueError("text became empty after normalization.")
+
+        text_max_len = (
+            self.default_text_max_len if req.max_text_len is None else int(req.max_text_len)
+        )
+        if text_max_len <= 0:
+            raise ValueError(f"max_text_len must be > 0, got {text_max_len}")
+        caption_max_len = (
+            self.default_caption_max_len
+            if req.max_caption_len is None
+            else int(req.max_caption_len)
+        )
+        if self.model_cfg.use_caption_condition and caption_max_len <= 0:
+            raise ValueError(f"max_caption_len must be > 0, got {caption_max_len}")
+        has_caption_text = bool(
+            self.model_cfg.use_caption_condition
+            and req.caption is not None
+            and str(req.caption).strip() != ""
+        )
+
+        truncation_factor = None if req.truncation_factor is None else float(req.truncation_factor)
+        rescale_k = None if req.rescale_k is None else float(req.rescale_k)
+        rescale_sigma = None if req.rescale_sigma is None else float(req.rescale_sigma)
+        if truncation_factor is not None and truncation_factor <= 0:
+            raise ValueError(f"truncation_factor must be > 0, got {truncation_factor}")
+        if (rescale_k is None) != (rescale_sigma is None):
+            raise ValueError("rescale_k and rescale_sigma must be set together.")
+        if rescale_k is not None and rescale_k <= 0:
+            raise ValueError(f"rescale_k must be > 0, got {rescale_k}")
+        if rescale_sigma is not None and rescale_sigma <= 0:
+            raise ValueError(f"rescale_sigma must be > 0, got {rescale_sigma}")
+
+        speaker_kv_scale = None if req.speaker_kv_scale is None else float(req.speaker_kv_scale)
+        speaker_kv_min_t = None
+        speaker_kv_max_layers = (
+            None if req.speaker_kv_max_layers is None else int(req.speaker_kv_max_layers)
+        )
+        if speaker_kv_scale is not None:
+            if not self.model_cfg.use_speaker_condition:
+                messages.append(
+                    "info: speaker conditioning is disabled for this checkpoint; ignoring speaker_kv_scale."
+                )
+                speaker_kv_scale = None
+            else:
+                if speaker_kv_scale <= 0:
+                    raise ValueError(f"speaker_kv_scale must be > 0, got {speaker_kv_scale}")
+                speaker_kv_min_t = (
+                    0.9 if req.speaker_kv_min_t is None else float(req.speaker_kv_min_t)
+                )
+                if not (0.0 <= speaker_kv_min_t <= 1.0):
+                    raise ValueError(f"speaker_kv_min_t must be in [0, 1], got {speaker_kv_min_t}")
+                if speaker_kv_max_layers is not None and speaker_kv_max_layers < 0:
+                    raise ValueError(
+                        f"speaker_kv_max_layers must be >= 0 when specified, got {speaker_kv_max_layers}"
+                    )
+
+        cfg_mode = str(req.cfg_guidance_mode).strip().lower()
+        if cfg_mode not in {"independent", "joint", "alternating"}:
+            raise ValueError(
+                f"Unsupported cfg_guidance_mode={req.cfg_guidance_mode!r}. "
+                "Expected one of: independent, joint, alternating."
+            )
+
+        cfg_scale_text, cfg_scale_caption, cfg_scale_speaker, scale_messages = resolve_cfg_scales(
+            cfg_guidance_mode=cfg_mode,
+            cfg_scale_text=req.cfg_scale_text,
+            cfg_scale_caption=req.cfg_scale_caption,
+            cfg_scale_speaker=req.cfg_scale_speaker,
+            cfg_scale=req.cfg_scale,
+            use_caption_condition=has_caption_text,
+            use_speaker_condition=self.model_cfg.use_speaker_condition,
+        )
+        messages.extend(scale_messages)
+        for msg in scale_messages:
+            _log(msg)
+
+        stage_timings: list[tuple[str, float]] = []
+        if req.seed is None:
+            used_seed = int(secrets.randbits(63))
+            msg = f"info: seed not specified; using random seed {used_seed}."
+            messages.append(msg)
+            _log(msg)
+        else:
+            used_seed = int(req.seed)
+            _log(f"[runtime] using seed: {used_seed}")
+        post_load_t0 = _measure_start(self.model_device, self.codec_device)
+
+        with self._infer_lock, torch.inference_mode():
+            t0 = _measure_start(self.model_device)
+            text_ids, text_mask = self.tokenizer.batch_encode(
+                [normalized_text] * num_candidates,
+                max_length=text_max_len,
+            )
+            stage_sec = _measure_end(self.model_device, t0)
+            stage_timings.append(("tokenize_text", stage_sec))
+            _log(f"[runtime] tokenize_text: {stage_sec * 1000.0:.1f} ms")
+            text_ids = text_ids.to(self.model_device)
+            text_mask = text_mask.to(self.model_device)
+            caption_ids = None
+            caption_mask = None
+            if self.model_cfg.use_caption_condition:
+                if self.caption_tokenizer is None:
+                    raise RuntimeError(
+                        "Caption conditioning is enabled but caption tokenizer is not loaded."
+                    )
+                caption_text = "" if req.caption is None else str(req.caption).strip()
+                caption_ids, caption_mask = self.caption_tokenizer.batch_encode(
+                    [caption_text] * num_candidates,
+                    max_length=caption_max_len,
+                )
+                if caption_text == "":
+                    caption_mask.zero_()
+                caption_ids = caption_ids.to(self.model_device)
+                caption_mask = caption_mask.to(self.model_device)
+
+            t0 = _measure_start(self.model_device, self.codec_device)
+            msg_count_before_ref = len(messages)
+            ref_latent, ref_mask = self._load_reference_latent(
+                req=req,
+                batch_size=num_candidates,
+                messages=messages,
+            )
+            stage_sec = _measure_end(self.model_device, t0, self.codec_device)
+            stage_timings.append(("prepare_reference", stage_sec))
+            for msg in messages[msg_count_before_ref:]:
+                _log(msg)
+            _log(f"[runtime] prepare_reference: {stage_sec * 1000.0:.1f} ms")
+
+            hop_length = int(self.codec.model.hop_length)
+            if manual_seconds is not None:
+                clamped_seconds = min(max_seconds, max(min_seconds, manual_seconds))
+                if clamped_seconds != manual_seconds:
+                    duration_msg = (
+                        f"warning: manual duration {manual_seconds:.3f}s was clamped to "
+                        f"{clamped_seconds:.3f}s."
+                    )
+                    messages.append(duration_msg)
+                    _log(duration_msg)
+                target_samples = max(1, int(clamped_seconds * self.codec.sample_rate))
+                latent_steps = math.ceil(target_samples / hop_length)
+                duration_msg = f"info: using manual duration {clamped_seconds:.3f}s."
+                messages.append(duration_msg)
+                _log(duration_msg)
+            elif self.model_cfg.use_duration_predictor:
+                t0 = _measure_start(self.model_device)
+                has_speaker_duration = torch.zeros(
+                    (num_candidates,), dtype=torch.bool, device=self.model_device
+                )
+                if self.model_cfg.use_speaker_condition and ref_mask is not None:
+                    has_speaker_duration = ref_mask.any(dim=1)
+                duration_features = build_duration_features(
+                    [normalized_text] * num_candidates,
+                    token_counts=text_mask.sum(dim=1),
+                    max_text_len=text_max_len,
+                    has_speaker=has_speaker_duration,
+                ).to(self.model_device)
+                (
+                    duration_text_state,
+                    duration_text_mask,
+                    duration_speaker_state,
+                    _duration_speaker_mask,
+                    _duration_caption_state,
+                    _duration_caption_mask,
+                ) = self.model.encode_conditions(
+                    text_input_ids=text_ids,
+                    text_mask=text_mask,
+                    ref_latent=ref_latent,
+                    ref_mask=ref_mask,
+                    caption_input_ids=caption_ids,
+                    caption_mask=caption_mask,
+                )
+                pred_log_frames = self.model.predict_duration_log_frames(
+                    text_state=duration_text_state,
+                    text_mask=duration_text_mask,
+                    speaker_state=duration_speaker_state,
+                    speaker_mask=_duration_speaker_mask,
+                    duration_features=duration_features,
+                    has_speaker=has_speaker_duration,
+                )
+                pred_frames = torch.expm1(pred_log_frames).float().mean().item()
+                scaled_frames = pred_frames * duration_scale
+                min_frames = max(1, math.ceil(min_seconds * self.codec.sample_rate / hop_length))
+                max_frames = max(1, math.floor(max_seconds * self.codec.sample_rate / hop_length))
+                latent_steps = int(round(scaled_frames))
+                latent_steps = max(min_frames, min(max_frames, latent_steps))
+                target_samples = int(latent_steps * hop_length)
+                stage_sec = _measure_end(self.model_device, t0)
+                stage_timings.append(("predict_duration", stage_sec))
+                msg = (
+                    f"info: predicted duration frames={pred_frames:.1f}, "
+                    f"scale={duration_scale:.3f}, using_frames={latent_steps} "
+                    f"({target_samples / float(self.codec.sample_rate):.3f}s)."
+                )
+                messages.append(msg)
+                _log(msg)
+                _log(f"[runtime] predict_duration: {stage_sec * 1000.0:.1f} ms")
+            else:
+                fallback_seconds = 30.0
+                target_samples = int(fallback_seconds * self.codec.sample_rate)
+                latent_steps = math.ceil(target_samples / hop_length)
+                msg = "info: checkpoint has no duration predictor; falling back to 30.000s."
+                messages.append(msg)
+                _log(msg)
+            patched_steps = math.ceil(latent_steps / self.model_cfg.latent_patch_size)
+
+            if isinstance(self.train_cfg, dict):
+                fixed_steps = self.train_cfg.get("fixed_target_latent_steps")
+                if isinstance(fixed_steps, int) and fixed_steps > 0 and latent_steps > fixed_steps:
+                    msg = (
+                        f"warning: requested latent length ({latent_steps}) exceeds fixed_target_latent_steps ({fixed_steps}) "
+                        "used in training. Long-tail stability may degrade."
+                    )
+                    messages.append(msg)
+                    _log(msg)
+
+            t0 = _measure_start(self.model_device)
+            z_patched = sample_euler_rf_cfg(
+                model=self.model,
+                text_input_ids=text_ids,
+                text_mask=text_mask,
+                ref_latent=ref_latent,
+                ref_mask=ref_mask,
+                sequence_length=patched_steps,
+                caption_input_ids=caption_ids,
+                caption_mask=caption_mask,
+                num_steps=int(req.num_steps),
+                cfg_scale_text=cfg_scale_text,
+                cfg_scale_caption=cfg_scale_caption,
+                cfg_scale_speaker=cfg_scale_speaker,
+                cfg_guidance_mode=cfg_mode,
+                cfg_min_t=float(req.cfg_min_t),
+                cfg_max_t=float(req.cfg_max_t),
+                seed=used_seed,
+                truncation_factor=truncation_factor,
+                rescale_k=rescale_k,
+                rescale_sigma=rescale_sigma,
+                use_context_kv_cache=bool(req.context_kv_cache),
+                speaker_kv_scale=speaker_kv_scale,
+                speaker_kv_max_layers=speaker_kv_max_layers,
+                speaker_kv_min_t=speaker_kv_min_t,
+            )
+            stage_sec = _measure_end(self.model_device, t0)
+            stage_timings.append(("sample_rf", stage_sec))
+            _log(f"[runtime] sample_rf: {stage_sec * 1000.0:.1f} ms")
+
+            t0 = _measure_start(self.model_device)
+            z = unpatchify_latent(
+                z_patched,
+                patch_size=self.model_cfg.latent_patch_size,
+                latent_dim=self.model_cfg.latent_dim,
+            )
+            stage_sec = _measure_end(self.model_device, t0)
+            stage_timings.append(("unpatchify_latent", stage_sec))
+            _log(f"[runtime] unpatchify_latent: {stage_sec * 1000.0:.1f} ms")
+            z = z[:, :latent_steps]
+
+            t0 = _measure_start(self.model_device, self.codec_device)
+            trimmed_audios: list[torch.Tensor] = []
+            if decode_mode == "batch":
+                audio_batch = self.codec.decode_latent(z).cpu()
+                for i in range(num_candidates):
+                    audio_i = audio_batch[i]
+                    max_samples = target_samples
+                    if bool(req.trim_tail):
+                        flattening_point = find_flattening_point(
+                            z[i],
+                            window_size=max(1, int(req.tail_window_size)),
+                            std_threshold=float(req.tail_std_threshold),
+                            mean_threshold=float(req.tail_mean_threshold),
+                        )
+                        flattening_samples = int(
+                            flattening_point * int(self.codec.model.hop_length)
+                        )
+                        if flattening_samples > 0:
+                            max_samples = min(max_samples, flattening_samples)
+                    trimmed_audios.append(audio_i[:, :max_samples])
+            else:
+                for i in range(num_candidates):
+                    audio_i = self.codec.decode_latent(z[i : i + 1]).cpu()[0]
+                    max_samples = target_samples
+                    if bool(req.trim_tail):
+                        flattening_point = find_flattening_point(
+                            z[i],
+                            window_size=max(1, int(req.tail_window_size)),
+                            std_threshold=float(req.tail_std_threshold),
+                            mean_threshold=float(req.tail_mean_threshold),
+                        )
+                        flattening_samples = int(
+                            flattening_point * int(self.codec.model.hop_length)
+                        )
+                        if flattening_samples > 0:
+                            max_samples = min(max_samples, flattening_samples)
+                    trimmed_audios.append(audio_i[:, :max_samples])
+            stage_sec = _measure_end(self.model_device, t0, self.codec_device)
+            stage_timings.append(("decode_latent", stage_sec))
+            _log(f"[runtime] decode_latent ({decode_mode}): {stage_sec * 1000.0:.1f} ms")
+
+            if self.watermarker.ready:
+                t0 = _measure_start(self.codec_device)
+                trimmed_audios = self.watermarker.encode_batch(
+                    trimmed_audios,
+                    sample_rate=int(self.codec.sample_rate),
+                )
+                stage_sec = _measure_end(self.codec_device, t0)
+                stage_timings.append(("silentcipher_watermark", stage_sec))
+                _log(f"[runtime] silentcipher_watermark: {stage_sec * 1000.0:.1f} ms")
+            else:
+                msg = (
+                    "warning: SilentCipher watermark is unavailable; generated audio was not "
+                    "watermarked."
+                )
+                messages.append(msg)
+                _log(msg)
+
+            total_to_decode = _measure_end(self.model_device, post_load_t0, self.codec_device)
+            _log(f"[runtime] total_to_decode: {total_to_decode:.3f} s")
+
+        _log("[runtime] done synthesize")
+        return SamplingResult(
+            audio=trimmed_audios[0],
+            audios=trimmed_audios,
+            sample_rate=int(self.codec.sample_rate),
+            stage_timings=stage_timings,
+            total_to_decode=total_to_decode,
+            used_seed=used_seed,
+            messages=messages,
+        )
+
+    def unload(self) -> None:
+        del self.model
+        del self.tokenizer
+        del self.codec
+        gc.collect()
+        for device in (self.model_device, self.codec_device):
+            if device.type == "cuda":
+                torch.cuda.empty_cache()
+            elif device.type == "mps":
+                mps = getattr(torch, "mps", None)
+                if mps is not None and hasattr(mps, "empty_cache"):
+                    mps.empty_cache()
+
+
+_RUNTIME_CACHE_LOCK = threading.Lock()
+_RUNTIME_CACHE_KEY: RuntimeKey | None = None
+_RUNTIME_CACHE_VALUE: InferenceRuntime | None = None
+
+
+def get_cached_runtime(key: RuntimeKey) -> tuple[InferenceRuntime, bool]:
+    global _RUNTIME_CACHE_KEY, _RUNTIME_CACHE_VALUE
+    with _RUNTIME_CACHE_LOCK:
+        if _RUNTIME_CACHE_VALUE is not None and _RUNTIME_CACHE_KEY == key:
+            return _RUNTIME_CACHE_VALUE, False
+
+        old_runtime = _RUNTIME_CACHE_VALUE
+        runtime = InferenceRuntime.from_key(key)
+        _RUNTIME_CACHE_KEY = key
+        _RUNTIME_CACHE_VALUE = runtime
+
+    if old_runtime is not None:
+        old_runtime.unload()
+
+    return runtime, True
+
+
+def clear_cached_runtime() -> None:
+    global _RUNTIME_CACHE_KEY, _RUNTIME_CACHE_VALUE
+    with _RUNTIME_CACHE_LOCK:
+        runtime = _RUNTIME_CACHE_VALUE
+        _RUNTIME_CACHE_KEY = None
+        _RUNTIME_CACHE_VALUE = None
+
+    if runtime is not None:
+        runtime.unload()
+
+
+def _load_audio(path: str | Path) -> tuple[torch.Tensor, int]:
+    try:
+        return torchaudio.load(str(path))
+    except RuntimeError:
+        import soundfile as sf
+
+        data, sr = sf.read(str(path), dtype="float32")
+        wav = torch.from_numpy(data)
+        if wav.ndim == 1:
+            wav = wav.unsqueeze(0)
+        else:
+            wav = wav.T
+        return wav, sr
+
+
+def save_wav(path: str | Path, audio: torch.Tensor, sample_rate: int) -> Path:
+    out_path = Path(path)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        torchaudio.save(str(out_path), audio, sample_rate)
+    except RuntimeError:
+        import soundfile as sf
+
+        sf.write(str(out_path), audio.squeeze(0).numpy(), sample_rate)
+    return out_path

irodori_tts/model.py

@@ -0,0 +1,1491 @@
+from __future__ import annotations
+
+import math
+from dataclasses import asdict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import ModelConfig
+
+DURATION_SPEAKER_FUSIONS = {
+    "concat",
+    "adarn",
+    "adarn_zero",
+    "speaker_cross_attn",
+    "text_cross_attn",
+}
+DURATION_ARCHITECTURES = {"pooled", "token_sum_adarn_zero_no_aux"}
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor:
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+    t = torch.arange(end, dtype=torch.float32)
+    freqs = torch.outer(t, freqs)
+    return torch.complex(torch.cos(freqs), torch.sin(freqs))
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    # x: (B, S, H, Dh), Dh must be even.
+    x_ = torch.view_as_complex(x.float().reshape(*x.shape[:3], -1, 2))
+    x_ = x_ * freqs_cis[None, :, None, :]
+    x_ = torch.view_as_real(x_).reshape_as(x)
+    return x_.type_as(x)
+
+
+def get_timestep_embedding(timestep: torch.Tensor, dim: int) -> torch.Tensor:
+    assert dim % 2 == 0
+    half = dim // 2
+    freqs = 1000.0 * torch.exp(
+        -torch.log(torch.tensor(10000.0, device=timestep.device, dtype=torch.float32))
+        * torch.arange(half, device=timestep.device, dtype=torch.float32)
+        / half
+    )
+    args = timestep[:, None].float() * freqs[None, :]
+    return torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(timestep.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int | tuple[int, ...], eps: float = 1e-6):
+        super().__init__()
+        if isinstance(dim, int):
+            dim = (dim,)
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt((x * x).mean(dim=-1, keepdim=True) + self.eps)
+        return (x * self.weight).to(x_dtype)
+
+
+class LowRankAdaLN(nn.Module):
+    """
+    Echo-style low-rank AdaLN that returns both modulated activations and a residual gate.
+    """
+
+    def __init__(self, model_dim: int, rank: int, eps: float):
+        super().__init__()
+        rank = max(1, min(int(rank), int(model_dim)))
+        self.eps = eps
+        self.shift_down = nn.Linear(model_dim, rank, bias=False)
+        self.scale_down = nn.Linear(model_dim, rank, bias=False)
+        self.gate_down = nn.Linear(model_dim, rank, bias=False)
+        self.shift_up = nn.Linear(rank, model_dim, bias=True)
+        self.scale_up = nn.Linear(rank, model_dim, bias=True)
+        self.gate_up = nn.Linear(rank, model_dim, bias=True)
+        # Match Echo/JAX AdaLN behavior: zero-init output projections.
+        nn.init.zeros_(self.shift_up.weight)
+        nn.init.zeros_(self.scale_up.weight)
+        nn.init.zeros_(self.gate_up.weight)
+        if self.shift_up.bias is not None:
+            nn.init.zeros_(self.shift_up.bias)
+        if self.scale_up.bias is not None:
+            nn.init.zeros_(self.scale_up.bias)
+        if self.gate_up.bias is not None:
+            nn.init.zeros_(self.gate_up.bias)
+
+    def forward(
+        self, x: torch.Tensor, cond_embed: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        shift, scale, gate = cond_embed.chunk(3, dim=-1)
+        shift = self.shift_up(self.shift_down(F.silu(shift))) + shift
+        scale = self.scale_up(self.scale_down(F.silu(scale))) + scale
+        gate = self.gate_up(self.gate_down(F.silu(gate))) + gate
+
+        x_dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt((x * x).mean(dim=-1, keepdim=True) + self.eps)
+        x = x * (1.0 + scale) + shift
+        gate = torch.tanh(gate)
+        return x.to(x_dtype), gate
+
+
+def patch_sequence_with_mask(
+    seq: torch.Tensor,
+    mask: torch.Tensor,
+    patch_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Patch along sequence axis:
+      seq: (B, S, D) -> (B, S//patch, D*patch)
+      mask: (B, S) -> (B, S//patch) with all() over patch window.
+
+    Note:
+      For speaker conditioning in this project, `seq` is already in
+      latent-patched space (D = latent_dim * latent_patch_size).
+      This helper applies an additional sequence patching for
+      `speaker_patch_size`.
+    """
+    if patch_size <= 1:
+        return seq, mask
+    if seq.ndim != 3 or mask.ndim != 2:
+        raise ValueError(
+            f"Expected seq=(B,S,D), mask=(B,S), got seq={tuple(seq.shape)} mask={tuple(mask.shape)}"
+        )
+    if seq.shape[0] != mask.shape[0] or seq.shape[1] != mask.shape[1]:
+        raise ValueError(
+            f"Sequence/mask shape mismatch: seq={tuple(seq.shape)}, mask={tuple(mask.shape)}. "
+            "Expected matching (B,S)."
+        )
+    bsz, seq_len, dim = seq.shape
+    usable = (seq_len // patch_size) * patch_size
+    if usable <= 0:
+        raise ValueError(
+            f"Reference sequence too short for speaker_patch_size={patch_size}: seq_len={seq_len}"
+        )
+    seq = seq[:, :usable].reshape(bsz, usable // patch_size, dim * patch_size)
+    mask = mask[:, :usable].reshape(bsz, usable // patch_size, patch_size).all(dim=-1)
+    return seq, mask
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, heads: int, norm_eps: float):
+        super().__init__()
+        if dim % heads != 0:
+            raise ValueError(f"dim={dim} must be divisible by heads={heads}")
+        if (dim // heads) % 2 != 0:
+            raise ValueError("head_dim must be even for RoPE")
+        self.dim = dim
+        self.heads = heads
+        self.head_dim = dim // heads
+
+        self.wq = nn.Linear(dim, dim, bias=False)
+        self.wk = nn.Linear(dim, dim, bias=False)
+        self.wv = nn.Linear(dim, dim, bias=False)
+        self.wo = nn.Linear(dim, dim, bias=False)
+        self.gate = nn.Linear(dim, dim, bias=False)
+
+        self.q_norm = RMSNorm((self.heads, self.head_dim), eps=norm_eps)
+        self.k_norm = RMSNorm((self.heads, self.head_dim), eps=norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        key_mask: torch.Tensor | None,
+        freqs_cis: torch.Tensor,
+    ) -> torch.Tensor:
+        bsz, seq_len, _ = x.shape
+        q = self.wq(x).reshape(bsz, seq_len, self.heads, self.head_dim)
+        k = self.wk(x).reshape(bsz, seq_len, self.heads, self.head_dim)
+        v = self.wv(x).reshape(bsz, seq_len, self.heads, self.head_dim)
+        gate = self.gate(x)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        q = apply_rotary_emb(q, freqs_cis[:seq_len])
+        k = apply_rotary_emb(k, freqs_cis[:seq_len])
+
+        attn_mask = None
+        if key_mask is not None:
+            attn_mask = key_mask[:, None, None, :]
+
+        y = F.scaled_dot_product_attention(
+            q.transpose(1, 2),
+            k.transpose(1, 2),
+            v.transpose(1, 2),
+            attn_mask=attn_mask,
+            is_causal=False,
+        ).transpose(1, 2)
+        y = y.reshape(bsz, seq_len, self.dim)
+        y = y * torch.sigmoid(gate)
+        return self.wo(y)
+
+
+class JointAttention(nn.Module):
+    """
+    Echo-style joint attention over latent self tokens + conditioning contexts.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        heads: int,
+        text_ctx_dim: int,
+        speaker_ctx_dim: int | None,
+        caption_ctx_dim: int | None,
+        norm_eps: float,
+    ):
+        super().__init__()
+        if dim % heads != 0:
+            raise ValueError(f"dim={dim} must be divisible by heads={heads}")
+        if (dim // heads) % 2 != 0:
+            raise ValueError("head_dim must be even for RoPE")
+        self.dim = dim
+        self.heads = heads
+        self.head_dim = dim // heads
+
+        self.wq = nn.Linear(dim, dim, bias=False)
+        self.wk = nn.Linear(dim, dim, bias=False)
+        self.wv = nn.Linear(dim, dim, bias=False)
+        self.wk_text = nn.Linear(text_ctx_dim, dim, bias=False)
+        self.wv_text = nn.Linear(text_ctx_dim, dim, bias=False)
+        self.has_speaker_condition = speaker_ctx_dim is not None
+        if self.has_speaker_condition:
+            self.wk_speaker = nn.Linear(int(speaker_ctx_dim), dim, bias=False)
+            self.wv_speaker = nn.Linear(int(speaker_ctx_dim), dim, bias=False)
+        self.has_caption_condition = caption_ctx_dim is not None
+        if self.has_caption_condition:
+            self.wk_caption = nn.Linear(int(caption_ctx_dim), dim, bias=False)
+            self.wv_caption = nn.Linear(int(caption_ctx_dim), dim, bias=False)
+        self.gate = nn.Linear(dim, dim, bias=False)
+        self.wo = nn.Linear(dim, dim, bias=False)
+
+        self.q_norm = RMSNorm((self.heads, self.head_dim), eps=norm_eps)
+        self.k_norm = RMSNorm((self.heads, self.head_dim), eps=norm_eps)
+
+    def _apply_rotary_half(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+        x_rot, x_passthrough = x.chunk(2, dim=-2)
+        x_rot = apply_rotary_emb(x_rot, freqs_cis)
+        return torch.cat([x_rot, x_passthrough], dim=-2)
+
+    def project_context_kv(
+        self,
+        text_context: torch.Tensor,
+        speaker_context: torch.Tensor | None,
+        caption_context: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, ...]:
+        """
+        Precompute conditioning KV projections for static conditioning.
+        """
+        bsz = text_context.shape[0]
+        k_text = self.wk_text(text_context).reshape(
+            bsz, text_context.shape[1], self.heads, self.head_dim
+        )
+        v_text = self.wv_text(text_context).reshape(
+            bsz, text_context.shape[1], self.heads, self.head_dim
+        )
+        k_text = self.k_norm(k_text)
+        projected: list[torch.Tensor] = [k_text, v_text]
+        if self.has_speaker_condition:
+            if speaker_context is None:
+                raise ValueError(
+                    "speaker_context is required when speaker conditioning is enabled."
+                )
+            if speaker_context.shape[0] != bsz:
+                raise ValueError(
+                    "Batch mismatch for context projection: "
+                    f"text={tuple(text_context.shape)} speaker={tuple(speaker_context.shape)}"
+                )
+            k_speaker = self.wk_speaker(speaker_context).reshape(
+                bsz, speaker_context.shape[1], self.heads, self.head_dim
+            )
+            v_speaker = self.wv_speaker(speaker_context).reshape(
+                bsz, speaker_context.shape[1], self.heads, self.head_dim
+            )
+            k_speaker = self.k_norm(k_speaker)
+            projected.extend([k_speaker, v_speaker])
+        elif speaker_context is not None and speaker_context.shape[0] != bsz:
+            raise ValueError(
+                "Batch mismatch for ignored speaker context: "
+                f"text={tuple(text_context.shape)} speaker={tuple(speaker_context.shape)}"
+            )
+        if not self.has_caption_condition:
+            return tuple(projected)
+        if caption_context is None:
+            raise ValueError("caption_context is required when caption conditioning is enabled.")
+        if caption_context.shape[0] != bsz:
+            raise ValueError(
+                "Batch mismatch for caption context projection: "
+                f"text={tuple(text_context.shape)} caption={tuple(caption_context.shape)}"
+            )
+        k_caption = self.wk_caption(caption_context).reshape(
+            bsz, caption_context.shape[1], self.heads, self.head_dim
+        )
+        v_caption = self.wv_caption(caption_context).reshape(
+            bsz, caption_context.shape[1], self.heads, self.head_dim
+        )
+        k_caption = self.k_norm(k_caption)
+        projected.extend([k_caption, v_caption])
+        return tuple(projected)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        text_context: torch.Tensor,
+        text_mask: torch.Tensor | None,
+        speaker_context: torch.Tensor | None,
+        speaker_mask: torch.Tensor | None,
+        caption_context: torch.Tensor | None,
+        caption_mask: torch.Tensor | None,
+        freqs_cis: torch.Tensor,
+        self_mask: torch.Tensor | None = None,
+        context_kv: tuple[torch.Tensor, ...] | None = None,
+    ) -> torch.Tensor:
+        bsz, seq_len, _ = x.shape
+        q = self.wq(x).reshape(bsz, seq_len, self.heads, self.head_dim)
+        k_self = self.wk(x).reshape(bsz, seq_len, self.heads, self.head_dim)
+        v_self = self.wv(x).reshape(bsz, seq_len, self.heads, self.head_dim)
+        if context_kv is None:
+            projected = self.project_context_kv(
+                text_context=text_context,
+                speaker_context=speaker_context,
+                caption_context=caption_context,
+            )
+        else:
+            projected = context_kv
+        if projected is None:
+            raise RuntimeError("JointAttention projected context unexpectedly missing.")
+        offset = 0
+        k_text, v_text = projected[offset], projected[offset + 1]
+        offset += 2
+        k_speaker = None
+        v_speaker = None
+        if self.has_speaker_condition:
+            k_speaker, v_speaker = projected[offset], projected[offset + 1]
+            offset += 2
+        k_caption = None
+        v_caption = None
+        if self.has_caption_condition:
+            k_caption, v_caption = projected[offset], projected[offset + 1]
+
+        q = self.q_norm(q)
+        k_self = self.k_norm(k_self)
+        q = self._apply_rotary_half(q, freqs_cis[:seq_len])
+        k_self = self._apply_rotary_half(k_self, freqs_cis[:seq_len])
+
+        if self_mask is None:
+            self_mask = torch.ones((bsz, seq_len), dtype=torch.bool, device=x.device)
+        if text_mask is None:
+            text_mask = torch.ones(
+                (bsz, text_context.shape[1]),
+                dtype=torch.bool,
+                device=x.device,
+            )
+        context_k = [k_self, k_text]
+        context_v = [v_self, v_text]
+        context_masks = [self_mask, text_mask]
+        if self.has_speaker_condition:
+            if speaker_context is None or k_speaker is None or v_speaker is None:
+                raise ValueError(
+                    "speaker_context is required when speaker conditioning is enabled."
+                )
+            if speaker_mask is None:
+                speaker_mask = torch.ones(
+                    (bsz, speaker_context.shape[1]),
+                    dtype=torch.bool,
+                    device=x.device,
+                )
+            context_k.append(k_speaker)
+            context_v.append(v_speaker)
+            context_masks.append(speaker_mask)
+        if self.has_caption_condition:
+            if caption_context is None:
+                raise ValueError(
+                    "caption_context is required when caption conditioning is enabled."
+                )
+            if caption_mask is None:
+                caption_mask = torch.ones(
+                    (bsz, caption_context.shape[1]),
+                    dtype=torch.bool,
+                    device=x.device,
+                )
+            if k_caption is None or v_caption is None:
+                raise RuntimeError(
+                    "Caption projections are missing despite enabled caption conditioning."
+                )
+            context_k.append(k_caption)
+            context_v.append(v_caption)
+            context_masks.append(caption_mask)
+
+        k = torch.cat(context_k, dim=1)
+        v = torch.cat(context_v, dim=1)
+        attn_mask = torch.cat(context_masks, dim=1)
+        attn_mask = attn_mask[:, None, None, :]
+
+        y = F.scaled_dot_product_attention(
+            q.transpose(1, 2),
+            k.transpose(1, 2),
+            v.transpose(1, 2),
+            attn_mask=attn_mask,
+            is_causal=False,
+        ).transpose(1, 2)
+        y = y.reshape(bsz, seq_len, self.dim)
+        y = y * torch.sigmoid(self.gate(x))
+        return self.wo(y)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+def _safe_attention_mask(
+    x: torch.Tensor,
+    mask: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if mask.ndim != 2 or mask.shape[0] != x.shape[0] or mask.shape[1] != x.shape[1]:
+        raise ValueError(
+            f"mask must have shape (B, S) matching x, got x={tuple(x.shape)} "
+            f"mask={tuple(mask.shape)}"
+        )
+    mask = mask.to(device=x.device, dtype=torch.bool)
+    has_any = mask.any(dim=1)
+    if bool(has_any.all()):
+        return x, mask
+    if x.shape[1] <= 0:
+        raise ValueError("Cannot attention-pool an empty sequence.")
+    x = x.clone()
+    mask = mask.clone()
+    x[~has_any] = 0
+    mask[~has_any, 0] = True
+    return x, mask
+
+
+class AttentionPooling(nn.Module):
+    def __init__(self, dim: int, heads: int, norm_eps: float):
+        super().__init__()
+        if dim % heads != 0:
+            raise ValueError(f"dim={dim} must be divisible by heads={heads}")
+        self.dim = int(dim)
+        self.heads = int(heads)
+        self.head_dim = int(dim) // int(heads)
+        self.query = nn.Parameter(torch.empty(1, 1, int(dim)))
+        nn.init.normal_(self.query, mean=0.0, std=0.02)
+        self.q_norm = RMSNorm(dim, eps=norm_eps)
+        self.k_norm = RMSNorm(dim, eps=norm_eps)
+        self.wq = nn.Linear(dim, dim, bias=False)
+        self.wk = nn.Linear(dim, dim, bias=False)
+        self.wv = nn.Linear(dim, dim, bias=False)
+        self.wo = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 3 or x.shape[-1] != self.dim:
+            raise ValueError(f"x must have shape (B, S, {self.dim}), got {tuple(x.shape)}")
+        x, mask = _safe_attention_mask(x, mask)
+        bsz, seq_len, _ = x.shape
+        q = self.query.to(dtype=x.dtype).expand(bsz, -1, -1)
+        q = self.wq(self.q_norm(q)).reshape(bsz, 1, self.heads, self.head_dim)
+        k = self.wk(self.k_norm(x)).reshape(bsz, seq_len, self.heads, self.head_dim)
+        v = self.wv(x).reshape(bsz, seq_len, self.heads, self.head_dim)
+        y = F.scaled_dot_product_attention(
+            q.transpose(1, 2),
+            k.transpose(1, 2),
+            v.transpose(1, 2),
+            attn_mask=mask[:, None, None, :],
+            is_causal=False,
+        )
+        y = y.transpose(1, 2).reshape(bsz, 1, self.dim)
+        return self.wo(y).squeeze(1)
+
+
+class CrossAttentionPooling(nn.Module):
+    def __init__(
+        self,
+        *,
+        query_dim: int,
+        context_dim: int,
+        output_dim: int,
+        heads: int,
+        norm_eps: float,
+    ):
+        super().__init__()
+        if output_dim % heads != 0:
+            raise ValueError(f"output_dim={output_dim} must be divisible by heads={heads}")
+        self.query_dim = int(query_dim)
+        self.context_dim = int(context_dim)
+        self.output_dim = int(output_dim)
+        self.heads = int(heads)
+        self.head_dim = int(output_dim) // int(heads)
+        self.q_norm = RMSNorm(query_dim, eps=norm_eps)
+        self.k_norm = RMSNorm(context_dim, eps=norm_eps)
+        self.wq = nn.Linear(query_dim, output_dim, bias=False)
+        self.wk = nn.Linear(context_dim, output_dim, bias=False)
+        self.wv = nn.Linear(context_dim, output_dim, bias=False)
+        self.wo = nn.Linear(output_dim, output_dim, bias=False)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        context: torch.Tensor,
+        context_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        if query.ndim != 2 or query.shape[-1] != self.query_dim:
+            raise ValueError(
+                f"query must have shape (B, {self.query_dim}), got {tuple(query.shape)}"
+            )
+        if context.ndim != 3 or context.shape[-1] != self.context_dim:
+            raise ValueError(
+                f"context must have shape (B, S, {self.context_dim}), got {tuple(context.shape)}"
+            )
+        context, context_mask = _safe_attention_mask(context, context_mask)
+        bsz, seq_len, _ = context.shape
+        q = query[:, None, :]
+        q = self.wq(self.q_norm(q)).reshape(bsz, 1, self.heads, self.head_dim)
+        k = self.wk(self.k_norm(context)).reshape(bsz, seq_len, self.heads, self.head_dim)
+        v = self.wv(context).reshape(bsz, seq_len, self.heads, self.head_dim)
+        y = F.scaled_dot_product_attention(
+            q.transpose(1, 2),
+            k.transpose(1, 2),
+            v.transpose(1, 2),
+            attn_mask=context_mask[:, None, None, :],
+            is_causal=False,
+        )
+        y = y.transpose(1, 2).reshape(bsz, 1, self.output_dim)
+        return self.wo(y).squeeze(1)
+
+
+class DurationSwiGLUBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim: int,
+        hidden_dim: int,
+        dropout: float,
+        norm_eps: float,
+        cond_dim: int | None = None,
+    ):
+        super().__init__()
+        self.norm = RMSNorm(dim, eps=norm_eps)
+        self.mlp = SwiGLU(dim, hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.cond_dim = cond_dim
+        self.modulation = None
+        if cond_dim is not None:
+            self.modulation = nn.Linear(cond_dim, dim * 3, bias=True)
+            nn.init.zeros_(self.modulation.weight)
+            nn.init.zeros_(self.modulation.bias)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor | None = None) -> torch.Tensor:
+        h = self.norm(x)
+        if self.modulation is not None:
+            if cond is None:
+                raise ValueError("cond is required for AdaRN-Zero duration blocks.")
+            shift, scale, gate = self.modulation(F.silu(cond)).chunk(3, dim=-1)
+            if h.ndim == 3 and shift.ndim == 2:
+                shift = shift.unsqueeze(1)
+                scale = scale.unsqueeze(1)
+                gate = gate.unsqueeze(1)
+            h = h * (1.0 + scale) + shift
+            return x + self.dropout(torch.tanh(gate) * self.mlp(h))
+        return x + self.dropout(self.mlp(h))
+
+
+class TextBlock(nn.Module):
+    def __init__(self, dim: int, heads: int, mlp_ratio: float, norm_eps: float, dropout: float):
+        super().__init__()
+        self.attention_norm = RMSNorm(dim, eps=norm_eps)
+        self.attention = SelfAttention(dim, heads, norm_eps=norm_eps)
+        self.mlp_norm = RMSNorm(dim, eps=norm_eps)
+        self.mlp = SwiGLU(dim, int(dim * mlp_ratio))
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+        x = x + self.dropout(
+            self.attention(self.attention_norm(x), key_mask=mask, freqs_cis=freqs_cis)
+        )
+        x = x + self.dropout(self.mlp(self.mlp_norm(x)))
+        return x
+
+
+class TextEncoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        vocab_size: int,
+        dim: int,
+        layers: int,
+        heads: int,
+        mlp_ratio: float,
+        norm_eps: float,
+        dropout: float,
+    ):
+        super().__init__()
+        self.text_embedding = nn.Embedding(vocab_size, dim)
+        self.blocks = nn.ModuleList(
+            TextBlock(
+                dim=dim,
+                heads=heads,
+                mlp_ratio=mlp_ratio,
+                norm_eps=norm_eps,
+                dropout=dropout,
+            )
+            for _ in range(layers)
+        )
+        self.head_dim = dim // heads
+        self.register_buffer(
+            "_freqs_cis_cache", torch.empty(0, 0, dtype=torch.complex64), persistent=False
+        )
+
+    def _rope_freqs(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        cache = self._freqs_cis_cache
+        if cache.device != device or cache.shape[0] < seq_len:
+            cache = precompute_freqs_cis(self.head_dim, seq_len).to(device)
+            self._freqs_cis_cache = cache
+        return cache[:seq_len]
+
+    def forward(self, input_ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        x = self.text_embedding(input_ids)
+        # Hard-mask invalid tokens so fully-masked conditioning becomes truly unconditional.
+        mask_f = mask.unsqueeze(-1).to(dtype=x.dtype)
+        x = x * mask_f
+        freqs = self._rope_freqs(input_ids.shape[1], x.device)
+        for block in self.blocks:
+            x = block(x, mask=mask, freqs_cis=freqs)
+            x = x * mask_f
+        return x * mask_f
+
+
+class ReferenceLatentEncoder(nn.Module):
+    """
+    Encoder for reference latents used as speaker/style conditioning.
+    """
+
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.in_proj = nn.Linear(cfg.speaker_patched_latent_dim, cfg.speaker_dim, bias=True)
+        speaker_mlp_ratio = cfg.speaker_mlp_ratio_resolved
+        self.blocks = nn.ModuleList(
+            TextBlock(
+                dim=cfg.speaker_dim,
+                heads=cfg.speaker_heads,
+                mlp_ratio=speaker_mlp_ratio,
+                norm_eps=cfg.norm_eps,
+                dropout=cfg.dropout,
+            )
+            for _ in range(cfg.speaker_layers)
+        )
+        self.head_dim = cfg.speaker_dim // cfg.speaker_heads
+        self.register_buffer(
+            "_freqs_cis_cache", torch.empty(0, 0, dtype=torch.complex64), persistent=False
+        )
+
+    def _rope_freqs(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        cache = self._freqs_cis_cache
+        if cache.device != device or cache.shape[0] < seq_len:
+            cache = precompute_freqs_cis(self.head_dim, seq_len).to(device)
+            self._freqs_cis_cache = cache
+        return cache[:seq_len]
+
+    def forward(self, latent: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        x = self.in_proj(latent)
+        x = x / 6.0
+        # Keep masked reference positions strictly zero across residual/MLP paths.
+        mask_f = mask.unsqueeze(-1).to(dtype=x.dtype)
+        x = x * mask_f
+        freqs = self._rope_freqs(x.shape[1], x.device)
+        for block in self.blocks:
+            x = block(x, mask=mask, freqs_cis=freqs)
+            x = x * mask_f
+        return x * mask_f
+
+
+class DiffusionBlock(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.attention = JointAttention(
+            cfg.model_dim,
+            cfg.num_heads,
+            cfg.text_dim,
+            cfg.speaker_dim if cfg.use_speaker_condition else None,
+            cfg.caption_dim_resolved if cfg.use_caption_condition else None,
+            norm_eps=cfg.norm_eps,
+        )
+        self.mlp = SwiGLU(cfg.model_dim, int(cfg.model_dim * cfg.mlp_ratio))
+        adaln_rank = max(1, min(int(cfg.adaln_rank), int(cfg.model_dim)))
+        self.attention_adaln = LowRankAdaLN(
+            model_dim=cfg.model_dim,
+            rank=adaln_rank,
+            eps=cfg.norm_eps,
+        )
+        self.mlp_adaln = LowRankAdaLN(
+            model_dim=cfg.model_dim,
+            rank=adaln_rank,
+            eps=cfg.norm_eps,
+        )
+        self.dropout = nn.Dropout(cfg.dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond_embed: torch.Tensor,
+        text_state: torch.Tensor,
+        text_mask: torch.Tensor,
+        speaker_state: torch.Tensor | None,
+        speaker_mask: torch.Tensor | None,
+        caption_state: torch.Tensor | None,
+        caption_mask: torch.Tensor | None,
+        freqs_cis: torch.Tensor,
+        self_mask: torch.Tensor | None = None,
+        context_kv: tuple[torch.Tensor, ...] | None = None,
+    ) -> torch.Tensor:
+        h, attention_gate = self.attention_adaln(x, cond_embed)
+        x = x + self.dropout(
+            attention_gate
+            * self.attention(
+                x=h,
+                text_context=text_state,
+                text_mask=text_mask,
+                speaker_context=speaker_state,
+                speaker_mask=speaker_mask,
+                caption_context=caption_state,
+                caption_mask=caption_mask,
+                freqs_cis=freqs_cis,
+                self_mask=self_mask,
+                context_kv=context_kv,
+            )
+        )
+
+        h, mlp_gate = self.mlp_adaln(x, cond_embed)
+        x = x + self.dropout(mlp_gate * self.mlp(h))
+        return x
+
+
+class DurationPredictor(nn.Module):
+    def __init__(
+        self,
+        *,
+        text_dim: int,
+        aux_dim: int,
+        hidden_dim: int,
+        layers: int,
+        dropout: float,
+        speaker_dim: int | None = None,
+        speaker_fusion: str = "concat",
+        attention_heads: int = 8,
+        norm_eps: float = 1e-5,
+        architecture: str = "pooled",
+        token_init_frames: float = 6.3,
+    ):
+        super().__init__()
+        if text_dim <= 0:
+            raise ValueError(f"duration predictor text_dim must be > 0, got {text_dim}")
+        if aux_dim <= 0:
+            raise ValueError(f"duration predictor aux_dim must be > 0, got {aux_dim}")
+        if hidden_dim <= 0:
+            raise ValueError(f"duration predictor hidden_dim must be > 0, got {hidden_dim}")
+        if layers <= 0:
+            raise ValueError(f"duration predictor layers must be > 0, got {layers}")
+        if speaker_dim is not None and speaker_dim <= 0:
+            raise ValueError(f"duration predictor speaker_dim must be > 0, got {speaker_dim}")
+        speaker_fusion = str(speaker_fusion).strip().lower()
+        if speaker_fusion not in DURATION_SPEAKER_FUSIONS:
+            raise ValueError(
+                f"duration speaker fusion must be one of {sorted(DURATION_SPEAKER_FUSIONS)}, "
+                f"got {speaker_fusion!r}"
+            )
+        architecture = str(architecture).strip().lower()
+        if architecture not in DURATION_ARCHITECTURES:
+            raise ValueError(
+                "duration architecture must be one of "
+                f"{sorted(DURATION_ARCHITECTURES)}, got {architecture!r}"
+            )
+        if attention_heads <= 0:
+            raise ValueError(
+                f"duration predictor attention_heads must be > 0, got {attention_heads}"
+            )
+        if token_init_frames <= 0:
+            raise ValueError(
+                f"duration token_init_frames must be > 0, got {token_init_frames}"
+            )
+        if speaker_dim is None and speaker_fusion != "concat":
+            raise ValueError(
+                f"duration speaker fusion {speaker_fusion!r} requires speaker_dim."
+            )
+        if architecture == "token_sum_adarn_zero_no_aux" and speaker_dim is None:
+            raise ValueError("token_sum_adarn_zero_no_aux requires speaker_dim.")
+        if architecture == "token_sum_adarn_zero_no_aux" and speaker_fusion != "adarn_zero":
+            raise ValueError(
+                "token_sum_adarn_zero_no_aux uses block-level speaker AdaRN-Zero and "
+                "requires speaker_fusion='adarn_zero'."
+            )
+
+        self.text_dim = int(text_dim)
+        self.aux_dim = int(aux_dim)
+        self.hidden_dim = int(hidden_dim)
+        self.speaker_dim = None if speaker_dim is None else int(speaker_dim)
+        self.speaker_fusion = speaker_fusion
+        self.duration_architecture = architecture
+        self.text_pool = None
+        self.null_speaker = (
+            nn.Parameter(torch.zeros(int(speaker_dim))) if speaker_dim is not None else None
+        )
+        self.text_adarn_norm = None
+        self.text_adarn = None
+        self.speaker_cross_attn = None
+        self.text_cross_attn = None
+        self.token_input_proj = None
+        self.token_blocks = None
+        self.token_out_norm = None
+        self.token_out_proj = None
+
+        if architecture == "token_sum_adarn_zero_no_aux":
+            self.token_input_proj = nn.Linear(int(text_dim), int(hidden_dim))
+            self.token_blocks = nn.ModuleList(
+                DurationSwiGLUBlock(
+                    dim=int(hidden_dim),
+                    hidden_dim=int(hidden_dim),
+                    dropout=float(dropout),
+                    norm_eps=float(norm_eps),
+                    cond_dim=int(speaker_dim),
+                )
+                for _ in range(int(layers))
+            )
+            self.token_out_norm = RMSNorm(int(hidden_dim), eps=float(norm_eps))
+            self.token_out_proj = nn.Linear(int(hidden_dim), 1)
+            nn.init.zeros_(self.token_out_proj.weight)
+            nn.init.constant_(
+                self.token_out_proj.bias,
+                float(math.log(math.expm1(float(token_init_frames)))),
+            )
+            return
+
+        self.text_pool = AttentionPooling(
+            dim=int(text_dim),
+            heads=int(attention_heads),
+            norm_eps=float(norm_eps),
+        )
+
+        if speaker_dim is not None:
+            if speaker_fusion == "concat":
+                input_dim = int(text_dim) + int(speaker_dim) + int(aux_dim)
+            elif speaker_fusion == "adarn":
+                input_dim = int(text_dim) + int(aux_dim)
+                self.text_adarn_norm = RMSNorm(int(text_dim), eps=float(norm_eps))
+                self.text_adarn = nn.Linear(int(speaker_dim), int(text_dim) * 2)
+                nn.init.zeros_(self.text_adarn.weight)
+                nn.init.zeros_(self.text_adarn.bias)
+            elif speaker_fusion == "adarn_zero":
+                input_dim = int(text_dim) + int(aux_dim)
+            elif speaker_fusion == "speaker_cross_attn":
+                input_dim = int(text_dim) * 2 + int(aux_dim)
+                self.speaker_cross_attn = CrossAttentionPooling(
+                    query_dim=int(text_dim),
+                    context_dim=int(speaker_dim),
+                    output_dim=int(text_dim),
+                    heads=int(attention_heads),
+                    norm_eps=float(norm_eps),
+                )
+            elif speaker_fusion == "text_cross_attn":
+                input_dim = int(text_dim) + int(speaker_dim) + int(aux_dim)
+                self.text_cross_attn = CrossAttentionPooling(
+                    query_dim=int(speaker_dim),
+                    context_dim=int(text_dim),
+                    output_dim=int(text_dim),
+                    heads=int(attention_heads),
+                    norm_eps=float(norm_eps),
+                )
+            else:
+                raise RuntimeError(f"Unsupported duration speaker fusion: {speaker_fusion!r}")
+        else:
+            input_dim = int(text_dim) + int(aux_dim)
+
+        self.input_proj = nn.Linear(int(input_dim), int(hidden_dim))
+        block_cond_dim = int(speaker_dim) if speaker_fusion == "adarn_zero" else None
+        self.blocks = nn.ModuleList(
+            DurationSwiGLUBlock(
+                dim=int(hidden_dim),
+                hidden_dim=int(hidden_dim),
+                dropout=float(dropout),
+                norm_eps=float(norm_eps),
+                cond_dim=block_cond_dim,
+            )
+            for _ in range(int(layers))
+        )
+        self.out_norm = RMSNorm(int(hidden_dim), eps=float(norm_eps))
+        self.out_proj = nn.Linear(int(hidden_dim), 1)
+
+    def _speaker_vec(
+        self,
+        *,
+        batch_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        speaker_state: torch.Tensor | None,
+        has_speaker: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.null_speaker is None or self.speaker_dim is None:
+            raise RuntimeError("Duration speaker modules are missing.")
+        null_vec = self.null_speaker.to(device=device, dtype=dtype)[None, :].expand(
+            batch_size, -1
+        )
+        if speaker_state is None:
+            return null_vec
+        if speaker_state.ndim != 3 or speaker_state.shape[0] != batch_size:
+            raise ValueError(
+                f"speaker_state must have shape (B, S, D), got {tuple(speaker_state.shape)}"
+            )
+        if speaker_state.shape[-1] != self.speaker_dim:
+            raise ValueError(
+                f"speaker_state last dim must be {self.speaker_dim}, got {speaker_state.shape[-1]}"
+            )
+        speaker_vec = speaker_state[:, 0].to(device=device, dtype=dtype)
+        return torch.where(has_speaker[:, None], speaker_vec, null_vec)
+
+    def _speaker_sequence(
+        self,
+        *,
+        batch_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        speaker_state: torch.Tensor | None,
+        speaker_mask: torch.Tensor | None,
+        has_speaker: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.null_speaker is None or self.speaker_dim is None:
+            raise RuntimeError("Duration speaker modules are missing.")
+        null_token = self.null_speaker.to(device=device, dtype=dtype)[None, None, :].expand(
+            batch_size, 1, -1
+        )
+        if speaker_state is None:
+            return null_token, torch.ones((batch_size, 1), dtype=torch.bool, device=device)
+        if speaker_state.ndim != 3 or speaker_state.shape[0] != batch_size:
+            raise ValueError(
+                f"speaker_state must have shape (B, S, D), got {tuple(speaker_state.shape)}"
+            )
+        if speaker_state.shape[-1] != self.speaker_dim:
+            raise ValueError(
+                f"speaker_state last dim must be {self.speaker_dim}, got {speaker_state.shape[-1]}"
+            )
+        speaker_state = speaker_state.to(device=device, dtype=dtype)
+        if speaker_mask is None:
+            speaker_mask = torch.ones(
+                (batch_size, speaker_state.shape[1]), dtype=torch.bool, device=device
+            )
+        elif speaker_mask.ndim != 2 or speaker_mask.shape[:2] != speaker_state.shape[:2]:
+            raise ValueError(
+                "speaker_mask must have shape matching speaker_state (B, S), "
+                f"got speaker_state={tuple(speaker_state.shape)} mask={tuple(speaker_mask.shape)}"
+            )
+        speaker_mask = speaker_mask.to(device=device, dtype=torch.bool)
+        real_mask = speaker_mask & has_speaker[:, None]
+        fallback_mask = ~real_mask.any(dim=1, keepdim=True)
+        context = torch.cat([speaker_state, null_token], dim=1)
+        context_mask = torch.cat([real_mask, fallback_mask], dim=1)
+        return context, context_mask
+
+    def forward(
+        self,
+        *,
+        text_state: torch.Tensor,
+        text_mask: torch.Tensor,
+        aux_features: torch.Tensor,
+        speaker_state: torch.Tensor | None = None,
+        speaker_mask: torch.Tensor | None = None,
+        has_speaker: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if text_state.ndim != 3 or text_state.shape[-1] != self.text_dim:
+            raise ValueError(
+                f"text_state must have shape (B, S, {self.text_dim}), "
+                f"got {tuple(text_state.shape)}"
+            )
+        if aux_features.ndim != 2 or aux_features.shape[1] != self.aux_dim:
+            raise ValueError(
+                f"aux_features must have shape (B, {self.aux_dim}), "
+                f"got {tuple(aux_features.shape)}"
+            )
+        if aux_features.shape[0] != text_state.shape[0]:
+            raise ValueError(
+                "Batch mismatch for duration predictor: "
+                f"text_state={tuple(text_state.shape)} aux_features={tuple(aux_features.shape)}"
+            )
+        text_state, text_mask = _safe_attention_mask(text_state, text_mask)
+        aux_features = aux_features.to(device=text_state.device, dtype=text_state.dtype)
+
+        if self.duration_architecture == "token_sum_adarn_zero_no_aux":
+            if self.speaker_dim is None:
+                raise RuntimeError("Token-sum duration architecture requires speaker modules.")
+            if has_speaker is None:
+                raise ValueError(
+                    "has_speaker is required for speaker-conditioned duration prediction."
+                )
+            has_speaker = has_speaker.to(device=text_state.device, dtype=torch.bool)
+            if has_speaker.ndim != 1 or has_speaker.shape[0] != text_state.shape[0]:
+                raise ValueError(
+                    f"has_speaker must have shape (B,), got {tuple(has_speaker.shape)}"
+                )
+            speaker_vec = self._speaker_vec(
+                batch_size=text_state.shape[0],
+                device=text_state.device,
+                dtype=text_state.dtype,
+                speaker_state=speaker_state,
+                has_speaker=has_speaker,
+            )
+            if (
+                self.token_input_proj is None
+                or self.token_blocks is None
+                or self.token_out_norm is None
+                or self.token_out_proj is None
+            ):
+                raise RuntimeError("Token-sum duration modules are missing.")
+            h = self.token_input_proj(text_state)
+            for block in self.token_blocks:
+                h = block(h, cond=speaker_vec)
+            token_logits = self.token_out_proj(self.token_out_norm(h)).squeeze(-1)
+            token_frames = F.softplus(token_logits.float())
+            total_frames = (token_frames * text_mask.to(dtype=token_frames.dtype)).sum(dim=1)
+            return torch.log1p(total_frames.clamp_min(0.0))
+
+        if self.text_pool is None:
+            raise RuntimeError("Pooled duration modules are missing.")
+        text_vec = self.text_pool(text_state, text_mask)
+        if self.speaker_dim is None:
+            x = torch.cat([text_vec, aux_features], dim=-1)
+            h = self.input_proj(x)
+            for block in self.blocks:
+                h = block(h)
+            return self.out_proj(self.out_norm(h)).squeeze(-1)
+
+        if has_speaker is None:
+            raise ValueError("has_speaker is required for speaker-conditioned duration prediction.")
+        has_speaker = has_speaker.to(device=text_vec.device, dtype=torch.bool)
+        if has_speaker.ndim != 1 or has_speaker.shape[0] != text_vec.shape[0]:
+            raise ValueError(
+                f"has_speaker must have shape (B,), got {tuple(has_speaker.shape)}"
+            )
+        speaker_vec = self._speaker_vec(
+            batch_size=text_vec.shape[0],
+            device=text_vec.device,
+            dtype=text_vec.dtype,
+            speaker_state=speaker_state,
+            has_speaker=has_speaker,
+        )
+
+        if self.speaker_fusion == "concat":
+            x = torch.cat([text_vec, speaker_vec, aux_features], dim=-1)
+            cond = None
+        elif self.speaker_fusion == "adarn":
+            if self.text_adarn_norm is None or self.text_adarn is None:
+                raise RuntimeError("AdaRN duration speaker modules are missing.")
+            scale, shift = self.text_adarn(speaker_vec).chunk(2, dim=-1)
+            text_vec = (self.text_adarn_norm(text_vec) * (1.0 + scale)) + shift
+            x = torch.cat([text_vec, aux_features], dim=-1)
+            cond = None
+        elif self.speaker_fusion == "adarn_zero":
+            x = torch.cat([text_vec, aux_features], dim=-1)
+            cond = speaker_vec
+        elif self.speaker_fusion == "speaker_cross_attn":
+            if self.speaker_cross_attn is None:
+                raise RuntimeError("speaker_cross_attn duration module is missing.")
+            speaker_context, speaker_context_mask = self._speaker_sequence(
+                batch_size=text_vec.shape[0],
+                device=text_vec.device,
+                dtype=text_vec.dtype,
+                speaker_state=speaker_state,
+                speaker_mask=speaker_mask,
+                has_speaker=has_speaker,
+            )
+            context_vec = self.speaker_cross_attn(
+                query=text_vec,
+                context=speaker_context,
+                context_mask=speaker_context_mask,
+            )
+            x = torch.cat([text_vec, context_vec, aux_features], dim=-1)
+            cond = None
+        elif self.speaker_fusion == "text_cross_attn":
+            if self.text_cross_attn is None:
+                raise RuntimeError("text_cross_attn duration module is missing.")
+            context_vec = self.text_cross_attn(
+                query=speaker_vec,
+                context=text_state,
+                context_mask=text_mask,
+            )
+            x = torch.cat([context_vec, speaker_vec, aux_features], dim=-1)
+            cond = None
+        else:
+            raise RuntimeError(f"Unsupported duration speaker fusion: {self.speaker_fusion!r}")
+
+        h = self.input_proj(x)
+        for block in self.blocks:
+            h = block(h, cond=cond)
+        return self.out_proj(self.out_norm(h)).squeeze(-1)
+
+
+class TextToLatentRFDiT(nn.Module):
+    """
+    Text + reference-latent conditioned RF diffusion model over patched DACVAE latent sequences.
+
+    Input x_t shape: (B, S, latent_dim * latent_patch_size)
+    Output v_pred shape: same as input.
+    """
+
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.text_encoder = TextEncoder(
+            vocab_size=cfg.text_vocab_size,
+            dim=cfg.text_dim,
+            layers=cfg.text_layers,
+            heads=cfg.text_heads,
+            mlp_ratio=cfg.text_mlp_ratio_resolved,
+            norm_eps=cfg.norm_eps,
+            dropout=cfg.dropout,
+        )
+        self.caption_encoder = None
+        self.caption_norm = None
+        if cfg.use_caption_condition:
+            self.caption_encoder = TextEncoder(
+                vocab_size=cfg.caption_vocab_size_resolved,
+                dim=cfg.caption_dim_resolved,
+                layers=cfg.caption_layers_resolved,
+                heads=cfg.caption_heads_resolved,
+                mlp_ratio=cfg.caption_mlp_ratio_resolved,
+                norm_eps=cfg.norm_eps,
+                dropout=cfg.dropout,
+            )
+            self.caption_norm = RMSNorm(cfg.caption_dim_resolved, eps=cfg.norm_eps)
+        self.speaker_encoder = None
+        if cfg.use_speaker_condition:
+            self.speaker_encoder = ReferenceLatentEncoder(cfg)
+        self.text_norm = RMSNorm(cfg.text_dim, eps=cfg.norm_eps)
+        self.speaker_norm = None
+        if cfg.use_speaker_condition:
+            self.speaker_norm = RMSNorm(cfg.speaker_dim, eps=cfg.norm_eps)
+        self.duration_predictor = None
+        if cfg.use_duration_predictor:
+            duration_speaker_dim = None
+            if cfg.use_speaker_condition:
+                duration_speaker_dim = int(cfg.speaker_dim)
+            self.duration_predictor = DurationPredictor(
+                text_dim=cfg.text_dim,
+                aux_dim=cfg.duration_aux_dim,
+                hidden_dim=cfg.duration_hidden_dim,
+                layers=cfg.duration_layers,
+                dropout=cfg.duration_dropout,
+                speaker_dim=duration_speaker_dim,
+                speaker_fusion=cfg.duration_speaker_fusion,
+                attention_heads=cfg.duration_attention_heads,
+                norm_eps=cfg.norm_eps,
+                architecture=cfg.duration_architecture,
+                token_init_frames=cfg.duration_token_init_frames,
+            )
+
+        self.cond_module = nn.Sequential(
+            nn.Linear(cfg.timestep_embed_dim, cfg.model_dim, bias=False),
+            nn.SiLU(),
+            nn.Linear(cfg.model_dim, cfg.model_dim, bias=False),
+            nn.SiLU(),
+            nn.Linear(cfg.model_dim, cfg.model_dim * 3, bias=False),
+        )
+
+        self.in_proj = nn.Linear(cfg.patched_latent_dim, cfg.model_dim)
+        self.blocks = nn.ModuleList(DiffusionBlock(cfg) for _ in range(cfg.num_layers))
+        self.out_norm = RMSNorm(cfg.model_dim, eps=cfg.norm_eps)
+        self.out_proj = nn.Linear(cfg.model_dim, cfg.patched_latent_dim)
+        # Echo/JAX training initializes decoder out projection to zero for stable early training.
+        nn.init.zeros_(self.out_proj.weight)
+        if self.out_proj.bias is not None:
+            nn.init.zeros_(self.out_proj.bias)
+
+        self.head_dim = cfg.model_dim // cfg.num_heads
+        if self.head_dim % 2 != 0:
+            raise ValueError("model head_dim must be even for RoPE")
+        self.register_buffer(
+            "_freqs_cis_cache", torch.empty(0, 0, dtype=torch.complex64), persistent=False
+        )
+
+    def _rope_freqs(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        cache = self._freqs_cis_cache
+        if cache.device != device or cache.shape[0] < seq_len:
+            cache = precompute_freqs_cis(self.head_dim, seq_len).to(device)
+            self._freqs_cis_cache = cache
+        return cache[:seq_len]
+
+    @staticmethod
+    def _prepend_masked_mean_token(
+        state: torch.Tensor,
+        mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Prepend one global summary token computed as masked mean over time.
+        """
+        mask_f = mask.unsqueeze(-1).to(dtype=state.dtype)
+        denom = mask_f.sum(dim=1, keepdim=True).clamp_min(1.0)
+        mean_token = (state * mask_f).sum(dim=1, keepdim=True) / denom
+        has_any = mask.any(dim=1, keepdim=True)
+        state = torch.cat([mean_token, state], dim=1)
+        mask = torch.cat([has_any, mask], dim=1)
+        return state, mask
+
+    def encode_conditions(
+        self,
+        text_input_ids: torch.Tensor,
+        text_mask: torch.Tensor,
+        ref_latent: torch.Tensor | None,
+        ref_mask: torch.Tensor | None,
+        caption_input_ids: torch.Tensor | None = None,
+        caption_mask: torch.Tensor | None = None,
+        text_condition_dropout: torch.Tensor | None = None,
+        speaker_condition_dropout: torch.Tensor | None = None,
+        caption_condition_dropout: torch.Tensor | None = None,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+    ]:
+        if text_condition_dropout is not None:
+            text_mask = text_mask.clone()
+            text_mask[text_condition_dropout] = False
+        if self.cfg.use_speaker_condition:
+            if self.speaker_encoder is None or self.speaker_norm is None:
+                raise RuntimeError(
+                    "Speaker conditioning is enabled but speaker modules are missing."
+                )
+            if ref_latent is None or ref_mask is None:
+                raise ValueError(
+                    "ref_latent and ref_mask are required when speaker conditioning is enabled."
+                )
+            if speaker_condition_dropout is not None:
+                ref_mask = ref_mask.clone()
+                ref_mask[speaker_condition_dropout] = False
+        if self.cfg.use_caption_condition:
+            if self.caption_encoder is None or self.caption_norm is None:
+                raise RuntimeError(
+                    "Caption conditioning is enabled but caption modules are missing."
+                )
+            if caption_input_ids is None or caption_mask is None:
+                raise ValueError(
+                    "caption_input_ids and caption_mask are required when caption conditioning is enabled."
+                )
+            if caption_condition_dropout is not None:
+                caption_mask = caption_mask.clone()
+                caption_mask[caption_condition_dropout] = False
+
+        text_state = self.text_encoder(text_input_ids, text_mask)
+        text_state = self.text_norm(text_state)
+        ref_state = None
+        if self.cfg.use_speaker_condition:
+            ref_latent, ref_mask = patch_sequence_with_mask(
+                seq=ref_latent,
+                mask=ref_mask,
+                patch_size=self.cfg.speaker_patch_size,
+            )
+            ref_state = self.speaker_encoder(ref_latent, ref_mask)
+            ref_state = self.speaker_norm(ref_state)
+            ref_state, ref_mask = self._prepend_masked_mean_token(ref_state, ref_mask)
+        caption_state = None
+        if self.cfg.use_caption_condition:
+            caption_state = self.caption_encoder(caption_input_ids, caption_mask)
+            caption_state = self.caption_norm(caption_state)
+        return text_state, text_mask, ref_state, ref_mask, caption_state, caption_mask
+
+    def forward_with_encoded_conditions(
+        self,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+        text_state: torch.Tensor,
+        text_mask: torch.Tensor,
+        speaker_state: torch.Tensor | None,
+        speaker_mask: torch.Tensor | None,
+        caption_state: torch.Tensor | None = None,
+        caption_mask: torch.Tensor | None = None,
+        latent_mask: torch.Tensor | None = None,
+        context_kv_cache: list[tuple[torch.Tensor, ...]] | None = None,
+    ) -> torch.Tensor:
+        t_embed = get_timestep_embedding(t, self.cfg.timestep_embed_dim).to(dtype=x_t.dtype)
+        cond_embed = self.cond_module(t_embed)
+        cond_embed = cond_embed[:, None, :]
+
+        x = self.in_proj(x_t)
+        freqs = self._rope_freqs(x.shape[1], x.device)
+        for i, block in enumerate(self.blocks):
+            x = block(
+                x=x,
+                cond_embed=cond_embed,
+                text_state=text_state,
+                text_mask=text_mask,
+                speaker_state=speaker_state,
+                speaker_mask=speaker_mask,
+                caption_state=caption_state,
+                caption_mask=caption_mask,
+                freqs_cis=freqs,
+                self_mask=latent_mask,
+                context_kv=context_kv_cache[i] if context_kv_cache is not None else None,
+            )
+
+        x = self.out_norm(x)
+        x = self.out_proj(x)
+        return x.to(dtype=x_t.dtype)
+
+    def forward(
+        self,
+        x_t: torch.Tensor | None,
+        t: torch.Tensor | None,
+        text_input_ids: torch.Tensor,
+        text_mask: torch.Tensor,
+        ref_latent: torch.Tensor | None,
+        ref_mask: torch.Tensor | None,
+        caption_input_ids: torch.Tensor | None = None,
+        caption_mask: torch.Tensor | None = None,
+        latent_mask: torch.Tensor | None = None,
+        text_condition_dropout: torch.Tensor | None = None,
+        speaker_condition_dropout: torch.Tensor | None = None,
+        caption_condition_dropout: torch.Tensor | None = None,
+        duration_features: torch.Tensor | None = None,
+        duration_has_speaker: torch.Tensor | None = None,
+        duration_only: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if duration_features is not None:
+            (
+                text_state,
+                text_mask_full,
+                speaker_state,
+                speaker_mask_full,
+                caption_state,
+                caption_mask_full,
+            ) = self.encode_conditions(
+                text_input_ids=text_input_ids,
+                text_mask=text_mask,
+                ref_latent=ref_latent,
+                ref_mask=ref_mask,
+                caption_input_ids=caption_input_ids,
+                caption_mask=caption_mask,
+            )
+            if duration_only:
+                return self.predict_duration_log_frames(
+                    text_state=text_state,
+                    text_mask=text_mask_full,
+                    speaker_state=speaker_state,
+                    speaker_mask=speaker_mask_full,
+                    duration_features=duration_features,
+                    has_speaker=duration_has_speaker,
+                )
+
+            if x_t is None or t is None:
+                raise ValueError("x_t and t are required unless duration_only=True.")
+            text_mask_dit = text_mask_full
+            speaker_mask_dit = speaker_mask_full
+            caption_mask_dit = caption_mask_full
+            if text_condition_dropout is not None:
+                text_mask_dit = text_mask_dit.clone()
+                text_mask_dit[text_condition_dropout] = False
+            if speaker_condition_dropout is not None and speaker_mask_dit is not None:
+                speaker_mask_dit = speaker_mask_dit.clone()
+                speaker_mask_dit[speaker_condition_dropout] = False
+            if caption_condition_dropout is not None and caption_mask_dit is not None:
+                caption_mask_dit = caption_mask_dit.clone()
+                caption_mask_dit[caption_condition_dropout] = False
+
+            v_pred = self.forward_with_encoded_conditions(
+                x_t=x_t,
+                t=t,
+                text_state=text_state,
+                text_mask=text_mask_dit,
+                speaker_state=speaker_state,
+                speaker_mask=speaker_mask_dit,
+                caption_state=caption_state,
+                caption_mask=caption_mask_dit,
+                latent_mask=latent_mask,
+            )
+            duration_pred = self.predict_duration_log_frames(
+                text_state=text_state,
+                text_mask=text_mask_full,
+                speaker_state=speaker_state,
+                speaker_mask=speaker_mask_full,
+                duration_features=duration_features,
+                has_speaker=duration_has_speaker,
+            )
+            return v_pred, duration_pred
+
+        if duration_only:
+            raise ValueError("duration_features is required when duration_only=True.")
+        if x_t is None or t is None:
+            raise ValueError("x_t and t are required for RF forward.")
+
+        (
+            text_state,
+            text_mask,
+            speaker_state,
+            speaker_mask,
+            caption_state,
+            caption_mask,
+        ) = self.encode_conditions(
+            text_input_ids=text_input_ids,
+            text_mask=text_mask,
+            ref_latent=ref_latent,
+            ref_mask=ref_mask,
+            caption_input_ids=caption_input_ids,
+            caption_mask=caption_mask,
+            text_condition_dropout=text_condition_dropout,
+            speaker_condition_dropout=speaker_condition_dropout,
+            caption_condition_dropout=caption_condition_dropout,
+        )
+        return self.forward_with_encoded_conditions(
+            x_t=x_t,
+            t=t,
+            text_state=text_state,
+            text_mask=text_mask,
+            speaker_state=speaker_state,
+            speaker_mask=speaker_mask,
+            caption_state=caption_state,
+            caption_mask=caption_mask,
+            latent_mask=latent_mask,
+        )
+
+    def build_context_kv_cache(
+        self,
+        text_state: torch.Tensor,
+        speaker_state: torch.Tensor | None,
+        caption_state: torch.Tensor | None = None,
+    ) -> list[tuple[torch.Tensor, ...]]:
+        """
+        Build per-layer projected conditioning KV tensors for faster repeated sampling steps.
+        """
+        return [
+            block.attention.project_context_kv(
+                text_context=text_state,
+                speaker_context=speaker_state,
+                caption_context=caption_state,
+            )
+            for block in self.blocks
+        ]
+
+    @staticmethod
+    def masked_mean(state: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        mask_f = mask.unsqueeze(-1).to(dtype=state.dtype)
+        denom = mask_f.sum(dim=1).clamp_min(1.0)
+        return (state * mask_f).sum(dim=1) / denom
+
+    def predict_duration_log_frames(
+        self,
+        *,
+        text_state: torch.Tensor,
+        text_mask: torch.Tensor,
+        speaker_state: torch.Tensor | None,
+        speaker_mask: torch.Tensor | None,
+        duration_features: torch.Tensor,
+        has_speaker: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if self.duration_predictor is None:
+            raise RuntimeError("Duration predictor is disabled for this model.")
+        if duration_features.ndim != 2:
+            raise ValueError(
+                f"duration_features must have shape (B, D), got {tuple(duration_features.shape)}"
+            )
+        if duration_features.shape[1] != self.cfg.duration_aux_dim:
+            raise ValueError(
+                "duration_features dim mismatch: "
+                f"expected {self.cfg.duration_aux_dim}, got {duration_features.shape[1]}"
+            )
+
+        pred = self.duration_predictor(
+            text_state=text_state.detach(),
+            text_mask=text_mask,
+            aux_features=duration_features,
+            speaker_state=None if speaker_state is None else speaker_state.detach(),
+            speaker_mask=speaker_mask,
+            has_speaker=has_speaker,
+        )
+        return pred.float()
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    def as_dict(self) -> dict:
+        return asdict(self.cfg)

irodori_tts/rf.py

@@ -0,0 +1,542 @@
+from __future__ import annotations
+
+import torch
+
+from .model import TextToLatentRFDiT
+
+
+def _make_rng(seed: int, device: torch.device) -> tuple[torch.Generator, torch.device]:
+    # MPS generators are not available on some PyTorch builds; use CPU generator as fallback.
+    try:
+        return torch.Generator(device=device).manual_seed(seed), device
+    except RuntimeError:
+        return torch.Generator(device="cpu").manual_seed(seed), torch.device("cpu")
+
+
+def sample_logit_normal_t(
+    batch_size: int,
+    device: torch.device,
+    mean: float = 0.0,
+    std: float = 1.0,
+    t_min: float = 1e-3,
+    t_max: float = 0.999,
+) -> torch.Tensor:
+    z = torch.randn(batch_size, device=device) * std + mean
+    t = torch.sigmoid(z)
+    return t.clamp(min=t_min, max=t_max)
+
+
+def sample_stratified_logit_normal_t(
+    batch_size: int,
+    device: torch.device,
+    mean: float = 0.0,
+    std: float = 1.0,
+    t_min: float = 1e-3,
+    t_max: float = 0.999,
+) -> torch.Tensor:
+    """
+    Stratified sampling for logit-normal timesteps.
+
+    u ~ stratified U(0, 1), z = mean + std * Phi^{-1}(u), t = sigmoid(z)
+    """
+    if batch_size <= 0:
+        return torch.empty((0,), device=device)
+    u = (
+        torch.arange(batch_size, device=device, dtype=torch.float32)
+        + torch.rand(batch_size, device=device)
+    ) / float(batch_size)
+    u = u.clamp(1e-6, 1.0 - 1e-6)
+    # Phi^{-1}(u) = sqrt(2) * erfinv(2u - 1)
+    z = torch.erfinv(2.0 * u - 1.0) * (2.0**0.5)
+    z = z * std + mean
+    t = torch.sigmoid(z)
+    # Randomize assignment order so dataset ordering does not correlate with t bins.
+    t = t[torch.randperm(batch_size, device=device)]
+    return t.clamp(min=t_min, max=t_max)
+
+
+def rf_interpolate(x0: torch.Tensor, noise: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    # Straight line interpolation: x_t = (1-t) x0 + t z.
+    return (1.0 - t[:, None, None]) * x0 + t[:, None, None] * noise
+
+
+def rf_velocity_target(x0: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+    # For x_t = (1-t) x0 + t z, velocity is d/dt x_t = z - x0.
+    return noise - x0
+
+
+def rf_predict_x0(x_t: torch.Tensor, v_pred: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    # x_t = x0 + t * v  =>  x0 = x_t - t * v
+    return x_t - t[:, None, None] * v_pred
+
+
+def temporal_score_rescale(
+    v_pred: torch.Tensor,
+    x_t: torch.Tensor,
+    t: float | torch.Tensor,
+    rescale_k: float,
+    rescale_sigma: float,
+) -> torch.Tensor:
+    """
+    Temporal score rescaling from https://arxiv.org/pdf/2510.01184.
+    """
+    t_value = float(t.item()) if isinstance(t, torch.Tensor) else float(t)
+    if t_value >= 1.0:
+        return v_pred
+    one_minus_t = 1.0 - t_value
+    snr = (one_minus_t * one_minus_t) / (t_value * t_value)
+    sigma_sq = float(rescale_sigma) * float(rescale_sigma)
+    ratio = (snr * sigma_sq + 1.0) / (snr * sigma_sq / float(rescale_k) + 1.0)
+    return (ratio * (one_minus_t * v_pred + x_t) - x_t) / one_minus_t
+
+
+def scale_speaker_kv_cache(
+    context_kv_cache: list[tuple[torch.Tensor, ...]],
+    scale: float,
+    max_layers: int | None = None,
+) -> None:
+    """
+    In-place scaling of speaker K/V tensors in precomputed context cache.
+    """
+    if max_layers is None:
+        n_layers = len(context_kv_cache)
+    else:
+        n_layers = max(0, min(int(max_layers), len(context_kv_cache)))
+    for i in range(n_layers):
+        layer_kv = context_kv_cache[i]
+        if len(layer_kv) < 4:
+            raise ValueError(
+                f"Expected at least 4 tensors in context KV cache entry, got {len(layer_kv)}"
+            )
+        k_speaker = layer_kv[2]
+        v_speaker = layer_kv[3]
+        k_speaker.mul_(scale)
+        v_speaker.mul_(scale)
+
+
+@torch.inference_mode()
+def sample_euler_rf_cfg(
+    model: TextToLatentRFDiT,
+    text_input_ids: torch.Tensor,
+    text_mask: torch.Tensor,
+    ref_latent: torch.Tensor | None,
+    ref_mask: torch.Tensor | None,
+    sequence_length: int,
+    caption_input_ids: torch.Tensor | None = None,
+    caption_mask: torch.Tensor | None = None,
+    num_steps: int = 40,
+    cfg_scale_text: float = 3.0,
+    cfg_scale_caption: float = 3.0,
+    cfg_scale_speaker: float = 5.0,
+    cfg_guidance_mode: str = "independent",
+    cfg_min_t: float = 0.5,
+    cfg_max_t: float = 1.0,
+    seed: int = 0,
+    cfg_scale: float | None = None,
+    truncation_factor: float | None = None,
+    rescale_k: float | None = None,
+    rescale_sigma: float | None = None,
+    use_context_kv_cache: bool = True,
+    speaker_kv_scale: float | None = None,
+    speaker_kv_max_layers: int | None = None,
+    speaker_kv_min_t: float | None = None,
+) -> torch.Tensor:
+    """
+    Euler sampling over RF ODE with text/reference/caption conditioning CFG.
+
+    Returns:
+      latent sequence in patched space, shape (B, sequence_length, patched_latent_dim)
+    """
+    device = model.device
+    dtype = model.dtype
+    batch_size = text_input_ids.shape[0]
+    latent_dim = model.cfg.patched_latent_dim
+
+    rng, rng_device = _make_rng(seed=seed, device=device)
+    x_t = torch.randn(
+        (batch_size, sequence_length, latent_dim), device=rng_device, dtype=dtype, generator=rng
+    )
+    if rng_device != device:
+        x_t = x_t.to(device=device)
+    if truncation_factor is not None:
+        x_t = x_t * float(truncation_factor)
+
+    if cfg_scale is not None:
+        # Backward compatibility for old single-scale caller.
+        cfg_scale_text = float(cfg_scale)
+        cfg_scale_caption = float(cfg_scale)
+        cfg_scale_speaker = float(cfg_scale)
+    if not model.cfg.use_speaker_condition:
+        cfg_scale_speaker = 0.0
+        speaker_kv_scale = None
+
+    cfg_guidance_mode = str(cfg_guidance_mode).strip().lower()
+    if cfg_guidance_mode not in {"independent", "joint", "alternating"}:
+        raise ValueError(
+            f"Unsupported cfg_guidance_mode={cfg_guidance_mode!r}. "
+            "Expected one of: independent, joint, alternating."
+        )
+
+    init_scale = 0.999
+    t_schedule = torch.linspace(1.0, 0.0, num_steps + 1, device=device) * init_scale
+    use_independent_cfg = cfg_guidance_mode == "independent"
+    use_joint_cfg = cfg_guidance_mode == "joint"
+    use_alternating_cfg = cfg_guidance_mode == "alternating"
+
+    (
+        text_state_cond,
+        text_mask_cond,
+        speaker_state_cond,
+        speaker_mask_cond,
+        caption_state_cond,
+        caption_mask_cond,
+    ) = model.encode_conditions(
+        text_input_ids=text_input_ids,
+        text_mask=text_mask,
+        ref_latent=ref_latent,
+        ref_mask=ref_mask,
+        caption_input_ids=caption_input_ids,
+        caption_mask=caption_mask,
+    )
+    text_state_uncond = torch.zeros_like(text_state_cond)
+    text_mask_uncond = torch.zeros_like(text_mask_cond)
+    speaker_state_uncond = None
+    speaker_mask_uncond = None
+    if model.cfg.use_speaker_condition:
+        if speaker_state_cond is None or speaker_mask_cond is None:
+            raise RuntimeError(
+                "Speaker conditioning is enabled but encoded speaker state is missing."
+            )
+        speaker_state_uncond = torch.zeros_like(speaker_state_cond)
+        speaker_mask_uncond = torch.zeros_like(speaker_mask_cond)
+    caption_state_uncond = None
+    caption_mask_uncond = None
+    if model.cfg.use_caption_condition:
+        if caption_state_cond is None or caption_mask_cond is None:
+            raise RuntimeError(
+                "Caption conditioning is enabled but encoded caption state is missing."
+            )
+        caption_state_uncond = torch.zeros_like(caption_state_cond)
+        caption_mask_uncond = torch.zeros_like(caption_mask_cond)
+
+    has_text_cfg = cfg_scale_text > 0
+    has_caption_cfg = (
+        model.cfg.use_caption_condition
+        and cfg_scale_caption > 0
+        and caption_mask_cond is not None
+        and bool(caption_mask_cond.any().item())
+    )
+    has_speaker_cfg = cfg_scale_speaker > 0
+
+    def _bundle(
+        *,
+        text_state: torch.Tensor,
+        text_mask_val: torch.Tensor,
+        speaker_state: torch.Tensor | None,
+        speaker_mask_val: torch.Tensor | None,
+        caption_state: torch.Tensor | None,
+        caption_mask_val: torch.Tensor | None,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+    ]:
+        return (
+            text_state,
+            text_mask_val,
+            speaker_state,
+            speaker_mask_val,
+            caption_state,
+            caption_mask_val,
+        )
+
+    cond_bundle = _bundle(
+        text_state=text_state_cond,
+        text_mask_val=text_mask_cond,
+        speaker_state=speaker_state_cond,
+        speaker_mask_val=speaker_mask_cond,
+        caption_state=caption_state_cond,
+        caption_mask_val=caption_mask_cond,
+    )
+    enabled_cfg_names: list[str] = []
+    cfg_scales: dict[str, float] = {}
+    if has_text_cfg:
+        enabled_cfg_names.append("text")
+        cfg_scales["text"] = float(cfg_scale_text)
+    if has_speaker_cfg:
+        enabled_cfg_names.append("speaker")
+        cfg_scales["speaker"] = float(cfg_scale_speaker)
+    if has_caption_cfg:
+        enabled_cfg_names.append("caption")
+        cfg_scales["caption"] = float(cfg_scale_caption)
+
+    independent_bundles = [cond_bundle]
+    independent_names = ["cond"]
+    if use_independent_cfg:
+        for name in enabled_cfg_names:
+            independent_names.append(name)
+            independent_bundles.append(
+                _bundle(
+                    text_state=text_state_uncond if name == "text" else text_state_cond,
+                    text_mask_val=text_mask_uncond if name == "text" else text_mask_cond,
+                    speaker_state=(
+                        speaker_state_uncond if name == "speaker" else speaker_state_cond
+                    ),
+                    speaker_mask_val=(
+                        speaker_mask_uncond if name == "speaker" else speaker_mask_cond
+                    ),
+                    caption_state=(
+                        caption_state_uncond if name == "caption" else caption_state_cond
+                    ),
+                    caption_mask_val=(
+                        caption_mask_uncond if name == "caption" else caption_mask_cond
+                    ),
+                )
+            )
+    cfg_batch_mult = len(independent_bundles)
+
+    def _cat_optional_tensors(values: list[torch.Tensor | None]) -> torch.Tensor | None:
+        present = [value for value in values if value is not None]
+        if not present:
+            return None
+        if len(present) != len(values):
+            raise ValueError("Cannot concatenate optional condition tensors with mixed presence.")
+        return torch.cat(present, dim=0)
+
+    independent_text_state = torch.cat([bundle[0] for bundle in independent_bundles], dim=0)
+    independent_text_mask = torch.cat([bundle[1] for bundle in independent_bundles], dim=0)
+    independent_speaker_state = _cat_optional_tensors([bundle[2] for bundle in independent_bundles])
+    independent_speaker_mask = _cat_optional_tensors([bundle[3] for bundle in independent_bundles])
+    independent_caption_state = _cat_optional_tensors([bundle[4] for bundle in independent_bundles])
+    independent_caption_mask = _cat_optional_tensors([bundle[5] for bundle in independent_bundles])
+
+    joint_uncond_bundle = _bundle(
+        text_state=text_state_uncond,
+        text_mask_val=text_mask_uncond,
+        speaker_state=speaker_state_uncond,
+        speaker_mask_val=speaker_mask_uncond,
+        caption_state=caption_state_uncond,
+        caption_mask_val=caption_mask_uncond,
+    )
+
+    alternating_bundles: dict[
+        str,
+        tuple[
+            torch.Tensor,
+            torch.Tensor,
+            torch.Tensor | None,
+            torch.Tensor | None,
+            torch.Tensor | None,
+            torch.Tensor | None,
+        ],
+    ] = {
+        "text": _bundle(
+            text_state=text_state_uncond,
+            text_mask_val=text_mask_uncond,
+            speaker_state=speaker_state_cond,
+            speaker_mask_val=speaker_mask_cond,
+            caption_state=caption_state_cond,
+            caption_mask_val=caption_mask_cond,
+        ),
+        "caption": _bundle(
+            text_state=text_state_cond,
+            text_mask_val=text_mask_cond,
+            speaker_state=speaker_state_cond,
+            speaker_mask_val=speaker_mask_cond,
+            caption_state=caption_state_uncond,
+            caption_mask_val=caption_mask_uncond,
+        ),
+    }
+    if has_speaker_cfg:
+        alternating_bundles["speaker"] = _bundle(
+            text_state=text_state_cond,
+            text_mask_val=text_mask_cond,
+            speaker_state=speaker_state_uncond,
+            speaker_mask_val=speaker_mask_uncond,
+            caption_state=caption_state_cond,
+            caption_mask_val=caption_mask_cond,
+        )
+
+    # Force-speaker scaling operates on projected speaker K/V, so it requires context KV caches.
+    effective_use_context_kv_cache = bool(use_context_kv_cache or (speaker_kv_scale is not None))
+
+    context_kv_cond = None
+    context_kv_cfg = None
+    context_kv_joint_uncond = None
+    context_kv_alternating: dict[str, list[tuple[torch.Tensor, ...]]] = {}
+    if effective_use_context_kv_cache:
+        context_kv_cond = model.build_context_kv_cache(
+            text_state=text_state_cond,
+            speaker_state=speaker_state_cond,
+            caption_state=caption_state_cond,
+        )
+        if use_independent_cfg and cfg_batch_mult > 1:
+            context_kv_cfg = model.build_context_kv_cache(
+                text_state=independent_text_state,
+                speaker_state=independent_speaker_state,
+                caption_state=independent_caption_state,
+            )
+        elif use_joint_cfg:
+            if enabled_cfg_names:
+                context_kv_joint_uncond = model.build_context_kv_cache(
+                    text_state=joint_uncond_bundle[0],
+                    speaker_state=joint_uncond_bundle[2],
+                    caption_state=joint_uncond_bundle[4],
+                )
+        elif use_alternating_cfg:
+            for name in enabled_cfg_names:
+                bundle = alternating_bundles[name]
+                context_kv_alternating[name] = model.build_context_kv_cache(
+                    text_state=bundle[0],
+                    speaker_state=bundle[2],
+                    caption_state=bundle[4],
+                )
+    if speaker_kv_scale is not None:
+        scale_speaker_kv_cache(
+            context_kv_cache=context_kv_cond,
+            scale=float(speaker_kv_scale),
+            max_layers=speaker_kv_max_layers,
+        )
+        if context_kv_cfg is not None:
+            scale_speaker_kv_cache(
+                context_kv_cache=context_kv_cfg,
+                scale=float(speaker_kv_scale),
+                max_layers=speaker_kv_max_layers,
+            )
+        for cache in context_kv_alternating.values():
+            scale_speaker_kv_cache(
+                context_kv_cache=cache,
+                scale=float(speaker_kv_scale),
+                max_layers=speaker_kv_max_layers,
+            )
+    speaker_kv_active = speaker_kv_scale is not None
+
+    for i in range(num_steps):
+        t = t_schedule[i]
+        t_next = t_schedule[i + 1]
+        tt = torch.full((batch_size,), t, device=device, dtype=dtype)
+
+        use_cfg = bool(enabled_cfg_names) and (cfg_min_t <= t.item() <= cfg_max_t)
+        if use_cfg:
+            if use_independent_cfg:
+                x_t_cfg = torch.cat([x_t] * cfg_batch_mult, dim=0).to(dtype)
+                tt_cfg = tt.repeat(cfg_batch_mult)
+                v_out = model.forward_with_encoded_conditions(
+                    x_t=x_t_cfg,
+                    t=tt_cfg,
+                    text_state=independent_text_state,
+                    text_mask=independent_text_mask,
+                    speaker_state=independent_speaker_state,
+                    speaker_mask=independent_speaker_mask,
+                    caption_state=independent_caption_state,
+                    caption_mask=independent_caption_mask,
+                    context_kv_cache=context_kv_cfg,
+                )
+                chunks = v_out.chunk(cfg_batch_mult, dim=0)
+                v = chunks[0]
+                for name, chunk in zip(independent_names[1:], chunks[1:], strict=True):
+                    v = v + cfg_scales[name] * (chunks[0] - chunk)
+            else:
+                v_cond = model.forward_with_encoded_conditions(
+                    x_t=x_t.to(dtype),
+                    t=tt,
+                    text_state=text_state_cond,
+                    text_mask=text_mask_cond,
+                    speaker_state=speaker_state_cond,
+                    speaker_mask=speaker_mask_cond,
+                    caption_state=caption_state_cond,
+                    caption_mask=caption_mask_cond,
+                    context_kv_cache=context_kv_cond,
+                )
+                if use_joint_cfg:
+                    if len(enabled_cfg_names) > 1:
+                        joint_scales = [cfg_scales[name] for name in enabled_cfg_names]
+                        if max(joint_scales) - min(joint_scales) > 1e-6:
+                            raise ValueError(
+                                "cfg_guidance_mode='joint' expects equal enabled guidance scales; "
+                                "set matching text/speaker/caption scales or use --cfg-scale."
+                            )
+                    joint_scale = cfg_scales[enabled_cfg_names[0]]
+                    v_uncond_joint = model.forward_with_encoded_conditions(
+                        x_t=x_t.to(dtype),
+                        t=tt,
+                        text_state=joint_uncond_bundle[0],
+                        text_mask=joint_uncond_bundle[1],
+                        speaker_state=joint_uncond_bundle[2],
+                        speaker_mask=joint_uncond_bundle[3],
+                        caption_state=joint_uncond_bundle[4],
+                        caption_mask=joint_uncond_bundle[5],
+                        context_kv_cache=context_kv_joint_uncond,
+                    )
+                    v = v_cond + joint_scale * (v_cond - v_uncond_joint)
+                elif use_alternating_cfg:
+                    alt_name = enabled_cfg_names[i % len(enabled_cfg_names)]
+                    alt_bundle = alternating_bundles[alt_name]
+                    v_uncond_alt = model.forward_with_encoded_conditions(
+                        x_t=x_t.to(dtype),
+                        t=tt,
+                        text_state=alt_bundle[0],
+                        text_mask=alt_bundle[1],
+                        speaker_state=alt_bundle[2],
+                        speaker_mask=alt_bundle[3],
+                        caption_state=alt_bundle[4],
+                        caption_mask=alt_bundle[5],
+                        context_kv_cache=context_kv_alternating.get(alt_name),
+                    )
+                    v = v_cond + cfg_scales[alt_name] * (v_cond - v_uncond_alt)
+                else:
+                    raise RuntimeError(f"Unexpected cfg_guidance_mode: {cfg_guidance_mode}")
+        else:
+            v = model.forward_with_encoded_conditions(
+                x_t=x_t.to(dtype),
+                t=tt,
+                text_state=text_state_cond,
+                text_mask=text_mask_cond,
+                speaker_state=speaker_state_cond,
+                speaker_mask=speaker_mask_cond,
+                caption_state=caption_state_cond,
+                caption_mask=caption_mask_cond,
+                context_kv_cache=context_kv_cond,
+            )
+
+        if rescale_k is not None and rescale_sigma is not None:
+            v = temporal_score_rescale(
+                v_pred=v,
+                x_t=x_t,
+                t=t,
+                rescale_k=float(rescale_k),
+                rescale_sigma=float(rescale_sigma),
+            )
+
+        if (
+            speaker_kv_active
+            and speaker_kv_min_t is not None
+            and (t_next < speaker_kv_min_t)
+            and (t >= speaker_kv_min_t)
+        ):
+            inv_scale = 1.0 / float(speaker_kv_scale)
+            scale_speaker_kv_cache(
+                context_kv_cache=context_kv_cond,
+                scale=inv_scale,
+                max_layers=speaker_kv_max_layers,
+            )
+            if context_kv_cfg is not None:
+                scale_speaker_kv_cache(
+                    context_kv_cache=context_kv_cfg,
+                    scale=inv_scale,
+                    max_layers=speaker_kv_max_layers,
+                )
+            for cache in context_kv_alternating.values():
+                scale_speaker_kv_cache(
+                    context_kv_cache=cache,
+                    scale=inv_scale,
+                    max_layers=speaker_kv_max_layers,
+                )
+            speaker_kv_active = False
+
+        x_t = x_t + v * (t_next - t)
+
+    return x_t

irodori_tts/text_normalization.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import re
+import unicodedata
+
+SIMPLE_REPLACE_MAP: dict[str, str] = {
+    "\t": "",
+    "[n]": "",
+    r"\[n\]": "",
+    "　": "",
+    "？": "?",
+    "！": "!",
+    "♥": "♡",
+    "●": "○",
+    "◯": "○",
+    "〇": "○",
+}
+
+REGEX_REPLACE_MAP = {
+    re.compile(r"[;▼♀♂《》≪≫①②③④⑤⑥]"): "",
+    re.compile(r"[\u02d7\u2010-\u2015\u2043\u2212\u23af\u23e4\u2500\u2501\u2e3a\u2e3b]"): "",
+    re.compile(r"[\uff5e\u301C]"): "ー",
+    re.compile(r"…{3,}"): "……",
+}
+
+
+def strip_outer_brackets(text: str) -> str:
+    pairs = {"「": "」", "『": "』", "（": "）", "【": "】", "(": ")"}
+
+    while True:
+        if len(text) < 2:
+            break
+
+        start_char = text[0]
+        end_char = text[-1]
+
+        if start_char in pairs and pairs[start_char] == end_char:
+            depth = 0
+            is_enclosing_all = True
+
+            for i, char in enumerate(text):
+                if char == start_char:
+                    depth += 1
+                elif char == end_char:
+                    depth -= 1
+
+                if depth == 0 and i < len(text) - 1:
+                    is_enclosing_all = False
+                    break
+
+            if is_enclosing_all and depth == 0:
+                text = text[1:-1]
+                continue
+
+        break
+
+    return text
+
+
+def normalize_text(text: str) -> str:
+    for old, new in SIMPLE_REPLACE_MAP.items():
+        text = text.replace(old, new)
+
+    for pattern, replacement in REGEX_REPLACE_MAP.items():
+        text = pattern.sub(replacement, text)
+
+    text = strip_outer_brackets(text)
+
+    text = unicodedata.normalize("NFKC", text)
+
+    text = text.replace("...", "…")
+    text = text.replace("..", "…")
+
+    return text

irodori_tts/tokenizer.py

@@ -0,0 +1,144 @@
+from collections.abc import Iterable
+
+import torch
+
+
+class ByteTokenizer:
+    """Simple byte-level tokenizer for text-to-speech."""
+
+    def __init__(self, bos_token: int = 256) -> None:
+        if bos_token < 0:
+            raise ValueError(f"bos_token must be >= 0, got {bos_token}")
+        self.bos_token = int(bos_token)
+
+    @classmethod
+    def for_vocab_size(cls, text_vocab_size: int) -> "ByteTokenizer":
+        if text_vocab_size < 256:
+            raise ValueError(
+                f"text_vocab_size must be >= 256 for byte-level tokenization, got {text_vocab_size}"
+            )
+        # Reserve a dedicated BOS token outside UTF-8 byte range when possible.
+        if text_vocab_size == 256:
+            return cls(bos_token=0)
+        return cls(bos_token=text_vocab_size - 1)
+
+    def encode(self, text: str, add_bos: bool = True) -> torch.Tensor:
+        tokens = list(text.encode("utf-8"))
+        if add_bos:
+            tokens.insert(0, self.bos_token)
+        return torch.tensor(tokens, dtype=torch.long)
+
+    def batch_encode(
+        self,
+        texts: Iterable[str],
+        max_length: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        encoded = [self.encode(t) for t in texts]
+        if max_length is None:
+            max_length = max(x.numel() for x in encoded)
+
+        batch = torch.zeros((len(encoded), max_length), dtype=torch.long)
+        mask = torch.zeros((len(encoded), max_length), dtype=torch.bool)
+        for i, seq in enumerate(encoded):
+            n = min(max_length, seq.numel())
+            batch[i, :n] = seq[:n]
+            mask[i, :n] = True
+        return batch, mask
+
+
+class PretrainedTextTokenizer:
+    """
+    Hugging Face tokenizer wrapper for text conditioning.
+    - right-padding for stable positional behavior
+    - optional explicit BOS prepend
+    """
+
+    def __init__(self, tokenizer, add_bos: bool = True) -> None:
+        self.tokenizer = tokenizer
+        self.add_bos = bool(add_bos)
+        # TTS collator uses fixed-length right-padding; enforce this regardless of pretrained defaults.
+        self.tokenizer.padding_side = "right"
+
+        if self.tokenizer.pad_token_id is None:
+            if self.tokenizer.eos_token_id is not None and self.tokenizer.eos_token is not None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            else:
+                raise ValueError(
+                    "Tokenizer has no pad_token_id (and no eos_token fallback). "
+                    "Set a pad token before training/inference."
+                )
+
+        if self.add_bos and self.tokenizer.bos_token_id is None:
+            raise ValueError("Tokenizer has no bos_token_id but add_bos=True.")
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id: str,
+        add_bos: bool = True,
+        local_files_only: bool = False,
+    ) -> "PretrainedTextTokenizer":
+        try:
+            from transformers import AutoTokenizer
+        except ImportError as exc:
+            raise RuntimeError(
+                "transformers is required for pretrained text tokenization. "
+                "Install with `pip install transformers sentencepiece`."
+            ) from exc
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            repo_id,
+            use_fast=True,
+            trust_remote_code=False,
+            local_files_only=local_files_only,
+        )
+        return cls(tokenizer=tokenizer, add_bos=add_bos)
+
+    @property
+    def vocab_size(self) -> int:
+        return int(len(self.tokenizer))
+
+    @property
+    def bos_token_id(self) -> int | None:
+        return self.tokenizer.bos_token_id
+
+    @property
+    def pad_token_id(self) -> int:
+        pad_id = self.tokenizer.pad_token_id
+        if pad_id is None:
+            raise RuntimeError("pad_token_id is unexpectedly None.")
+        return int(pad_id)
+
+    def encode(self, text: str, add_bos: bool | None = None) -> torch.Tensor:
+        token_ids = self.tokenizer.encode(text, add_special_tokens=False)
+        use_bos = self.add_bos if add_bos is None else bool(add_bos)
+        if use_bos:
+            bos_id = self.bos_token_id
+            if bos_id is None:
+                raise ValueError("Tokenizer has no bos_token_id but BOS prepend was requested.")
+            token_ids.insert(0, int(bos_id))
+        return torch.tensor(token_ids, dtype=torch.long)
+
+    def batch_encode(
+        self,
+        texts: Iterable[str],
+        max_length: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        encoded = [self.encode(t) for t in texts]
+        if max_length is None:
+            max_length = max(max(x.numel(), 1) for x in encoded)
+        if max_length <= 0:
+            raise ValueError(f"max_length must be > 0, got {max_length}")
+
+        batch = torch.full(
+            (len(encoded), max_length),
+            fill_value=self.pad_token_id,
+            dtype=torch.long,
+        )
+        mask = torch.zeros((len(encoded), max_length), dtype=torch.bool)
+        for i, seq in enumerate(encoded):
+            n = min(max_length, seq.numel())
+            if n > 0:
+                batch[i, :n] = seq[:n]
+                mask[i, :n] = True
+        return batch, mask

irodori_tts/watermark.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterable
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+IRODORI_WATERMARK_PAYLOAD = (73, 82, 68, 84, 83)  # "IRDTS"
+
+
+def _as_single_channel_vector(audio: torch.Tensor) -> torch.Tensor | None:
+    squeezed = audio.detach().float().squeeze()
+    if squeezed.ndim == 0 or squeezed.numel() == 0:
+        return None
+    if squeezed.ndim == 1:
+        return squeezed
+    return squeezed.reshape(-1)
+
+
+def _match_original_rank(audio: torch.Tensor, *, reference: torch.Tensor) -> torch.Tensor:
+    if reference.ndim == 2:
+        return audio.reshape(1, -1)
+    return audio.reshape(-1)
+
+
+class SilentCipherWatermarker:
+    def __init__(self, *, device: str, model_type: str = "44.1k") -> None:
+        self.model = self._load_backend(device=device, model_type=model_type)
+
+    @staticmethod
+    def _load_backend(*, device: str, model_type: str):
+        try:
+            import silentcipher
+        except ImportError:
+            logger.warning(
+                "SilentCipher package is unavailable; generated audio will not be watermarked."
+            )
+            return None
+
+        try:
+            return silentcipher.get_model(model_type=model_type, device=device)
+        except Exception as exc:
+            logger.warning(
+                "SilentCipher model could not be loaded (%s); generated audio will not be "
+                "watermarked.",
+                exc,
+            )
+            return None
+
+    @property
+    def ready(self) -> bool:
+        return self.model is not None
+
+    def encode_one(
+        self,
+        audio: torch.Tensor,
+        *,
+        sample_rate: int,
+        payload: Iterable[int] = IRODORI_WATERMARK_PAYLOAD,
+    ) -> torch.Tensor:
+        if self.model is None:
+            return audio
+
+        vector = _as_single_channel_vector(audio)
+        if vector is None:
+            return audio
+
+        encoded, _ = self.model.encode_wav(
+            vector.to(self.model.device),
+            int(sample_rate),
+            list(payload),
+            calc_sdr=False,
+        )
+        encoded_audio = torch.as_tensor(encoded, dtype=torch.float32, device="cpu")
+        return _match_original_rank(encoded_audio, reference=audio)
+
+    def encode_batch(self, audios: list[torch.Tensor], *, sample_rate: int) -> list[torch.Tensor]:
+        if self.model is None:
+            return audios
+        return [self.encode_one(audio, sample_rate=sample_rate) for audio in audios]

requirements.txt

@@ -0,0 +1,12 @@
+torch>=2.5.1
+torchaudio>=2.5.1
+transformers<5
+sentencepiece>=0.1.99,<0.2
+safetensors>=0.7.0
+soundfile>=0.12.0
+huggingface-hub>=0.34.0,<1.0
+gradio>=5.0.0
+numpy
+dacvae @ git+https://github.com/facebookresearch/dacvae
+torchcodec>=0.10.0
+silentcipher @ git+https://github.com/SesameAILabs/silentcipher.git@d46d7d0893a583d8968ab3a6626e2289faec9152