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