Spaces:
Running on Zero
Running on Zero
| """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() | |
| 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 = """ | |
| <div id="hero"> | |
| <h1>💎 Diamond</h1> | |
| <p>Speech restoration — turn degraded audio (podcast · phone · streamed · | |
| over-compressed) into near-studio <b>44.1 kHz</b>.</p> | |
| <div class="badge-row"> | |
| <img src="https://img.shields.io/badge/44.1_kHz-output-1d4ed8"> | |
| <img src="https://img.shields.io/badge/license-Apache_2.0-3b82f6"> | |
| </div> | |
| </div> | |
| """ | |
| # ── 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() | |