Spaces:
Running on Zero
Running on Zero
| from kani_tts import KaniTTS | |
| from kani_tts import SpeakerEmbedder | |
| import os | |
| import json | |
| import torch | |
| from omegaconf import OmegaConf | |
| def load_config(config_path: str): | |
| """Load configuration from a YAML file using OmegaConf. | |
| Args: | |
| config_path (str): Path to the YAML configuration file. | |
| Returns: | |
| Any: The loaded OmegaConf DictConfig. | |
| """ | |
| resolved_path = os.path.abspath(config_path) | |
| if not os.path.exists(resolved_path): | |
| raise FileNotFoundError(f"Config file not found: {resolved_path}") | |
| config = OmegaConf.load(resolved_path) | |
| return config | |
| class InitModels: | |
| """ | |
| Lazy initializer that constructs a map of model name -> KaniTTS. | |
| Parameters | |
| ---------- | |
| models_configs : OmegaConf | DictConfig | |
| The `models` section from `model_config.yaml` describing one or | |
| more HF model checkpoints and their options (device_map, use_bematts, etc.). | |
| Returns | |
| ------- | |
| dict | |
| When called, returns a dictionary `{model_name: KaniTTS}`. | |
| Notes | |
| ----- | |
| - All models are loaded immediately in `__call__` so the UI can list | |
| them and switch between them without extra latency. | |
| - Each KaniTTS instance is initialized with its config directly. | |
| """ | |
| def __init__(self, models_configs: OmegaConf): | |
| self.models_configs = models_configs | |
| def __call__(self): | |
| models = {} | |
| for model_name, config in self.models_configs.items(): | |
| print(f"Loading {model_name}...") | |
| # Convert OmegaConf to dict to access parameters | |
| cfg_dict = dict(config) | |
| models[model_name] = KaniTTS( | |
| model_name=cfg_dict.get('model_name'), | |
| device_map=cfg_dict.get('device_map'), | |
| ) | |
| print(f"{model_name} loaded!") | |
| print("All models loaded!") | |
| return models | |
| class SpeakerManager: | |
| """ | |
| Manages speaker embeddings for the TTS application. | |
| Supports two modes: | |
| 1. Select speaker: Load pre-saved speaker embeddings from speaker_map.json | |
| 2. Generate embedding: Generate speaker embedding from uploaded audio using SpeakerEmbedder | |
| Parameters | |
| ---------- | |
| speaker_map_path : str | |
| Path to speaker_map.json file | |
| Methods | |
| ------- | |
| get_speaker_emb(mode, speaker_name=None) -> str | torch.Tensor | None | |
| Returns speaker embedding based on mode: | |
| - "select": Returns path to .pt file from speaker_map | |
| - "generate": Returns cached generated embedding tensor or None | |
| generate_embedding(audio_data, sample_rate) -> torch.Tensor | |
| Generates speaker embedding from audio using SpeakerEmbedder. | |
| Expects audio at 16kHz. Caches the result internally. | |
| clean() | |
| Clears cached generated embedding. | |
| get_speaker_names() -> list[str] | |
| Returns list of available speaker names from speaker_map.json. | |
| """ | |
| def __init__(self, speaker_map_path: str = "./speakers/speaker_map.json"): | |
| self.speaker_map_path = speaker_map_path | |
| self.speaker_map = self._load_speaker_map() | |
| self.cached_embedding = None | |
| self.embedder = None | |
| def _load_speaker_map(self): | |
| """Load speaker map from JSON file.""" | |
| if not os.path.exists(self.speaker_map_path): | |
| return {} | |
| with open(self.speaker_map_path, "r") as f: | |
| return json.load(f) | |
| def get_speaker_names(self): | |
| """Get list of available speaker names.""" | |
| return list(self.speaker_map.keys()) | |
| def get_speaker_emb(self, mode: str, speaker_name: str = None): | |
| """ | |
| Get speaker embedding based on mode. | |
| Parameters | |
| ---------- | |
| mode : str | |
| Either "select" or "generate" | |
| speaker_name : str, optional | |
| Name of speaker from speaker_map (only used in "select" mode) | |
| Returns | |
| ------- | |
| str | torch.Tensor | None | |
| Path to .pt file (select mode) or embedding tensor (generate mode) | |
| """ | |
| if mode == "select": | |
| if speaker_name and speaker_name in self.speaker_map: | |
| path = self.speaker_map[speaker_name] | |
| print(f"[SpeakerManager] Returning speaker path: {path}") | |
| return path | |
| return None | |
| elif mode == "generate": | |
| print(f"[SpeakerManager] Cached embedding: {self.cached_embedding}") | |
| print(f"[SpeakerManager] Cached embedding type: {type(self.cached_embedding)}") | |
| if self.cached_embedding is not None: | |
| print(f"[SpeakerManager] Cached embedding shape: {self.cached_embedding.shape}") | |
| print(f"[SpeakerManager] Cached embedding device: {self.cached_embedding.device}") | |
| return self.cached_embedding | |
| return None | |
| def generate_embedding(self, audio_data, sample_rate: int): | |
| """ | |
| Generate speaker embedding from audio data. | |
| Parameters | |
| ---------- | |
| audio_data : tuple | np.ndarray | |
| Either (sample_rate, audio_array) tuple from Gradio or numpy array | |
| sample_rate : int | |
| Sample rate of the audio (default: 16000) | |
| Returns | |
| ------- | |
| torch.Tensor | |
| Generated speaker embedding [1, 128] | |
| """ | |
| # Initialize embedder lazily | |
| if self.embedder is None: | |
| print("[SpeakerManager] Initializing SpeakerEmbedder...") | |
| self.embedder = SpeakerEmbedder() | |
| # Handle Gradio audio format (sr, audio) tuple | |
| if isinstance(audio_data, tuple): | |
| sample_rate, audio_array = audio_data | |
| print(f"[SpeakerManager] Audio tuple: sr={sample_rate}, shape={audio_array.shape}") | |
| else: | |
| audio_array = audio_data | |
| print(f"[SpeakerManager] Audio array shape: {audio_array.shape}") | |
| # Generate embedding | |
| print("[SpeakerManager] Generating embedding...") | |
| embedding = self.embedder.embed_audio(audio_array, sample_rate=sample_rate) | |
| print(f"[SpeakerManager] Generated embedding shape: {embedding.shape}, device: {embedding.device}") | |
| # Cache the result | |
| self.cached_embedding = embedding | |
| print(f"[SpeakerManager] Cached embedding (id={id(self.cached_embedding)})") | |
| return embedding | |
| def clean(self): | |
| """Clear cached generated embedding.""" | |
| self.cached_embedding = None | |
| return "Embedding cleared" | |
| def get_status(self): | |
| """Get current status of generated embedding.""" | |
| if self.cached_embedding is not None: | |
| return "✅ Embedding ready" | |
| return "No embedding generated" | |
| class Examples: | |
| """ | |
| Adapter that converts YAML examples into Gradio `gr.Examples` rows. | |
| Parameters | |
| ---------- | |
| exam_cfg : OmegaConf | DictConfig | |
| Parsed contents of `examples.yaml`. Expected structure: | |
| `examples: [ {text, model, speaker?, temperature?, top_p?, repetition_penalty?}, ... ]`. | |
| Behavior | |
| -------- | |
| - Produces a list-of-lists whose order must match the `inputs` order | |
| used when constructing `gr.Examples` in `app.py`. | |
| - Current order: `[text, model_dropdown, speaker_mode, speaker_dropdown, temp, top_p, rp]`. | |
| Why this exists | |
| --------------- | |
| - Keeps format and defaults centralized, so changing the UI inputs | |
| order only requires a single change here and in `app.py`. | |
| """ | |
| def __init__(self, exam_cfg: OmegaConf): | |
| self.exam_cfg = exam_cfg | |
| def __call__(self) -> list[list]: | |
| rows = [] | |
| for e in self.exam_cfg.examples: | |
| text = e.get("text") | |
| model = e.get("model") | |
| speaker_mode = e.get("speaker_mode", "select") # Default to "select" mode | |
| speaker = e.get("speaker", "Kore (en)") | |
| temperature = e.get("temperature", 1.0) | |
| top_p = e.get("top_p", 0.95) | |
| repetition_penalty = e.get("repetition_penalty", 1.1) | |
| # Order must match gr.Examples inputs: [text, model_dropdown, speaker_mode, speaker_dropdown, temp, top_p, rp] | |
| rows.append([text, model, speaker_mode, speaker, temperature, top_p, repetition_penalty]) | |
| return rows | |