feat: implement cached audio cloning functionality

- Added a new module `cheched_vallex.py` for cached audio cloning, allowing users to save and infer from audio prompts.
- Integrated the cached functionality into the main application, providing a new tab for zero-shot audio cloning with cached prompts.
- Enhanced the `infer_from_audio` function to include timing metrics for better performance tracking.

apps/audio_cloning/cheched_vallex.py

@@ -0,0 +1,363 @@
+import logging
+import os
+import re
+import shutil
+import time
+from typing import List, Optional, Tuple
+
+import gradio as gr
+import numpy as np
+import torch
+
+from .vallex import main as vallex
+from .vallex.descriptions import infer_from_audio_ja_md, top_ja_md
+from .vallex.examples import infer_from_audio_examples
+from .vallex.macros import code2lang, lang2token, langdropdown2token, token2lang
+
+logger = logging.getLogger(__name__)
+
+PROMPTS_DIR = "./models/prompts"
+PROMPT_ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]+$")
+
+
+def _ensure_prompt_dir() -> str:
+    os.makedirs(PROMPTS_DIR, exist_ok=True)
+    return PROMPTS_DIR
+
+
+def _list_saved_prompts() -> List[str]:
+    directory = _ensure_prompt_dir()
+    files = [f for f in os.listdir(directory) if f.endswith(".npz")]
+    return sorted(files)
+
+
+def _format_prompt_list() -> str:
+    prompts = _list_saved_prompts()
+    return "\n".join(prompts) if prompts else "保存済みプロンプトはありません。"
+
+
+def save_prompt_to_cache(
+    prompt_id: str,
+    upload_audio_prompt: Optional[Tuple[int, np.ndarray]],
+    record_audio_prompt: Optional[Tuple[int, np.ndarray]],
+    transcript_content: str,
+):
+    prompt_id = prompt_id.strip()
+    if prompt_id.lower().endswith(".npz"):
+        prompt_id = prompt_id[:-4]
+
+    if not prompt_id:
+        return (
+            "プロンプト ID を入力してください。",
+            None,
+            gr.update(choices=_list_saved_prompts(), value=None),
+            gr.update(value=_format_prompt_list()),
+        )
+
+    if not PROMPT_ID_PATTERN.match(prompt_id):
+        return (
+            "プロンプト ID には英数字・ハイフン・アンダースコアのみ使用できます。",
+            None,
+            gr.update(choices=_list_saved_prompts(), value=None),
+            gr.update(value=_format_prompt_list()),
+        )
+
+    audio_prompt = (
+        upload_audio_prompt if upload_audio_prompt is not None else record_audio_prompt
+    )
+    if audio_prompt is None:
+        return (
+            "音声をアップロードするか録音してください。",
+            None,
+            gr.update(choices=_list_saved_prompts(), value=None),
+            gr.update(value=_format_prompt_list()),
+        )
+
+    try:
+        message, temp_path = vallex.make_npz_prompt(
+            prompt_id,
+            upload_audio_prompt,
+            record_audio_prompt,
+            transcript_content,
+        )
+    except Exception as err:  # pylint: disable=broad-except
+        logger.exception("Failed to create prompt", exc_info=err)
+        return (
+            f"プロンプト作成に失敗しました: {err}",
+            None,
+            gr.update(choices=_list_saved_prompts(), value=None),
+            gr.update(value=_format_prompt_list()),
+        )
+
+    _ensure_prompt_dir()
+    cached_filename = f"{prompt_id}.npz"
+    cached_path = os.path.join(PROMPTS_DIR, cached_filename)
+
+    try:
+        shutil.copy(temp_path, cached_path)
+    except OSError as err:
+        logger.exception("Failed to copy prompt to cache", exc_info=err)
+        return (
+            f"プロンプトの保存に失敗しました: {err}",
+            None,
+            gr.update(choices=_list_saved_prompts(), value=None),
+            gr.update(value=_format_prompt_list()),
+        )
+    finally:
+        try:
+            os.remove(temp_path)
+        except OSError:
+            pass
+
+    choices = _list_saved_prompts()
+    message = (
+        f"{message}\nSaved cached prompt to {cached_path}"
+        if message
+        else f"Saved cached prompt to {cached_path}"
+    )
+    return (
+        message,
+        cached_path,
+        gr.update(choices=choices, value=cached_filename),
+        gr.update(value=_format_prompt_list()),
+    )
+
+
+def refresh_prompt_choices():
+    choices = _list_saved_prompts()
+    value = choices[0] if choices else None
+    return (
+        gr.update(choices=choices, value=value),
+        gr.update(value=_format_prompt_list()),
+    )
+
+
+def infer_from_cached_prompt(
+    text: str,
+    language: str,
+    accent: str,
+    prompt_filename: Optional[str],
+):
+    if not text:
+        return "テキストを入力してください。", None
+
+    if not prompt_filename:
+        return "プロンプトを選択してください。", None
+
+    prompt_path = os.path.join(_ensure_prompt_dir(), prompt_filename)
+    if not os.path.exists(prompt_path):
+        return f"プロンプトが見つかりません: {prompt_path}", None
+
+    timings: List[Tuple[str, float]] = []
+    start_time = time.perf_counter()
+    try:
+        logger.info("Loading cached prompt from: %s", prompt_path)
+        prompt_data = np.load(prompt_path)
+        audio_tokens = torch.from_numpy(prompt_data["audio_tokens"]).to(
+            dtype=torch.long
+        )
+        text_prompts = torch.from_numpy(prompt_data["text_tokens"]).to(dtype=torch.long)
+        lang_code = (
+            int(prompt_data["lang_code"])
+            if prompt_data["lang_code"].shape == ()
+            else int(prompt_data["lang_code"][0])
+        )
+    except Exception as err:  # pylint: disable=broad-except
+        logger.exception("Failed to load cached prompt", exc_info=err)
+        return (f"プロンプトの読み込みに失敗しました: {err}", None)
+    timings.append(("プロンプト読込", time.perf_counter() - start_time))
+
+    lang_pr = code2lang.get(lang_code, "en")
+
+    start_time = time.perf_counter()
+    if language == "auto-detect":
+        detected_lang = vallex.langid.classify(text)[0]
+        lang_token = lang2token.get(detected_lang, "[EN]")
+    else:
+        lang_token = langdropdown2token[language]
+
+    conditioned_text = f"{lang_token}{text}{lang_token}"
+
+    phone_tokens, langs = vallex.text_tokenizer.tokenize(
+        text=f"_{conditioned_text}".strip()
+    )
+    text_tokens, text_tokens_lens = vallex.text_collater([phone_tokens])
+
+    enroll_x_lens = torch.IntTensor([text_prompts.shape[-1]])
+    text_tokens = torch.cat([text_prompts, text_tokens], dim=-1)
+    text_tokens_lens += enroll_x_lens
+    timings.append(("テキスト準備", time.perf_counter() - start_time))
+
+    vallex.model.to(vallex.device)
+
+    audio_prompts = audio_tokens.to(vallex.device)
+    if audio_prompts.dim() == 2:
+        audio_prompts = audio_prompts.unsqueeze(0)
+
+    start_time = time.perf_counter()
+    logger.info("Start inferring from cached prompt: %s", prompt_path)
+    encoded_frames = vallex.model.inference(
+        text_tokens.to(vallex.device),
+        text_tokens_lens.to(vallex.device),
+        audio_prompts,
+        enroll_x_lens=enroll_x_lens.to(vallex.device),
+        top_k=-100,
+        temperature=1,
+        prompt_language=lang_pr,
+        text_language=langs
+        if accent == "no-accent"
+        else token2lang[langdropdown2token[accent]],
+        best_of=5,
+    )
+    timings.append(("モデル推論", time.perf_counter() - start_time))
+    logger.info("Inference completed")
+
+    start_time = time.perf_counter()
+    logger.info("Decoding with Vocos...")
+    frames = encoded_frames.permute(2, 0, 1)
+    features = vallex.vocos.codes_to_features(frames)
+    samples = vallex.vocos.decode(
+        features, bandwidth_id=torch.tensor([2], device=vallex.device)
+    )
+    timings.append(("ボコーダ復号", time.perf_counter() - start_time))
+    logger.info("Decoding completed")
+
+    message = (
+        f"Loaded cached prompt: {prompt_filename}\n"
+        f"Prompt language: {lang_pr}\n"
+        f"Synthesized text: {conditioned_text}"
+    )
+
+    timing_report = "\n↓\n".join(
+        f"{step}：{duration:.4f} sec" for step, duration in timings
+    )
+    logger.info("推論ステップ計測結果\n%s", timing_report)
+
+    return message, (24000, samples.squeeze(0).cpu().numpy())
+
+
+def main():
+    prompt_choices = _list_saved_prompts()
+
+    gr.Markdown(top_ja_md)
+    gr.Markdown(infer_from_audio_ja_md)
+    gr.Markdown("[Cached] Zero-shot 音声クローニング")
+    with gr.Row():
+        with gr.Column():
+            textbox = gr.TextArea(
+                label="音声合成で喋らせたいテキスト",
+                placeholder="ここに音声合成で喋らせたいテキストを入力してください。",
+                value="Welcome back, Master. What can I do for you today?",
+                elem_id="tts-input-cached",
+            )
+            language_dropdown = gr.Dropdown(
+                choices=["auto-detect", "English", "中文", "日本語"],
+                value="auto-detect",
+                label="language",
+            )
+            accent_dropdown = gr.Dropdown(
+                choices=["no-accent", "English", "中文", "日本語"],
+                value="no-accent",
+                label="accent",
+            )
+            textbox_transcript = gr.TextArea(
+                label="Transcript",
+                placeholder="アップロードした音声、または録音した音声のテキストを入力してください。(whisper を使用する場合は空のままにしてください。)",
+                value="",
+            )
+            upload_audio_prompt = gr.Audio(
+                label="音声アップロード",
+                sources=["upload"],
+                interactive=True,
+            )
+            record_audio_prompt = gr.Audio(
+                label="音声を録音する",
+                sources=["microphone"],
+                interactive=True,
+            )
+            prompt_id_box = gr.Textbox(
+                label="Prompt ID",
+                placeholder="例: my_speaker01",
+                value="",
+            )
+            cached_prompt_dropdown = gr.Dropdown(
+                label="Cached prompts",
+                choices=prompt_choices,
+                value=prompt_choices[0] if prompt_choices else None,
+                interactive=True,
+            )
+            prompt_list_box = gr.Textbox(
+                label="保存済みプロンプト一���",
+                value=_format_prompt_list(),
+                interactive=False,
+                lines=6,
+            )
+            refresh_btn = gr.Button("キャッシュ一覧を更新")
+
+        with gr.Column():
+            text_output = gr.Textbox(label="Message")
+            audio_output = gr.Audio(label="Output Audio", elem_id="tts-audio-cached")
+            btn_infer = gr.Button("音声合成を開始する")
+            btn_infer.click(
+                vallex.infer_from_audio,
+                inputs=[
+                    textbox,
+                    language_dropdown,
+                    accent_dropdown,
+                    upload_audio_prompt,
+                    record_audio_prompt,
+                    textbox_transcript,
+                ],
+                outputs=[text_output, audio_output],
+            )
+
+            prompt_output = gr.File(label="Generated prompt", interactive=False)
+            btn_save = gr.Button("./models/prompts に保存")
+            btn_save.click(
+                save_prompt_to_cache,
+                inputs=[
+                    prompt_id_box,
+                    upload_audio_prompt,
+                    record_audio_prompt,
+                    textbox_transcript,
+                ],
+                outputs=[
+                    text_output,
+                    prompt_output,
+                    cached_prompt_dropdown,
+                    prompt_list_box,
+                ],
+            )
+
+            btn_cached_infer = gr.Button("キャッシュしたプロンプトで合成")
+            btn_cached_infer.click(
+                infer_from_cached_prompt,
+                inputs=[
+                    textbox,
+                    language_dropdown,
+                    accent_dropdown,
+                    cached_prompt_dropdown,
+                ],
+                outputs=[text_output, audio_output],
+            )
+
+    refresh_btn.click(
+        refresh_prompt_choices,
+        inputs=None,
+        outputs=[cached_prompt_dropdown, prompt_list_box],
+    )
+
+    gr.Examples(
+        examples=infer_from_audio_examples,
+        inputs=[
+            textbox,
+            language_dropdown,
+            accent_dropdown,
+            upload_audio_prompt,
+            record_audio_prompt,
+            textbox_transcript,
+        ],
+        outputs=[text_output, audio_output],
+        fn=vallex.infer_from_audio,
+        cache_examples=False,
+    )

apps/audio_cloning/main.py

@@ -4,6 +4,7 @@ import gradio as gr
 
 from logger import setup_logger
 
+from .cheched_vallex import main as cached_vallex
 from .vallex.main import main as vallex
 
 logger = logging.getLogger(__name__)
@@ -18,6 +19,9 @@ def main():
     gr.Markdown("# Charamix Audio Cloning Prototype")
 
     # zero-shot audio cloning
+    with gr.Tab("[Cached] Zero-shot Audio Cloning"):
+        cached_vallex()
+
     with gr.Tab("Zero-shot Audio Cloning with VALL-E-X"):
         vallex()
 

apps/audio_cloning/vallex/main.py

@@ -359,6 +359,9 @@ def infer_from_audio(
     text, language, accent, audio_prompt, record_audio_prompt, transcript_content
 ):
     global model, text_collater, text_tokenizer, audio_tokenizer
+    timings = []
+
+    start_time = time.perf_counter()
     audio_prompt = audio_prompt if audio_prompt is not None else record_audio_prompt
     sr, wav_pr = audio_prompt
     if not isinstance(wav_pr, torch.FloatTensor):
@@ -370,28 +373,36 @@ def infer_from_audio(
     if wav_pr.ndim == 1:
         wav_pr = wav_pr.unsqueeze(0)
     assert wav_pr.ndim and wav_pr.size(0) == 1
+    timings.append(("音声前処理", time.perf_counter() - start_time))
 
+    start_time = time.perf_counter()
     if transcript_content == "":
         text_pr, lang_pr = make_prompt("dummy", wav_pr, sr, save=False)
     else:
         lang_pr = langid.classify(str(transcript_content))[0]
         lang_token = lang2token[lang_pr]
         text_pr = f"{lang_token}{str(transcript_content)}{lang_token}"
+    timings.append(("プロンプト生成", time.perf_counter() - start_time))
 
+    start_time = time.perf_counter()
     if language == "auto-detect":
         lang_token = lang2token[langid.classify(text)[0]]
     else:
         lang_token = langdropdown2token[language]
     lang = token2lang[lang_token]
     text = lang_token + text + lang_token
+    timings.append(("言語設定", time.perf_counter() - start_time))
 
     # onload model
     model.to(device)
 
+    start_time = time.perf_counter()
     # tokenize audio
     encoded_frames = tokenize_audio(audio_tokenizer, (wav_pr, sr))
     audio_prompts = encoded_frames[0][0].transpose(2, 1).to(device)
+    timings.append(("音声トークナイズ", time.perf_counter() - start_time))
 
+    start_time = time.perf_counter()
     # tokenize text
     logging.info(f"synthesize text: {text}")
     phone_tokens, langs = text_tokenizer.tokenize(text=f"_{text}".strip())
@@ -404,6 +415,9 @@ def infer_from_audio(
     text_tokens = torch.cat([text_prompts, text_tokens], dim=-1)
     text_tokens_lens += enroll_x_lens
     lang = lang if accent == "no-accent" else token2lang[langdropdown2token[accent]]
+    timings.append(("テキストトークナイズ", time.perf_counter() - start_time))
+
+    start_time = time.perf_counter()
     encoded_frames = model.inference(
         text_tokens.to(device),
         text_tokens_lens.to(device),
@@ -415,14 +429,18 @@ def infer_from_audio(
         text_language=langs if accent == "no-accent" else lang,
         best_of=5,
     )
+    timings.append(("モデル推論", time.perf_counter() - start_time))
     # Decode with Vocos
+    start_time = time.perf_counter()
     frames = encoded_frames.permute(2, 0, 1)
     features = vocos.codes_to_features(frames)
     samples = vocos.decode(features, bandwidth_id=torch.tensor([2], device=device))
+    timings.append(("ボコーダ復号", time.perf_counter() - start_time))
 
-    # offload model
-    model.to("cpu")
-    torch.cuda.empty_cache()
+    timing_report = "\n↓\n".join(
+        f"{step}：{duration:.4f} sec" for step, duration in timings
+    )
+    logger.info("推論ステップ計測結果\n%s", timing_report)
 
     message = f"text prompt: {text_pr}\nsythesized text: {text}"
     return message, (24000, samples.squeeze(0).cpu().numpy())