Update app.py

app.py

@@ -1,520 +1,28 @@
-import gc
-import html
-import io
-import os
-import queue
-import wave
-from argparse import ArgumentParser
-from functools import partial
-from pathlib import Path
-
-import gradio as gr
-import librosa
-import numpy as np
-import pyrootutils
-import torch
-from loguru import logger
-from transformers import AutoTokenizer
-
-pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
-
-from fish_speech.i18n import i18n
-from fish_speech.text.chn_text_norm.text import Text as ChnNormedText
-from fish_speech.utils import autocast_exclude_mps
-from tools.api import decode_vq_tokens, encode_reference
-from tools.auto_rerank import batch_asr, calculate_wer, is_chinese, load_model
-from tools.llama.generate import (
-    GenerateRequest,
-    GenerateResponse,
-    WrappedGenerateResponse,
-    launch_thread_safe_queue,
-)
-from tools.vqgan.inference import load_model as load_decoder_model
-
-# Make einx happy
-os.environ["EINX_FILTER_TRACEBACK"] = "false"
-
-HEADER_MD = f"""# Fish Speech
-
-{i18n("A text-to-speech model based on VQ-GAN and Llama developed by [Fish Audio](https://fish.audio).")}  
-
-{i18n("You can find the source code [here](https://github.com/fishaudio/fish-speech) and models [here](https://huggingface.co/fishaudio/fish-speech-1.4).")}  
-
-{i18n("Related code and weights are released under CC BY-NC-SA 4.0 License.")}  
-
-{i18n("We are not responsible for any misuse of the model, please consider your local laws and regulations before using it.")}  
-"""
-
-TEXTBOX_PLACEHOLDER = i18n("Put your text here.")
-SPACE_IMPORTED = False
-
-# 定义参数变量
-llama_checkpoint_path = "checkpoints/fish-speech-1.4-sft-yth-lora"
-decoder_checkpoint_path = "checkpoints/fish-speech-1.4/firefly-gan-vq-fsq-8x1024-21hz-generator.pth"
-decoder_config_name = "firefly_gan_vq"
-device = "cpu"
-
-
-def build_html_error_message(error):
-    return f"""
-    <div style="color: red; 
-    font-weight: bold;">
-        {html.escape(str(error))}
-    </div>
-    """
-
-
-@torch.inference_mode()
-def inference(
-    text,
-    enable_reference_audio,
-    reference_audio,
-    reference_text,
-    max_new_tokens,
-    chunk_length,
-    top_p,
-    repetition_penalty,
-    temperature,
-    streaming=False,
-):
-    if args.max_gradio_length > 0 and len(text) > args.max_gradio_length:
-        return (
-            None,
-            None,
-            i18n("Text is too long, please keep it under {} characters.").format(
-                args.max_gradio_length
-            ),
-        )
-
-    # Parse reference audio aka prompt
-    prompt_tokens = encode_reference(
-        decoder_model=decoder_model,
-        reference_audio=reference_audio,
-        enable_reference_audio=enable_reference_audio,
-    )
-
-    # LLAMA Inference
-    request = dict(
-        device=device,  # 使用指定的设备
-        max_new_tokens=max_new_tokens,
-        text=text,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        temperature=temperature,
-        compile=args.compile,
-        iterative_prompt=chunk_length > 0,
-        chunk_length=chunk_length,
-        max_length=2048,
-        prompt_tokens=prompt_tokens if enable_reference_audio else None,
-        prompt_text=reference_text if enable_reference_audio else None,
-    )
-
-    response_queue = queue.Queue()
-    llama_queue.put(
-        GenerateRequest(
-            request=request,
-            response_queue=response_queue,
-        )
-    )
-
-    if streaming:
-        yield wav_chunk_header(), None, None
-
-    segments = []
-
-    while True:
-        result: WrappedGenerateResponse = response_queue.get()
-        if result.status == "error":
-            yield None, None, build_html_error_message(result.response)
-            break
-
-        result: GenerateResponse = result.response
-        if result.action == "next":
-            break
-
-        with autocast_exclude_mps(device_type=device, dtype=args.precision):  # 使用指定的设备
-            fake_audios = decode_vq_tokens(
-                decoder_model=decoder_model,
-                codes=result.codes,
-            )
-
-        fake_audios = fake_audios.float().cpu().numpy()
-        segments.append(fake_audios)
-
-        if streaming:
-            yield (fake_audios * 32768).astype(np.int16).tobytes(), None, None
-
-    if len(segments) == 0:
-        return (
-            None,
-            None,
-            build_html_error_message(
-                i18n("No audio generated, please check the input text.")
-            ),
-        )
-
-    # No matter streaming or not, we need to return the final audio
-    audio = np.concatenate(segments, axis=0)
-    yield None, (decoder_model.spec_transform.sample_rate, audio), None
-
-    if torch.cuda.is_available():  # 如果没有 GPU，则不执行此部分
-        torch.cuda.empty_cache()
-        gc.collect()
-
-
-def inference_with_auto_rerank(
-    text,
-    enable_reference_audio,
-    reference_audio,
-    reference_text,
-    max_new_tokens,
-    chunk_length,
-    top_p,
-    repetition_penalty,
-    temperature,
-    use_auto_rerank,
-    streaming=False,
-):
-
-    max_attempts = 2 if use_auto_rerank else 1
-    best_wer = float("inf")
-    best_audio = None
-    best_sample_rate = None
-
-    for attempt in range(max_attempts):
-        audio_generator = inference(
-            text,
-            enable_reference_audio,
-            reference_audio,
-            reference_text,
-            max_new_tokens,
-            chunk_length,
-            top_p,
-            repetition_penalty,
-            temperature,
-            streaming=False,
-        )
-
-        # 获取音频数据
-        for _ in audio_generator:
-            pass
-        _, (sample_rate, audio), message = _
-
-        if audio is None:
-            return None, None, message
-
-        if not use_auto_rerank:
-            return None, (sample_rate, audio), None
-
-        asr_result = batch_asr(asr_model, [audio], sample_rate)[0]
-        wer = calculate_wer(text, asr_result["text"])
-        if wer <= 0.3 and not asr_result["huge_gap"]:
-            return None, (sample_rate, audio), None
-
-        if wer < best_wer:
-            best_wer = wer
-            best_audio = audio
-            best_sample_rate = sample_rate
-
-        if attempt == max_attempts - 1:
-            break
-
-    return None, (best_sample_rate, best_audio), None
-
-
-inference_stream = partial(inference, streaming=True)
-
-n_audios = 4
-
-global_audio_list = []
-global_error_list = []
-
-
-def inference_wrapper(
-    text,
-    enable_reference_audio,
-    reference_audio,
-    reference_text,
-    max_new_tokens,
-    chunk_length,
-    top_p,
-    repetition_penalty,
-    temperature,
-    batch_infer_num,
-    if_load_asr_model,
-):
-    audios = []
-    errors = []
-
-    for _ in range(batch_infer_num):
-        result = inference_with_auto_rerank(
-            text,
-            enable_reference_audio,
-            reference_audio,
-            reference_text,
-            max_new_tokens,
-            chunk_length,
-            top_p,
-            repetition_penalty,
-            temperature,
-            if_load_asr_model,
-        )
-
-        _, audio_data, error_message = result
-
-        audios.append(
-            gr.Audio(value=audio_data if audio_data else None, visible=True),
-        )
-        errors.append(
-            gr.HTML(value=error_message if error_message else None, visible=True),
-        )
-
-    for _ in range(batch_infer_num, n_audios):
-        audios.append(
-            gr.Audio(value=None, visible=False),
-        )
-        errors.append(
-            gr.HTML(value=None, visible=False),
-        )
-
-    return None, *audios, *errors
-
-
-def wav_chunk_header(sample_rate=44100, bit_depth=16, channels=1):
-    buffer = io.BytesIO()
-
-    with wave.open(buffer, "wb") as wav_file:
-        wav_file.setnchannels(channels)
-        wav_file.setsampwidth(bit_depth // 8)
-        wav_file.setframerate(sample_rate)
-
-    wav_header_bytes = buffer.getvalue()
-    buffer.close()
-    return wav_header_bytes
-
-
-def normalize_text(user_input, use_normalization):
-    if use_normalization:
-        return ChnNormedText(raw_text=user_input).normalize()
-    else:
-        return user_input
-
-
-asr_model = None
-
-
-def change_if_load_asr_model(if_load):
-    global asr_model
-
-    if if_load:
-        gr.Warning("Loading faster whisper model...")
-        if asr_model is None:
-            asr_model = load_model()
-        return gr.Checkbox(label="Unload faster whisper model", value=if_load)
-
-    if if_load is False:
-        gr.Warning("Unloading faster whisper model...")
-        del asr_model
-        asr_model = None
-        if torch.cuda.is_available():  # 如果没有 GPU，则不执行此部分
-            torch.cuda.empty_cache()
-            gc.collect()
-        return gr.Checkbox(label="Load faster whisper model", value=if_load)
-
-
-def change_if_auto_label(if_load, if_auto_label, enable_ref, ref_audio, ref_text):
-    if if_load and asr_model is not None:
-        if (
-            if_auto_label
-            and enable_ref
-            and ref_audio
-            and ref_text.strip() == ""
-        ):
-            ref_text = batch_asr(asr_model, [ref_audio])[0]["text"]
-            return ref_text
-    return ref_text
-
-
-def setup_gradio_interface():
-    with gr.Blocks() as app:
-        gr.Markdown(HEADER_MD)
-
-        with gr.Row():
-            with gr.Column(scale=2):
-                text_box = gr.Textbox(
-                    label=i18n("Input Text"),
-                    placeholder=TEXTBOX_PLACEHOLDER,
-                    max_lines=6,
-                )
-                normalization_checkbox = gr.Checkbox(
-                    label=i18n("Enable Text Normalization"),
-                    value=False,
-                )
-
-                reference_audio_file = gr.Audio(
-                    label=i18n("Reference Audio"),
-                    type="filepath",
-                    source="upload",
-                    interactive=True,
-                )
-
-                reference_text_box = gr.Textbox(
-                    label=i18n("Reference Text"),
-                    placeholder=i18n("Put your reference text here."),
-                    max_lines=3,
-                )
-
-                with gr.Row():
-                    max_new_tokens_input = gr.Slider(
-                        label=i18n("Max New Tokens"),
-                        minimum=1,
-                        maximum=200,
-                        value=60,
-                        step=1,
-                    )
-                    chunk_length_input = gr.Slider(
-                        label=i18n("Chunk Length"),
-                        minimum=0,
-                        maximum=20,
-                        value=0,
-                        step=1,
-                    )
-
-                with gr.Row():
-                    temperature_input = gr.Slider(
-                        label=i18n("Temperature"),
-                        minimum=0,
-                        maximum=1,
-                        value=0.7,
-                        step=0.01,
-                    )
-                    repetition_penalty_input = gr.Slider(
-                        label=i18n("Repetition Penalty"),
-                        minimum=0,
-                        maximum=2,
-                        value=1,
-                        step=0.01,
-                    )
-                    top_p_input = gr.Slider(
-                        label=i18n("Top P"),
-                        minimum=0,
-                        maximum=1,
-                        value=0.9,
-                        step=0.01,
-                    )
-
-                with gr.Row():
-                    load_asr_model_checkbox = gr.Checkbox(
-                        label=i18n("Load ASR Model"),
-                        value=False,
-                    )
-                    auto_label_checkbox = gr.Checkbox(
-                        label=i18n("Auto Labeling"),
-                        value=False,
-                    )
-
-            with gr.Column(scale=1):
-                submit_btn = gr.Button(i18n("Submit"))
-
-        output_audio = gr.Audio(label=i18n("Generated Audio"))
-        output_error = gr.HTML(label=i18n("Error Message"))
-
-        submit_btn.click(
-            inference_wrapper,
-            inputs=[
-                text_box,
-                load_asr_model_checkbox,
-                reference_audio_file,
-                reference_text_box,
-                max_new_tokens_input,
-                chunk_length_input,
-                top_p_input,
-                repetition_penalty_input,
-                temperature_input,
-                gr.Slider(value=n_audios, visible=False),
-            ],
-            outputs=[output_error, output_audio],
-        )
-
-        # Interface to reload ASR model
-        load_asr_model_checkbox.change(
-            change_if_load_asr_model,
-            inputs=[load_asr_model_checkbox],
-            outputs=[load_asr_model_checkbox],
-        )
-
-        # Interface for auto labeling
-        auto_label_checkbox.change(
-            change_if_auto_label,
-            inputs=[
-                auto_label_checkbox,
-                load_asr_model_checkbox,
-                reference_audio_file,
-                reference_text_box,
-            ],
-            outputs=[reference_text_box],
-        )
-
-    app.launch()
-
-
-if __name__ == "__main__":
-    parser = ArgumentParser()
-    parser.add_argument(
-        "--max-gradio-length",
-        type=int,
-        default=2048,
-        help="Maximum length of input text for Gradio.",
-    )
-    parser.add_argument(
-        "--compile",
-        action="store_true",
-        help="Compile the model.",
-    )
-    parser.add_argument(
-        "--precision",
-        type=str,
-        default="float32",
-        help="Model precision, one of ['float16', 'float32', 'bfloat16'].",
-    )
-    parser.add_argument(
-        "--llama-checkpoint-path",
-        type=str,
-        required=True,
-        help="Path to the Llama checkpoint.",
-    )
-    parser.add_argument(
-        "--decoder-checkpoint-path",
-        type=str,
-        required=True,
-        help="Path to the decoder checkpoint.",
-    )
-    parser.add_argument(
-        "--decoder-config-name",
-        type=str,
-        required=True,
-        help="Name of the decoder config.",
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="cpu",
-        help="Device to run the model on, one of ['cpu', 'cuda'].",
-    )
+import argparse
+import subprocess
+
+def main():
+    # 创建命令行参数解析器
+    parser = argparse.ArgumentParser(description="启动 Fish Speech 应用")
+    parser.add_argument("--llama-checkpoint-path", type=str, required=True, help="Llama 检查点路径")
+    parser.add_argument("--decoder-checkpoint-path", type=str, required=True, help="解码器检查点路径")
+    parser.add_argument("--decoder-config-name", type=str, required=True, help="解码器配置名称")
+    parser.add_argument("--device", type=str, default="cpu", help="设备类型（cpu 或 cuda）")
+
+    # 解析参数
     args = parser.parse_args()
 
-    llama_checkpoint_path = args.llama_checkpoint_path
-    decoder_checkpoint_path = args.decoder_checkpoint_path
-    decoder_config_name = args.decoder_config_name
-    device = args.device
+    # 构造命令
+    command = [
+        "python", "tools/webui.py",
+        "--llama-checkpoint-path", args.llama_checkpoint_path,
+        "--decoder-checkpoint-path", args.decoder_checkpoint_path,
+        "--decoder-config-name", args.decoder_config_name,
+        "--device", args.device
+    ]
 
-    logger.info("Loading decoder model...")
-    decoder_model = load_decoder_model()
+    # 运行命令
+    subprocess.run(command)
 
-    # Initialize Llama and ASR models
-    llama_queue = launch_thread_safe_queue()
-    logger.info("Loading Llama model...")
-    load_model(0)
-
-    # Setup the Gradio interface
-    setup_gradio_interface()
+if __name__ == "__main__":
+    main()