refactor(backend): 启动环境初始化与资源项目化（零跨项目依赖）
- main.py：uv run main.py 一键启动——读取 backend/.env、TORCH_HOME/NLTK_DATA 指向项目内缓存、certifi SSL + NO_PROXY 网络直连、默认开启 s2s，启动横幅 - vendor/：补丁后 s2s 云化栈移入项目（backend/vendor/s2s-cloud），pyproject sources 改相对路径 - .gitignore：忽略 .torch-cache / nltk_data 缓存 - AGENTS.md：进程管理铁律（AI 不启动/重启服务，由用户操作）
This commit is contained in:
Pine
+11
@@ -0,0 +1,11 @@
|
||||
# Archived Models
|
||||
|
||||
This directory stores sunset model implementations that are kept in-repo but are no longer wired into `s2s_pipeline.py`.
|
||||
|
||||
- STT: `moonshine` -> `archive/STT/moonshine_handler.py`
|
||||
- TTS: `parler` -> `archive/TTS/parler_handler.py`
|
||||
- TTS: `melo` -> `archive/TTS/melo_handler.py`
|
||||
- Legacy args: `archive/arguments_classes/parler_tts_arguments.py`
|
||||
- Legacy args: `archive/arguments_classes/melo_tts_arguments.py`
|
||||
|
||||
These models are also removed from default requirements. If you want to run them manually, install their dependencies separately.
|
||||
@@ -0,0 +1,72 @@
|
||||
import os
|
||||
|
||||
os.environ['KERAS_BACKEND'] = 'torch'
|
||||
|
||||
import logging
|
||||
|
||||
import moonshine
|
||||
import torch
|
||||
from rich.console import Console
|
||||
|
||||
from speech_to_speech.baseHandler import BaseHandler
|
||||
from speech_to_speech.pipeline.messages import VADAudio
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
console = Console()
|
||||
|
||||
|
||||
class MoonshineSTTHandler(BaseHandler[VADAudio]):
|
||||
"""
|
||||
Handles the Speech To Text generation using a Moonshine model.
|
||||
"""
|
||||
|
||||
def setup(
|
||||
self,
|
||||
model_name="moonshine/base",
|
||||
torch_dtype="float16",
|
||||
gen_kwargs={},
|
||||
):
|
||||
self.torch_dtype = getattr(torch, torch_dtype)
|
||||
self.gen_kwargs = gen_kwargs
|
||||
|
||||
self.tokenizer = moonshine.load_tokenizer()
|
||||
self.model = moonshine.load_model(model_name)
|
||||
|
||||
self.warmup()
|
||||
|
||||
def warmup(self):
|
||||
logger.info(f"Warming up {self.__class__.__name__}")
|
||||
|
||||
n_steps = 2
|
||||
dummy_input = torch.randn(
|
||||
(1, 16000),
|
||||
dtype=self.torch_dtype,
|
||||
)
|
||||
|
||||
if torch.cuda.is_available():
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
torch.cuda.synchronize()
|
||||
start_event.record()
|
||||
|
||||
for _ in range(n_steps):
|
||||
_ = self.model.generate(dummy_input)
|
||||
|
||||
if torch.cuda.is_available():
|
||||
end_event.record()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
logger.info(
|
||||
f"{self.__class__.__name__}: warmed up! time: {start_event.elapsed_time(end_event) * 1e-3:.3f} s"
|
||||
)
|
||||
|
||||
def process(self, vad_audio: VADAudio):
|
||||
logger.debug("infering moonshine...")
|
||||
|
||||
pred_ids = self.model.generate(vad_audio.audio[None, :])
|
||||
pred_text = self.tokenizer.decode_batch(pred_ids)[0]
|
||||
|
||||
logger.debug("finished whisper inference")
|
||||
console.print(f"[yellow]USER: {pred_text}")
|
||||
|
||||
yield (pred_text, "en")
|
||||
@@ -0,0 +1,130 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from threading import Event
|
||||
from typing import Any, Iterator
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import torch
|
||||
from melo.api import TTS
|
||||
from rich.console import Console
|
||||
|
||||
from speech_to_speech.baseHandler import BaseHandler
|
||||
from speech_to_speech.pipeline.cancel_scope import CancelScope
|
||||
from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut
|
||||
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse
|
||||
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
console = Console()
|
||||
|
||||
WHISPER_LANGUAGE_TO_MELO_LANGUAGE = {
|
||||
"en": "EN",
|
||||
"fr": "FR",
|
||||
"es": "ES",
|
||||
"zh": "ZH",
|
||||
"ja": "JP",
|
||||
"ko": "KR",
|
||||
}
|
||||
|
||||
WHISPER_LANGUAGE_TO_MELO_SPEAKER = {
|
||||
"en": "EN-BR",
|
||||
"fr": "FR",
|
||||
"es": "ES",
|
||||
"zh": "ZH",
|
||||
"ja": "JP",
|
||||
"ko": "KR",
|
||||
}
|
||||
|
||||
|
||||
class MeloTTSHandler(BaseHandler[TTSIn, TTSOut]):
|
||||
def setup(
|
||||
self,
|
||||
should_listen: Event,
|
||||
device: str = "mps",
|
||||
language: str = "en",
|
||||
speaker_to_id: str = "en",
|
||||
gen_kwargs: dict[str, Any] = {}, # Unused
|
||||
blocksize: int = 512,
|
||||
cancel_scope: CancelScope | None = None,
|
||||
speculative_turns: SpeculativeTurnTracker | None = None,
|
||||
) -> None:
|
||||
self.should_listen = should_listen
|
||||
self.cancel_scope = cancel_scope
|
||||
self.speculative_turns = speculative_turns
|
||||
self.device = device
|
||||
self.language = language
|
||||
self.model = TTS(language=WHISPER_LANGUAGE_TO_MELO_LANGUAGE[self.language], device=device)
|
||||
self.speaker_id = self.model.hps.data.spk2id[WHISPER_LANGUAGE_TO_MELO_SPEAKER[speaker_to_id]]
|
||||
self.blocksize = blocksize
|
||||
self._initial_language = self.language
|
||||
self.warmup()
|
||||
|
||||
def warmup(self) -> None:
|
||||
logger.info(f"Warming up {self.__class__.__name__}")
|
||||
_ = self.model.tts_to_file("text", self.speaker_id, quiet=True)
|
||||
|
||||
def process(self, tts_input: TTSIn) -> Iterator[TTSOut]:
|
||||
if isinstance(tts_input, EndOfResponse):
|
||||
yield AUDIO_RESPONSE_DONE
|
||||
return
|
||||
|
||||
speculative_turns = getattr(self, "speculative_turns", None)
|
||||
if speculative_turns and not speculative_turns.is_latest(
|
||||
tts_input.turn_id,
|
||||
tts_input.turn_revision,
|
||||
):
|
||||
logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision)
|
||||
return
|
||||
|
||||
gen = self.cancel_scope.generation if self.cancel_scope else None
|
||||
language_code = tts_input.language_code
|
||||
text = tts_input.text
|
||||
|
||||
console.print(f"[green]ASSISTANT: {text}")
|
||||
|
||||
if language_code is not None and self.language != language_code:
|
||||
try:
|
||||
self.model = TTS(
|
||||
language=WHISPER_LANGUAGE_TO_MELO_LANGUAGE[language_code],
|
||||
device=self.device,
|
||||
)
|
||||
self.speaker_id = self.model.hps.data.spk2id[WHISPER_LANGUAGE_TO_MELO_SPEAKER[language_code]]
|
||||
self.language = language_code
|
||||
except KeyError:
|
||||
console.print(f"[red]Language {language_code} not supported by Melo. Using {self.language} instead.")
|
||||
|
||||
if self.device == "mps":
|
||||
import time
|
||||
|
||||
start = time.time()
|
||||
torch.mps.synchronize() # Waits for all kernels in all streams on the MPS device to complete.
|
||||
torch.mps.empty_cache() # Frees all memory allocated by the MPS device.
|
||||
_ = time.time() - start # Removing this line makes it fail more often. I'm looking into it.
|
||||
|
||||
try:
|
||||
audio_chunk = self.model.tts_to_file(text, self.speaker_id, quiet=True)
|
||||
except (AssertionError, RuntimeError) as e:
|
||||
logger.error(f"Error in MeloTTSHandler: {e}")
|
||||
audio_chunk = np.array([])
|
||||
if len(audio_chunk) == 0:
|
||||
return
|
||||
audio_chunk = librosa.resample(audio_chunk, orig_sr=44100, target_sr=16000)
|
||||
audio_chunk = (audio_chunk * 32768).astype(np.int16)
|
||||
for i in range(0, len(audio_chunk), self.blocksize):
|
||||
if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen):
|
||||
logger.info("TTS generation cancelled (interruption)")
|
||||
return
|
||||
yield np.pad(
|
||||
audio_chunk[i : i + self.blocksize],
|
||||
(0, self.blocksize - len(audio_chunk[i : i + self.blocksize])),
|
||||
)
|
||||
|
||||
def on_session_end(self) -> None:
|
||||
if self.language != self._initial_language:
|
||||
self.language = self._initial_language
|
||||
self.model = TTS(language=WHISPER_LANGUAGE_TO_MELO_LANGUAGE[self.language], device=self.device)
|
||||
self.speaker_id = self.model.hps.data.spk2id[WHISPER_LANGUAGE_TO_MELO_SPEAKER[self.language]]
|
||||
logger.debug("Melo TTS session state reset")
|
||||
@@ -0,0 +1,244 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from threading import Event, Thread
|
||||
from time import perf_counter
|
||||
from typing import Any, Optional
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import torch
|
||||
from parler_tts import ParlerTTSForConditionalGeneration, ParlerTTSStreamer
|
||||
from rich.console import Console
|
||||
from transformers import (
|
||||
AutoTokenizer,
|
||||
)
|
||||
from transformers.utils.import_utils import (
|
||||
is_flash_attn_2_available,
|
||||
)
|
||||
|
||||
from speech_to_speech.baseHandler import BaseHandler
|
||||
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse, TTSInput
|
||||
from speech_to_speech.utils.utils import next_power_of_2
|
||||
|
||||
torch._inductor.config.fx_graph_cache = True
|
||||
# mind about this parameter ! should be >= 2 * number of padded prompt sizes for TTS
|
||||
torch._dynamo.config.cache_size_limit = 15
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
if not is_flash_attn_2_available() and torch.cuda.is_available():
|
||||
logger.warn(
|
||||
"""Parler TTS works best with flash attention 2, but is not installed
|
||||
Given that CUDA is available in this system, you can install flash attention 2 with `uv pip install flash-attn --no-build-isolation`"""
|
||||
)
|
||||
|
||||
|
||||
WHISPER_LANGUAGE_TO_PARLER_SPEAKER = {
|
||||
"en": "Jason",
|
||||
"fr": "Christine",
|
||||
"es": "Steven",
|
||||
"de": "Nicole",
|
||||
"pt": "Sophia",
|
||||
"pl": "Alex",
|
||||
"it": "Richard",
|
||||
"nl": "Mark",
|
||||
}
|
||||
|
||||
|
||||
class ParlerTTSHandler(BaseHandler[TTSInput | EndOfResponse]):
|
||||
def setup(
|
||||
self,
|
||||
should_listen,
|
||||
model_name="parler-tts/parler-mini-v1-jenny",
|
||||
device="cuda",
|
||||
torch_dtype="float16",
|
||||
compile_mode=None,
|
||||
gen_kwargs={},
|
||||
max_prompt_pad_length=8,
|
||||
description=(
|
||||
"Jenny speaks at a slightly slow pace with an animated delivery with clear audio quality."
|
||||
),
|
||||
play_steps_s=1,
|
||||
blocksize=512,
|
||||
use_default_speakers_list=True,
|
||||
cancel_response: Event | None = None,
|
||||
):
|
||||
self.should_listen = should_listen
|
||||
self.cancel_response = cancel_response
|
||||
self.device = device
|
||||
self.torch_dtype = getattr(torch, torch_dtype)
|
||||
self.gen_kwargs = gen_kwargs
|
||||
self.compile_mode = compile_mode
|
||||
self.max_prompt_pad_length = max_prompt_pad_length
|
||||
self.use_default_speakers_list = use_default_speakers_list
|
||||
if self.use_default_speakers_list:
|
||||
description = description.replace("Jenny", "")
|
||||
|
||||
self.speaker = "Jason"
|
||||
self.description = description
|
||||
|
||||
self.model = ParlerTTSForConditionalGeneration.from_pretrained(
|
||||
model_name, torch_dtype=self.torch_dtype
|
||||
).to(device)
|
||||
|
||||
self.description_tokenizer = AutoTokenizer.from_pretrained(self.model.config.text_encoder._name_or_path)
|
||||
self.prompt_tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
|
||||
|
||||
framerate = self.model.audio_encoder.config.frame_rate
|
||||
self.play_steps = int(framerate * play_steps_s)
|
||||
self.blocksize = blocksize
|
||||
|
||||
if self.compile_mode not in (None, "default"):
|
||||
logger.warning(
|
||||
"Torch compilation modes that captures CUDA graphs are not yet compatible with the TTS part. Reverting to 'default'"
|
||||
)
|
||||
self.compile_mode = "default"
|
||||
|
||||
if self.compile_mode:
|
||||
self.model.generation_config.cache_implementation = "static"
|
||||
self.model.forward = torch.compile(
|
||||
self.model.forward, mode=self.compile_mode, fullgraph=True
|
||||
)
|
||||
|
||||
self.warmup()
|
||||
|
||||
def prepare_model_inputs(
|
||||
self,
|
||||
prompt,
|
||||
max_length_prompt=50,
|
||||
pad=False,
|
||||
):
|
||||
pad_args_prompt = (
|
||||
{"padding": "max_length", "max_length": max_length_prompt} if pad else {}
|
||||
)
|
||||
|
||||
description = self.description
|
||||
if self.use_default_speakers_list:
|
||||
description = self.speaker + " " + self.description
|
||||
|
||||
tokenized_description = self.description_tokenizer(
|
||||
description, return_tensors="pt"
|
||||
).to(self.device)
|
||||
input_ids = tokenized_description.input_ids
|
||||
attention_mask = tokenized_description.attention_mask
|
||||
|
||||
tokenized_prompt = self.prompt_tokenizer(
|
||||
prompt, return_tensors="pt", **pad_args_prompt
|
||||
).to(self.device)
|
||||
prompt_input_ids = tokenized_prompt.input_ids
|
||||
prompt_attention_mask = tokenized_prompt.attention_mask
|
||||
|
||||
gen_kwargs = {
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": attention_mask,
|
||||
"prompt_input_ids": prompt_input_ids,
|
||||
"prompt_attention_mask": prompt_attention_mask,
|
||||
**self.gen_kwargs,
|
||||
}
|
||||
|
||||
return gen_kwargs
|
||||
|
||||
def warmup(self):
|
||||
logger.info(f"Warming up {self.__class__.__name__}")
|
||||
|
||||
if self.device == "cuda":
|
||||
start_event = torch.cuda.Event(enable_timing=True)
|
||||
end_event = torch.cuda.Event(enable_timing=True)
|
||||
|
||||
# 2 warmup steps for no compile or compile mode with CUDA graphs capture
|
||||
n_steps = 1 if self.compile_mode == "default" else 2
|
||||
|
||||
if self.device == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
start_event.record()
|
||||
if self.compile_mode:
|
||||
pad_lengths = [2**i for i in range(2, self.max_prompt_pad_length)]
|
||||
for pad_length in pad_lengths[::-1]:
|
||||
model_kwargs = self.prepare_model_inputs(
|
||||
"dummy prompt", max_length_prompt=pad_length, pad=True
|
||||
)
|
||||
for _ in range(n_steps):
|
||||
_ = self.model.generate(**model_kwargs)
|
||||
logger.info(f"Warmed up length {pad_length} tokens!")
|
||||
else:
|
||||
model_kwargs = self.prepare_model_inputs("dummy prompt")
|
||||
for _ in range(n_steps):
|
||||
_ = self.model.generate(**model_kwargs)
|
||||
|
||||
if self.device == "cuda":
|
||||
end_event.record()
|
||||
torch.cuda.synchronize()
|
||||
logger.info(
|
||||
f"{self.__class__.__name__}: warmed up! time: {start_event.elapsed_time(end_event) * 1e-3:.3f} s"
|
||||
)
|
||||
|
||||
def process(self, tts_input: TTSInput | EndOfResponse):
|
||||
if isinstance(tts_input, EndOfResponse):
|
||||
yield AUDIO_RESPONSE_DONE
|
||||
return
|
||||
|
||||
runtime_config = tts_input.runtime_config
|
||||
response = tts_input.response
|
||||
language_code = tts_input.language_code
|
||||
text = tts_input.text
|
||||
|
||||
voice: Optional[str] = None
|
||||
if response and response.audio and response.audio.output:
|
||||
voice = str(response.audio.output.voice) if response.audio.output.voice is not None else None
|
||||
if not voice and runtime_config:
|
||||
audio_cfg = runtime_config.session.audio
|
||||
audio_output = audio_cfg.output if audio_cfg is not None else None
|
||||
voice = str(audio_output.voice) if audio_output is not None and audio_output.voice else None
|
||||
if voice:
|
||||
self.speaker = voice
|
||||
elif language_code:
|
||||
self.speaker = WHISPER_LANGUAGE_TO_PARLER_SPEAKER.get(language_code, "Jason")
|
||||
|
||||
console.print(f"[green]ASSISTANT: {text}")
|
||||
nb_tokens = len(self.prompt_tokenizer(text).input_ids)
|
||||
|
||||
pad_args: dict[str, Any] = {}
|
||||
if self.compile_mode:
|
||||
# pad to closest upper power of two
|
||||
pad_length = next_power_of_2(nb_tokens)
|
||||
logger.debug(f"padding to {pad_length}")
|
||||
pad_args["pad"] = True
|
||||
pad_args["max_length_prompt"] = pad_length
|
||||
|
||||
tts_gen_kwargs = self.prepare_model_inputs(
|
||||
text,
|
||||
**pad_args,
|
||||
)
|
||||
|
||||
streamer = ParlerTTSStreamer(
|
||||
self.model, device=self.device, play_steps=self.play_steps
|
||||
)
|
||||
tts_gen_kwargs = {"streamer": streamer, **tts_gen_kwargs}
|
||||
torch.manual_seed(0)
|
||||
thread = Thread(target=self.model.generate, kwargs=tts_gen_kwargs)
|
||||
thread.start()
|
||||
|
||||
pipeline_start = perf_counter()
|
||||
for i, audio_chunk in enumerate(streamer):
|
||||
if self.cancel_response and self.cancel_response.is_set():
|
||||
logger.info("TTS generation cancelled (interruption)")
|
||||
return
|
||||
if i == 0:
|
||||
logger.info(
|
||||
f"Time to first audio: {perf_counter() - pipeline_start:.3f}s"
|
||||
)
|
||||
audio_chunk = librosa.resample(audio_chunk, orig_sr=44100, target_sr=16000)
|
||||
audio_chunk = (audio_chunk * 32768).astype(np.int16)
|
||||
for i in range(0, len(audio_chunk), self.blocksize):
|
||||
yield np.pad(
|
||||
audio_chunk[i : i + self.blocksize],
|
||||
(0, self.blocksize - len(audio_chunk[i : i + self.blocksize])),
|
||||
)
|
||||
|
||||
if not runtime_config:
|
||||
self.should_listen.set()
|
||||
@@ -0,0 +1,17 @@
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class MeloTTSHandlerArguments:
|
||||
melo_language: str = field(
|
||||
default="en",
|
||||
metadata={"help": "The language of the text to be synthesized. Default is 'EN_NEWEST'."},
|
||||
)
|
||||
melo_device: str = field(
|
||||
default="auto",
|
||||
metadata={"help": "The device to be used for speech synthesis. Default is 'auto'."},
|
||||
)
|
||||
melo_speaker_to_id: str = field(
|
||||
default="en",
|
||||
metadata={"help": "Mapping of speaker names to speaker IDs. Default is ['EN-Newest']."},
|
||||
)
|
||||
@@ -0,0 +1,68 @@
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParlerTTSHandlerArguments:
|
||||
tts_model_name: str = field(
|
||||
default="parler-tts/parler-mini-v1-jenny",
|
||||
metadata={
|
||||
"help": "The pretrained TTS model to use. Default is 'parler-tts/parler-mini-v1-jenny'."
|
||||
},
|
||||
)
|
||||
tts_device: str = field(
|
||||
default="cuda",
|
||||
metadata={
|
||||
"help": "The device type on which the model will run. Default is 'cuda' for GPU acceleration."
|
||||
},
|
||||
)
|
||||
tts_torch_dtype: str = field(
|
||||
default="float16",
|
||||
metadata={
|
||||
"help": "The PyTorch data type for the model and input tensors. One of `float32` (full-precision), `float16` or `bfloat16` (both half-precision)."
|
||||
},
|
||||
)
|
||||
tts_compile_mode: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "Compile mode for torch compile. Either 'default', 'reduce-overhead' and 'max-autotune'. Default is None (no compilation)"
|
||||
},
|
||||
)
|
||||
tts_gen_min_new_tokens: int = field(
|
||||
default=64,
|
||||
metadata={
|
||||
"help": "Maximum number of new tokens to generate in a single completion. Default is 64, which corresponds to ~0.74 secs"
|
||||
},
|
||||
)
|
||||
tts_gen_max_new_tokens: int = field(
|
||||
default=1024,
|
||||
metadata={
|
||||
"help": "Maximum number of new tokens to generate in a single completion. Default is 1024, which corresponds to ~12 secs"
|
||||
},
|
||||
)
|
||||
description: str = field(
|
||||
default=(
|
||||
"Jenny speaks at a slightly slow pace with an animated delivery with clear audio quality."
|
||||
),
|
||||
metadata={
|
||||
"help": "Description of the speaker's voice and speaking style to guide the TTS model."
|
||||
},
|
||||
)
|
||||
play_steps_s: float = field(
|
||||
default=1.0,
|
||||
metadata={
|
||||
"help": "The time interval in seconds for playing back the generated speech in steps. Default is 1.0 seconds."
|
||||
},
|
||||
)
|
||||
max_prompt_pad_length: int = field(
|
||||
default=8,
|
||||
metadata={
|
||||
"help": "When using compilation, the prompt as to be padded to closest power of 2. This parameters sets the maximun power of 2 possible."
|
||||
},
|
||||
)
|
||||
use_default_speakers_list: bool = field(
|
||||
default=False,
|
||||
metadata={
|
||||
"help": "Whether to use the default list of speakers or not."
|
||||
},
|
||||
)
|
||||
Reference in New Issue
Block a user