app.py

"""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 = """
<div id="hero">
  <h1>💎 Diamond</h1>
  <p>Speech restoration — turn degraded audio (podcast · phone · streamed ·
     over-compressed) into near-studio <b>44.1&nbsp;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()