"""Diamond — speech restoration demo (Hugging Face Space, ZeroGPU). Degraded speech in → near-studio 44.1 kHz out. The model is an RQ-Transformer (encoder → time-transformer → depth-transformer over 9 DAC RVQ codebooks → frozen DAC decoder). Decoding is chunked greedy with a repetition penalty. The heavy lifting runs inside `@spaces.GPU`, so the model sits on CPU between requests and is moved to the GPU only for the duration of a restoration call. """ from __future__ import annotations import hashlib import os import tempfile import gradio as gr import numpy as np import soundfile as sf import spaces import torch from huggingface_hub import hf_hub_download from diamond.infer import Diamond, SampleConfig, _load_audio HERE = os.path.dirname(os.path.abspath(__file__)) EX = os.path.join(HERE, "examples") # ── Checkpoint ──────────────────────────────────────────────────────────────── # Public release. Override with the DIAMOND_HF_REPO env var on the Space. HF_REPO = os.environ.get("DIAMOND_HF_REPO", "nineninesix/diamond-1.0") WEIGHTS = os.environ.get("DIAMOND_WEIGHTS", "diamond.safetensors") CONFIG = os.environ.get("DIAMOND_CONFIG", "diamond.json") # Pull weights + config sidecar into the same dir so load_checkpoint finds both. _ckpt_path = hf_hub_download(HF_REPO, WEIGHTS) hf_hub_download(HF_REPO, CONFIG) # sibling .json, same snapshot dir # Build on CPU at import time (ZeroGPU forbids CUDA init outside @spaces.GPU). # The DAC codec weights download automatically here on first run. DIA = Diamond.from_pretrained(_ckpt_path, device="cpu") def _move(device: str) -> None: """Move the whole pipeline (model + mel + DAC) to `device`.""" DIA.model.to(device) DIA.mel.to(device) DIA.dac.model.to(device) DIA.dac.device = device # Restorations are cached to disk keyed by (audio bytes + params) so clicking an # example only hits the GPU the first time; repeat clicks replay the cached file. CACHE_DIR = os.path.join(tempfile.gettempdir(), "diamond_cache") os.makedirs(CACHE_DIR, exist_ok=True) def _cache_key(audio_path, chunk_sec, overlap_sec, rep_penalty, trim_leadin) -> str: h = hashlib.sha256() with open(audio_path, "rb") as f: h.update(f.read()) h.update(repr((float(chunk_sec), float(overlap_sec), float(rep_penalty), bool(trim_leadin))).encode()) return h.hexdigest() @spaces.GPU(duration=120) def _restore_gpu(audio_path, chunk_sec, overlap_sec, rep_penalty, trim_leadin, out_path): torch.manual_seed(42) np.random.seed(42) wav, sr = _load_audio(audio_path) if wav.size == 0: raise gr.Error("Empty audio.") _move("cuda") try: restored, out_sr = DIA.restore( wav, sr, chunk_sec=float(chunk_sec), overlap_sec=float(overlap_sec), warmup_sec=1.0, tail_pad_sec=1.0, normalize=True, trim_leadin=bool(trim_leadin), recover_collapse=True, sample=SampleConfig(rep_penalty=float(rep_penalty)), ) finally: _move("cpu") # release the GPU; keep clean state between ZeroGPU calls sf.write(out_path, restored, out_sr, subtype="PCM_16") return out_path def restore(audio_path, chunk_sec, overlap_sec, rep_penalty, trim_leadin): """Cached restore: return the stored result on a hit, else run on GPU. Skipping the ``@spaces.GPU`` call on a cache hit means example replays are instant and don't burn a ZeroGPU allocation. """ if audio_path is None: raise gr.Error("Upload or record some audio first.") out_path = os.path.join( CACHE_DIR, _cache_key(audio_path, chunk_sec, overlap_sec, rep_penalty, trim_leadin) + ".wav", ) if os.path.exists(out_path): return out_path return _restore_gpu(audio_path, chunk_sec, overlap_sec, rep_penalty, trim_leadin, out_path) # ── Theme ───────────────────────────────────────────────────────────────────── THEME = gr.themes.Soft( primary_hue="blue", secondary_hue="blue", neutral_hue="slate", font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"], ).set( body_background_fill="*neutral_50", block_radius="14px", button_primary_shadow="*shadow_drop_lg", ) CSS = """ .gradio-container {max-width: 980px !important; margin: 0 auto;} #hero {text-align:center; padding: 8px 0 2px;} #hero h1 {font-size: 2.2rem; margin: 0; letter-spacing:-0.5px;} #hero p {color: var(--body-text-color-subdued); margin: 6px 0 0;} .badge-row {display:flex; gap:8px; justify-content:center; flex-wrap:wrap; margin-top:10px;} footer {display:none !important;} """ HERO = """

💎 Diamond

Speech restoration — turn degraded audio (podcast · phone · streamed · over-compressed) into near-studio 44.1 kHz.

""" # ── UI ──────────────────────────────────────────────────────────────────────── with gr.Blocks(theme=THEME, css=CSS, title="Diamond — Speech Restoration") as demo: gr.HTML(HERO) with gr.Row(equal_height=True): with gr.Column(): inp = gr.Audio( label="Degraded input — upload or record", type="filepath", sources=["upload", "microphone"], ) with gr.Row(): chunk = gr.Slider(1.5, 4.0, value=2.5, step=0.5, label="Chunk length (s)") overlap = gr.Slider(0.0, 1.0, value=0.4, step=0.1, label="Crossfade overlap (s)") with gr.Accordion("Advanced", open=False): rep = gr.Slider(1.0, 1.6, value=1.3, step=0.1, label="Repetition penalty (suppresses AR looping)") lead = gr.Checkbox(value=True, label="Trim non-speech intro") btn = gr.Button("✨ Restore", variant="primary", size="lg") with gr.Column(): out = gr.Audio(label="Restored — 44.1 kHz", type="filepath") btn.click(restore, inputs=[inp, chunk, overlap, rep, lead], outputs=out) gr.Markdown("### Examples — click to restore (generated once, then cached)") gr.Examples( examples=[ [os.path.join(EX, "01_degraded.wav"), 2.5, 0.4, 1.3, True], [os.path.join(EX, "02_degraded.wav"), 2.5, 0.4, 1.3, True], [os.path.join(EX, "03_degraded.wav"), 2.5, 0.4, 1.3, True], [os.path.join(EX, "04_degraded.wav"), 2.5, 0.4, 1.3, True], [os.path.join(EX, "05_degraded.wav"), 2.5, 0.4, 1.3, True], [os.path.join(EX, "06_degraded.wav"), 2.5, 0.4, 1.3, True], ], inputs=[inp, chunk, overlap, rep, lead], outputs=out, fn=restore, run_on_click=True, # Gradio's own example cache doesn't play well with ZeroGPU (it can't warm # a @spaces.GPU fn at build and leaves the dataset non-clickable). We keep # the examples live and cache inside `restore` ourselves instead. cache_examples=False, label="Degraded input · Chunk (s) · Crossfade overlap (s) · Rep penalty · Trim intro", ) gr.Markdown( "Model: [`nineninesix/diamond-1.0`]" "(https://huggingface.co/nineninesix/diamond-1.0) · " "Training: [`diamond-train`](https://github.com/nineninesix-ai/diamond-train.git) · " "Inference: [`diamond-inference`](https://github.com/nineninesix-ai/diamond-inference.git)" ) if __name__ == "__main__": demo.queue().launch()