src/dots_tts/runtime_double_streaming.py

from __future__ import annotations

from pathlib import Path

import torch
from loguru import logger

from dots_tts.data.pipelines.tts_pipeline import TTS_INTERLEAVE_PREFIX
from dots_tts.runtime import DotsTtsRuntime
from dots_tts.utils.util import get_dtype


class DoubleStreamingSession:
    """Incremental interleave session for text-token to audio-chunk generation."""

    def __init__(
        self,
        runtime: DotsTtsRuntime,
        *,
        prompt_audio_path: str | None = None,
        prompt_text: str | None = None,
        ode_method: str = "euler",
        num_steps: int = 10,
        guidance_scale: float = 1.2,
        speaker_scale: float = 1.5,
        eos_threshold: float = 0.8,
        initial_silence_audio_tokens: int = 1,
    ) -> None:
        normalized_prompt_text = runtime._process_prompt_text(prompt_text)
        if normalized_prompt_text:
            raise ValueError("Double streaming does not support prompt_text.")

        self.runtime = runtime
        self.model = runtime.model
        self.device = runtime.device
        self.ode_method = ode_method
        self.num_steps = int(num_steps)
        self.guidance_scale = float(guidance_scale)
        self.speaker_scale = float(speaker_scale)
        self.eos_threshold = float(eos_threshold)
        self.max_generate_length = runtime.max_generate_length
        self._initial_silence_audio_tokens = max(
            0,
            min(10, int(initial_silence_audio_tokens or 0)),
        )

        self._dtype = get_dtype(runtime.precision)
        self._use_amp = self.device.type == "cuda" and self._dtype in {
            torch.float16,
            torch.bfloat16,
        }
        self._prefix_token_ids = tuple(
            self.model.tokenizer.encode(
                TTS_INTERLEAVE_PREFIX,
                add_special_tokens=False,
            )
        )
        self._state = self.model._allocate_generate_state(
            max_audio_patch_count=self.max_generate_length,
            device=self.device,
            dtype=self._dtype,
        )
        self._vocoder_state = self.model.vocoder.init_stream_state(
            batch_size=1,
            chunk_size=self.model.core.latent_patch_size,
        )
        self._g_cond = None
        self._started = False
        self._text_finished = False
        self._closed = False
        self._decoded_patch_count = 0

        if prompt_audio_path is not None:
            cache = getattr(self.runtime, "_double_streaming_prompt_g_cond_cache", None)
            if cache is None:
                cache = {}
                setattr(self.runtime, "_double_streaming_prompt_g_cond_cache", cache)
            prompt_cache_key = (
                str(Path(prompt_audio_path).expanduser().resolve()),
                str(self.device),
                str(self._dtype),
                self.speaker_scale,
            )
            cached_g_cond = cache.get(prompt_cache_key)
            if cached_g_cond is None:
                prompt_audio = self.runtime._load_prompt_audio(prompt_audio_path)
                with torch.no_grad():
                    with torch.autocast(
                        device_type=self.device.type,
                        dtype=self._dtype,
                        enabled=self._use_amp,
                    ):
                        prompt_conditioning = self.model._prepare_prompt_conditioning(
                            prompt_audio,
                            use_prompt_prefill=False,
                            speaker_scale=self.speaker_scale,
                        )
                cached_g_cond = prompt_conditioning.g_cond.detach()
                cache[prompt_cache_key] = cached_g_cond
                logger.info(
                    "Double streaming prompt conditioning cached: path={} device={} "
                    "dtype={} speaker_scale={}",
                    prompt_cache_key[0],
                    self.device,
                    self._dtype,
                    self.speaker_scale,
                )
            else:
                logger.info(
                    "Double streaming prompt conditioning cache hit: path={} device={} "
                    "dtype={} speaker_scale={}",
                    prompt_cache_key[0],
                    self.device,
                    self._dtype,
                    self.speaker_scale,
                )
            self._g_cond = cached_g_cond

        logger.info(
            "Double streaming session started: prefix_token_count={} precision={} "
            "ode_method={} num_steps={} guidance_scale={} speaker_scale={} max_audio_patch_count={} "
            "initial_silence_audio_tokens={} has_ref_audio_only={}",
            len(self._prefix_token_ids),
            runtime.precision,
            self.ode_method,
            self.num_steps,
            self.guidance_scale,
            self.speaker_scale,
            self.max_generate_length,
            self._initial_silence_audio_tokens,
            self._g_cond is not None,
        )

    @property
    def is_finished(self) -> bool:
        return self._closed

    def push_text_token(self, text_token: int) -> torch.Tensor | None:
        self._ensure_active()
        if self._text_finished:
            raise RuntimeError("Cannot push text tokens after finish_text().")
        if self._state.end_flag:
            raise RuntimeError(
                "Double streaming generation has already reached EOS. "
                "Call finish_text() to flush the remaining audio tail."
            )

        token_id = int(text_token)
        if not self._started:
            chunk_token_ids = [*self._prefix_token_ids, token_id]
            self._started = True
        else:
            chunk_token_ids = [token_id]

        self._consume_text_chunk(chunk_token_ids)
        return self._decode_audio_chunk()

    def finish_text(self):
        self._ensure_active()

        if not self._state.end_flag:
            if not self._text_finished:
                text_end_chunk = [self.model.core.text_cond_end_id]
                if not self._started:
                    text_end_chunk = [*self._prefix_token_ids, *text_end_chunk]
                    self._started = True
                self._consume_text_chunk(text_end_chunk)
                self._text_finished = True

            while not self._state.end_flag:
                audio_chunk = self._decode_audio_chunk(continue_audio_span=True)
                if audio_chunk is not None:
                    yield audio_chunk
        else:
            self._text_finished = True

        final_chunk = self.model.vocoder.stream_flush(self._vocoder_state)
        self._closed = True
        logger.info(
            "Double streaming session finished: decoded_patch_count={}",
            self._decoded_patch_count,
        )
        if final_chunk.size(-1) > 0:
            yield final_chunk

    def _ensure_active(self) -> None:
        if self._closed:
            raise RuntimeError("Double streaming session is already closed.")

    def _consume_text_chunk(self, token_ids: list[int]) -> None:
        schedule = torch.tensor(
            [token_ids],
            dtype=torch.long,
            device=self.device,
        )
        with torch.no_grad():
            with torch.autocast(
                device_type=self.device.type,
                dtype=self._dtype,
                enabled=self._use_amp,
            ):
                self.model._consume_text_schedule(
                    schedule,
                    position=0,
                    next_audio_position=schedule.size(1),
                    state=self._state,
                )

    def _get_initial_silence_audio_patch(
        self,
        patch_index: int,
        audio_patch: torch.Tensor,
    ) -> torch.Tensor:
        cache = getattr(self.runtime, "_double_streaming_silence_audio_patch_cache", None)
        if cache is None:
            cache = {}
            setattr(self.runtime, "_double_streaming_silence_audio_patch_cache", cache)

        cache_count = 10
        patch_size = int(self.model.core.latent_patch_size)
        key = (
            str(self.device),
            str(self._dtype),
            patch_size,
            int(audio_patch.size(-1)),
            cache_count,
        )
        cached_patches = cache.get(key)
        if cached_patches is None:
            hop_size = int(getattr(self.model.vocoder, "hop_size", 1))
            zero_samples = cache_count * patch_size * hop_size
            zero_audio = torch.zeros(
                (1, 1, zero_samples),
                device=self.device,
                dtype=torch.float32,
            )
            silence_latents = self.model.vocoder.extract_latents(zero_audio)
            silence_latents, _ = torch.split(
                silence_latents,
                int(audio_patch.size(-1)),
                dim=1,
            )
            silence_latents = silence_latents.transpose(1, 2)
            target_frames = cache_count * patch_size
            if silence_latents.size(1) < target_frames:
                silence_latents = torch.cat(
                    [
                        silence_latents,
                        silence_latents.new_zeros(
                            (
                                silence_latents.size(0),
                                target_frames - silence_latents.size(1),
                                silence_latents.size(2),
                            )
                        ),
                    ],
                    dim=1,
                )
            silence_latents = silence_latents[:, :target_frames, :]
            cached_patches = self.model.core.io_helper.normalize(silence_latents)
            cached_patches = cached_patches.to(device=self.device, dtype=audio_patch.dtype)
            cached_patches = cached_patches.reshape(
                1,
                cache_count,
                patch_size,
                int(audio_patch.size(-1)),
            ).detach()
            cache[key] = cached_patches
            logger.info(
                "Double streaming initial silence cache built: patches={} patch_size={} "
                "hop_size={} device={} dtype={}",
                cache_count,
                patch_size,
                hop_size,
                self.device,
                audio_patch.dtype,
            )
        return cached_patches[:, int(patch_index)].clone()

    def _consume_audio_patch(self, audio_patch: torch.Tensor) -> None:
        self.model._consume_audio_patch(self._state, audio_patch=audio_patch)

    def _decode_audio_chunk(self, *, continue_audio_span: bool = False) -> torch.Tensor | None:
        if self._decoded_patch_count >= self.max_generate_length:
            raise RuntimeError(
                "Double streaming exceeded max_generate_length before reaching EOS."
            )

        with torch.no_grad():
            with torch.autocast(
                device_type=self.device.type,
                dtype=self._dtype,
                enabled=self._use_amp,
            ):
                stop_after_current_audio = self.model._should_stop_after_current_audio(
                    self._state,
                    eos_threshold=self.eos_threshold,
                )
                audio_patch = self.model._decode_next_audio(
                    self._state,
                    device=self.device,
                    g_cond=self._g_cond,
                    ode_method=self.ode_method,
                    num_steps=self.num_steps,
                    guidance_scale=self.guidance_scale,
                )
                if self._decoded_patch_count < self._initial_silence_audio_tokens:
                    audio_patch = self._get_initial_silence_audio_patch(
                        self._decoded_patch_count,
                        audio_patch,
                    )
                self._consume_audio_patch(audio_patch)
                if continue_audio_span:
                    self.model._append_hidden_chunk(self._state, self._state.llm_hiddens)
                self._decoded_patch_count += 1
                latent_patch = self.model.core.io_helper.denormalize(audio_patch)
                audio_chunk = self.model.vocoder.stream_step(
                    latent_patch.transpose(1, 2),
                    self._vocoder_state,
                )
                if stop_after_current_audio:
                    self._state.end_flag = True

        if audio_chunk.size(-1) == 0:
            return None
        return audio_chunk


class DotsTtsRuntimeDoubleStreaming(DotsTtsRuntime):
    def start_double_streaming(
        self,
        *,
        prompt_audio_path: str | None = None,
        prompt_text: str | None = None,
        ode_method: str = "euler",
        num_steps: int = 10,
        guidance_scale: float = 1.2,
        speaker_scale: float = 1.5,
        eos_threshold: float = 0.8,
        initial_silence_audio_tokens: int = 1,
    ) -> DoubleStreamingSession:
        return DoubleStreamingSession(
            self,
            prompt_audio_path=prompt_audio_path,
            prompt_text=prompt_text,
            ode_method=ode_method,
            num_steps=num_steps,
            guidance_scale=guidance_scale,
            speaker_scale=speaker_scale,
            eos_threshold=eos_threshold,
            initial_silence_audio_tokens=initial_silence_audio_tokens,
        )


__all__ = ["DotsTtsRuntimeDoubleStreaming", "DoubleStreamingSession"]