Update app.py

app.py

@@ -1,63 +1,230 @@
+import gc
+import html
+import io
 import os
-import torch
+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 tools.llama.generate import launch_thread_safe_queue
+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
-import gradio as gr  # 导入 Gradio
 
-def parse_args():
-    parser = ArgumentParser()
-    parser.add_argument(
-        "--llama-checkpoint-path",
-        type=str,
-        default="checkpoints/fish-speech-1.4-sft-yth-lora",
-        help="Path to the Llama checkpoint"
-    )
-    parser.add_argument(
-        "--decoder-checkpoint-path",
-        type=str,
-        default="checkpoints/fish-speech-1.4/firefly-gan-vq-fsq-8x1024-21hz-generator.pth",
-        help="Path to the VQ-GAN checkpoint"
-    )
-    parser.add_argument(
-        "--decoder-config-name",
-        type=str,
-        default="firefly_gan_vq",
-        help="VQ-GAN config name"
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="cpu",
-        help="Device to run on (cpu or cuda)"
-    )
-    parser.add_argument(
-        "--half", 
-        action="store_true", 
-        help="Use half precision"
-    )
-    parser.add_argument(
-        "--compile",
-        action="store_true",
-        default=True,
-        help="Compile the model for optimized inference"
+# 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
+
+
+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,
     )
-    parser.add_argument(
-        "--max-gradio-length",
-        type=int,
-        default=0,
-        help="Maximum length for Gradio input"
+
+    # LLAMA Inference
+    request = dict(
+        device="cpu",  # 设置为 CPU
+        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,
     )
-    parser.add_argument(
-        "--theme",
-        type=str,
-        default="light",
-        help="Theme for the Gradio app"
+
+    response_queue = queue.Queue()
+    llama_queue.put(
+        GenerateRequest(
+            request=request,
+            response_queue=response_queue,
+        )
     )
-    return parser.parse_args()
 
-def inference(
+    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="cpu", dtype=args.precision):  # 设置为 CPU
+            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,
@@ -67,66 +234,253 @@ def inference(
     top_p,
     repetition_penalty,
     temperature,
+    batch_infer_num,
+    if_load_asr_model,
 ):
-    logger.info(f"Running inference on: {text}")
-    # 模拟推理过程
-    result = f"Processed text: {text}"
-    return result
+    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
 
-def inference_function(text):
-    return f"Processed: {text}"
 
-def build_app(args):
+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 is not None
+            and ref_text.strip() == ""
+        ):
+            data, sample_rate = librosa.load(ref_audio)
+            res = batch_asr(asr_model, [data], sample_rate)[0]
+            ref_text = res["text"]
+    return ref_text
+
+
+def setup_gradio_interface():
     with gr.Blocks() as app:
-        gr.Markdown(f"# Fish Speech Inference - Theme: {args.theme}")
-        text_input = gr.Textbox(label="Input Text")
-        output = gr.Textbox(label="Output Text")
-        submit_button = gr.Button("Submit")
+        gr.Markdown(HEADER_MD)
 
-        submit_button.click(fn=inference_function, inputs=text_input, outputs=output)
-    return app
+        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,
+                )
 
-def main():
-    args = parse_args()
+                reference_audio_file = gr.Audio(
+                    label=i18n("Reference Audio"),
+                    type="filepath",
+                    source="upload",
+                    interactive=True,
+                )
 
-    args.precision = torch.half if args.half else torch.bfloat16
+                reference_text_box = gr.Textbox(
+                    label=i18n("Reference Text"),
+                    placeholder=i18n("Put your reference text here."),
+                    max_lines=3,
+                )
 
-    logger.info("Loading Llama model...")
-    llama_queue = launch_thread_safe_queue(
-        checkpoint_path=args.llama_checkpoint_path,
-        device=args.device,
-        precision=args.precision,
-        compile=args.compile,
-    )
-    logger.info("Llama model loaded, loading VQ-GAN model...")
+                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,
+                    )
 
-    decoder_model = load_decoder_model(
-        config_name=args.decoder_config_name,
-        checkpoint_path=args.decoder_checkpoint_path,
-        device=args.device,
-    )
+                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,
+                    )
 
-    logger.info("Decoder model loaded, warming up...")
-
-    # Perform a dry run to warm up the model
-    inference(
-        text="Hello, world!",
-        enable_reference_audio=False,
-        reference_audio=None,
-        reference_text="",
-        max_new_tokens=0,
-        chunk_length=100,
-        top_p=0.7,
-        repetition_penalty=1.2,
-        temperature=0.7,
-    )
+                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],
+        )
 
-    logger.info("Warming up done, launching the web UI...")
+        # 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],
+        )
 
-    # Launch the Gradio app, passing args to build_app
-    app = build_app(args)
-    app.launch(show_api=True)
+    app.launch()
 
 
 if __name__ == "__main__":
-    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'].",
+    )
+    args = parser.parse_args()
+
+    logger.info("Loading decoder model...")
+    decoder_model = load_decoder_model()
+
+    # 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()