first commit
This commit is contained in:
@@ -0,0 +1,355 @@
|
||||
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"]
|
||||
Reference in New Issue
Block a user