Mirror from aufklarer/Stable-Audio-3-DiT-Medium-MLX-8bit

README.md

@@ -0,0 +1,99 @@
+---
+license: other
+license_name: stability-ai-community-license
+license_link: https://stability.ai/license
+language:
+  - en
+tags:
+  - mlx
+  - audio
+  - text-to-audio
+  - music
+  - sfx
+  - latent-diffusion
+  - inpainting
+  - dit
+  - stable-audio-3
+  - quantized
+  - int8
+base_model: stabilityai/stable-audio-3-optimized
+library_name: mlx
+pipeline_tag: text-to-audio
+---
+
+# Stable-Audio-3-DiT-Medium-MLX-8bit
+
+- [speech-swift](https://github.com/soniqo/speech-swift) — Apple SDK
+- [soniqo.audio](https://soniqo.audio) — website
+- [blog](https://soniqo.audio/blog) — blog
+
+MLX port of [Stability AI Stable Audio 3 (optimized)](https://huggingface.co/stabilityai/stable-audio-3-optimized).
+Latent-diffusion text-to-audio with **mask-based inpainting** and continuation, with the DiT denoiser quantized to **INT8** for Apple Silicon.
+
+## What's in this bundle
+
+| Component | Format | Notes |
+|---|---|---|
+| DiT (Medium, 1.4B) | **INT8** | Diffusion Transformer denoiser, group-size 64 |
+| SAME-L encoder | FP32 | Audio → latents (codec is precision-sensitive — differential attention cancels in FP16) |
+| SAME-L decoder | FP32 | Latents → 44.1 kHz stereo waveform |
+| T5Gemma text encoder | FP16 | Prompt conditioning |
+
+Codec stays FP32 because the SAME differential attention catastrophically cancels in FP16 (per Stability's own MLX runtime). T5Gemma stays FP16 — it's small relative to the DiT and quantization gives no speed-up on the short prompt encode pass.
+
+## Files
+
+| File | Size | Format |
+|---|---|---|
+| `dit_medium/model.safetensors` | 1 GB | int8 |
+| `same_l_encoder/model.safetensors` | 2 GB | fp32 |
+| `same_l_decoder/model.safetensors` | 2 GB | fp32 |
+| `t5gemma/model.safetensors` | 541 MB | fp16 |
+
+## Capabilities
+
+- **Text-to-audio** generation (music + SFX depending on DiT specialisation)
+- **Inpainting / region editing** via masked latent diffusion
+- **Audio continuation** from a short prompt clip
+- **Variable-length** generation up to several minutes
+
+The `DiT-Small-Music-*` variant is music-specialised; `DiT-Small-SFX-*` is sound-effects specialised; `DiT-Medium-*` is the higher-quality general model.
+
+## Usage
+
+This bundle is the quantized weights only — inference uses Stability AI's official pure-MLX runtime at [`stable-audio-3/optimized/mlx`](https://github.com/Stability-AI/stable-audio-3/tree/main/optimized/mlx). At load time, each `(base.weight, base.scales, base.biases)` triplet is dequantized via `mlx.core.dequantize` back to FP16; codec and T5Gemma load as-is.
+
+```python
+from huggingface_hub import snapshot_download
+import mlx.core as mx
+
+bundle = snapshot_download("aufklarer/Stable-Audio-3-DiT-Medium-MLX-8bit")
+
+def load_component(comp_dir):
+    w = dict(mx.load(f"{comp_dir}/model.safetensors"))
+    bases = {k[:-7] for k in w if k.endswith(".scales")
+              if f"{k[:-7]}.weight" in w and f"{k[:-7]}.biases" in w}
+    out = {}
+    for k, v in w.items():
+        if k.endswith((".scales", ".biases")) and k.rsplit(".", 1)[0] in bases:
+            continue
+        if k.endswith(".weight") and k[:-7] in bases:
+            base = k[:-7]
+            out[k] = mx.dequantize(w[f"{base}.weight"], w[f"{base}.scales"],
+                                   w[f"{base}.biases"], group_size=64, bits=8)
+        else:
+            out[k] = v
+    return out
+```
+
+Plug the rehydrated dict into the matching model class from `stable-audio-3/optimized/mlx/models/defs/`.
+
+## Source
+
+- Upstream weights: [stabilityai/stable-audio-3-optimized](https://huggingface.co/stabilityai/stable-audio-3-optimized) (Stability AI Community License)
+- Text encoder lineage: [google/t5gemma-b-b-ul2](https://huggingface.co/google/t5gemma-b-b-ul2) (Gemma Terms of Use)
+- Paper: [Stable Audio 3 — Stability AI](https://arxiv.org/abs/2605.17991)
+
+## License
+
+**Stability AI Community License** — free for non-commercial research and for commercial use up to the revenue threshold defined by Stability AI; see the [license text](https://stability.ai/license). T5Gemma component additionally inherits the **Gemma Terms of Use**.

config.json

@@ -0,0 +1,3 @@
+{
+  "model_type": "stable_audio_dit"
+}

dit_medium/manifest.json

@@ -0,0 +1,21 @@
+{
+  "source_npz": "dit_medium_f16.npz",
+  "format": "int8",
+  "group_size": 64,
+  "stats": {
+    "quantized": 202,
+    "float": 323,
+    "skipped_shape": 0,
+    "by_rank": {
+      "1": 295,
+      "2": 228,
+      "3": 2
+    },
+    "size_in": 2907132960,
+    "size_out": 1555714112,
+    "bits": 8,
+    "group_size": 64
+  },
+  "size_bytes": 1555822118,
+  "sha256": "0d145d800c28e7ac68208375b9ae5ab923ade7ffc00588d30566e2fbc80b94fc"
+}

dit_medium/model.safetensors

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

manifest.json

@@ -0,0 +1,25 @@
+{
+  "bundle": "Stable-Audio-3-DiT-Medium-MLX-8bit",
+  "components": {
+    "dit_medium": {
+      "format": "int8",
+      "size_bytes": 1555822118,
+      "sha256": "0d145d800c28e7ac68208375b9ae5ab923ade7ffc00588d30566e2fbc80b94fc"
+    },
+    "same_l_encoder": {
+      "format": "fp32",
+      "size_bytes": 1704272119,
+      "sha256": "9f1ba713acf75034216cfbdd5f765edca81fcaaf7719231922cbd14cb7852680"
+    },
+    "same_l_decoder": {
+      "format": "fp32",
+      "size_bytes": 1704270957,
+      "sha256": "685bf46f7d3f7b2d10ff8d48e5757fb8da80dbc3912e4ddc8f415ca559be7679"
+    },
+    "t5gemma": {
+      "format": "fp16",
+      "size_bytes": 567416487,
+      "sha256": "3c6ba187587a2a124ba4224fdc3dea85e6dfae33745fc5b735326f80e2c31dff"
+    }
+  }
+}

same_l_decoder/manifest.json

@@ -0,0 +1,21 @@
+{
+  "source_npz": "same_l_decoder_f32.npz",
+  "format": "fp32",
+  "group_size": null,
+  "stats": {
+    "quantized": 0,
+    "float": 222,
+    "skipped_shape": 0,
+    "by_rank": {
+      "1": 171,
+      "2": 49,
+      "3": 2
+    },
+    "size_in": 1704249540,
+    "size_out": 1704249540,
+    "bits": -1,
+    "group_size": 64
+  },
+  "size_bytes": 1704270957,
+  "sha256": "685bf46f7d3f7b2d10ff8d48e5757fb8da80dbc3912e4ddc8f415ca559be7679"
+}

same_l_decoder/model.safetensors

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

same_l_encoder/manifest.json

@@ -0,0 +1,21 @@
+{
+  "source_npz": "same_l_encoder_f32.npz",
+  "format": "fp32",
+  "group_size": null,
+  "stats": {
+    "quantized": 0,
+    "float": 224,
+    "skipped_shape": 0,
+    "by_rank": {
+      "3": 3,
+      "1": 171,
+      "2": 50
+    },
+    "size_in": 1704250564,
+    "size_out": 1704250564,
+    "bits": -1,
+    "group_size": 64
+  },
+  "size_bytes": 1704272119,
+  "sha256": "9f1ba713acf75034216cfbdd5f765edca81fcaaf7719231922cbd14cb7852680"
+}

same_l_encoder/model.safetensors

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

t5gemma/manifest.json

@@ -0,0 +1,20 @@
+{
+  "source_npz": "t5gemma_f16.npz",
+  "format": "fp16",
+  "group_size": null,
+  "stats": {
+    "quantized": 0,
+    "float": 137,
+    "skipped_shape": 0,
+    "by_rank": {
+      "1": 52,
+      "2": 85
+    },
+    "size_in": 567402094,
+    "size_out": 567402094,
+    "bits": -1,
+    "group_size": 64
+  },
+  "size_bytes": 567416487,
+  "sha256": "3c6ba187587a2a124ba4224fdc3dea85e6dfae33745fc5b735326f80e2c31dff"
+}

t5gemma/model.safetensors

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