Upload folder using huggingface_hub

README.md

@@ -0,0 +1,48 @@
+---
+license: mit
+library_name: mlx
+tags:
+  - audio
+  - audio-super-resolution
+  - mlx
+base_model: haoheliu/audiosr_basic
+---
+
+# AudioSR basic, trimmed
+
+The released [AudioSR](https://github.com/haoheliu/versatile_audio_super_resolution)
+`basic` checkpoint, cut down to what super-resolution runs and saved as
+safetensors. It is used by IntelligenceKit's `MLXAudioSR` (Swift, MLX) and
+by `afm upscale`.
+
+`basic.safetensors`, 2.69 GB, float32, with PyTorch names and layouts. It holds:
+
+- `model.`: the diffusion UNet
+- `first_stage_model.`: the VAE and its 48 kHz HiFi-GAN vocoder
+- the scale factor and the noise schedule's buffers
+
+These are left out:
+
+- CLAP, which super-resolution doesn't use;
+- the EMA copy of the UNet, since inference uses the live weights;
+- the condition encoder's VAE, which is identical to `first_stage_model`.
+
+The MLX port casts to float16 at load time.
+
+`convert.py`, the script that made the file:
+
+```sh
+python convert.py basic basic.safetensors
+```
+
+## Parity
+
+The MLX port matches PyTorch stage by stage, from 76 to 137 dB SNR. It was
+tested on a clip low-passed at 6 kHz and scored by LSD against the original
+(lower is better):
+
+| | LSD |
+|---|---|
+| input | 3.88 |
+| PyTorch | 1.16 |
+| MLX | 1.14 |

convert.py

@@ -0,0 +1,35 @@
+"""The released AudioSR checkpoint, trimmed to what super-resolution uses, as safetensors.
+
+    python convert.py basic audiosr-basic.safetensors
+
+Kept: the diffusion UNet (`model.`), the VAE with its HiFi-GAN vocoder
+(`first_stage_model.`) and the scalar/schedule buffers. Dropped: CLAP
+(`clap.`, unused for super-resolution), the EMA copy of the UNet (`model_ema.`,
+inference uses the live weights) and the condition encoder's VAE
+(`cond_stage_models.0.vae.`), after checking it equals `first_stage_model`.
+"""
+
+import sys
+
+import torch
+from audiosr.utils import download_checkpoint
+from safetensors.torch import save_file
+
+name, out = sys.argv[1], sys.argv[2]
+state = torch.load(download_checkpoint(name), map_location="cpu", weights_only=False)["state_dict"]
+
+cond_prefix = "cond_stage_models.0.vae."
+mismatched = [
+    k for k in state
+    if k.startswith(cond_prefix) and not torch.equal(state[k], state["first_stage_model." + k[len(cond_prefix):]])
+]
+assert not mismatched, f"condition VAE differs from first_stage_model: {mismatched[:5]}"
+
+dropped = ("clap.", "model_ema.", "cond_stage_models.")
+kept = {k: v.contiguous() for k, v in state.items() if not k.startswith(dropped)}
+total = sum(v.numel() for v in kept.values())
+print(f"kept {len(kept)} of {len(state)} tensors, {total / 1e6:.0f}M values, {total * 4 / 1e9:.2f} GB at fp32")
+for prefix in ("model.", "first_stage_model.encoder", "first_stage_model.decoder", "first_stage_model.vocoder"):
+    n = sum(v.numel() for k, v in kept.items() if k.startswith(prefix))
+    print(f"  {prefix:30} {n / 1e6:7.1f}M")
+save_file(kept, out, metadata={"source": f"audiosr {name}", "dtype": "float32"})