Create tts.py

tts.py

@@ -0,0 +1,512 @@
+import argparse
+from dataclasses import dataclass, field
+from enum import Enum
+import huggingface_hub
+from pathlib import Path
+import random
+import time
+import typing as tp
+
+import numpy as np
+from safetensors.torch import load_file
+import torch
+
+from moshi.conditioners import ConditionAttributes, dropout_all_conditions, TensorCondition
+from moshi.models import loaders
+from moshi.models.lm import _LMGenState, LMGen
+from moshi.models.tts import TTSModel, Entry, State, StateMachine, DEFAULT_DSM_TTS_REPO
+from moshi.modules.transformer import StreamingMultiheadAttention
+from pydantic import BaseModel
+
+
+class MaskFlags(Enum):
+    # Output PCM is ready
+    HAS_PCM = 1
+    # Generation is done, no need to step again.
+    IS_EOS = 2
+    # One word was consumed in the text stream.
+    WORD_FINISHED = 4
+    # One AR step was performed.
+    AR_STEP = 8
+    # AR step was skipped because the client is not sending words fast enough.
+    MISSING_WORDS = 16
+
+
+def flags_out_from_mask_(flags_out: np.ndarray, mask: torch.Tensor, value: int):
+    flags_out[mask.numpy()] |= value
+
+
+def split_at_specific_separator(text: str, separator: str, index_of_separator: int) -> tuple[str, str]:
+    """ kyutai/tts-voices/unmute-prod-website/*.safetensors
+    becomes
+    ('kyutai/tts-voices', 'unmute-prod-website/*.safetensors)
+    with index_of_separator=1.
+    """
+    if text.count(separator) <= index_of_separator:
+        raise ValueError(f"Separator '{separator}' not found {index_of_separator + 1} times in `{text}`.")
+    parts = text.split(separator, index_of_separator + 1)
+    return separator.join(parts[:-1]), parts[-1]
+
+
+class Config(BaseModel):
+    log_folder: Path = Path.home() / 'tmp/tts-service'
+    hf_repo: str = DEFAULT_DSM_TTS_REPO
+    mimi_weight: Path | None = None
+    moshi_weight: Path | None = None
+    config_path: Path | None = None
+    tokenizer: Path | None = None
+    device: str = 'cuda'
+
+    n_q: int = 4
+    # This can have multiple formats:
+    # - A path to a folder with voices, e.g. `models/tts`
+    # - A huggingface snapshot, e.g. `hf-snapshot://kyutai/tts-voices`
+    # - A huggingface snapshot with a pattern,
+    #     e.g. `hf-snapshot://kyutai/tts-voices/unmute-prod-website/*.safetensors`
+    voice_folder: str = str(Path.home() / 'models/tts-voices')
+    default_voice: str = "barack_demo.wav"
+
+    temp: float = 0.6
+    cfg_coef: float = 2.
+
+    max_padding: int = 8
+    initial_padding: int = 2
+    final_padding: int = 4
+    padding_between: int = 1
+
+    interleaved_text_only: int = 2
+    debug: bool = False
+
+
+def init(batch_size: int, config_override: dict) -> 'TTSService':
+    config = Config(**config_override)
+    config.log_folder.mkdir(parents=True, exist_ok=True)
+    print("ADJUSTED CODE")
+    print("retrieving checkpoint")
+    checkpoint_info = loaders.CheckpointInfo.from_hf_repo(
+        config.hf_repo, moshi_weights=config.moshi_weight, mimi_weights=config.mimi_weight,
+        config_path=config.config_path, tokenizer=config.tokenizer)
+
+    cfg_condition = None
+    tts_model = TTSModel.from_checkpoint_info(
+        checkpoint_info, n_q=config.n_q, temp=config.temp, cfg_coef=config.cfg_coef,
+        max_padding=config.max_padding, initial_padding=config.initial_padding, final_padding=config.final_padding,
+        device=config.device, dtype=torch.float16)
+    if tts_model.valid_cfg_conditionings:
+        # Model was trained with CFG distillation.
+        cfg_condition = tts_model.cfg_coef
+        tts_model.cfg_coef = 1.
+        cfg_is_no_text = False
+    else:
+        cfg_is_no_text = True
+
+    voice_suffix = tts_model.voice_suffix
+    print(f"loading voices from {config.voice_folder}, with suffix {voice_suffix}.")
+    all_attributes = {}
+    voice_folder = config.voice_folder
+    if voice_folder.startswith("hf-snapshot://"):
+        voice_folder = voice_folder.removeprefix("hf-snapshot://")
+        # We detect if there is a pattern in the voice folder.
+        if voice_folder.count("/") > 1:
+            voice_folder, pattern = split_at_specific_separator(voice_folder, '/', 1)
+        else:
+            pattern = None
+        print(f"retrieving voices from {voice_folder}")
+        voice_folder = huggingface_hub.snapshot_download(voice_folder, allow_patterns=pattern)
+    voice_folder = Path(voice_folder)
+
+    for file in voice_folder.glob(f'**/*{voice_suffix}'):
+        relative = file.relative_to(voice_folder)
+        name = str(relative.with_name(relative.name.removesuffix(voice_suffix)))
+        try:
+            attributes = tts_model.make_condition_attributes([file, file], cfg_coef=cfg_condition)
+        except Exception:
+            print(f"[WARNING] failed to load voice {name}")
+        else:
+            all_attributes[name] = attributes
+
+    if not all_attributes:
+        raise RuntimeError(
+            "No voices found, please check your voice folder. "
+            f"Searched for files matching {voice_folder}/**/*{voice_suffix}"
+        )
+
+    if config.default_voice not in all_attributes:
+        raise RuntimeError(
+            f"Default voice {config.default_voice}, please check your voice folder. "
+            f"Expected {voice_folder}/{config.default_voice}{voice_suffix} to exist"
+        )
+
+    service = TTSService(
+        batch_size=batch_size, default_attribute_name=config.default_voice,
+        all_attributes=all_attributes,
+        tts_model=tts_model,
+        cfg_condition=cfg_condition,
+        cfg_is_no_text=cfg_is_no_text,
+        padding_between=config.padding_between,
+        debug=config.debug,
+        interleaved_text_only=config.interleaved_text_only)
+
+    return service
+
+
+@dataclass
+class ClientState:
+    is_complete: bool = False
+    state: State | None = None
+    offset: int = 0
+
+    def reset(self, state_machine: StateMachine) -> None:
+        self.is_complete = False
+        self.offset = 0
+        self.state = state_machine.new_state([])
+
+
+@dataclass
+class TTSService:
+    batch_size: int
+    default_attribute_name: str
+    all_attributes: dict[str, ConditionAttributes]
+
+    tts_model: TTSModel
+
+    cfg_is_no_text: bool = True
+    cfg_condition: float | None = None
+    padding_between: int = 1
+    n_q: int = 4
+    debug: bool = False
+    interleaved_text_only: int = 0
+
+    flags_out: np.ndarray | None = None
+    clients: list[ClientState] = field(default_factory=list)
+    cross_attention_cache: dict[str, torch.Tensor] = field(default_factory=dict)
+    cross_attentions: list[StreamingMultiheadAttention] = field(default_factory=list)
+
+    def __post_init__(self):
+        print(self.n_q)
+        print("NQ AMOUNT")
+
+        lm = self.tts_model.lm
+        tts_model = self.tts_model
+        mimi = self.tts_model.mimi
+        machine = self.tts_model.machine
+
+        self.device = lm.device
+        self.dtype = lm.dtype
+        self.remaining_text_only = self.interleaved_text_only
+
+        for _ in range(self.batch_size):
+            client = ClientState()
+            self.clients.append(client)
+
+        print("Filling cross attention cache.")
+        for name, attributes in self.all_attributes.items():
+            self.cross_attention_cache[name] = self._get_cross_attention_source([attributes])
+
+        assert lm.condition_provider is not None
+
+        cas = [self.all_attributes[self.default_attribute_name]] * self.batch_size
+        if self.tts_model.cfg_coef != 1.0:
+            nulled = make_null(cas)
+            cas = cas + nulled
+        prepared = lm.condition_provider.prepare(cas)
+        condition_tensors = lm.condition_provider(prepared)
+
+        for module in lm.modules():
+            if isinstance(module, StreamingMultiheadAttention) and module.cross_attention:
+                self.cross_attentions.append(module)
+
+        self.lm_gen = LMGen(
+            lm, temp=tts_model.temp, temp_text=tts_model.temp, cfg_coef=tts_model.cfg_coef,
+            condition_tensors=condition_tensors, on_text_hook=self._on_text_hook,
+            on_audio_hook=self._on_audio_hook, cfg_is_no_text=self.cfg_is_no_text,
+            support_out_of_sync=True)
+        self.lm_gen.streaming_forever(self.batch_size)
+        mimi.streaming_forever(self.batch_size)
+
+        missing = lm.n_q - lm.dep_q
+        self.input_tokens = torch.full(
+            (self.batch_size, missing, 1), machine.token_ids.zero,
+            dtype=torch.long, device=self.device)
+        self.no_depformer_tokens = torch.full(
+            (self.batch_size, lm.dep_q, 1), machine.token_ids.zero,
+            dtype=torch.long, device=self.device)
+        self.last_actives: list[bool] = [False] * self.batch_size
+        print("warming up.")
+        for _ in range(3):
+            mimi.set_exec_mask(torch.ones(self.batch_size, dtype=torch.bool))
+            self.lm_gen.set_exec_mask(torch.ones(self.batch_size, dtype=torch.bool))
+            frame = self.lm_gen.step(self.input_tokens)
+            assert frame is not None
+            mimi.decode(frame[:, 1:].clamp(min=0))
+        print("ready to roll.")
+
+    def _get_cross_attention_source(self, all_attributes: list[ConditionAttributes]) -> torch.Tensor:
+        lm = self.tts_model.lm
+        assert lm.condition_provider is not None
+        assert lm.fuser is not None
+        prepared = lm.condition_provider.prepare(all_attributes)
+        condition_tensors = lm.condition_provider(prepared)
+        cross = lm.fuser.get_cross(condition_tensors)
+        assert cross is not None
+        return cross.to(device=self.device, dtype=self.dtype)
+
+    @property
+    def _lm_gen_state(self) -> _LMGenState:
+        assert self.lm_gen._streaming_state is not None
+        return self.lm_gen._streaming_state
+
+    def _on_audio_hook(self, audio_tokens: torch.Tensor) -> None:
+        delays = self.lm_gen.delays_cuda[1: 1 + self.tts_model.lm.dep_q]
+        mask = self._lm_gen_state.offsets[:, None] < delays + self.tts_model.delay_steps
+        audio_tokens.masked_fill_(mask, self.tts_model.machine.token_ids.zero)
+
+    def _on_text_hook(self, text_tokens) -> None:
+        tokens = text_tokens.tolist()
+        out_tokens = []
+        for b, (token, client) in enumerate(zip(tokens, self.clients)):
+            if not self.last_actives[b]:
+                out_tokens.append(token)
+                continue
+            assert client.state is not None
+            out_token, consumed_new_word = self.tts_model.machine.process(client.offset, client.state, token)
+
+            if self.flags_out is not None and consumed_new_word:
+                self.flags_out[b] |= MaskFlags.WORD_FINISHED.value
+            out_tokens.append(out_token)
+        text_tokens[:] = torch.tensor(out_tokens, dtype=torch.long, device=text_tokens.device)
+
+    def _print(self, *args, **kwargs):
+        if self.debug:
+            print(*args, **kwargs)
+
+    @torch.no_grad()
+    def step(self, updates: list[tuple[int, list[int], np.ndarray | str | None]], pcm_out: np.ndarray,
+             flags_out: np.ndarray, code_out: np.ndarray) -> None:
+        mimi = self.tts_model.mimi
+        machine = self.tts_model.machine
+        delay_steps = self.tts_model.delay_steps
+
+        self.flags_out = flags_out
+        flags_out[:] = 0
+
+        reset_mask = torch.zeros(self.batch_size, dtype=torch.bool)
+        # List of pre computed cross attention values.
+        new_cross_sources: list[torch.Tensor] = []
+        new_cross_indexes: list[int] = []
+        # List of new dynamic conditioning that we need to compute.
+        new_voice_indexes: list[int] = []
+        new_voice_sources: list[torch.Tensor] = []
+        for b, new_entry, voice in updates:
+            client = self.clients[b]
+            if not new_entry:
+                self._print(f"[{b}] NO TOKENS REALLY LAURENT.")
+            if new_entry[0] == -1:
+                client.reset(machine)
+                reset_mask[b] = True
+                new_entry = new_entry[1:]
+                if isinstance(voice, np.ndarray):
+                    new_voice_indexes.append(b)
+                    new_voice_sources.append(torch.from_numpy(voice))
+                else:
+                    cross_source = self.cross_attention_cache.get(voice or '', None)
+                    if cross_source is None:
+                        cross_source = self.cross_attention_cache[self.default_attribute_name]
+                    new_cross_sources.append(cross_source)
+                    new_cross_indexes.append(b)
+                self._print(f"[{b}] Reset, voice is {voice}.")
+            if client.state is None:
+                self._print(f"[{b}] Trying to push {new_entry}, but not assigned.")
+            elif new_entry == [-2]:
+                self._print(f"[{b}] Done.")
+                client.is_complete = True
+            else:
+                self._print(f"[{b}] Pushing {new_entry}.")
+                padding = 0
+                if self.padding_between > 0:
+                    padding = max(0, self.padding_between + len(new_entry) - 1)
+                client.state.entries.append(Entry(new_entry, '', padding=padding))
+
+        actives = []
+        mimi_actives = []
+        in_text_onlys = []
+        for b, client in enumerate(self.clients):
+            if client.state is None:
+                # client is not currently assigned.
+                active = False
+            elif client.is_complete:
+                # We got all the words from the client and are wrapping up.
+                active = True
+            elif client.state.forced_padding > 0:
+                # We are sure we won't try to consume a word at this point.
+                active = True
+            elif len(client.state.entries) > self.tts_model.machine.second_stream_ahead:
+                # We have some words ready to be consumed.
+                active = True
+            else:
+                flags_out[b] |= MaskFlags.MISSING_WORDS.value
+                active = False
+            actives.append(active)
+
+            real_offset = client.offset - self.lm_gen.max_delay
+
+            mimi_active = active and (real_offset >= delay_steps)
+            mimi_actives.append(mimi_active)
+
+            in_text_only = active and (client.offset < delay_steps)
+            in_text_onlys.append(in_text_only)
+
+        in_text_only_mask = torch.tensor(in_text_onlys, dtype=torch.bool)
+        run_in_text_only = self.remaining_text_only > 0 and in_text_only_mask.any()
+
+        if run_in_text_only:
+            self.remaining_text_only -= 1
+            mimi_exec_mask = torch.zeros(self.batch_size, dtype=torch.bool)
+            exec_mask = in_text_only_mask
+            actives = in_text_onlys
+        else:
+            self.remaining_text_only = self.interleaved_text_only
+            exec_mask = torch.tensor(actives, dtype=torch.bool)
+            mimi_exec_mask = torch.tensor(mimi_actives, dtype=torch.bool)
+        del mimi_actives
+        self.last_actives = actives
+
+        flags_out_from_mask_(flags_out, exec_mask, MaskFlags.AR_STEP.value)
+        flags_out_from_mask_(flags_out, mimi_exec_mask, MaskFlags.HAS_PCM.value)
+
+        # We check on exec_mask whether we actually need to run anything, before we move it to CUDA.
+        # However, we still need to perform the reset and update of cross attention for models
+        # with a text lookahead stream.
+        skip_exec = not exec_mask.any()
+
+        exec_mask = exec_mask.to(self.device)
+        mimi_exec_mask = mimi_exec_mask.to(self.device)
+        need_reset = reset_mask.any()
+        reset_mask = reset_mask.to(self.device)
+
+        if new_voice_sources:
+            all_attributes = [make_condition_attributes([voice_source], cfg_condition=self.cfg_condition)
+                              for voice_source in new_voice_sources]
+            new_cross_sources += self._get_cross_attention_source(all_attributes).split(1)
+            new_cross_indexes += new_voice_indexes
+        if new_cross_sources:
+            cross_source = torch.cat(new_cross_sources)
+            cross_indexes = torch.tensor(new_cross_indexes, dtype=torch.long, device=self.device)
+            for attention in self.cross_attentions:
+                k, v = attention._compute_cross_attention(cross_source, cross_source)
+                state = attention._streaming_state
+                assert state is not None
+                assert state.k_cross is not None
+                assert state.v_cross is not None
+                state.k_cross.index_copy_(0, cross_indexes, k)
+                state.v_cross.index_copy_(0, cross_indexes, v)
+
+        if need_reset:
+            self.lm_gen.reset_streaming(reset_mask=reset_mask)
+            mimi.reset_streaming(reset_mask=reset_mask)
+
+        if skip_exec:
+            return
+
+        self.lm_gen.set_exec_mask(exec_mask)
+        mimi.set_exec_mask(mimi_exec_mask)
+
+        depformer_replace_tokens = self.no_depformer_tokens if run_in_text_only else None
+        frame = self.lm_gen.step(self.input_tokens, depformer_replace_tokens=depformer_replace_tokens)
+        assert frame is not None
+        audio_frame = frame[:, 1:]
+        audio_frame.clamp_(min=0)
+
+        if run_in_text_only:
+            pcm = None
+        else:
+            pcm = mimi.decode(audio_frame)
+            pcm.clamp_(-0.99, 0.99)
+
+        for b, client in enumerate(self.clients):
+            if actives[b]:
+                assert client.state is not None
+                client.offset += 1
+                self._print(f"[{b}] Offset {client.offset: 3d}, pendings={len(client.state.entries): 3d}.")
+                if client.is_complete and client.state.end_step is not None:
+                    # We were waiting for the end of the generation.
+                    real_end = (
+                        client.state.end_step + delay_steps + self.tts_model.final_padding + self.lm_gen.max_delay)
+                    if client.offset >= real_end:
+                        self._print(f"[{b}] Done.")
+                        client.reset(machine)
+                        flags_out[b] |= MaskFlags.IS_EOS.value
+        if pcm is not None:
+            pcm_out[:] = pcm[:, 0].cpu().numpy()
+        code_out[:, :frame.shape[1]] = frame[:, :, 0].int().cpu().numpy()
+        code_out[:, frame.shape[1]:] = 0
+        self.flags_out = None
+
+
+class Profiler:
+    """Context manager wrapper for xformers profiler.
+    """
+    def __init__(self, enabled: bool = False):
+        self.profiler: tp.Optional[tp.Any] = None
+        if enabled:
+            from xformers.profiler import profile
+            from xformers.profiler.api import PyTorchProfiler
+            output_dir = './profiler_data'
+            schedule = (
+                (PyTorchProfiler, 6, 12),
+            )
+            self.profiler = profile(output_dir=output_dir, schedule=schedule)
+
+    def step(self):
+        if self.profiler is not None:
+            self.profiler.step()  # type: ignore
+
+    def __enter__(self):
+        if self.profiler is not None:
+            return self.profiler.__enter__()  # type: ignore
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        if self.profiler is not None:
+            return self.profiler.__exit__(exc_type, exc_value, exc_tb)  # type: ignore
+
+
+def make_condition_attributes(voices: list[Path | torch.Tensor],
+                              max_speakers: int = 5,
+                              cfg_condition: float | None = None) -> ConditionAttributes:
+    assert voices
+    voice_tensor = None
+    mask = None
+    for idx in range(5):
+        if idx < len(voices):
+            voice = voices[idx]
+            if isinstance(voice, Path):
+                emb = load_file(voice, device='cuda')['speaker_wavs']
+            else:
+                emb = voice
+            assert emb.dim() == 3
+            if voice_tensor is None:
+                voice_tensor = torch.zeros(1, max_speakers, emb.shape[2], emb.shape[1], device='cuda')
+            if mask is None:
+                mask = torch.zeros(1, max_speakers, emb.shape[2], dtype=torch.bool, device='cuda')
+            voice_tensor[:, idx, :, :] = emb.transpose(1, 2)
+            mask[:, idx, :] = True
+    assert voice_tensor is not None
+    assert mask is not None
+    voice_tensor = voice_tensor.view(1, -1, voice_tensor.shape[-1])
+    mask = mask.view(1, -1)
+    tensors = {
+        'speaker_wavs': TensorCondition(voice_tensor, mask)
+    }
+    text: dict[str, str | None] = {
+        'control': 'ok',
+    }
+    if cfg_condition is None:
+        text['cfg'] = None
+    else:
+        text['cfg'] = format(cfg_condition, '.1f')
+    return ConditionAttributes(text=dict(text), tensor=tensors)
+
+
+def make_null(all_attributes: tp.Sequence[ConditionAttributes]) -> list[ConditionAttributes]:
+    return dropout_all_conditions(all_attributes)
+