Upload app.py with huggingface_hub

app.py

@@ -1,110 +1,225 @@
+import json
 import os
-import shutil
+import sys
 import threading
 import time
-import sys
-import torch
 
-from huggingface_hub import snapshot_download
+import warnings
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", category=UserWarning)
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(current_dir)
 sys.path.append(os.path.join(current_dir, "indextts"))
 
+import argparse
+parser = argparse.ArgumentParser(description="IndexTTS WebUI")
+parser.add_argument("--verbose", action="store_true", default=False, help="Enable verbose mode")
+parser.add_argument("--port", type=int, default=7860, help="Port to run the web UI on")
+parser.add_argument("--host", type=str, default="127.0.0.1", help="Host to run the web UI on")
+parser.add_argument("--model_dir", type=str, default="checkpoints", help="Model checkpoints directory")
+cmd_args = parser.parse_args()
+MODE = 'local'
+snapshot_download("IndexTeam/IndexTTS-1.5", local_dir="checkpoints")
+
+if not os.path.exists(cmd_args.model_dir):
+    print(f"Model directory {cmd_args.model_dir} does not exist. Please download the model first.")
+    sys.exit(1)
+
+for file in [
+    "bigvgan_generator.pth",
+    "bpe.model",
+    "gpt.pth",
+    "config.yaml",
+]:
+    file_path = os.path.join(cmd_args.model_dir, file)
+    if not os.path.exists(file_path):
+        print(f"Required file {file_path} does not exist. Please download it.")
+        sys.exit(1)
+
 import gradio as gr
+
 from indextts.infer import IndexTTS
 from tools.i18n.i18n import I18nAuto
 
-# 设置多语言
-i18n = I18nAuto(language="en")
-
-# 下载模型
+i18n = I18nAuto(language="zh_CN")
 MODE = 'local'
-snapshot_download("IndexTeam/IndexTTS-1.5", local_dir="checkpoints")
+tts = IndexTTS(model_dir=cmd_args.model_dir, cfg_path=os.path.join(cmd_args.model_dir, "config.yaml"),)
+
+
+os.makedirs("outputs/tasks",exist_ok=True)
+os.makedirs("prompts",exist_ok=True)
 
-# 自动选择设备：优先 GPU，没有就 CPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"🔥 Using device: {device}")
-
-# 初始化 TTS (不传 device)
-tts = IndexTTS(model_dir="checkpoints", cfg_path="checkpoints/config.yaml")
-
-# 如果 IndexTTS 支持 to() 或 to_device()，切换设备
-if hasattr(tts, "to"):
-    tts.to(device)
-elif hasattr(tts, "to_device"):
-    tts.to_device(device)
-else:
-    print("⚠️ IndexTTS 没有 to()/to_device() 方法，可能内部已自动处理设备。")
-
-# 确保必要的目录存在
-os.makedirs("outputs/tasks", exist_ok=True)
-os.makedirs("prompts", exist_ok=True)
-
-# 推理函数
-def infer(voice, text, output_path=None):
-    if not tts:
-        raise Exception("Model not loaded")
+with open("tests/cases.jsonl", "r", encoding="utf-8") as f:
+    example_cases = []
+    for line in f:
+        line = line.strip()
+        if not line:
+            continue
+        example = json.loads(line)
+        example_cases.append([os.path.join("tests", example.get("prompt_audio", "sample_prompt.wav")),
+                              example.get("text"), ["普通推理", "批次推理"][example.get("infer_mode", 0)]])
+
+def gen_single(prompt, text, infer_mode, max_text_tokens_per_sentence=120, sentences_bucket_max_size=4,
+                *args, progress=gr.Progress()):
+    output_path = None
     if not output_path:
         output_path = os.path.join("outputs", f"spk_{int(time.time())}.wav")
-    tts.infer(voice, text, output_path)
-    return output_path
-
-# API 调用
-def tts_api(voice, text):
-    try:
-        output_path = infer(voice, text)
-        with open(output_path, "rb") as f:
-            audio_bytes = f.read()
-        return (200, {}, audio_bytes)
-    except Exception as e:
-        return (500, {"error": str(e)}, None)
-
-# 单次生成
-def gen_single(prompt, text):
-    output_path = infer(prompt, text)
-    return gr.update(value=output_path, visible=True)
-
-# 上传后启用按钮
+    # set gradio progress
+    tts.gr_progress = progress
+    do_sample, top_p, top_k, temperature, \
+        length_penalty, num_beams, repetition_penalty, max_mel_tokens = args
+    kwargs = {
+        "do_sample": bool(do_sample),
+        "top_p": float(top_p),
+        "top_k": int(top_k) if int(top_k) > 0 else None,
+        "temperature": float(temperature),
+        "length_penalty": float(length_penalty),
+        "num_beams": num_beams,
+        "repetition_penalty": float(repetition_penalty),
+        "max_mel_tokens": int(max_mel_tokens),
+        # "typical_sampling": bool(typical_sampling),
+        # "typical_mass": float(typical_mass),
+    }
+    if infer_mode == "普通推理":
+        output = tts.infer(prompt, text, output_path, verbose=cmd_args.verbose,
+                           max_text_tokens_per_sentence=int(max_text_tokens_per_sentence),
+                           **kwargs)
+    else:
+        # 批次推理
+        output = tts.infer_fast(prompt, text, output_path, verbose=cmd_args.verbose,
+            max_text_tokens_per_sentence=int(max_text_tokens_per_sentence),
+            sentences_bucket_max_size=(sentences_bucket_max_size),
+            **kwargs)
+    return gr.update(value=output,visible=True)
+
 def update_prompt_audio():
     update_button = gr.update(interactive=True)
     return update_button
 
-# Gradio 界面
-with gr.Blocks() as demo:
+with gr.Blocks(title="IndexTTS Demo") as demo:
     mutex = threading.Lock()
-    gr.HTML(f'''
-    <h2><center>IndexTTS WebUI</center></h2>
-    <p align="center">当前设备: <b>{device}</b></p>
+    gr.HTML('''
+    <h2><center>IndexTTS: An Industrial-Level Controllable and Efficient Zero-Shot Text-To-Speech System</h2>
+    <h2><center>(一款工业级可控且高效的零样本文本转语音系统)</h2>
+<p align="center">
+<a href='https://arxiv.org/abs/2502.05512'><img src='https://img.shields.io/badge/ArXiv-2502.05512-red'></a>
+</p>
     ''')
     with gr.Tab("音频生成"):
         with gr.Row():
-            os.makedirs("prompts", exist_ok=True)
-            prompt_audio = gr.Audio(
-                label="请上传参考音频",
-                key="prompt_audio",
-                sources=["upload", "microphone"],
-                type="filepath"
-            )
+            os.makedirs("prompts",exist_ok=True)
+            prompt_audio = gr.Audio(label="参考音频",key="prompt_audio",
+                                    sources=["upload","microphone"],type="filepath")
             prompt_list = os.listdir("prompts")
             default = ''
             if prompt_list:
                 default = prompt_list[0]
-            input_text_single = gr.Textbox(label="请输入目标文本", key="input_text_single")
-            gen_button = gr.Button("生成语音", key="gen_button", interactive=True)
-            output_audio = gr.Audio(label="生成结果", visible=False, key="output_audio")
+            with gr.Column():
+                input_text_single = gr.TextArea(label="文本",key="input_text_single", placeholder="请输入目标文本", info="当前模型版本{}".format(tts.model_version or "1.0"))
+                infer_mode = gr.Radio(choices=["普通推理", "批次推理"], label="推理模式",info="批次推理：更适合长句，性能翻倍",value="普通推理")        
+                gen_button = gr.Button("生成语音", key="gen_button",interactive=True)
+            output_audio = gr.Audio(label="生成结果", visible=True,key="output_audio")
+        with gr.Accordion("高级生成参数设置", open=False):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown("**GPT2 采样设置** _参数会影响音频多样性和生成速度详见[Generation strategies](https://huggingface.co/docs/transformers/main/en/generation_strategies)_")
+                    with gr.Row():
+                        do_sample = gr.Checkbox(label="do_sample", value=True, info="是否进行采样")
+                        temperature = gr.Slider(label="temperature", minimum=0.1, maximum=2.0, value=1.0, step=0.1)
+                    with gr.Row():
+                        top_p = gr.Slider(label="top_p", minimum=0.0, maximum=1.0, value=0.8, step=0.01)
+                        top_k = gr.Slider(label="top_k", minimum=0, maximum=100, value=30, step=1)
+                        num_beams = gr.Slider(label="num_beams", value=3, minimum=1, maximum=10, step=1)
+                    with gr.Row():
+                        repetition_penalty = gr.Number(label="repetition_penalty", precision=None, value=10.0, minimum=0.1, maximum=20.0, step=0.1)
+                        length_penalty = gr.Number(label="length_penalty", precision=None, value=0.0, minimum=-2.0, maximum=2.0, step=0.1)
+                    max_mel_tokens = gr.Slider(label="max_mel_tokens", value=600, minimum=50, maximum=tts.cfg.gpt.max_mel_tokens, step=10, info="生成Token最大数量，过小导致音频被截断", key="max_mel_tokens")
+                    # with gr.Row():
+                    #     typical_sampling = gr.Checkbox(label="typical_sampling", value=False, info="不建议使用")
+                    #     typical_mass = gr.Slider(label="typical_mass", value=0.9, minimum=0.0, maximum=1.0, step=0.1)
+                with gr.Column(scale=2):
+                    gr.Markdown("**分句设置** _参数会影响音频质量和生成速度_")
+                    with gr.Row():
+                        max_text_tokens_per_sentence = gr.Slider(
+                            label="分句最大Token数", value=120, minimum=20, maximum=tts.cfg.gpt.max_text_tokens, step=2, key="max_text_tokens_per_sentence",
+                            info="建议80~200之间，值越大，分句越长；值越小，分句越碎；过小过大都可能导致音频质量不高",
+                        )
+                        sentences_bucket_max_size = gr.Slider(
+                            label="分句分桶的最大容量（批次推理生效）", value=4, minimum=1, maximum=16, step=1, key="sentences_bucket_max_size",
+                            info="建议2-8之间，值越大，一批次推理包含的分句数越多，过大可能导致内存溢出",
+                        )
+                    with gr.Accordion("预览分句结果", open=True) as sentences_settings:
+                        sentences_preview = gr.Dataframe(
+                            headers=["序号", "分句内容", "Token数"],
+                            key="sentences_preview",
+                            wrap=True,
+                        )
+            advanced_params = [
+                do_sample, top_p, top_k, temperature,
+                length_penalty, num_beams, repetition_penalty, max_mel_tokens,
+                # typical_sampling, typical_mass,
+            ]
+        
+        if len(example_cases) > 0:
+            gr.Examples(
+                examples=example_cases,
+                inputs=[prompt_audio, input_text_single, infer_mode],
+            )
 
+    def on_input_text_change(text, max_tokens_per_sentence):
+        if text and len(text) > 0:
+            text_tokens_list = tts.tokenizer.tokenize(text)
+
+            sentences = tts.tokenizer.split_sentences(text_tokens_list, max_tokens_per_sentence=int(max_tokens_per_sentence))
+            data = []
+            for i, s in enumerate(sentences):
+                sentence_str = ''.join(s)
+                tokens_count = len(s)
+                data.append([i, sentence_str, tokens_count])
+            
+            return {
+                sentences_preview: gr.update(value=data, visible=True, type="array"),
+            }
+        else:
+            df = pd.DataFrame([], columns=["序号", "分句内容", "Token数"])
+            return {
+                sentences_preview: gr.update(value=df)
+            }
+
+    input_text_single.change(
+        on_input_text_change,
+        inputs=[input_text_single, max_text_tokens_per_sentence],
+        outputs=[sentences_preview]
+    )
+    max_text_tokens_per_sentence.change(
+        on_input_text_change,
+        inputs=[input_text_single, max_text_tokens_per_sentence],
+        outputs=[sentences_preview]
+    )
     prompt_audio.upload(update_prompt_audio,
                          inputs=[],
                          outputs=[gen_button])
 
     gen_button.click(gen_single,
-                     inputs=[prompt_audio, input_text_single],
+                     inputs=[prompt_audio, input_text_single, infer_mode,
+                             max_text_tokens_per_sentence, sentences_bucket_max_size,
+                             *advanced_params,
+                     ],
                      outputs=[output_audio])
 
+
+
+
+# if __name__ == "__main__":
+#     demo.queue(20)
+#     demo.launch(server_name=cmd_args.host, server_port=cmd_args.port)
+
+
+
 def main():
     tts.load_normalizer()
     demo.launch(server_name="0.0.0.0", server_port=7860)
 
 if __name__ == "__main__":
-    main()
+    main()