Blame for app.py

"""WaveGen on Hugging Face ZeroGPU. See README for an unchanged-environment local check."""

from __future__ import annotations

import argparse
import os
from pathlib import Path
import sys
import time

ROOT = Path(__file__).resolve().parent
EXAMPLE_DIR = ROOT / "src/wavegen_v_spec/infer/examples"
sys.path.insert(0, str(ROOT / "src"))

# No torch, model, or Gradio imports here: spaces must patch CUDA first.
ENGINES = {}
OUTPUT_DIR = None


def gpu_duration(ref_audio, ref_text, text, model_variant, seed, speed, mel_temperature):
    import gradio as gr
    from runtime import validate_request
    try:
        request = validate_request(ref_audio, ref_text, text, model_variant, seed, speed, mel_temperature)
    except (ValueError, OSError) as exc:
        raise gr.Error(str(exc)) from exc
    seconds = request["predicted_seconds"]
    # Conservative allocation budgets, not a promise of synthesis latency.
    return min(120, max(30, int(20 + 3 * seconds))) if model_variant == "V-Spec" else min(40, max(15, int(10 + seconds)))


def generate(ref_audio, ref_text, text, model_variant="V-Spec", seed=0, speed=1.0, mel_temperature=0.0):
    """This exact top-level function is GPU-decorated and registered with Gradio."""
    import gradio as gr
    import torch
    from runtime import save_audio, synthesize_request, validate_request
    try:
        request = validate_request(ref_audio, ref_text, text, model_variant, seed, speed, mel_temperature)
        if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
            raise RuntimeError("A BF16-capable GPU is required for this demo.")
        if model_variant not in ENGINES or OUTPUT_DIR is None:
            raise RuntimeError("Models are not ready yet. Please retry after the Space has started.")
        started = time.perf_counter()
        sample_rate, waveform = synthesize_request(ENGINES, request)
        path = save_audio(sample_rate, waveform, OUTPUT_DIR)
        elapsed = time.perf_counter() - started
        steps = 4 if model_variant == "Flash" else 50
        return path, f"WaveGen {model_variant} · {steps} steps · Mel temp {request['mel_temperature']:g} · seed {request['seed']} · {elapsed:.1f}s"
    except ValueError as exc:
        raise gr.Error(str(exc)) from exc
    except (RuntimeError, OSError) as exc:
        # Avoid exposing server paths or model-cache details to public clients.
        print(f"Generation failed ({type(exc).__name__})", flush=True)
        raise gr.Error("Generation failed. Try a shorter sentence or retry when a GPU is available.") from exc


def create_app(generate_fn):
    import gradio as gr
    from runtime import MODEL_CHOICES, model_description
    global OUTPUT_DIR
    with gr.Blocks(title="WaveGen", analytics_enabled=False, delete_cache=(3600, 86400)) as demo:
        gr.Markdown("# WaveGen\n**NeurIPS 2026** · Sang-Hoon Lee and Ha-Yeong Choi · Ajou University and KT\n\n"
                    "Generate speech from a reference voice and your text. English and Mandarin Chinese are supported.\n\n"
                    "[Source Code](https://github.com/sh-lee-prml/WaveGen) · "
                    "[Models](https://huggingface.co/LeeSangHoon/WaveGen) · "
                    "[Demo Page](https://wavegen26.github.io/demo/) · "
                    "[Benchmark Demo](https://sh-lee-prml.github.io/wavegen-tts-demo/)")
        model_variant = gr.Dropdown(MODEL_CHOICES, value="V-Spec", label="Model")
        details = gr.Markdown(model_description("V-Spec"))
        gr.Markdown("Reference and generated speech must fit within **30 seconds combined**. "
                    "Use a short, clear reference and enter its exact transcript. Long text is rejected rather than cut off.")
        with gr.Row():
            with gr.Column():
                ref_audio = gr.Audio(label="Reference speech", sources=["upload", "microphone"], type="filepath")
                ref_text = gr.Textbox(label="Reference transcript", lines=3)
                text = gr.Textbox(label="Text to generate", lines=4, max_lines=10)
            with gr.Column():
                output = gr.Audio(label="Generated speech", type="filepath", format="wav",
                                  interactive=False, buttons=["download"])
                status = gr.Textbox(label="Generation details", interactive=False)
        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                seed = gr.Number(value=0, precision=0, minimum=0, maximum=2**32 - 1, label="Seed")
                speed = gr.Slider(0.5, 2, value=1, step=0.05, label="Speed")
                mel_temperature = gr.Slider(0, 1, value=0, step=0.05, label="Mel temperature")
        button = gr.Button("Generate speech", variant="primary")
        gr.Examples(
            examples=[[str(EXAMPLE_DIR / "basic_ref_en.wav"),
                       "Some call me nature, others call me mother nature.",
                       "Here we generate something just for test."]],
            inputs=[ref_audio, ref_text, text], cache_examples=False)
        gr.Markdown("Use reference speech only with the speaker's permission. Do not use this demo for "
                    "deceptive impersonation. Clearly identify shared outputs as AI-generated audio. "
                    "The model weights are licensed under [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/).")
        model_variant.change(model_description, inputs=model_variant, outputs=details, queue=False)
        OUTPUT_DIR = Path(output.GRADIO_CACHE)
        button.click(generate_fn, inputs=[ref_audio, ref_text, text, model_variant, seed, speed, mel_temperature],
                     outputs=[output, status], concurrency_limit=1, concurrency_id="wavegen", api_name="generate")
    return demo.queue(default_concurrency_limit=1, max_size=16)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--local", action="store_true", help="Use an existing local CUDA environment without ZeroGPU scheduling.")
    parser.add_argument("--model-dir", type=Path, help="Optional existing ckpts directory for offline local verification.")
    parser.add_argument("--server-name", default=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"))
    parser.add_argument("--server-port", type=int, default=int(os.getenv("GRADIO_SERVER_PORT", "7860")))
    args = parser.parse_args()
    if not args.local:
        import spaces  # MUST precede torch and any model imports.
    from runtime import MODEL_REPO, MODEL_REVISION, load_models
    global ENGINES
    ENGINES = load_models(model_dir=args.model_dir, repo_id=os.getenv("WAVEGEN_MODEL_REPO", MODEL_REPO),
                          revision=os.getenv("WAVEGEN_MODEL_REVISION", MODEL_REVISION))
    generate_fn = generate if args.local else spaces.GPU(duration=gpu_duration)(generate)
    demo = create_app(generate_fn)
    demo.launch(server_name=args.server_name, server_port=args.server_port,
                allowed_paths=[str(EXAMPLE_DIR / "basic_ref_en.wav")],
                max_file_size="20mb", share=False, show_error=False, ssr_mode=False)


if __name__ == "__main__":
    main()