Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/251-118436-0003.wav filter=lfs diff=lfs merge=lfs -text
+assets/5_1812_20170628135856.wav filter=lfs diff=lfs merge=lfs -text
+assets/p226_006.wav filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,50 @@
 ---
-title: Focalcodec Quantization
-emoji: 🏆
-colorFrom: indigo
+title: FocalCodec Quantization Lab
+emoji: ⚡
+colorFrom: blue
 colorTo: purple
 sdk: gradio
 sdk_version: 6.30.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: FocalCodec weights at int4–int32 vs fp4–fp32
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ⚡ FocalCodec 12.5 Hz — weight quantization lab
+
+Interactive demo of what happens to [`lucadellalib/focalcodec_12_5hz`](https://huggingface.co/lucadellalib/focalcodec_12_5hz)
+when its **145 M weights** are stored at different numeric precisions:
+
+`int4` · `int8` · `int16` · `int32` · `fp4` · `fp8` · `fp16` · `fp32`
+
+The token stream is untouched, so the codec still runs at **0.1625 kbps**
+(12.5 tokens/s × 13 bits) at every level — only the size of the network changes
+(72.6 MB at 4-bit down to 581.1 MB at fp32).
+
+## Three views
+
+| Tab | What it shows |
+| --- | --- |
+| **Compare** | One reconstructed clip + input/reconstruction spectrograms, side by side for all 8 levels |
+| **Playground** | Encode → decode a single clip at one chosen level, with its metrics row |
+| **Benchmark** | SI-SDR, STOI, LSD, stored model size, bitrate and real-time factor for all 8 levels, plus a summary chart |
+
+## Method
+
+Post-training **fake quantization** of every floating-point tensor: per-output-channel
+absmax scale, round-to-nearest in the target format, dequantize back to fp32, then run
+the untouched FocalCodec forward pass. `fp4` is IEEE-style E2M1, `fp8` is E4M3, `fp16`
+and `fp32` are exact casts. Nothing is fine-tuned or calibrated on data.
+
+⚠️ **On SI-SDR:** it is deeply negative at every level, including `fp32`. That is
+expected — at 0.16 kbps FocalCodec does not preserve waveform phase, and the
+[FocalCodec paper](https://arxiv.org/abs/2502.04465) reports SI-SNR ≈ −37 dB for the
+untouched model. Use **STOI** (intelligibility, higher is better) and **LSD**
+(log-spectral distance, lower is better) to compare levels.
+
+## Credits
+
+Model, method and example audio by
+[Luca Della Libera](https://huggingface.co/lucadellalib) et al. — Apache-2.0.
+[FocalCodec: Low-Bitrate Speech Coding via Focal Modulation Networks](https://arxiv.org/abs/2502.04465), NeurIPS 2025.

app.py

@@ -0,0 +1,193 @@
+"""FocalCodec 12.5 Hz weight-quantization lab (Gradio / ZeroGPU)."""
+
+import hashlib
+import time
+
+import spaces  # must precede torch: patches torch.cuda for ZeroGPU
+
+import gradio as gr
+import soundfile as sf
+import torch
+
+import core
+from core import QUANT_LEVELS
+
+MODEL_REPO = "lucadellalib/focalcodec_12_5hz"
+HUB_REPO = "lucadellalib/focalcodec"
+MAX_SECONDS = 20
+HEADERS = ["Level", "SI-SDR (dB)", "STOI", "LSD (dB)", "Model size (MB)", "Bitrate (kbps)", "RTF"]
+DATATYPES = ["str", "number", "number", "number", "number", "number", "number"]
+EXAMPLES = [
+    ["assets/251-118436-0003.wav"],
+    ["assets/p226_006.wav"],
+    ["assets/5_1812_20170628135856.wav"],
+]
+
+codec = torch.hub.load(HUB_REPO, "focalcodec", config=MODEL_REPO, trust_repo=True)
+codec.eval().requires_grad_(False)
+codec.to("cuda")
+BASE_STATE = {k: v.detach().cpu().clone() for k, v in codec.state_dict().items()}
+SIZES = {level: core.model_size_mb(BASE_STATE, level) for level in QUANT_LEVELS}
+SAMPLE_RATE = codec.sample_rate_input
+assert codec.sample_rate_output == SAMPLE_RATE, "demo assumes a 16/16 kHz codec"
+
+
+def _estimate_duration(audio_path: str, levels: tuple) -> float:
+    try:
+        seconds = min(sf.info(audio_path).duration, MAX_SECONDS)
+    except Exception:
+        seconds = float(MAX_SECONDS)
+    return min(240.0, 5.0 + 2.0 * len(levels) + 2.5 * seconds)
+
+
+@spaces.GPU(duration=_estimate_duration)
+def _encode_decode(audio_path: str, levels: tuple) -> dict:
+    """Reconstruct `audio_path` once per weight-quantization level, scoring each result."""
+    ref, sr = core.load_audio(audio_path, SAMPLE_RATE, MAX_SECONDS)
+    sig = torch.from_numpy(ref).unsqueeze(0).cuda()
+    rows, reconstructed = [], {}
+    for level in levels:
+        codec.load_state_dict(BASE_STATE if level == "fp32" else core.quantize_state(BASE_STATE, level))
+        torch.cuda.synchronize()
+        start = time.perf_counter()
+        with torch.inference_mode():
+            rec = codec.sig_to_sig(sig).squeeze(0).float().cpu().numpy()
+        torch.cuda.synchronize()
+        rtf = (time.perf_counter() - start) / (len(ref) / sr)
+        rows.append(core.evaluate(ref, rec, sr, rtf, SIZES[level], level))
+        reconstructed[level] = (sr, rec)
+    codec.load_state_dict(BASE_STATE)
+    return {"input": (sr, ref), "rows": rows, "audio": reconstructed}
+
+
+_CACHE: dict[str, dict] = {}
+
+
+def run(audio_path: str, levels: list | tuple) -> dict:
+    """Encode and decode the audio at every requested weight-quantization level.
+
+    Results are cached by audio content and level set, so the Compare and
+    Benchmark tabs share a single GPU pass.
+    """
+    if not audio_path:
+        raise gr.Error("Provide audio first: upload, record, or pick an example.")
+    levels = tuple(levels)
+    with open(audio_path, "rb") as fh:
+        key = hashlib.md5(fh.read()).hexdigest() + "|" + ",".join(levels)
+    if key not in _CACHE:
+        if len(_CACHE) >= 6:
+            _CACHE.clear()
+        _CACHE[key] = _encode_decode(audio_path, levels)
+    return _CACHE[key]
+
+
+def compare(audio_path: str):
+    """Rebuild the audio at all eight levels and return one player + spectrogram per level."""
+    result = run(audio_path, QUANT_LEVELS)
+    audio = [result["audio"][level] for level in QUANT_LEVELS]
+    plots = [
+        core.spectrogram_figure(result["input"][1], result["audio"][level][1], SAMPLE_RATE, level)
+        for level in QUANT_LEVELS
+    ]
+    return audio + plots
+
+
+def playground(audio_path: str, level: str):
+    """Rebuild the audio at a single quantization level, with its metrics row."""
+    result = run(audio_path, [level])
+    rec = result["audio"][level]
+    return (
+        rec,
+        core.spectrogram_figure(result["input"][1], rec[1], SAMPLE_RATE, level),
+        [result["rows"][0]],
+    )
+
+
+def benchmark(audio_path: str):
+    """Score every quantization level on one input: metrics table plus summary chart."""
+    result = run(audio_path, QUANT_LEVELS)
+    return result["rows"], core.metrics_figure(result["rows"])
+
+
+HEADER = f"""# ⚡ FocalCodec 12.5 Hz — weight quantization lab
+
+Same checkpoint ([`lucadellalib/focalcodec_12_5hz`](https://huggingface.co/lucadellalib/focalcodec_12_5hz)),
+eight numeric formats for its **{sum(v.numel() for v in BASE_STATE.values()) / 1e6:.0f} M weights**: `int4 … int32` and `fp4 … fp32`.
+The bitrate never changes — **{core.BITRATE_KBPS:g} kbps** (12.5 tokens/s × 13 bits) — only the storage precision of the network itself.
+
+**Method:** post-training fake quantization, per-output-channel absmax scaling, round-to-nearest, dequantized back to fp32 for inference. `fp4` = E2M1, `fp8` = E4M3, `fp16`/`fp32` are exact casts. No fine-tuning.
+
+⚠️ *SI-SDR sits far below 0 dB at every level: at 0.16 kbps the codec does not preserve waveform phase, and the FocalCodec paper reports SI-SNR ≈ −37 dB for the untouched model. It is an anchor, not a quality score — read **STOI** (intelligibility) and **LSD** (log-spectral distance, lower is better) instead.*
+"""
+
+FOOTER = """Audio examples from the [FocalCodec repository](https://github.com/lucadellalib/focalcodec)
+(LibriSpeech, VoiceBank, MagicData). Model and method by Luca Della Libera et al., Apache-2.0."""
+
+with gr.Blocks(title="FocalCodec quantization lab", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(HEADER)
+
+    with gr.Tab("Compare"):
+        with gr.Row():
+            with gr.Column(scale=1, min_width=320):
+                cmp_input = gr.Audio(
+                    sources=["upload", "microphone"], type="filepath", label="Input audio"
+                )
+                gr.Examples(examples=EXAMPLES, inputs=cmp_input, label="Examples")
+                cmp_button = gr.Button("Reconstruct at all 8 levels", variant="primary")
+                gr.Markdown(f"Input capped at {MAX_SECONDS} s.")
+            with gr.Column(scale=4):
+                cmp_audio, cmp_plot = {}, {}
+                for chunk in (QUANT_LEVELS[:4], QUANT_LEVELS[4:]):
+                    with gr.Row():
+                        for level in chunk:
+                            with gr.Column():
+                                gr.Markdown(f"**{level}**")
+                                cmp_audio[level] = gr.Audio(interactive=False, label=level)
+                                cmp_plot[level] = gr.Plot(label=f"{level} spectrogram")
+        cmp_button.click(
+            compare,
+            inputs=cmp_input,
+            outputs=[cmp_audio[level] for level in QUANT_LEVELS]
+            + [cmp_plot[level] for level in QUANT_LEVELS],
+        )
+
+    with gr.Tab("Playground"):
+        with gr.Row():
+            with gr.Column(scale=1, min_width=320):
+                pg_input = gr.Audio(
+                    sources=["upload", "microphone"], type="filepath", label="Input audio"
+                )
+                gr.Examples(examples=EXAMPLES, inputs=pg_input, label="Examples")
+                pg_level = gr.Dropdown(
+                    choices=list(QUANT_LEVELS), value="int8", label="Weight format"
+                )
+                pg_button = gr.Button("Encode + decode", variant="primary")
+            with gr.Column(scale=3):
+                pg_output = gr.Audio(interactive=False, label="Reconstruction")
+                pg_plot = gr.Plot(label="Input vs reconstruction")
+        pg_table = gr.Dataframe(headers=HEADERS, datatype=DATATYPES, interactive=False)
+        pg_button.click(
+            playground, inputs=[pg_input, pg_level], outputs=[pg_output, pg_plot, pg_table]
+        )
+
+    with gr.Tab("Benchmark"):
+        with gr.Row():
+            with gr.Column(scale=1, min_width=320):
+                bm_input = gr.Audio(
+                    sources=["upload", "microphone"], type="filepath", label="Input audio"
+                )
+                gr.Examples(examples=EXAMPLES, inputs=bm_input, label="Examples")
+                bm_button = gr.Button("Run all 8 levels", variant="primary")
+            with gr.Column(scale=3):
+                gr.Markdown(
+                    "**STOI** ↑ intelligibility · **LSD** ↓ spectral error · "
+                    "**RTF** = encode+decode time ÷ audio duration · size assumes "
+                    "the stored weight format."
+                )
+        bm_table = gr.Dataframe(headers=HEADERS, datatype=DATATYPES, interactive=False)
+        bm_plot = gr.Plot(label="Quality and size per level")
+        bm_button.click(benchmark, inputs=bm_input, outputs=[bm_table, bm_plot])
+
+    gr.Markdown(FOOTER)
+
+demo.launch(mcp_server=True)

assets/251-118436-0003.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:340ee98cdd138b743dcf12d82bc2d83b39b72a413bb6f1d8a975d7d1ae8d7217
+size 705360

assets/5_1812_20170628135856.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dc7aaa044b83ba0a88b4089d2914440f82490456ce3b294297847b60a8aa7050
+size 151624

assets/p226_006.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53b9be9b754b28c4907469333a1d938bc3d6ae7e513e7b371c999f337720d095
+size 489980

core.py

@@ -0,0 +1,158 @@
+"""Quantization, audio IO, objective metrics and plots for the FocalCodec lab.
+
+Kept free of Gradio / torch.hub so it can be exercised standalone.
+"""
+
+from __future__ import annotations
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+import matplotlib.pyplot as plt
+import numpy as np
+import soundfile as sf
+import torch
+from pystoi import stoi
+from scipy.signal import resample_poly, stft
+
+QUANT_LEVELS = ("int4", "int8", "int16", "int32", "fp4", "fp8", "fp16", "fp32")
+
+TOKEN_RATE_HZ = 12.5
+BITS_PER_TOKEN = 13
+BITRATE_KBPS = TOKEN_RATE_HZ * BITS_PER_TOKEN / 1000.0
+
+_BITS = {"int4": 4, "int8": 8, "int16": 16, "int32": 32, "fp4": 4, "fp8": 8, "fp16": 16, "fp32": 32}
+_FP_FORMATS = {"fp4": (2, 1), "fp8": (4, 3)}  # (exponent bits, mantissa bits) -> E2M1 / E4M3
+
+
+def _ref_scale(w: torch.Tensor) -> torch.Tensor:
+    """Per-output-channel absmax, per-tensor for 1-D params."""
+    if w.dim() > 1:
+        return w.abs().amax(dim=tuple(range(1, w.dim())), keepdim=True)
+    return w.abs().amax()
+
+
+def _quant_int(w: torch.Tensor, bits: int) -> torch.Tensor:
+    if bits >= 32:
+        return w
+    qmax = 1 << (bits - 1)
+    scale = (_ref_scale(w) / qmax).clamp_min(1e-12)
+    return (w / scale).round().clamp(-qmax, qmax - 1) * scale
+
+
+def _quant_fp(w: torch.Tensor, e_bits: int, m_bits: int) -> torch.Tensor:
+    bias = (1 << (e_bits - 1)) - 1
+    max_exp = (1 << e_bits) - 1 - bias
+    max_norm = (2.0 - 2.0 ** -m_bits) * 2.0**max_exp
+    scale = (_ref_scale(w) / max_norm).clamp_min(1e-12)
+    x = w / scale
+    ax = x.abs()
+    exp = torch.floor(torch.log2(ax.clamp_min(1e-30))).clamp(1 - bias, max_exp)
+    step = torch.pow(2.0, exp - m_bits)
+    return torch.sign(x) * ((ax / step).round() * step) * scale
+
+
+def quantize_state(state: dict, level: str) -> dict:
+    """Fake-quantize every float tensor in `state` to `level` (per-channel absmax scaling)."""
+    if level == "fp32":
+        return dict(state)
+    out = {}
+    for key, value in state.items():
+        if not value.is_floating_point():
+            out[key] = value
+        elif level == "fp16":
+            out[key] = value.half().float()
+        elif level.startswith("int"):
+            out[key] = _quant_int(value, int(level[3:]))
+        else:
+            out[key] = _quant_fp(value, *_FP_FORMATS[level])
+    return out
+
+
+def model_size_mb(state: dict, level: str) -> float:
+    return sum(v.numel() for v in state.values()) * _BITS[level] / 8 / 1e6
+
+
+def load_audio(path: str, target_sr: int, max_seconds: float | None = None) -> tuple[np.ndarray, int]:
+    """Read `path` as mono float32, resample to `target_sr`, optionally truncate."""
+    x, sr = sf.read(path, dtype="float32", always_2d=True)
+    x = x.mean(axis=1)
+    if sr != target_sr:
+        g = np.gcd(sr, target_sr)
+        x = resample_poly(x, target_sr // g, sr // g).astype(np.float32)
+        sr = target_sr
+    if max_seconds is not None and len(x) > int(max_seconds * sr):
+        x = x[: int(max_seconds * sr)]
+    return x, sr
+
+
+def si_sdr(ref: np.ndarray, rec: np.ndarray) -> float:
+    ref = ref - ref.mean()
+    rec = rec - rec.mean()
+    denom = float(ref @ ref) + 1e-12
+    target = (float(rec @ ref) / denom) * ref
+    noise = rec - target
+    return float(10 * np.log10((float(target @ target) + 1e-12) / (float(noise @ noise) + 1e-12)))
+
+
+def log_spectral_distance(ref: np.ndarray, rec: np.ndarray, sr: int) -> float:
+    """Frame-averaged log-spectral distance (dB) between reference and reconstruction."""
+    _, _, R = stft(ref, fs=sr, nperseg=1024, noverlap=768)
+    _, _, S = stft(rec, fs=sr, nperseg=1024, noverlap=768)
+    diff = 20 * np.log10(np.abs(R) + 1e-8) - 20 * np.log10(np.abs(S) + 1e-8)
+    return float(np.mean(np.sqrt(np.mean(diff**2, axis=0))))
+
+
+def evaluate(ref: np.ndarray, rec: np.ndarray, sr: int, rtf: float, size_mb: float, level: str) -> dict:
+    n = min(len(ref), len(rec))
+    ref, rec = ref[:n], rec[:n]
+    return {
+        "Level": level,
+        "SI-SDR (dB)": round(si_sdr(ref, rec), 2),
+        "STOI": round(float(stoi(ref, rec, sr, extended=False)), 3),
+        "LSD (dB)": round(log_spectral_distance(ref, rec, sr), 2),
+        "Model size (MB)": round(size_mb, 1),
+        "Bitrate (kbps)": round(BITRATE_KBPS, 4),
+        "RTF": round(rtf, 3),
+    }
+
+
+def spectrogram_figure(ref: np.ndarray, rec: np.ndarray, sr: int, title: str):
+    fig, axes = plt.subplots(1, 2, figsize=(9.0, 2.6), sharey=True)
+    for ax, sig, label in zip(axes, (ref, rec), ("Input", "Reconstruction")):
+        ax.specgram(sig, NFFT=1024, noverlap=768, Fs=sr, cmap="magma")
+        ax.set_title(label, fontsize=9)
+        ax.set_xlabel("Time (s)")
+    axes[0].set_ylabel("Frequency (Hz)")
+    fig.suptitle(title, fontsize=10)
+    fig.tight_layout()
+    return fig
+
+
+def metrics_figure(rows: list[dict]):
+    levels = [r["Level"] for r in rows]
+    panels = (("SI-SDR (dB)", "#4c72b0"), ("STOI", "#55a868"), ("Model size (MB)", "#c44e52"))
+    fig, axes = plt.subplots(1, len(panels), figsize=(11.5, 2.9))
+    for ax, (key, color) in zip(axes, panels):
+        ax.bar(levels, [r[key] for r in rows], color=color)
+        ax.set_title(key, fontsize=9)
+        ax.tick_params(axis="x", rotation=45)
+        ax.grid(axis="y", alpha=0.3)
+    fig.tight_layout()
+    return fig
+
+
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    w = (torch.randn(64, 32) * 0.05).requires_grad_(False)
+    for level in QUANT_LEVELS:
+        q = quantize_state({"w": w}, level)["w"]
+        assert torch.isfinite(q).all(), f"{level} not finite"
+        assert q.abs().max() <= w.abs().max() * 1.01 + 1e-9, f"{level} broke the magnitude bound"
+    x = np.random.randn(16000).astype(np.float32)
+    assert si_sdr(x, x) > 60, "SI-SDR of identical signals should be large"
+    assert si_sdr(x, 3 * x) > 60, "SI-SDR must be scale invariant"
+    assert abs(evaluate(x, x, 16000, 1.0, 1.0, "fp32")["STOI"] - 1.0) < 1e-3
+    assert evaluate(x, x, 16000, 1.0, 1.0, "fp32")["Bitrate (kbps)"] == round(BITRATE_KBPS, 4)
+    print("core self-check ok")

requirements.txt

@@ -0,0 +1,5 @@
+matplotlib
+pystoi
+safetensors
+scipy
+soundfile