Upload folder using huggingface_hub

README.md

@@ -26,20 +26,22 @@ The token stream is untouched, so the codec still runs at **0.1625 kbps**
 
 | 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 an `lscodec-25hz` baseline row and a summary chart |
+| **Compare** | One reconstructed clip + input/reconstruction spectrograms, side by side for all 8 FocalCodec levels |
+| **Playground** | Encode → decode a single clip at one chosen codec + weight format, with its metrics row |
+| **Benchmark** | SI-SDR, STOI, LSD, stored model size, bitrate and real-time factor for all 16 levels — 8 FocalCodec + 8 LSCodec — 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`
+the untouched codec forward pass. `fp4` is IEEE-style E2M1, `fp8` is E4M3, `fp16`
 and `fp32` are exact casts. Nothing is fine-tuned or calibrated on data.
 
-The Benchmark tab also runs [LSCodec-25 Hz](https://github.com/X-LANCE/LSCodec-Inference)
-(self-prompted with WavLM-Large) as a second codec for reference: it is a different model
-at a different bitrate (0.25 kbps vs 0.1625 kbps), so read it as a floor, not a control.
+The same sweep runs on a second codec, [LSCodec-25 Hz](https://github.com/X-LANCE/LSCodec-Inference),
+whose rows are prefixed `lscodec-`: its encoder, vocoder and codebook are quantized
+identically, while the WavLM-Large features it self-prompts with are not. It is a
+different model at a different bitrate (0.25 kbps vs 0.1625 kbps), so read it as a
+floor, not a control.
 
 ⚠️ **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

app.py

@@ -16,9 +16,7 @@ from core import QUANT_LEVELS
 MODEL_REPO = "lucadellalib/focalcodec_12_5hz"
 HUB_REPO = "lucadellalib/focalcodec"
 LS_REPO = "cantabile-kwok/lscodec_25hz"
-LS_LEVEL = "lscodec-25hz"
-LS_SIZE_MB = 196.3  # encoder 70.4 + vocoder 125.6 + codebook 0.3 MB
-LS_BITRATE_KBPS = 25 * 10 / 1000  # 25 tokens/s x 10-bit codebook
+LS_PREFIX = "lscodec-"
 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"]
@@ -27,6 +25,7 @@ EXAMPLES = [
     ["assets/p226_006.wav"],
     ["assets/5_1812_20170628135856.wav"],
 ]
+LS_BITRATE_KBPS = 25 * 10 / 1000  # 25 tokens/s x 10-bit codebook
 
 codec = torch.hub.load(HUB_REPO, "focalcodec", config=MODEL_REPO, trust_repo=True)
 codec.eval().requires_grad_(False)
@@ -84,6 +83,9 @@ def _boot_lscodec() -> dict | None:
         "dec": dec,
         "books": books,
         "wavlm": wavlm,
+        "enc_state": {k: v.detach().cpu().clone() for k, v in enc.state_dict().items()},
+        "dec_state": {k: v.detach().cpu().clone() for k, v in dec.state_dict().items()},
+        "books_state": {k: v.detach().cpu().clone() for k, v in books.state_dict().items()},
         "repeat": bool(vocoder_config.get("repeat_input_tokens")),
         "out_sr": int(vocoder_config["sampling_rate"]),
     }
@@ -95,15 +97,20 @@ except Exception as exc:  # noqa: BLE001  a third-party dep breaking must not ta
     LS = None
     print(f"[lscodec] baseline disabled: {exc!r}")
 
-BENCH_LEVELS = QUANT_LEVELS + ((LS_LEVEL,) if LS is not None else ())
+LS_LEVELS = tuple(LS_PREFIX + level for level in QUANT_LEVELS) if LS is not None else ()
+BENCH_LEVELS = QUANT_LEVELS + LS_LEVELS
 
 
-def _lscodec_decode(ref: np.ndarray, sr: int) -> np.ndarray:
-    """Self-prompted LSCodec-25 Hz reconstruction, resampled back to `sr`."""
-    audio = torch.from_numpy(ref).float().view(1, -1).cuda()
+def _lscodec_prompt(audio: torch.Tensor) -> torch.Tensor:
+    """WavLM-Large layers 0..5 features: a function of the input only, so compute it once per clip."""
     with torch.inference_mode():
         prompt_input = torch.nn.functional.layer_norm(audio, audio.shape)  # fairseq WavLM cfg.normalize
-        prompt = LS["wavlm"](prompt_input, output_hidden_states=True).hidden_states[6]
+        return LS["wavlm"](prompt_input, output_hidden_states=True).hidden_states[6]
+
+
+def _lscodec_decode(prompt: torch.Tensor, audio: torch.Tensor, sr: int) -> np.ndarray:
+    """Reconstruct `audio` with LSCodec at the currently loaded weight level, resampled back to `sr`."""
+    with torch.inference_mode():
         _, _, index = LS["enc"].encode(audio.unsqueeze(0))
         if LS["repeat"]:
             index = index.repeat_interleave(2, dim=0)
@@ -116,13 +123,25 @@ def _lscodec_decode(ref: np.ndarray, sr: int) -> np.ndarray:
     return rec
 
 
+def _load_lscodec(level: str) -> None:
+    """Swap the LSCodec encoder, vocoder and codebook onto `level` (any QUANT_LEVELS entry)."""
+    for key in ("enc", "dec", "books"):
+        state = LS[key + "_state"]
+        LS[key].load_state_dict(state if level == "fp32" else core.quantize_state(state, level))
+
+
+def _lscodec_size_mb(level: str) -> float:
+    """Stored size of the codec at `level`; the WavLM prompt encoder is excluded (see the tab note)."""
+    return sum(core.model_size_mb(LS[key + "_state"], level) for key in ("enc", "dec", "books"))
+
+
 def _estimate_duration(audio_path: str, levels: tuple) -> float:
     try:
         seconds = min(sf.info(audio_path).duration, MAX_SECONDS)
     except Exception:  # noqa: BLE001  unreadable input falls back to the full cap
         seconds = float(MAX_SECONDS)
     # ZeroGPU rejects a request above the visitor's tier cap, so never ask for more than 60 s.
-    return min(60.0, 4.0 + 2.0 * len(levels) + 2.0 * seconds + (8.0 if LS_LEVEL in levels else 0.0))
+    return min(60.0, 3.0 + len(levels) + 1.5 * seconds)
 
 
 @spaces.GPU(duration=_estimate_duration)
@@ -132,6 +151,7 @@ def _encode_decode(audio_path: str, levels: tuple) -> dict:
     seconds = len(ref) / sr
     rows, reconstructed = [], {}
     focal_levels = tuple(level for level in levels if level in QUANT_LEVELS)
+    ls_levels = tuple(level for level in levels if level.startswith(LS_PREFIX))
     if focal_levels:
         sig = torch.from_numpy(ref).unsqueeze(0).cuda()
         with torch.inference_mode():  # untimed warmup so the first row's RTF isn't inflated
@@ -147,15 +167,22 @@ def _encode_decode(audio_path: str, levels: tuple) -> dict:
             rows.append(core.evaluate(ref, rec, sr, rtf, SIZES[level], level))
             reconstructed[level] = (sr, rec)
         codec.load_state_dict(BASE_STATE)
-    if LS_LEVEL in levels:
-        with torch.inference_mode():  # untimed warmup, as above
-            _lscodec_decode(ref, sr)
-        torch.cuda.synchronize()
-        start = time.perf_counter()
-        rec = _lscodec_decode(ref, sr)
-        torch.cuda.synchronize()
-        rtf = (time.perf_counter() - start) / seconds
-        rows.append(core.evaluate(ref, rec, sr, rtf, LS_SIZE_MB, LS_LEVEL, LS_BITRATE_KBPS))
+    if ls_levels:
+        audio = torch.from_numpy(ref).float().view(1, -1).cuda()
+        with torch.inference_mode():  # one prompt for every level, plus an untimed warmup pass
+            prompt = _lscodec_prompt(audio)
+            _lscodec_decode(prompt, audio, sr)
+        for level in ls_levels:
+            _load_lscodec(level[len(LS_PREFIX) :])
+            torch.cuda.synchronize()
+            start = time.perf_counter()
+            rec = _lscodec_decode(prompt, audio, sr)
+            torch.cuda.synchronize()
+            rtf = (time.perf_counter() - start) / seconds
+            base = level[len(LS_PREFIX) :]
+            rows.append(core.evaluate(ref, rec, sr, rtf, _lscodec_size_mb(base), level, LS_BITRATE_KBPS))
+            reconstructed[level] = (sr, rec)
+        _load_lscodec("fp32")  # leave the baseline weights in place for the next request
     return {"input": (sr, ref), "rows": rows, "audio": reconstructed}
 
 
@@ -164,15 +191,10 @@ def _table(rows: list[dict]) -> list[list]:
 
 
 _CACHE: dict[str, dict] = {}
-_LS_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.
-    """
+    """Encode and decode the audio at every requested level (either codec), cached per level set."""
     if not audio_path:
         raise gr.Error("Provide audio first: upload, record, or pick an example.")
     levels = tuple(levels)
@@ -207,22 +229,9 @@ def playground(audio_path: str, level: str):
     )
 
 
-def _lscodec_row(audio_path: str) -> dict:
-    """LSCodec-25 Hz score for one input, cached separately from the quantization sweep."""
-    with open(audio_path, "rb") as fh:
-        key = hashlib.md5(fh.read()).hexdigest()
-    if key not in _LS_CACHE:
-        if len(_LS_CACHE) >= 6:
-            _LS_CACHE.clear()
-        _LS_CACHE[key] = _encode_decode(audio_path, (LS_LEVEL,))["rows"][0]
-    return _LS_CACHE[key]
-
-
 def benchmark(audio_path: str):
-    """Score every quantization level on one input: metrics table plus summary chart."""
-    rows = run(audio_path, QUANT_LEVELS)["rows"]
-    if LS_LEVEL in BENCH_LEVELS:
-        rows = rows + [_lscodec_row(audio_path)]
+    """Score every level of both codecs on one input: metrics table plus summary chart."""
+    rows = run(audio_path, BENCH_LEVELS)["rows"]
     return _table(rows), core.metrics_figure(rows)
 
 
@@ -239,8 +248,9 @@ The bitrate never changes — **{core.BITRATE_KBPS:g} kbps** (12.5 tokens/s × 1
 
 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.
-The **lscodec-25hz** baseline is [LSCodec](https://github.com/X-LANCE/LSCodec-Inference) by
-Yiwei Guo et al. (MIT), self-prompted with WavLM-Large from [`microsoft/wavlm-large`](https://huggingface.co/microsoft/wavlm-large)."""
+The **`lscodec-…`** rows are [LSCodec](https://github.com/X-LANCE/LSCodec-Inference) by
+Yiwei Guo et al. (MIT), quantized identically and self-prompted with WavLM-Large from
+[`microsoft/wavlm-large`](https://huggingface.co/microsoft/wavlm-large)."""
 
 with gr.Blocks(title="FocalCodec quantization lab") as demo:
     gr.Markdown(HEADER)
@@ -278,7 +288,9 @@ with gr.Blocks(title="FocalCodec quantization lab") as demo:
                 )
                 gr.Examples(examples=EXAMPLES, inputs=pg_input, label="Examples")
                 pg_level = gr.Dropdown(
-                    choices=list(QUANT_LEVELS), value="int8", label="Weight format"
+                    choices=list(QUANT_LEVELS) + list(LS_LEVELS),
+                    value="int8",
+                    label="Codec / weight format",
                 )
                 pg_button = gr.Button("Encode + decode", variant="primary")
             with gr.Column(scale=3):
@@ -303,12 +315,13 @@ with gr.Blocks(title="FocalCodec quantization lab") as demo:
                     "**RTF** = encode+decode time ÷ audio duration · size assumes "
                     "the stored weight format."
                 )
-                if LS_LEVEL in BENCH_LEVELS:
+                if LS_LEVELS:
                     gr.Markdown(
-                        f"**{LS_LEVEL}** is a second codec, not a weight format: it runs at "
-                        f"**{LS_BITRATE_KBPS:g} kbps** (25 tokens/s × 10 bits) versus FocalCodec's "
-                        f"**{core.BITRATE_KBPS:g} kbps**, and its size counts only the codec "
-                        f"({LS_SIZE_MB:g} MB) — the WavLM-Large prompt encoder it shares is 1.2 GB."
+                        f"Rows prefixed `lscodec-` are the second codec, quantized the same way: "
+                        f"it runs at **{LS_BITRATE_KBPS:g} kbps** (25 tokens/s × 10 bits) versus "
+                        f"FocalCodec's **{core.BITRATE_KBPS:g} kbps**, and its size covers the encoder, "
+                        "vocoder and codebook only — the WavLM-Large prompt encoder it needs is 1.2 GB "
+                        "and stays excluded."
                     )
         bm_table = gr.Dataframe(headers=HEADERS, datatype=DATATYPES, interactive=False)
         bm_plot = gr.Plot(label="Quality and size per level")

core.py

@@ -88,7 +88,11 @@ def quantize_state(state: dict, level: str) -> dict:
 
 
 def model_size_mb(state: dict, level: str) -> float:
-    return sum(v.numel() for v in state.values()) * _BITS[level] / 8 / 1e6
+    """Stored size of `state` at `level`; non-float buffers keep their own width."""
+    return sum(
+        v.numel() * (v.element_size() if not v.is_floating_point() else _BITS[level] / 8)
+        for v in state.values()
+    ) / 1e6
 
 
 def load_audio(path: str, target_sr: int, max_seconds: float | None = None) -> tuple[np.ndarray, int]:
@@ -180,6 +184,8 @@ if __name__ == "__main__":
     assert torch.equal(grid8.to(torch.float8_e4m3fn).float(), grid8), "fp8 grid escapes E4M3-FN"
     grid4 = fp_levels(w, "fp4")
     assert grid4.abs().max() <= _FP_MAX_NORM["fp4"] * 1.0000001, "fp4 grid escapes E2M1"
+    assert model_size_mb({"w": w}, "fp32") == w.numel() * 4 / 1e6
+    assert model_size_mb({"i": torch.zeros(8, dtype=torch.int64)}, "int4") == 8 * 8 / 1e6  # buffers keep 8 bytes
     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"