Blame for app.py

## IMPORTS ##
import os
import tempfile
import time
from pathlib import Path

import gradio as gr
import numpy as np
import spaces

import mfn_hf  # shared look + usage analytics; import before @spaces.GPU
import torch
import torchaudio
import soundfile as sf
import librosa
from cached_path import cached_path
from huggingface_hub import hf_hub_download
from transformers import pipeline



def _load_audio(path, *args, **kwargs):
    # torchaudio>=2.9 routes load() through torchcodec; read files with soundfile/librosa instead
    try:
        data, sr = sf.read(path, dtype="float32", always_2d=True)
        wav = torch.from_numpy(data.T.copy())
    except Exception:
        y, sr = librosa.load(path, sr=None, mono=False)
        wav = torch.from_numpy(np.atleast_2d(y).astype(np.float32))
    return wav, sr


torchaudio.load = _load_audio

from infer import DMOInference

## CUDA DEVICE ##
device = "cuda" if torch.cuda.is_available() else "cpu"

## LOAD MODELS ##
asr_pipe = pipeline(
    "automatic-speech-recognition", model="openai/whisper-large-v3-turbo", device=device
)
model = DMOInference(
    student_checkpoint_path=str(cached_path("hf://yl4579/DMOSpeech2/model_85000.pt")),
    duration_predictor_path=str(cached_path("hf://yl4579/DMOSpeech2/model_1500.pt")),
    device=device,
    model_type="F5TTS_Base",
)


def transcribe(ref_audio, language=None):
    """Transcribe audio using the pre-loaded ASR pipeline."""
    return asr_pipe(
        ref_audio,
        chunk_length_s=30,
        batch_size=128,
        generate_kwargs=(
            {"task": "transcribe", "language": language}
            if language
            else {"task": "transcribe"}
        ),
        return_timestamps=False,
    )["text"].strip()

MODES = {
    "Student Only (4 steps)": {
        "teacher_steps": 0,
        "teacher_stopping_time": 1.0,
        "student_start_step": 0,
        "description": "Fastest (4 steps), good quality"
    },
    "Teacher-Guided (8 steps)": {
        "teacher_steps": 16,
        "teacher_stopping_time": 0.07,
        "student_start_step": 1,
        "description": "Best balance (8 steps), recommended"
    },
    "High Diversity (16 steps)": {
        "teacher_steps": 24,
        "teacher_stopping_time": 0.3,
        "student_start_step": 2,
        "description": "More natural prosody (16 steps)"
    },
    "Custom": {
        "teacher_steps": None,
        "teacher_stopping_time": None,
        "student_start_step": None,
        "description": "Fine-tune all parameters"
    }
}

@spaces.GPU(duration=120)
def _generate_speech(
    prompt_audio,
    prompt_text,
    target_text,
    mode,
    temperature,
    custom_teacher_steps,
    custom_teacher_stopping_time,
    custom_student_start_step,
    verbose,
):
    if prompt_audio is None:
        raise gr.Error("Please upload a reference audio!")

    if not target_text:
        raise gr.Error("Please enter text to generate!")

    if not prompt_text and prompt_text != "":
        prompt_text = transcribe(prompt_audio)


    if mode == "Custom":
        teacher_steps, teacher_stopping_time, student_start_step = custom_teacher_steps, custom_teacher_stopping_time, custom_student_start_step
    else:
        teacher_steps = MODES[mode]["teacher_steps"]
        teacher_stopping_time = MODES[mode]["teacher_stopping_time"]
        student_start_step = MODES[mode]["student_start_step"]

    generated_audio = model.generate(
        gen_text=target_text,
        audio_path=prompt_audio,
        prompt_text=prompt_text if prompt_text else None,
        teacher_steps=teacher_steps,
        teacher_stopping_time=teacher_stopping_time,
        student_start_step=student_start_step,
        temperature=temperature,
        verbose=verbose,
    )


    if isinstance(generated_audio, torch.Tensor):
        audio_np = generated_audio.cpu().numpy()
    else:
        audio_np = generated_audio

    # Ensure audio is properly normalized and in the correct format
    if audio_np.ndim == 2 and audio_np.shape[0] == 1:
        audio_np = audio_np.squeeze(0)  # Remove batch dimension if present
    
    # Normalize audio to [-1, 1] range if needed
    if np.abs(audio_np).max() > 1.0:
        audio_np = audio_np / np.abs(audio_np).max()
    
    # Ensure audio is in float32 format
    audio_np = audio_np.astype(np.float32)

    return (24000, audio_np)


def generate_speech(
    prompt_audio,
    prompt_text,
    target_text,
    mode,
    temperature,
    custom_teacher_steps,
    custom_teacher_stopping_time,
    custom_student_start_step,
    verbose,
):
    t0 = time.time()
    audio = _generate_speech(
        prompt_audio,
        prompt_text,
        target_text,
        mode,
        temperature,
        custom_teacher_steps,
        custom_teacher_stopping_time,
        custom_student_start_step,
        verbose,
    )
    return audio, f"Took {time.time() - t0:.1f} s"


# Create Gradio interface
with gr.Blocks(title="DMOSpeech 2 - Zero-Shot TTS") as demo:
    mfn_hf.header(
        "DMOSpeech 2",
        model="yl4579/DMOSpeech2",
        description="Speak any text in any voice from a short sample of that voice.",
        flow_fn="/generate_speech",
    )
    gr.Markdown("Code and paper: [github.com/yl4579/DMOSpeech2](https://github.com/yl4579/DMOSpeech2)")

    with gr.Row():
        with gr.Column(scale=1):
            # Reference audio input
            prompt_audio = gr.Audio(
                label="Reference Audio",
                type="filepath",
                sources=["upload", "microphone"],
            )

            prompt_text = gr.Textbox(
                label="Reference Text (leave empty for auto-transcription)",
                placeholder="The text spoken in the reference audio...",
                lines=2,
            )

            target_text = gr.Textbox(
                label="Text to Generate",
                placeholder="Enter the text you want to synthesize...",
                lines=4,
            )

            # Generation mode
            mode = gr.Radio(
                choices=[
                    "Student Only (4 steps)",
                    "Teacher-Guided (8 steps)",
                    "High Diversity (16 steps)",
                    "Custom",
                ],
                value="Teacher-Guided (8 steps)",
                label="Generation Mode",
                info="Choose speed vs quality/diversity tradeoff",
            )

            # Advanced settings (collapsible)
            with gr.Accordion("Advanced Settings", open=False):
                temperature = gr.Slider(
                    minimum=0.0,
                    maximum=2.0,
                    value=0.0,
                    step=0.1,
                    label="Duration Temperature",
                    info="0 = deterministic, >0 = more variation in speech rhythm",
                )

                with gr.Group(visible=False) as custom_settings:
                    gr.Markdown("### Custom Mode Settings")
                    custom_teacher_steps = gr.Slider(
                        minimum=0,
                        maximum=32,
                        value=16,
                        step=1,
                        label="Teacher Steps",
                        info="More steps = higher quality",
                    )

                    custom_teacher_stopping_time = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.07,
                        step=0.01,
                        label="Teacher Stopping Time",
                        info="When to switch to student",
                    )

                    custom_student_start_step = gr.Slider(
                        minimum=0,
                        maximum=4,
                        value=1,
                        step=1,
                        label="Student Start Step",
                        info="Which student step to start from",
                    )

                verbose = gr.Checkbox(
                    value=False,
                    label="Verbose Output",
                    info="Show detailed generation steps",
                )

            generate_btn = gr.Button("Generate Speech", variant="primary", size="lg")

        with gr.Column(scale=1):
            # Output
            output_audio = gr.Audio(
                label="Generated Speech", type="filepath", autoplay=True
            )
            took = mfn_hf.took()

            # Tips
            gr.Markdown(
                """
            ### Tips
            
            - **Auto-transcription**: Leave reference text empty to auto-transcribe
            - **Student Only**: Fastest (4 steps), good quality
            - **Teacher-Guided**: Best balance (8 steps), recommended
            - **High Diversity**: More natural prosody (16 steps)
            - **Custom Mode**: Fine-tune all parameters
            
            ### Speed
            - Student Only: ~0.05x (20x faster than real-time)
            - Teacher-Guided: ~0.10x (10x faster)
            - High Diversity: ~0.20x (5x faster)
            """
            )

    # Event handler
    generate_btn.click(
        generate_speech,
        inputs=[
            prompt_audio,
            prompt_text,
            target_text,
            mode,
            temperature,
            custom_teacher_steps,
            custom_teacher_stopping_time,
            custom_student_start_step,
            verbose,
        ],
        outputs=[output_audio, took],
    )

    mode.change(lambda x: gr.update(visible=x == "Custom"), inputs=[mode], outputs=[custom_settings])

    mfn_hf.footer()

demo.queue().launch()