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