Spaces:
Running on Zero
Running on Zero
| """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() | |