app.py

#!/usr/bin/env python
# ruff: noqa: E402

import json
import re
import tempfile
import os

import click
import gradio as gr
import numpy as np
import soundfile as sf
import torchaudio

from importlib.resources import files
from groq import Groq
from cached_path import cached_path
from transformers import AutoModelForCausalLM, AutoTokenizer

# Try to import spaces; if available, set USING_SPACES to True so we can decorate functions for GPU support.
try:
    import spaces

    USING_SPACES = True
except ImportError:
    USING_SPACES = False


def gpu_decorator(func):
    """
    Decorator that wraps a function with GPU acceleration if running in a Spaces environment.
    """
    if USING_SPACES:
        return spaces.GPU(func)
    return func


# Local package imports
from f5_tts.model import DiT, UNetT
from f5_tts.infer.utils_infer import (
    load_vocoder,
    load_model,
    preprocess_ref_audio_text,
    infer_process,
    remove_silence_for_generated_wav,
    save_spectrogram,
)

DEFAULT_TTS_MODEL = "F5-TTS"
DEFAULT_TTS_MODEL_CFG = [
    "hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors",
    "hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt",
    json.dumps(dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)),
]

# Load vocoder and TTS model
vocoder = load_vocoder()


def load_f5tts(
    ckpt_path: str = str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors"))
):
    """
    Load the F5-TTS model from the given checkpoint path.
    """
    F5TTS_model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
    return load_model(DiT, F5TTS_model_cfg, ckpt_path)


F5TTS_ema_model = load_f5tts()


def load_groq(apikey):
    client = Groq(api_key=apikey)
    return client = Groq(api_key=apikey)

@gpu_decorator
def generate_response(messages):
    """
    Generate a chat response using the Groq API.
    If messages is a string, wrap it as a user message.
    """
    if isinstance(messages, str):
        messages_payload = [{"role": "user", "content": messages}]
    else:
        messages_payload = messages

    chat_completion = client.chat.completions.create(
        messages=messages_payload,
        model="deepseek-r1-distill-llama-70b",
        stream=False,
    )
    # Check that we got a valid response.
    if chat_completion.choices and hasattr(chat_completion.choices[0].message, "content"):
        return chat_completion.choices[0].message.content
    return ""


@gpu_decorator
def process_audio_input(audio_path, text, history, conv_state):
    """
    Process audio and/or text input from the user:
      - If an audio file is provided, its transcript is obtained.
      - The conversation state and history are updated.
    """
    if not audio_path and not text.strip():
        return history, conv_state, ""

    if audio_path:
        # preprocess_ref_audio_text returns a tuple (audio, transcript).
        _, text = preprocess_ref_audio_text(audio_path, text)

    if not text.strip():
        return history, conv_state, ""

    conv_state.append({"role": "user", "content": text})
    history.append((text, None))
    response = generate_response(conv_state)
    conv_state.append({"role": "assistant", "content": response})
    history[-1] = (text, response)
    return history, conv_state, ""


@gpu_decorator
def infer(
    ref_audio_orig,
    ref_text,
    gen_text,
    remove_silence,
    cross_fade_duration: float = 0.15,
    nfe_step: int = 32,
    speed: float = 1,
    show_info=print,
):
    """
    Generate speech audio using the F5-TTS system based on a reference audio/text and generated text.
    """
    if not ref_audio_orig:
        gr.Warning("Please provide reference audio.")
        return gr.update(), gr.update(), ref_text

    if not gen_text.strip():
        gr.Warning("Please enter text to generate.")
        return gr.update(), gr.update(), ref_text

    # Preprocess the reference audio and text.
    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=show_info)
    ema_model = F5TTS_ema_model  # Use the default F5-TTS model.

    final_wave, final_sample_rate, combined_spectrogram = infer_process(
        ref_audio,
        ref_text,
        gen_text,
        ema_model,
        vocoder,
        cross_fade_duration=cross_fade_duration,
        nfe_step=nfe_step,
        speed=speed,
        show_info=show_info,
        progress=gr.Progress(),
    )

    if remove_silence:
        # Write the generated waveform to a temporary file.
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
            temp_audio_path = f.name
            sf.write(temp_audio_path, final_wave, final_sample_rate)
        # Process the file to remove silence.
        remove_silence_for_generated_wav(temp_audio_path)
        final_wave_tensor, _ = torchaudio.load(temp_audio_path)
        final_wave = final_wave_tensor.squeeze().cpu().numpy()
        os.unlink(temp_audio_path)  # Clean up the temporary file.

    # Save the spectrogram as a temporary PNG file.
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
        spectrogram_path = tmp_spectrogram.name
        save_spectrogram(combined_spectrogram, spectrogram_path)

    return (final_sample_rate, final_wave), spectrogram_path, ref_text


with gr.Blocks() as app_chat:
    gr.Markdown(
        """
# Voice Chat
Have a conversation with an AI using your reference voice! 
1. Upload a reference audio clip and optionally its transcript.
2. Load the chat model.
3. Record your message through your microphone.
4. The AI will respond using the reference voice.
"""
    )

    # The chat interface container is visible only if running in a Spaces environment.
    chat_interface_container = gr.Column(visible=USING_SPACES)

    with chat_interface_container:
        with gr.Row():

            with gr.Column():
                ref_audio_chat = gr.Audio(label="Reference Audio", type="filepath")
            with gr.Column():
                with gr.Accordion("Advanced Settings", open=False):
                    remove_silence_chat = gr.Checkbox(label="Remove Silences", value=True)
                    ref_text_chat = gr.Textbox(
                        label="Reference Text",
                        info="Optional: Leave blank to auto-transcribe",
                        lines=2,
                    )
                    system_prompt_chat = gr.Textbox(
                        label="System Prompt",
                        value=(
                            "You are not an AI assistant, you are whoever the user says you are. "
                            "You must stay in character. Keep your responses concise since they will be spoken out loud."
                        ),
                        lines=2,
                    )

        chatbot_interface = gr.Chatbot(label="Conversation")
        with gr.Row():
            with gr.Column():
                audio_input_chat = gr.Microphone(label="Speak your message", type="filepath")
                audio_output_chat = gr.Audio(autoplay=True)
            with gr.Column():
                with gr.Row():
                    groq_apikey = gr.Textbox(label="Your Groq API Key")
                    load_api = gr.Button("Load API")
                text_input_chat = gr.Textbox(label="Type your message", lines=1)
                send_btn_chat = gr.Button("Send Message")
                clear_btn_chat = gr.Button("Clear Conversation")

        # Initialize the conversation state with the system prompt.
        conversation_state = gr.State(
            value=[
                {
                    "role": "system",
                    "content": (
                        "You are not an AI assistant, you are whoever the user says you are. "
                        "You must stay in character. Keep your responses concise since they will be spoken out loud."
                    ),
                }
            ]
        )

        @gpu_decorator
        def generate_audio_response(history, ref_audio, ref_text, remove_silence):
            """
            Generate an audio response from the last AI message in the conversation.
            """
            if not history or not ref_audio:
                return None, ref_text

            last_user_message, last_ai_response = history[-1]
            if not last_ai_response:
                return None, ref_text

            audio_result, _, ref_text_out = infer(
                ref_audio,
                ref_text,
                last_ai_response,
                remove_silence,
                cross_fade_duration=0.15,
                speed=1.0,
                show_info=print,
            )
            return audio_result, ref_text_out

        def clear_conversation():
            """
            Clear the chat conversation and reset the conversation state.
            """
            initial_state = [
                {
                    "role": "system",
                    "content": (
                        "You are not an AI assistant, you are whoever the user says you are. "
                        "You must stay in character. Keep your responses concise since they will be spoken out loud."
                    ),
                }
            ]
            return [], initial_state

        def update_system_prompt(new_prompt):
            """
            Update the system prompt and reset the conversation.
            """
            initial_state = [{"role": "system", "content": new_prompt}]
            return [], initial_state

        # Set up callbacks so that when recording stops, or text is submitted, the chain of processing is run.
        audio_input_chat.stop_recording(
            process_audio_input,
            inputs=[audio_input_chat, text_input_chat, chatbot_interface, conversation_state],
            outputs=[chatbot_interface, conversation_state],
        ).then(
            generate_audio_response,
            inputs=[chatbot_interface, ref_audio_chat, ref_text_chat, remove_silence_chat],
            outputs=[audio_output_chat, ref_text_chat],
        ).then(lambda: None, None, audio_input_chat)

        load_api.click(load_groq, inputs=groq_apikey, outputs=None)
        text_input_chat.submit(
            process_audio_input,
            inputs=[audio_input_chat, text_input_chat, chatbot_interface, conversation_state],
            outputs=[chatbot_interface, conversation_state],
        ).then(
            generate_audio_response,
            inputs=[chatbot_interface, ref_audio_chat, ref_text_chat, remove_silence_chat],
            outputs=[audio_output_chat, ref_text_chat],
        ).then(lambda: None, None, text_input_chat)

        send_btn_chat.click(
            process_audio_input,
            inputs=[audio_input_chat, text_input_chat, chatbot_interface, conversation_state],
            outputs=[chatbot_interface, conversation_state],
        ).then(
            generate_audio_response,
            inputs=[chatbot_interface, ref_audio_chat, ref_text_chat, remove_silence_chat],
            outputs=[audio_output_chat, ref_text_chat],
        ).then(lambda: None, None, text_input_chat)

        clear_btn_chat.click(clear_conversation, outputs=[chatbot_interface, conversation_state])
        system_prompt_chat.change(
            update_system_prompt,
            inputs=system_prompt_chat,
            outputs=[chatbot_interface, conversation_state],
        )


app = app_chat


@click.command()
@click.option("--port", "-p", default=None, type=int, help="Port to run the app on")
@click.option("--host", "-H", default=None, help="Host to run the app on")
@click.option("--share", "-s", default=False, is_flag=True, help="Share the app via Gradio share link")
@click.option("--api", "-a", default=True, is_flag=True, help="Allow API access")
@click.option("--root_path", "-r", default=None, type=str, help="Root path for the application")
def main(port, host, share, api, root_path):
    """
    Launch the Gradio app.
    """
    app.queue(api_open=api).launch(
        server_name=host,
        server_port=port,
        share=share,
        show_api=api,
        root_path=root_path,
    )


if __name__ == "__main__":
    main()