Spaces:
Running on Zero
Running on Zero
| """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 | |
| _FP_MAX_NORM = {"fp4": 6.0, "fp8": 448.0} # E2M1 max; E4M3-FN finite max (top encoding is NaN) | |
| 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 int_scale(w: torch.Tensor, bits: int) -> torch.Tensor: | |
| return (_ref_scale(w) / (1 << (bits - 1))).clamp_min(1e-12) | |
| def int_levels(w: torch.Tensor, bits: int) -> torch.Tensor: | |
| """Signed integer levels before rescaling (what an int-format export stores).""" | |
| qmax = 1 << (bits - 1) | |
| return (w / int_scale(w, bits)).round().clamp(-qmax, qmax - 1) | |
| def fp_scale(w: torch.Tensor, level: str) -> torch.Tensor: | |
| return (_ref_scale(w) / _FP_MAX_NORM[level]).clamp_min(1e-12) | |
| def fp_levels(w: torch.Tensor, level: str) -> torch.Tensor: | |
| """Values on the format's own grid (what a float-format export stores, unscaled).""" | |
| e_bits, m_bits = _FP_FORMATS[level] | |
| bias = (1 << (e_bits - 1)) - 1 | |
| max_exp = (1 << e_bits) - 1 - bias | |
| x = w / fp_scale(w, level) | |
| 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) | |
| def _quant_int(w: torch.Tensor, bits: int) -> torch.Tensor: | |
| if bits >= 32: | |
| return w | |
| return int_levels(w, bits) * int_scale(w, bits) | |
| def _quant_fp(w: torch.Tensor, level: str) -> torch.Tensor: | |
| return fp_levels(w, level) * fp_scale(w, level) | |
| 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, level) | |
| return out | |
| def model_size_mb(state: dict, level: str) -> float: | |
| """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]: | |
| """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, | |
| bitrate_kbps: float | None = None, | |
| ) -> dict: | |
| n = min(len(ref), len(rec)) | |
| ref, rec = ref[:n], rec[:n] | |
| if bitrate_kbps is None: | |
| bitrate_kbps = BITRATE_KBPS | |
| 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" | |
| # the packed levels a real export stores must reproduce the fake-quant demo exactly | |
| assert torch.equal(int_levels(w, 8) * int_scale(w, 8), quantize_state({"w": w}, "int8")["w"]) | |
| assert torch.equal(fp_levels(w, "fp8") * fp_scale(w, "fp8"), quantize_state({"w": w}, "fp8")["w"]) | |
| # fp8 grid values are exactly float8_e4m3fn-representable, so no precision is lost on store | |
| grid8 = fp_levels(w, "fp8") | |
| 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" | |
| 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) | |
| assert evaluate(x, x, 16000, 1.0, 1.0, "lscodec-25hz", 0.25)["Bitrate (kbps)"] == 0.25 | |
| print("core self-check ok") | |