Remove Faster Irodori TTS runtime and related files

Dockerfile

@@ -34,7 +34,6 @@ RUN python -m pip install --upgrade pip setuptools wheel \
 
 COPY app.py .
 COPY faster_qwen3_tts ./faster_qwen3_tts
-COPY faster_irodori_tts ./faster_irodori_tts
 COPY qwen_tts ./qwen_tts
 COPY irodori_tts ./irodori_tts
 

app.py

@@ -22,15 +22,14 @@ from huggingface_hub import hf_hub_download
 #from huggingface_hub import snapshot_download
 #from qwen_tts import Qwen3TTSModel
 from faster_qwen3_tts import FasterQwen3TTS
-#from irodori_tts.inference_runtime import InferenceRuntime, RuntimeKey, SamplingRequest
-from faster_irodori_tts import FasterIrodoriTTSRuntime, RuntimeKey, SamplingRequest
+from irodori_tts.inference_runtime import InferenceRuntime, RuntimeKey, SamplingRequest
 
 
 load_dotenv(verbose=False)
 
 #TTS_MODEL = Qwen3TTSModel.from_pretrained(snapshot_download('Qwen/Qwen3-TTS-12Hz-1.7B-Base', token=os.environ['HF_TOKEN']), device_map=torch.device('cuda' if torch.cuda.is_available() else 'cpu'), dtype=torch.bfloat16, token=os.environ['HF_TOKEN'], attn_implementation='kernels-community/flash-attn3')
 TTS_MODEL = FasterQwen3TTS.from_pretrained('Qwen/Qwen3-TTS-12Hz-1.7B-Base')
-IRODORI_TTS_RUNTIME: Optional[FasterIrodoriTTSRuntime] = None
+IRODORI_TTS_RUNTIME: Optional[InferenceRuntime] = None
 WHISPER_MODEL = whisper.load_model('turbo', device='cpu', download_root=os.environ.get('WHISPER_CACHE_DIR'))
 REFERENCE_AUDIO_TRANSCRIPTION_CACHE: dict[str, tuple[float, str, str]] = {}
 REFERENCE_AUDIO_TRANSCRIPTION_CACHE_LOCK = threading.Lock()
@@ -235,7 +234,7 @@ def generate_voice_clone(model: str | None, input_text: str, language: str | Non
         if IRODORI_TTS_RUNTIME is None:
             device = 'cuda' if torch.cuda.is_available() else 'cpu'
             precision = 'bf16' if device == 'cuda' else 'fp32'
-            IRODORI_TTS_RUNTIME = FasterIrodoriTTSRuntime.from_key(RuntimeKey(
+            IRODORI_TTS_RUNTIME = InferenceRuntime.from_key(RuntimeKey(
                 checkpoint=hf_hub_download(repo_id='Aratako/Irodori-TTS-500M-v2', filename='model.safetensors'),
                 model_device=device,
                 codec_repo='Aratako/Semantic-DACVAE-Japanese-32dim',
@@ -245,6 +244,9 @@ def generate_voice_clone(model: str | None, input_text: str, language: str | Non
                 enable_watermark=False,
             ))
 
+        if sample_rate != 48000:
+            reference_audio = (_resample(reference_audio[0], sample_rate, 48000), 48000)
+
         result = IRODORI_TTS_RUNTIME.synthesize(SamplingRequest(
             text=input_text,
             ref_wav=reference_audio,

faster_irodori_tts/__init__.py

@@ -1,17 +0,0 @@
-"""CUDA Graph accelerated runtime helpers for Irodori-TTS."""
-
-from .runtime import (
-    FasterInferenceRuntime,
-    FasterIrodoriTTSRuntime,
-    RuntimeKey,
-    SamplingRequest,
-    SamplingResult,
-)
-
-__all__ = [
-    "FasterInferenceRuntime",
-    "FasterIrodoriTTSRuntime",
-    "RuntimeKey",
-    "SamplingRequest",
-    "SamplingResult",
-]

faster_irodori_tts/rf_graph.py

@@ -1,511 +0,0 @@
-from __future__ import annotations
-
-from collections import OrderedDict
-from dataclasses import dataclass
-
-import torch
-
-from irodori_tts.rf import sample_euler_rf_cfg
-
-
-@dataclass(frozen=True)
-class RFGraphSignature:
-    batch_size: int
-    sequence_length: int
-    latent_dim: int
-    text_len: int
-    speaker_len: int
-    num_steps: int
-    cfg_scale_text: float
-    cfg_scale_speaker: float
-    cfg_min_t: float
-    cfg_max_t: float
-    dtype: str
-    device: str
-
-
-@dataclass
-class RFGraphSampleResult:
-    latent: torch.Tensor
-    graph_used: bool
-    fallback_reason: str | None = None
-
-
-def _device_key(device: torch.device) -> str:
-    index = 0 if device.index is None else int(device.index)
-    return f"{device.type}:{index}"
-
-
-def _pad_reference_to_bucket(
-    ref_latent: torch.Tensor,
-    ref_mask: torch.Tensor,
-    *,
-    speaker_patch_size: int,
-    bucket_multiple: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    if bucket_multiple <= 1:
-        return ref_latent, ref_mask
-
-    patch = max(1, int(speaker_patch_size))
-    current = int(ref_latent.shape[1])
-    after_patch = max(1, (current + patch - 1) // patch)
-    bucketed_after_patch = (
-        (after_patch + int(bucket_multiple) - 1) // int(bucket_multiple)
-    ) * int(bucket_multiple)
-    target = bucketed_after_patch * patch
-    if target <= current:
-        return ref_latent, ref_mask
-
-    pad_len = target - current
-    latent_pad = torch.zeros(
-        (ref_latent.shape[0], pad_len, ref_latent.shape[2]),
-        device=ref_latent.device,
-        dtype=ref_latent.dtype,
-    )
-    mask_pad = torch.zeros(
-        (ref_mask.shape[0], pad_len),
-        device=ref_mask.device,
-        dtype=ref_mask.dtype,
-    )
-    return torch.cat([ref_latent, latent_pad], dim=1), torch.cat([ref_mask, mask_pad], dim=1)
-
-
-def _copy_context_kv(
-    dst: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],
-    src: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],
-) -> None:
-    if len(dst) != len(src):
-        raise ValueError(f"Context KV layer count mismatch: graph={len(dst)} input={len(src)}")
-    for dst_layer, src_layer in zip(dst, src):
-        for dst_tensor, src_tensor in zip(dst_layer, src_layer):
-            if tuple(dst_tensor.shape) != tuple(src_tensor.shape):
-                raise ValueError(
-                    "Context KV shape mismatch: "
-                    f"graph={tuple(dst_tensor.shape)} input={tuple(src_tensor.shape)}"
-                )
-            dst_tensor.copy_(src_tensor)
-
-
-class IrodoriRFGraph:
-    """Captured RF Euler sampler for one fixed Irodori-TTS shape/configuration."""
-
-    def __init__(
-        self,
-        model,
-        signature: RFGraphSignature,
-        *,
-        num_warmup: int = 2,
-    ) -> None:
-        self.model = model
-        self.signature = signature
-        self.device = model.device
-        self.device_index = 0 if self.device.index is None else int(self.device.index)
-        self.dtype = model.dtype
-        self.num_warmup = int(num_warmup)
-        self.cfg_batch_mult = 3
-
-        bsz = signature.batch_size
-        seq_len = signature.sequence_length
-        latent_dim = signature.latent_dim
-        cfg_bsz = bsz * self.cfg_batch_mult
-        text_dim = model.cfg.text_dim
-        speaker_dim = model.cfg.speaker_dim
-
-        self.x_buf = torch.zeros((bsz, seq_len, latent_dim), device=self.device, dtype=self.dtype)
-        self.x_cfg_buf = torch.zeros(
-            (cfg_bsz, seq_len, latent_dim), device=self.device, dtype=self.dtype
-        )
-        self.v_buf = torch.zeros_like(self.x_buf)
-        self.latent_mask = torch.ones((bsz, seq_len), device=self.device, dtype=torch.bool)
-        self.latent_mask_cfg = torch.ones((cfg_bsz, seq_len), device=self.device, dtype=torch.bool)
-
-        self.text_state_cond = torch.zeros(
-            (bsz, signature.text_len, text_dim), device=self.device, dtype=self.dtype
-        )
-        self.text_mask_cond = torch.zeros(
-            (bsz, signature.text_len), device=self.device, dtype=torch.bool
-        )
-        self.speaker_state_cond = torch.zeros(
-            (bsz, signature.speaker_len, speaker_dim), device=self.device, dtype=self.dtype
-        )
-        self.speaker_mask_cond = torch.zeros(
-            (bsz, signature.speaker_len), device=self.device, dtype=torch.bool
-        )
-
-        self.text_state_cfg = torch.zeros(
-            (cfg_bsz, signature.text_len, text_dim), device=self.device, dtype=self.dtype
-        )
-        self.text_mask_cfg = torch.zeros(
-            (cfg_bsz, signature.text_len), device=self.device, dtype=torch.bool
-        )
-        self.speaker_state_cfg = torch.zeros(
-            (cfg_bsz, signature.speaker_len, speaker_dim), device=self.device, dtype=self.dtype
-        )
-        self.speaker_mask_cfg = torch.zeros(
-            (cfg_bsz, signature.speaker_len), device=self.device, dtype=torch.bool
-        )
-
-        self.context_kv_cond = self._make_context_kv_buffers(bsz)
-        self.context_kv_cfg = self._make_context_kv_buffers(cfg_bsz)
-
-        init_scale = 0.999
-        t_schedule = torch.linspace(
-            1.0,
-            0.0,
-            signature.num_steps + 1,
-            device=self.device,
-            dtype=torch.float32,
-        ) * init_scale
-        self.t_cond = [torch.full((bsz,), t_schedule[i], device=self.device, dtype=self.dtype)
-                       for i in range(signature.num_steps)]
-        self.t_cfg = [self.t_cond[i].repeat(self.cfg_batch_mult)
-                      for i in range(signature.num_steps)]
-        self.deltas = [
-            float((t_schedule[i + 1] - t_schedule[i]).detach().cpu())
-            for i in range(signature.num_steps)
-        ]
-        self.use_cfg = [
-            bool(signature.cfg_min_t <= float(t_schedule[i].detach().cpu()) <= signature.cfg_max_t)
-            for i in range(signature.num_steps)
-        ]
-
-        self.graph: torch.cuda.CUDAGraph | None = None
-        self.captured = False
-
-    def _make_context_kv_buffers(
-        self,
-        batch_size: int,
-    ) -> list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]:
-        buffers: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]] = []
-        for block in self.model.blocks:
-            attn = block.attention
-            k_text = torch.zeros(
-                (batch_size, self.signature.text_len, attn.heads, attn.head_dim),
-                device=self.device,
-                dtype=self.dtype,
-            )
-            v_text = torch.zeros_like(k_text)
-            k_speaker = torch.zeros(
-                (batch_size, self.signature.speaker_len, attn.heads, attn.head_dim),
-                device=self.device,
-                dtype=self.dtype,
-            )
-            v_speaker = torch.zeros_like(k_speaker)
-            buffers.append((k_text, v_text, k_speaker, v_speaker))
-        return buffers
-
-    def _copy_cfg_x(self) -> None:
-        bsz = self.signature.batch_size
-        self.x_cfg_buf[:bsz].copy_(self.x_buf)
-        self.x_cfg_buf[bsz : 2 * bsz].copy_(self.x_buf)
-        self.x_cfg_buf[2 * bsz : 3 * bsz].copy_(self.x_buf)
-
-    def _full_loop(self) -> None:
-        bsz = self.signature.batch_size
-        scale_text = float(self.signature.cfg_scale_text)
-        scale_speaker = float(self.signature.cfg_scale_speaker)
-        cond_weight = 1.0 + scale_text + scale_speaker
-
-        for i in range(self.signature.num_steps):
-            if self.use_cfg[i]:
-                self._copy_cfg_x()
-                v_out = self.model.forward_with_encoded_conditions(
-                    x_t=self.x_cfg_buf,
-                    t=self.t_cfg[i],
-                    text_state=self.text_state_cfg,
-                    text_mask=self.text_mask_cfg,
-                    speaker_state=self.speaker_state_cfg,
-                    speaker_mask=self.speaker_mask_cfg,
-                    latent_mask=self.latent_mask_cfg,
-                    context_kv_cache=self.context_kv_cfg,
-                )
-                v_cond = v_out[:bsz]
-                v_uncond_text = v_out[bsz : 2 * bsz]
-                v_uncond_speaker = v_out[2 * bsz : 3 * bsz]
-                self.v_buf.copy_(v_cond)
-                self.v_buf.mul_(cond_weight)
-                self.v_buf.add_(v_uncond_text, alpha=-scale_text)
-                self.v_buf.add_(v_uncond_speaker, alpha=-scale_speaker)
-            else:
-                v_out = self.model.forward_with_encoded_conditions(
-                    x_t=self.x_buf,
-                    t=self.t_cond[i],
-                    text_state=self.text_state_cond,
-                    text_mask=self.text_mask_cond,
-                    speaker_state=self.speaker_state_cond,
-                    speaker_mask=self.speaker_mask_cond,
-                    latent_mask=self.latent_mask,
-                    context_kv_cache=self.context_kv_cond,
-                )
-                self.v_buf.copy_(v_out)
-
-            self.x_buf.add_(self.v_buf, alpha=self.deltas[i])
-
-    @torch.inference_mode()
-    def capture(self) -> None:
-        if self.captured:
-            return
-
-        # Populate module-side RoPE caches and allocator pools before capture.
-        for _ in range(max(1, self.num_warmup)):
-            self._full_loop()
-        torch.cuda.synchronize(self.device)
-
-        with torch.cuda.device(self.device_index):
-            stream = torch.cuda.Stream()
-            stream.wait_stream(torch.cuda.current_stream())
-            with torch.cuda.stream(stream):
-                for _ in range(max(1, self.num_warmup)):
-                    self._full_loop()
-                torch.cuda.synchronize(self.device)
-
-                self.graph = torch.cuda.CUDAGraph()
-                with torch.cuda.graph(self.graph):
-                    self._full_loop()
-
-            torch.cuda.current_stream().wait_stream(stream)
-        torch.cuda.synchronize(self.device)
-        self.captured = True
-
-    @torch.inference_mode()
-    def run(
-        self,
-        *,
-        x_t: torch.Tensor,
-        text_state_cond: torch.Tensor,
-        text_mask_cond: torch.Tensor,
-        speaker_state_cond: torch.Tensor,
-        speaker_mask_cond: torch.Tensor,
-        text_state_cfg: torch.Tensor,
-        text_mask_cfg: torch.Tensor,
-        speaker_state_cfg: torch.Tensor,
-        speaker_mask_cfg: torch.Tensor,
-        context_kv_cond: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],
-        context_kv_cfg: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],
-    ) -> torch.Tensor:
-        if not self.captured or self.graph is None:
-            self.capture()
-
-        self.x_buf.copy_(x_t)
-        self.text_state_cond.copy_(text_state_cond)
-        self.text_mask_cond.copy_(text_mask_cond)
-        self.speaker_state_cond.copy_(speaker_state_cond)
-        self.speaker_mask_cond.copy_(speaker_mask_cond)
-        self.text_state_cfg.copy_(text_state_cfg)
-        self.text_mask_cfg.copy_(text_mask_cfg)
-        self.speaker_state_cfg.copy_(speaker_state_cfg)
-        self.speaker_mask_cfg.copy_(speaker_mask_cfg)
-        _copy_context_kv(self.context_kv_cond, context_kv_cond)
-        _copy_context_kv(self.context_kv_cfg, context_kv_cfg)
-
-        self.graph.replay()
-        return self.x_buf.clone()
-
-
-class FasterIrodoriRFSampler:
-    """Graph cache and safe fallback wrapper for Irodori RF sampling."""
-
-    def __init__(
-        self,
-        *,
-        max_graphs: int = 2,
-        speaker_bucket_multiple: int = 64,
-        num_warmup: int = 2,
-    ) -> None:
-        self.max_graphs = max(1, int(max_graphs))
-        self.speaker_bucket_multiple = max(1, int(speaker_bucket_multiple))
-        self.num_warmup = max(1, int(num_warmup))
-        self._graphs: OrderedDict[RFGraphSignature, IrodoriRFGraph] = OrderedDict()
-
-    def _unsupported_reason(
-        self,
-        *,
-        model,
-        cfg_guidance_mode: str,
-        cfg_scale_text: float,
-        cfg_scale_speaker: float,
-        rescale_k: float | None,
-        rescale_sigma: float | None,
-        use_context_kv_cache: bool,
-        speaker_kv_scale: float | None,
-    ) -> str | None:
-        if model.device.type != "cuda" or not torch.cuda.is_available():
-            return "CUDA Graph requires a CUDA device"
-        if str(cfg_guidance_mode).strip().lower() != "independent":
-            return "only cfg_guidance_mode='independent' is currently graphed"
-        if cfg_scale_text <= 0 or cfg_scale_speaker <= 0:
-            return "graph path currently expects both text and speaker CFG scales to be > 0"
-        if rescale_k is not None or rescale_sigma is not None:
-            return "rescale_k/rescale_sigma path is not graph-enabled"
-        if not use_context_kv_cache:
-            return "context_kv_cache=False is not graph-enabled"
-        if speaker_kv_scale is not None:
-            return "speaker_kv_scale path is not graph-enabled"
-        return None
-
-    def _get_graph(self, model, signature: RFGraphSignature) -> IrodoriRFGraph:
-        graph = self._graphs.get(signature)
-        if graph is not None:
-            self._graphs.move_to_end(signature)
-            return graph
-
-        graph = IrodoriRFGraph(model, signature, num_warmup=self.num_warmup)
-        graph.capture()
-        self._graphs[signature] = graph
-        self._graphs.move_to_end(signature)
-        while len(self._graphs) > self.max_graphs:
-            self._graphs.popitem(last=False)
-        return graph
-
-    @torch.inference_mode()
-    def sample(
-        self,
-        *,
-        model,
-        text_input_ids: torch.Tensor,
-        text_mask: torch.Tensor,
-        ref_latent: torch.Tensor,
-        ref_mask: torch.Tensor,
-        sequence_length: int,
-        num_steps: int = 40,
-        cfg_scale_text: 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,
-        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,
-    ) -> RFGraphSampleResult:
-        def fallback(reason: str) -> RFGraphSampleResult:
-            return RFGraphSampleResult(
-                latent=sample_euler_rf_cfg(
-                    model=model,
-                    text_input_ids=text_input_ids,
-                    text_mask=text_mask,
-                    ref_latent=ref_latent,
-                    ref_mask=ref_mask,
-                    sequence_length=sequence_length,
-                    num_steps=num_steps,
-                    cfg_scale_text=cfg_scale_text,
-                    cfg_scale_speaker=cfg_scale_speaker,
-                    cfg_guidance_mode=cfg_guidance_mode,
-                    cfg_min_t=cfg_min_t,
-                    cfg_max_t=cfg_max_t,
-                    seed=seed,
-                    truncation_factor=truncation_factor,
-                    rescale_k=rescale_k,
-                    rescale_sigma=rescale_sigma,
-                    use_context_kv_cache=use_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,
-                ),
-                graph_used=False,
-                fallback_reason=reason,
-            )
-
-        reason = self._unsupported_reason(
-            model=model,
-            cfg_guidance_mode=cfg_guidance_mode,
-            cfg_scale_text=float(cfg_scale_text),
-            cfg_scale_speaker=float(cfg_scale_speaker),
-            rescale_k=rescale_k,
-            rescale_sigma=rescale_sigma,
-            use_context_kv_cache=bool(use_context_kv_cache),
-            speaker_kv_scale=speaker_kv_scale,
-        )
-        if reason is not None:
-            return fallback(reason)
-
-        device = model.device
-        dtype = model.dtype
-        batch_size = int(text_input_ids.shape[0])
-        latent_dim = model.cfg.patched_latent_dim
-
-        ref_latent, ref_mask = _pad_reference_to_bucket(
-            ref_latent,
-            ref_mask,
-            speaker_patch_size=model.cfg.speaker_patch_size,
-            bucket_multiple=self.speaker_bucket_multiple,
-        )
-
-        rng = torch.Generator(device=device).manual_seed(int(seed))
-        x_t = torch.randn(
-            (batch_size, int(sequence_length), latent_dim),
-            device=device,
-            dtype=dtype,
-            generator=rng,
-        )
-        if truncation_factor is not None:
-            x_t = x_t * float(truncation_factor)
-
-        text_state_cond, text_mask_cond, speaker_state_cond, speaker_mask_cond = (
-            model.encode_conditions(
-                text_input_ids=text_input_ids,
-                text_mask=text_mask,
-                ref_latent=ref_latent,
-                ref_mask=ref_mask,
-            )
-        )
-        text_state_uncond = torch.zeros_like(text_state_cond)
-        text_mask_uncond = torch.zeros_like(text_mask_cond)
-        speaker_state_uncond = torch.zeros_like(speaker_state_cond)
-        speaker_mask_uncond = torch.zeros_like(speaker_mask_cond)
-
-        text_state_cfg = torch.cat([text_state_cond, text_state_uncond, text_state_cond], dim=0)
-        text_mask_cfg = torch.cat([text_mask_cond, text_mask_uncond, text_mask_cond], dim=0)
-        speaker_state_cfg = torch.cat(
-            [speaker_state_cond, speaker_state_cond, speaker_state_uncond], dim=0
-        )
-        speaker_mask_cfg = torch.cat(
-            [speaker_mask_cond, speaker_mask_cond, speaker_mask_uncond], dim=0
-        )
-
-        context_kv_cond = model.build_context_kv_cache(
-            text_state=text_state_cond,
-            speaker_state=speaker_state_cond,
-        )
-        context_kv_cfg = model.build_context_kv_cache(
-            text_state=text_state_cfg,
-            speaker_state=speaker_state_cfg,
-        )
-
-        signature = RFGraphSignature(
-            batch_size=batch_size,
-            sequence_length=int(sequence_length),
-            latent_dim=int(latent_dim),
-            text_len=int(text_state_cond.shape[1]),
-            speaker_len=int(speaker_state_cond.shape[1]),
-            num_steps=int(num_steps),
-            cfg_scale_text=float(cfg_scale_text),
-            cfg_scale_speaker=float(cfg_scale_speaker),
-            cfg_min_t=float(cfg_min_t),
-            cfg_max_t=float(cfg_max_t),
-            dtype=str(dtype),
-            device=_device_key(device),
-        )
-        try:
-            graph = self._get_graph(model, signature)
-            latent = graph.run(
-                x_t=x_t,
-                text_state_cond=text_state_cond,
-                text_mask_cond=text_mask_cond,
-                speaker_state_cond=speaker_state_cond,
-                speaker_mask_cond=speaker_mask_cond,
-                text_state_cfg=text_state_cfg,
-                text_mask_cfg=text_mask_cfg,
-                speaker_state_cfg=speaker_state_cfg,
-                speaker_mask_cfg=speaker_mask_cfg,
-                context_kv_cond=context_kv_cond,
-                context_kv_cfg=context_kv_cfg,
-            )
-        except Exception as exc:
-            self._graphs.pop(signature, None)
-            return fallback(f"CUDA Graph capture/replay failed: {exc}")
-        return RFGraphSampleResult(latent=latent, graph_used=True)

faster_irodori_tts/runtime.py

@@ -1,290 +0,0 @@
-from __future__ import annotations
-
-import math
-import secrets
-from collections.abc import Callable
-
-import torch
-
-from irodori_tts.codec import unpatchify_latent
-from irodori_tts.inference_runtime import (
-    InferenceRuntime,
-    RuntimeKey,
-    SamplingRequest,
-    SamplingResult,
-    _measure_end,
-    _measure_start,
-    find_flattening_point,
-    resolve_cfg_scales,
-)
-from irodori_tts.text_normalization import normalize_text
-
-from .rf_graph import FasterIrodoriRFSampler
-
-
-class FasterIrodoriTTSRuntime(InferenceRuntime):
-    """Irodori runtime that uses CUDA Graphs for supported RF sampling requests."""
-
-    def __init__(self, **kwargs) -> None:
-        super().__init__(**kwargs)
-        self.rf_sampler = FasterIrodoriRFSampler()
-
-    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(
-            (
-                "[faster_runtime] start synthesize "
-                "model_device={} model_precision={} codec_device={} codec_precision={} "
-                "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.codec.enable_watermark,
-                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,
-            )
-        )
-
-        if req.seconds <= 0:
-            raise ValueError(f"seconds must be > 0, got {req.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}")
-
-        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 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_speaker, scale_messages = resolve_cfg_scales(
-            cfg_guidance_mode=cfg_mode,
-            cfg_scale_text=req.cfg_scale_text,
-            cfg_scale_speaker=req.cfg_scale_speaker,
-            cfg_scale=req.cfg_scale,
-        )
-        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"[faster_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"[faster_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)
-
-            target_samples = int(float(req.seconds) * self.codec.sample_rate)
-            latent_steps = math.ceil(target_samples / int(self.codec.model.hop_length))
-            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, 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"[faster_runtime] prepare_reference: {stage_sec * 1000.0:.1f} ms")
-
-            t0 = _measure_start(self.model_device)
-            sample_result = self.rf_sampler.sample(
-                model=self.model,
-                text_input_ids=text_ids,
-                text_mask=text_mask,
-                ref_latent=ref_latent,
-                ref_mask=ref_mask,
-                sequence_length=patched_steps,
-                num_steps=int(req.num_steps),
-                cfg_scale_text=cfg_scale_text,
-                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,
-            )
-            z_patched = sample_result.latent
-            stage_sec = _measure_end(self.model_device, t0)
-            stage_timings.append(("sample_rf", stage_sec))
-            if sample_result.graph_used:
-                _log(f"[faster_runtime] sample_rf (cuda_graph): {stage_sec * 1000.0:.1f} ms")
-            else:
-                msg = f"info: RF CUDA Graph fallback: {sample_result.fallback_reason}"
-                messages.append(msg)
-                _log(msg)
-                _log(f"[faster_runtime] sample_rf (fallback): {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"[faster_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"[faster_runtime] decode_latent ({decode_mode}): {stage_sec * 1000.0:.1f} ms")
-
-            total_to_decode = _measure_end(self.model_device, post_load_t0, self.codec_device)
-            _log(f"[faster_runtime] total_to_decode: {total_to_decode:.3f} s")
-
-        _log("[faster_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,
-        )
-
-
-# Backward-friendly alias for callers that prefer an InferenceRuntime-like name.
-FasterInferenceRuntime = FasterIrodoriTTSRuntime
-
-__all__ = [
-    "FasterIrodoriTTSRuntime",
-    "FasterInferenceRuntime",
-    "RuntimeKey",
-    "SamplingRequest",
-    "SamplingResult",
-]