"""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()