Blame for runtime.py

"""Spaces adapter; the vendored WaveGen inference implementations are unchanged."""

from __future__ import annotations

import math
import hashlib
import os
from pathlib import Path
import re
import tempfile

import numpy as np
import soundfile as sf


MODEL_REPO = "LeeSangHoon/WaveGen"
MODEL_REVISION = "da53c835a39272c05e01526fe3c10b2b9d4ce60b"
MAX_SECONDS = 30
MAX_FILE_BYTES = 20 * 1024 * 1024
MAX_TEXT_CHARS = 2000
MODEL_CHOICES = [("WaveGen V-Spec (Recommended)", "V-Spec"),
                 ("WaveGen-Flash (4 Steps)", "Flash")]
ASSETS = {
    "V-Spec": ("WaveGen_VSpec/model_1400000.safetensors",
               "WaveGen_VSpec/WaveGen_VSpec.yaml", "WaveGen_VSpec/vocab.txt"),
    "Flash": ("WaveGen_Flash/model.safetensors",
              "WaveGen_Flash/WaveGen_Flash.yaml", "WaveGen_Flash/vocab.txt"),
}


def _number(value, name, low, high, *, integer=False):
    if (isinstance(value, bool) or not isinstance(value, (int, float))
            or not math.isfinite(value) or not low <= value <= high
            or (integer and value != int(value))):
        raise ValueError(f"{name} must be {'an integer' if integer else 'a number'} between {low} and {high}.")
    return int(value) if integer else float(value)


def validate_request(ref_audio, ref_text, text, model_variant="V-Spec", seed=0,
                     speed=1.0, mel_temperature=0.0):
    """Bound decoded input and mirror native UTF-8 duration prediction on CPU."""
    if model_variant not in ASSETS:
        raise ValueError("Select WaveGen V-Spec or WaveGen-Flash.")
    for value, name in ((ref_text, "Reference transcript"), (text, "Text to generate")):
        if not isinstance(value, str) or not value.strip():
            raise ValueError(f"{name} is required.")
        if len(value) > MAX_TEXT_CHARS:
            raise ValueError(f"{name} must be at most {MAX_TEXT_CHARS:,} characters.")
    seed = _number(seed, "Seed", 0, 2**32 - 1, integer=True)
    speed = _number(speed, "Speed", 0.5, 2)
    mel_temperature = _number(mel_temperature, "Mel temperature", 0, 1)
    if not isinstance(ref_audio, (str, os.PathLike)) or not Path(ref_audio).is_file():
        raise ValueError("Upload or record reference speech first.")
    path = Path(ref_audio)
    if path.stat().st_size > MAX_FILE_BYTES:
        raise ValueError("Reference audio must be at most 20 MB.")
    try:
        info = sf.info(str(path))
        if not 0 < info.duration <= MAX_SECONDS:
            raise ValueError("Reference speech must be nonempty and at most 30 seconds.")
        if not 8000 <= info.samplerate <= 192000 or info.channels not in (1, 2):
            raise ValueError("Use mono or stereo reference speech at 8–192 kHz.")
        audio, _ = sf.read(str(path), dtype="float32", always_2d=True)
    except (RuntimeError, OSError) as exc:
        raise ValueError("Could not decode the reference audio. Try a WAV file.") from exc
    mono = audio.mean(axis=1)
    if not np.isfinite(audio).all() or not len(mono) or np.sqrt(np.mean(mono**2)) <= 1e-8:
        raise ValueError("Reference audio must contain audible, finite speech.")
    # torchaudio.Resample rounds the output length up, then the native loader
    # truncates/pads the reference to a whole 480-sample waveform patch.
    resampled_length = (info.frames * 24000 + info.samplerate - 1) // info.samplerate
    prompt_length = max(resampled_length // 480 * 480, 480)
    duration_ref_text = ref_text + (" " if len(ref_text[-1].encode("utf-8")) == 1 else "")
    total_length = prompt_length + int(
        prompt_length / len(duration_ref_text.encode("utf-8")) * len(text.encode("utf-8")) / speed
    )
    if total_length <= prompt_length:
        raise ValueError("Predicted speech is empty. Enter more text or reduce speed.")
    if total_length > MAX_SECONDS * 24000:
        raise ValueError("Reference plus predicted speech exceeds 30 seconds. Shorten the text or reference audio.")
    return dict(ref_audio=str(path), ref_text=ref_text, text=text, model_variant=model_variant,
                seed=seed, speed=speed, mel_temperature=mel_temperature,
                predicted_seconds=total_length / 24000)


def model_description(model_variant):
    if model_variant == "Flash":
        return "**WaveGen-Flash:** 4 steps · CFG 0 · config dropout · final Mel head skipped."
    return "**WaveGen V-Spec:** 50 steps · Prompt Anchor · Scheduled K3 (CFG 1.25) · config dropout · final Mel head skipped."


def load_models(*, model_dir=None, repo_id=MODEL_REPO, revision=MODEL_REVISION):
    """Call after importing spaces, before Gradio launch / GPU worker creation.

    CPU loading avoids the native loader's CUDA/BF16 device probes. ZeroGPU
    captures startup model.to('cuda'); actual CUDA checks belong in generate().
    """
    import torch
    from wavegen_v_spec.infer.utils_infer import WaveGen
    from wavegen_flash import WaveGenFlash

    torch.set_num_threads(1)
    factories = {"V-Spec": WaveGen, "Flash": WaveGenFlash}
    engines = {}
    resolved_revision = revision
    if model_dir is None and not re.fullmatch(r"[0-9a-fA-F]{40}", revision):
        from huggingface_hub import model_info
        resolved_revision = model_info(repo_id, revision=revision).sha
        if not resolved_revision or not re.fullmatch(r"[0-9a-fA-F]{40}", resolved_revision):
            raise RuntimeError("The model revision did not resolve to an immutable commit.")
    if model_dir is None:
        from huggingface_hub import hf_hub_download

        # On Linux, the normal Hub snapshot cache uses symlinks to extensionless
        # blobs. The native loader resolves paths before checking the format, so
        # request actual named files instead of passing snapshot symlinks to it.
        repo_key = hashlib.sha256(repo_id.encode("utf-8")).hexdigest()[:16]
        download_root = Path(tempfile.gettempdir()) / "wavegen-hub" / repo_key / resolved_revision
    for variant, (checkpoint, config, vocabulary) in ASSETS.items():
        print(f"Loading {variant} for inference", flush=True)
        options = dict(device="cpu", dtype="fp32")
        if variant == "V-Spec":
            options["use_ema"] = True
        if model_dir is not None:
            root = Path(model_dir)
            engine = factories[variant].from_checkpoint(
                checkpoint=root / checkpoint, vocab_file=root / vocabulary,
                config=root / config, **options)
        else:
            checkpoint_path, config_path, vocabulary_path = [hf_hub_download(
                repo_id=repo_id, filename=filename, revision=resolved_revision,
                local_dir=download_root,
            ) for filename in (checkpoint, config, vocabulary)]
            engine = factories[variant].from_checkpoint(
                checkpoint=checkpoint_path, config=config_path,
                vocab_file=vocabulary_path, **options)
            engine.hub_repo_id = repo_id
            engine.hub_revision = resolved_revision
        if engine.sample_rate != 24000:
            raise ValueError("The demo requires the 24 kHz model bundles.")
        engine.model.requires_grad_(False).to("cuda")
        engine.device = torch.device("cuda")
        engine.dtype = torch.bfloat16
        engines[variant] = engine
    return engines


def synthesize_request(engines, request):
    variant = request["model_variant"]
    engine = engines[variant]
    options = {key: request[key] for key in
               ("ref_audio", "ref_text", "text", "seed", "speed", "mel_temperature")}
    options["compute_final_mel"] = False
    if variant == "V-Spec":
        options.update(steps=50, cfg_schedule="scheduled_k3", cfg_strength=1.25,
                       inference_method="Prompt_xt_Anchor", inference_dropout=None,
                       max_duration_frames=1500)
    else:
        options["max_duration_seconds"] = MAX_SECONDS
    return engine.synthesize(**options)


def save_audio(sample_rate, waveform, output_dir):
    """Preserve native levels: return PCM16 WAV, not a Gradio numpy-audio tuple."""
    waveform = np.asarray(waveform)
    if sample_rate != 24000 or waveform.ndim != 1 or not waveform.size or not np.isfinite(waveform).all():
        raise RuntimeError("The model returned invalid audio.")
    if waveform.size > MAX_SECONDS * sample_rate:
        raise RuntimeError("The generated audio exceeds the demo duration limit.")
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(prefix="wavegen-", suffix=".wav", dir=output_dir, delete=False) as f:
        path = Path(f.name)
    try:
        sf.write(str(path), waveform, sample_rate, subtype="PCM_16")
    except Exception:
        path.unlink(missing_ok=True)
        raise
    return str(path)