app.py

try:
    import spaces
except ImportError:
    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
    class spaces:
        class GPU:
            def __init__(self, func=None, duration=60):
                self.func = func

            def __call__(self, *args, **kwargs):
                if self.func is not None:
                    return self.func(*args, **kwargs)
                func = args[0]
                return func


import sys

sys.stdout.reconfigure(line_buffering=True)

import os
import re
import tempfile
import threading

import torch
import gradio as gr
import whisper
from huggingface_hub import snapshot_download
from pyharp import ModelCard, build_endpoint

from models.svc.vevo2.vevo2_utils import Vevo2InferencePipeline, save_audio

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

CKPT_DIR = "./ckpts/Vevo2"
# training-only artifacts (optimizer/scheduler/rng state, and the earlier
# "pretrained" AR checkpoint superseded by "posttrained") aren't needed for
# inference and would roughly double the download.
CKPT_IGNORE_PATTERNS = [
    "*/optimizer.pt",
    "*/optimizer.bin",
    "*/scheduler.pt",
    "*/scheduler.bin",
    "*/rng_state*.pth",
    "*/random_states_*.pkl",
    "*/trainer_state.json",
    "*/training_args.bin",
    "contentstyle_modeling/pretrained/*",
]

pipeline = None
download_ready = False
download_error = None


FMT_CONFIG_PATH = os.path.join(
    CKPT_DIR, "acoustic_modeling/fm_emilia101k_singnet7k_repa/config.json"
)


def _fix_whisper_stats_path():
    """The checkpoint's config.json hardcodes whisper_stats_path as
    models/svc/vevosing/config/whisper_stats.pt, relative to the full
    Amphion repo — a directory this deployment doesn't copy. The identical
    file ships right next to this config, so point it there instead."""
    correct_path = os.path.join(
        CKPT_DIR, "acoustic_modeling/fm_emilia101k_singnet7k_repa/whisper_stats.pt"
    )
    with open(FMT_CONFIG_PATH) as f:
        text = f.read()
    text = re.sub(
        r'"whisper_stats_path":\s*"[^"]*"',
        f'"whisper_stats_path": "{correct_path}"',
        text,
    )
    with open(FMT_CONFIG_PATH, "w") as f:
        f.write(text)


def load_checkpoint():
    """Download Vevo2's weights and pre-warm the Whisper cache. CPU-only —
    building the actual pipeline happens lazily on first process_fn call
    (see get_pipeline) instead of here, since Vevo2InferencePipeline moves
    every submodule onto `device` inside its own constructor and ZeroGPU
    only intercepts CUDA calls made inside an @spaces.GPU-decorated call,
    not from a background thread."""
    global download_ready, download_error
    try:
        snapshot_download(
            repo_id="RMSnow/Vevo2",
            local_dir=CKPT_DIR,
            ignore_patterns=CKPT_IGNORE_PATTERNS,
        )
        _fix_whisper_stats_path()
        whisper.load_model("medium", device="cpu")  # warms ~/.cache/whisper
        print("Checkpoint download complete.")
    except Exception as e:
        download_error = str(e)
        print(f"Download error: {e}")
    finally:
        download_ready = True


threading.Thread(target=load_checkpoint, daemon=True).start()


def get_pipeline():
    """Build the inference pipeline on first use, inside the @spaces.GPU
    call (see load_checkpoint for why this can't happen in the background
    thread)."""
    global pipeline
    if pipeline is None:
        pipeline = Vevo2InferencePipeline(
            prosody_tokenizer_ckpt_path=os.path.join(
                CKPT_DIR, "tokenizer/prosody_fvq512_6.25hz"
            ),
            content_style_tokenizer_ckpt_path=os.path.join(
                CKPT_DIR, "tokenizer/contentstyle_fvq16384_12.5hz"
            ),
            ar_cfg_path=os.path.join(
                CKPT_DIR, "contentstyle_modeling/posttrained/amphion_config.json"
            ),
            ar_ckpt_path=os.path.join(CKPT_DIR, "contentstyle_modeling/posttrained"),
            fmt_cfg_path=os.path.join(
                CKPT_DIR, "acoustic_modeling/fm_emilia101k_singnet7k_repa/config.json"
            ),
            fmt_ckpt_path=os.path.join(
                CKPT_DIR, "acoustic_modeling/fm_emilia101k_singnet7k_repa"
            ),
            vocoder_cfg_path=os.path.join(CKPT_DIR, "vocoder/config.json"),
            vocoder_ckpt_path=os.path.join(CKPT_DIR, "vocoder"),
            device=DEVICE,
        )
    return pipeline


model_card = ModelCard(
    name="Vevo2",
    description=(
        "Zero-shot speech and singing voice generation and conversion: "
        "voice/singing conversion, text-to-singing, singing editing, "
        "singing style conversion, and melody control.\n\n"
        "What each task needs, beyond picking it from the Task dropdown:\n"
        "Voice/Singing Conversion: Input Audio + Reference Voice.\n"
        "Text-to-Speech: Input Audio + Text/Lyrics (Reference Voice optional).\n"
        "Singing Editing: Input Audio + Text/Lyrics.\n"
        "Singing Style Conversion: Input Audio + Reference Voice + Text/Lyrics.\n"
        "Melody Control: Input Audio + Reference Voice + Text/Lyrics."
    ),
    author=(
        "Xueyao Zhang, Junan Zhang, Yuancheng Wang, Chaoren Wang, "
        "Yuanzhe Chen, Dongya Jia, Zhuo Chen, Zhizheng Wu"
    ),
    tags=["voice-conversion", "singing-synthesis", "text-to-speech"],
)

TASKS = [
    "Voice/Singing Conversion",
    "Text-to-Speech / Text-to-Singing",
    "Singing Editing",
    "Singing Style Conversion",
    "Melody Control (Humming/Instrument to Singing)",
]


@spaces.GPU(duration=120)
@torch.inference_mode()
def process_fn(
    task: str,
    input_audio_path: str,
    reference_audio_path: str,
    target_text: str,
    style_ref_text: str,
    flow_matching_steps: int,
) -> str:
    """Runs the Vevo2 task selected in the Task dropdown. Each task calls a
    different combination of Vevo2InferencePipeline.inference_fm /
    inference_ar_and_fm, mirroring the task wrapper functions in the
    original repo's infer_vevo2_fm.py / infer_vevo2_ar.py."""
    if not download_ready:
        raise gr.Error("Model is still downloading, please wait a moment and try again.")
    if download_error is not None:
        raise gr.Error(f"Model failed to download: {download_error}")

    pipe = get_pipeline()
    target_text = (target_text or "").strip() or None
    style_ref_text = style_ref_text or ""
    reference_audio_path = reference_audio_path or None

    if task == "Voice/Singing Conversion":
        if reference_audio_path is None:
            raise gr.Error(
                "Voice/Singing Conversion needs a reference voice (Reference Voice input)."
            )
        gen_audio = pipe.inference_fm(
            src_wav_path=input_audio_path,
            timbre_ref_wav_path=reference_audio_path,
            use_pitch_shift=True,
            flow_matching_steps=flow_matching_steps,
        )
    elif task == "Text-to-Speech / Text-to-Singing":
        if target_text is None:
            raise gr.Error("Text-to-Speech needs Text / Lyrics.")
        gen_audio = pipe.inference_ar_and_fm(
            target_text=target_text,
            style_ref_wav_path=input_audio_path,
            style_ref_wav_text=style_ref_text,
            timbre_ref_wav_path=reference_audio_path or input_audio_path,
            use_prosody_code=False,
            flow_matching_steps=flow_matching_steps,
        )
    elif task == "Singing Editing":
        if target_text is None:
            raise gr.Error("Singing Editing needs the edited Text / Lyrics.")
        gen_audio = pipe.inference_ar_and_fm(
            target_text=target_text,
            prosody_wav_path=input_audio_path,
            style_ref_wav_path=input_audio_path,
            style_ref_wav_text=style_ref_text,
            timbre_ref_wav_path=input_audio_path,
            use_prosody_code=True,
            flow_matching_steps=flow_matching_steps,
        )
    elif task == "Singing Style Conversion":
        if reference_audio_path is None:
            raise gr.Error(
                "Singing Style Conversion needs a style reference (Reference Voice input)."
            )
        if target_text is None:
            raise gr.Error(
                "Singing Style Conversion needs Text / Lyrics — use Input Audio's own "
                "transcript, since only the style changes, not the content."
            )
        gen_audio = pipe.inference_ar_and_fm(
            target_text=target_text,
            prosody_wav_path=input_audio_path,
            style_ref_wav_path=reference_audio_path,
            style_ref_wav_text=style_ref_text,
            timbre_ref_wav_path=input_audio_path,
            use_prosody_code=True,
            use_pitch_shift=True,
            flow_matching_steps=flow_matching_steps,
        )
    else:  # Melody Control
        if target_text is None:
            raise gr.Error("Melody Control needs the Text / Lyrics to sing.")
        if reference_audio_path is None:
            raise gr.Error(
                "Melody Control needs a reference voice (Reference Voice input) — "
                "the melody input alone (e.g. humming) has no usable vocal timbre."
            )
        gen_audio = pipe.inference_ar_and_fm(
            target_text=target_text,
            prosody_wav_path=input_audio_path,
            style_ref_wav_path=reference_audio_path,
            style_ref_wav_text=style_ref_text,
            timbre_ref_wav_path=reference_audio_path,
            use_prosody_code=True,
            use_pitch_shift=True,
            flow_matching_steps=flow_matching_steps,
        )

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        output_path = f.name
    save_audio(gen_audio, output_path=output_path)
    return output_path


with gr.Blocks() as demo:
    input_components = [
        gr.Dropdown(
            choices=TASKS,
            value=TASKS[0],
            label="Task",
            info="What to do with the input audio.",
        ),
        gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info(
            "Voice/Singing Conversion: the source to convert. "
            "Text-to-Speech: style reference. "
            "Singing Editing: the recording to edit. "
            "Singing Style Conversion: the source to convert. "
            "Melody Control: melody reference (e.g. humming or an instrument)."
        ),
        gr.Audio(type="filepath", label="Reference Voice")
        .harp_required(False)
        .set_info(
            "Voice/Singing Conversion: sets the output's voice. "
            "Text-to-Speech: optional, defaults to Input Audio's voice. "
            "Singing Editing: ignored. "
            "Singing Style Conversion: style/technique reference only — output keeps "
            "Input Audio's own voice. "
            "Melody Control: sets the output's voice."
        ),
        gr.Textbox(
            label="Text / Lyrics",
            info="Voice/Singing Conversion: ignored. "
            "Text-to-Speech: the text or lyrics to generate. "
            "Singing Editing: the edited lyrics. "
            "Singing Style Conversion: Input Audio's own transcript — content stays "
            "the same, only the style changes. "
            "Melody Control: the lyrics to sing.",
        ),
        gr.Textbox(
            label="Style Reference Transcript (optional)",
            info="Voice/Singing Conversion: not used. "
            "Text-to-Speech: transcript of Input Audio. "
            "Singing Editing: transcript of Input Audio. "
            "Singing Style Conversion: transcript of Reference Voice. "
            "Melody Control: transcript of Reference Voice.",
        ),
        gr.Slider(
            minimum=8,
            maximum=64,
            step=1,
            value=32,
            label="Generation Detail (Steps)",
            info="More steps trade generation speed for audio detail (default: 32, per repo).",
        ),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Output Audio").set_info(
            "Generated speech or singing voice."
        ),
    ]

    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

if __name__ == "__main__":
    demo.queue().launch(pwa=True)