Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ 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
+epic_orchestral.wav filter=lfs diff=lfs merge=lfs -text
+piano_strings.wav filter=lfs diff=lfs merge=lfs -text
+reference_epic_orchestral_clean.wav filter=lfs diff=lfs merge=lfs -text
+reference_piano_strings_clean.wav filter=lfs diff=lfs merge=lfs -text

LICENSE-REFINER

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, Any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,13 +1,47 @@
 ---
-title: Minimax Music3 Latent Refiner
-emoji: 🏆
-colorFrom: indigo
-colorTo: indigo
+title: MiniMax Music 3 Latent Refiner
+emoji: 🎵
+colorFrom: purple
+colorTo: gray
 sdk: gradio
 sdk_version: 6.26.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Restore damaged music in MiniMax Music 3 DAV latent space
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MiniMax Music 3 — Latent Refiner v0.10
+
+Restore damaged music while retaining the performance, timing, vocals, and arrangement.
+This Space runs the [MiniMax Music 3 latent refiner](https://huggingface.co/terminusresearch/minimax-music3-latent-refiner-v0.10) — a 137M-parameter bridge
+transformer operating directly in MiniMax Music 3's continuous DAV latent space, conditioned on
+MERT frame-aligned features (layer 7), pooled CLAP music embeddings, and an SR3-style degraded
+latent stream with in-context reference tokens.
+
+## How it works
+
+The refiner integrates a deterministic 32-step Euler **bridge** sampler from the damaged latent
+endpoint (t=1) back to the clean one (t=0), then decodes the refined latents through the DAV
+decoder. Inputs are resampled to 44.1 kHz stereo; clips longer than 30 s are processed in
+overlapping 30-second windows with a 2-second latent cross-fade (the training configuration).
+
+## Usage
+
+1. **Restore** — upload damaged/degraded music (bandwidth-limited, noisy, quantized, clipped
+   recordings) and press *Restore*. 32 bridge steps, 30 s windows by default.
+2. **Damage your own audio** — apply the training-style degradation chain (bandwidth reduction,
+   additive noise, bit-depth reduction, soft clipping) to a clean clip, then restore it in the
+   *Restore* tab and compare against the original.
+
+The bundled examples are public-domain instrumental music (FreePD/CC0, via
+[mrfakename/cc0-music-captioned](https://huggingface.co/datasets/mrfakename/cc0-music-captioned))
+with the degradation chain applied; the undamaged originals are included for side-by-side
+listening.
+
+## Notes
+
+- FP32 is the verified release precision for the refiner.
+- Inputs are trimmed to 120 s to keep ZeroGPU durations bounded.
+- The `minimax_music3_latent_refiner` inference package is vendored in `src/` (Apache-2.0, from
+  the model repository).

app.py

@@ -0,0 +1,235 @@
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import sys
+from pathlib import Path
+
+import spaces  # MUST come before any CUDA-touching import
+import torch
+
+sys.path.insert(0, str(Path(__file__).parent / "src"))
+
+import numpy as np
+import soundfile as sf
+import torchaudio.functional as AF
+import gradio as gr
+
+from minimax_music3_latent_refiner import MiniMaxMusic3RefinerPipeline
+
+MODEL_ID = "terminusresearch/minimax-music3-latent-refiner-v0.10"
+MAX_SECONDS = 120.0  # keep ZeroGPU duration honest; model is quadratic in frames
+
+print("Loading MiniMax Music 3 latent refiner pipeline (refiner + DAV + MERT + CLAP)...")
+REFINER = MiniMaxMusic3RefinerPipeline.from_pretrained(MODEL_ID, device="cuda")
+REFINER.refiner.to("cuda")
+REFINER.audio_vae.to("cuda")
+REFINER.mert.to("cuda")
+REFINER.clap.to("cuda")
+for _v in REFINER.normalization.values():
+    _v.to("cuda")
+print("Pipeline loaded.")
+
+
+def _load_waveform(path: str):
+    """Read any audio file as float32 [channels, samples] at its native rate."""
+    if not path:
+        raise gr.Error("Please upload an audio file or pick an example first.")
+    data, sr = sf.read(path, dtype="float32", always_2d=True)
+    if data.shape[1] > 2:
+        data = data[:, :2]
+    return torch.from_numpy(data.T.copy()), sr
+
+
+def _trim(waveform: torch.Tensor, sr: int, max_seconds: float):
+    limit = int(max_seconds * sr)
+    if waveform.shape[-1] > limit:
+        return waveform[..., :limit], True
+    return waveform, False
+
+
+@spaces.GPU(duration=120)
+def refine(
+    audio_path: str,
+    steps: int = 32,
+    window_seconds: float = 30.0,
+    overlap_seconds: float = 2.0,
+    direct: bool = False,
+):
+    """Restore a damaged music clip with the MiniMax Music 3 latent refiner.
+
+    Args:
+        audio_path: degraded/damaged music audio to restore.
+        steps: deterministic Euler bridge sampling steps (release default 32).
+        window_seconds: overlapping inference window in seconds (30 s matches training).
+        overlap_seconds: cross-fade overlap between windows in latent frames.
+        direct: process the whole clip as one dense sequence instead of windows.
+
+    Returns:
+        Tuple of (restored audio at 44.1 kHz stereo, status text).
+    """
+    if not audio_path:
+        raise gr.Error("Please upload an audio file or pick an example first.")
+    waveform, sr = _load_waveform(audio_path)
+    waveform, trimmed = _trim(waveform, sr, MAX_SECONDS)
+    steps = int(steps)
+    if steps < 1:
+        steps = 32
+    ws = None if direct else float(window_seconds)
+    result = REFINER(
+        waveform,
+        sr,
+        steps=steps,
+        window_seconds=ws,
+        overlap_seconds=float(overlap_seconds),
+    )
+    out = result.audio.squeeze(0).T.numpy()
+    out_path = Path("/tmp") / "refined.wav"
+    sf.write(out_path, out, result.sample_rate, subtype="PCM_16")
+    status = (
+        f"Restored {out.shape[0] / result.sample_rate:.1f}s of audio "
+        f"({steps} bridge steps, {'one dense sequence' if ws is None else f'{ws:.0f}s windows'})."
+    )
+    if trimmed:
+        status += f" Input was trimmed to the {MAX_SECONDS:.0f}s demo limit."
+    return str(out_path), status
+
+
+def degrade(audio_path: str, lowpass_hz: float, noise_db: float, bits: float):
+    """Apply a degradation chain (bandwidth cut + noise + bit depth + soft clip)
+    to a clean clip, so you can hear what the refiner restores.
+
+    Args:
+        audio_path: clean music audio to damage.
+        lowpass_hz: lowpass cutoff in Hz (bandwidth reduction).
+        noise_db: additive white noise level in dBFS.
+        bits: bit-depth reduction (bits per sample).
+
+    Returns:
+        Tuple of (damaged audio, status text).
+    """
+    if not audio_path:
+        raise gr.Error("Upload clean audio to damage first.")
+    import scipy.signal as sps
+
+    data, sr = sf.read(audio_path, dtype="float32", always_2d=True)
+    if data.shape[1] > 2:
+        data = data[:, :2]
+    limit = int(MAX_SECONDS * sr)
+    data = data[:limit]
+    y = data.copy()
+    nyq = sr / 2
+    if lowpass_hz < nyq:
+        y = sps.sosfilt(sps.butter(2, lowpass_hz / nyq, btype="low", output="sos"), y, axis=0)
+    y = sps.sosfilt(sps.butter(2, 60.0 / nyq, btype="high", output="sos"), y, axis=0)
+    if noise_db > -90:
+        rng = np.random.default_rng(0)
+        y = y + rng.normal(0.0, 1.0, y.shape).astype(np.float32) * (10 ** (noise_db / 20.0))
+    if bits >= 2:
+        levels = 2 ** int(bits)
+        y = np.round(y * levels) / levels
+    peak = max(float(np.abs(y).max()), 1e-6)
+    y = np.tanh(y * (0.95 / peak) * 2.0) * 0.9
+    out_path = str(Path("/tmp") / "damaged.wav")
+    sf.write(out_path, np.clip(y, -1, 1), sr, subtype="PCM_16")
+    return out_path, f"Damaged {y.shape[0] / sr:.1f}s — now press Restore to refine it."
+
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    gr.Markdown(
+        """
+        # 🎵 MiniMax Music 3 — Latent Refiner v0.10
+
+        Restore damaged music while keeping the performance, timing, and arrangement intact.
+        A 137M-parameter bridge transformer that works directly in MiniMax Music 3's continuous
+        DAV latent space, conditioned on MERT frame features and CLAP audio embeddings.
+
+        **Upload a degraded clip** (bandwidth-limited, noisy, quantized, clipped recordings), or
+        use the *Damage* tab to apply the training-style degradation chain to your own clean audio,
+        then press **Restore**.
+
+        Model: [`terminusresearch/minimax-music3-latent-refiner-v0.10`](https://huggingface.co/terminusresearch/minimax-music3-latent-refiner-v0.10)
+        """
+    )
+    with gr.Column(elem_id="col-container"):
+        with gr.Tab("Restore"):
+            with gr.Row():
+                with gr.Column():
+                    audio_in = gr.Audio(
+                        label="Damaged music (input)",
+                        type="filepath",
+                        sources=["upload", "microphone"],
+                    )
+                    restore_btn = gr.Button("Restore", variant="primary")
+                    status = gr.Textbox(label="Status", interactive=False, lines=2)
+                with gr.Column():
+                    audio_out = gr.Audio(label="Restored music (output)", type="filepath")
+            with gr.Accordion("Advanced settings", open=False):
+                steps = gr.Slider(1, 64, value=32, step=1, label="Bridge steps")
+                window_seconds = gr.Slider(
+                    5.0, 60.0, value=30.0, step=5.0,
+                    label="Window seconds (30 s matches training)",
+                )
+                overlap_seconds = gr.Slider(
+                    0.0, 10.0, value=2.0, step=0.5, label="Window overlap (seconds)"
+                )
+                direct = gr.Checkbox(
+                    value=False,
+                    label="One dense sequence (no windows; not the quality baseline)",
+                )
+        with gr.Tab("Damage your own audio"):
+            gr.Markdown(
+                "Apply the refiner's training-style degradation chain (bandwidth reduction, additive "
+                "noise, bit-depth reduction, soft clipping) to a clean clip, then restore it in the "
+                "**Restore** tab to compare. Public-domain (FreePD/CC0) source clips."
+            )
+            clean_in = gr.Audio(label="Clean music (input)", type="filepath", sources=["upload", "microphone"])
+            with gr.Row():
+                lowpass_hz = gr.Slider(500, 20000, value=3200, step=100, label="Lowpass cutoff (Hz)")
+                noise_db = gr.Slider(-90, -10, value=-34, step=1, label="Noise level (dBFS)")
+                bits = gr.Slider(2, 16, value=6, step=1, label="Bit depth")
+            damage_btn = gr.Button("Damage", variant="secondary")
+            damaged_out = gr.Audio(label="Damaged music (output)", type="filepath")
+            damage_status = gr.Textbox(label="Status", interactive=False, lines=1)
+
+        gr.Examples(
+            examples=[
+                ["piano_strings.wav"],
+                ["epic_orchestral.wav"],
+            ],
+            inputs=[audio_in],
+            outputs=[audio_out, status],
+            fn=refine,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Examples (public-domain music with the training-style degradation applied)",
+        )
+        gr.Markdown(
+            """
+            #### Reference clean sources
+            These are the undamaged originals of the two examples — listen side by side to judge the restoration.
+            """
+        )
+        with gr.Row():
+            gr.Audio(value="reference_piano_strings_clean.wav", label="Clean original — piano & strings", type="filepath", show_download_button=True)
+            gr.Audio(value="reference_epic_orchestral_clean.wav", label="Clean original — epic orchestral", type="filepath", show_download_button=True)
+
+    restore_btn.click(
+        refine,
+        inputs=[audio_in, steps, window_seconds, overlap_seconds, direct],
+        outputs=[audio_out, status],
+        api_name="restore",
+    )
+    damage_btn.click(
+        degrade,
+        inputs=[clean_in, lowpass_hz, noise_db, bits],
+        outputs=[damaged_out, damage_status],
+        api_name="damage",
+    )
+
+demo.launch(mcp_server=True)

epic_orchestral.wav

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

piano_strings.wav

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

reference_epic_orchestral_clean.wav

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

reference_piano_strings_clean.wav

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

requirements.txt

@@ -0,0 +1,8 @@
+diffusers>=0.38.0
+transformers>=4.52.0
+accelerate>=1.8.0
+torchaudio
+soundfile>=0.12.1
+safetensors>=0.5.0
+numpy>=1.26
+scipy

src/minimax_music3_latent_refiner/__init__.py

@@ -0,0 +1,13 @@
+from .dav import MiniMaxMusic3DAV
+from .diffusers_patch import attach_latent_refiner
+from .model import MiniMaxMusic3LatentRefiner, bridge_sample
+from .pipeline import MiniMaxMusic3RefinerPipeline, RefinerOutput
+
+__all__ = [
+    "MiniMaxMusic3DAV",
+    "MiniMaxMusic3LatentRefiner",
+    "MiniMaxMusic3RefinerPipeline",
+    "RefinerOutput",
+    "attach_latent_refiner",
+    "bridge_sample",
+]

src/minimax_music3_latent_refiner/cli.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import soundfile as sf
+import torch
+
+from .pipeline import DEFAULT_MODEL_ID, MiniMaxMusic3RefinerPipeline
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Refine audio with the MiniMax Music 3 bridge latent refiner")
+    parser.add_argument("input", type=Path)
+    parser.add_argument("output", type=Path)
+    parser.add_argument("--model", default=DEFAULT_MODEL_ID)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--steps", type=int, default=32)
+    parser.add_argument("--window-seconds", type=float, default=30.0)
+    parser.add_argument("--overlap-seconds", type=float, default=2.0)
+    parser.add_argument("--direct", action="store_true", help="use one dense sequence instead of overlapping windows")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    data, sample_rate = sf.read(args.input, dtype="float32", always_2d=True)
+    waveform = torch.from_numpy(data.T.copy())
+    pipeline = MiniMaxMusic3RefinerPipeline.from_pretrained(args.model, device=args.device)
+    output = pipeline(
+        waveform,
+        sample_rate,
+        steps=args.steps,
+        window_seconds=None if args.direct else args.window_seconds,
+        overlap_seconds=args.overlap_seconds,
+    )
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    sf.write(args.output, output.audio.squeeze(0).T.numpy(), output.sample_rate)
+
+
+if __name__ == "__main__":
+    main()

src/minimax_music3_latent_refiner/dav.py

@@ -0,0 +1,179 @@
+# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved.
+# Licensed under the Apache License, Version 2.0
+
+from __future__ import annotations
+
+import math
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from torch import nn
+from torch.nn.utils import weight_norm
+
+
+class MiniMaxMusic3Snake1d(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        shape = hidden_states.shape
+        hidden_states = hidden_states.reshape(shape[0], shape[1], -1)
+        hidden_states = hidden_states + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * hidden_states).pow(2)
+        return hidden_states.reshape(shape)
+
+
+class DAVResidualUnit(nn.Module):
+    def __init__(self, dim: int, dilation: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            MiniMaxMusic3Snake1d(dim),
+            weight_norm(nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=3 * dilation)),
+            MiniMaxMusic3Snake1d(dim),
+            weight_norm(nn.Conv1d(dim, dim, kernel_size=1)),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = self.block(hidden_states)
+        if residual.shape[-1] != hidden_states.shape[-1]:
+            padding = (hidden_states.shape[-1] - residual.shape[-1]) // 2
+            hidden_states = hidden_states[..., padding : hidden_states.shape[-1] - padding]
+        return hidden_states + residual
+
+
+class DAVEncoderBlock(nn.Module):
+    def __init__(self, dim: int, stride: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            DAVResidualUnit(dim // 2, dilation=1),
+            DAVResidualUnit(dim // 2, dilation=3),
+            DAVResidualUnit(dim // 2, dilation=9),
+            MiniMaxMusic3Snake1d(dim // 2),
+            weight_norm(nn.Conv1d(dim // 2, dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2))),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.block(hidden_states)
+
+
+class DAVEncoder(nn.Module):
+    def __init__(self, encoder_dim: int, encoder_rates: tuple[int, ...], latent_dim: int):
+        super().__init__()
+        layers: list[nn.Module] = [weight_norm(nn.Conv1d(1, encoder_dim, kernel_size=7, padding=3))]
+        for stride in encoder_rates:
+            encoder_dim *= 2
+            layers.append(DAVEncoderBlock(encoder_dim, stride=stride))
+        layers.extend(
+            (
+                MiniMaxMusic3Snake1d(encoder_dim),
+                weight_norm(nn.Conv1d(encoder_dim, latent_dim, kernel_size=3, padding=1)),
+            )
+        )
+        self.block = nn.Sequential(*layers)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.block(hidden_states)
+
+
+class DAVDecoderBlock(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, stride: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            MiniMaxMusic3Snake1d(input_dim),
+            weight_norm(
+                nn.ConvTranspose1d(
+                    input_dim,
+                    output_dim,
+                    kernel_size=2 * stride,
+                    stride=stride,
+                    padding=math.ceil(stride / 2),
+                )
+            ),
+            DAVResidualUnit(output_dim, dilation=1),
+            DAVResidualUnit(output_dim, dilation=3),
+            DAVResidualUnit(output_dim, dilation=9),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.block(hidden_states)
+
+
+class DAVDecoder(nn.Module):
+    def __init__(self, input_dim: int, hidden_dim: int, upsampling_ratios: tuple[int, ...]):
+        super().__init__()
+        layers: list[nn.Module] = [weight_norm(nn.Conv1d(input_dim, hidden_dim, kernel_size=7, padding=3))]
+        output_dim = hidden_dim
+        for index, stride in enumerate(upsampling_ratios):
+            input_channels = hidden_dim // (2**index)
+            output_dim = hidden_dim // (2 ** (index + 1))
+            layers.append(DAVDecoderBlock(input_channels, output_dim, stride=stride))
+        layers.extend(
+            (
+                MiniMaxMusic3Snake1d(output_dim),
+                weight_norm(nn.Conv1d(output_dim, 1, kernel_size=7, padding=3)),
+                nn.Tanh(),
+            )
+        )
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.model(hidden_states)
+
+
+class MiniMaxMusic3DAV(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        latent_channels: int = 128,
+        channel_latent_channels: int = 64,
+        encoder_dim: int = 64,
+        encoder_rates: tuple[int, ...] = (2, 4, 8, 8),
+        encoder_latent_dim: int = 1024,
+        decoder_input_dim: int = 1024,
+        decoder_hidden_dim: int = 1536,
+        upsampling_ratios: tuple[int, ...] = (8, 8, 4, 2),
+        sampling_rate: int = 44_100,
+    ):
+        super().__init__()
+        if channel_latent_channels * 2 != latent_channels:
+            raise ValueError("latent_channels must be twice channel_latent_channels")
+        self.hop_length = math.prod(encoder_rates)
+        self.encoder = DAVEncoder(encoder_dim, encoder_rates, encoder_latent_dim)
+        self.mean_proj = nn.Conv1d(encoder_latent_dim, channel_latent_channels, kernel_size=1)
+        self.logs_proj = nn.Conv1d(encoder_latent_dim, channel_latent_channels, kernel_size=1)
+        self.dec_in_proj = nn.Conv1d(channel_latent_channels, decoder_input_dim, kernel_size=1)
+        self.decoder = DAVDecoder(decoder_input_dim, decoder_hidden_dim, upsampling_ratios)
+
+    def _prepare_waveform(self, waveform: torch.Tensor) -> torch.Tensor:
+        if waveform.ndim == 1:
+            waveform = waveform[None, None]
+        elif waveform.ndim == 2:
+            waveform = waveform[None]
+        if waveform.ndim != 3:
+            raise ValueError("waveform must have shape [batch, channels, samples]")
+        if waveform.shape[1] == 1:
+            waveform = waveform.repeat(1, 2, 1)
+        elif waveform.shape[1] != 2:
+            raise ValueError("waveform must be mono or stereo")
+        remainder = waveform.shape[-1] % self.hop_length
+        if remainder:
+            waveform = torch.nn.functional.pad(waveform, (0, self.hop_length - remainder))
+        return waveform
+
+    def encode(self, waveform: torch.Tensor) -> torch.Tensor:
+        waveform = self._prepare_waveform(waveform)
+        batch_size = waveform.shape[0]
+        hidden_states = self.encoder(waveform.reshape(batch_size * 2, 1, -1))
+        latents = self.mean_proj(hidden_states)
+        return latents.reshape(batch_size, self.config.latent_channels, -1)
+
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        if latents.ndim != 3 or latents.shape[1] != self.config.latent_channels:
+            raise ValueError(f"latents must have shape [batch, {self.config.latent_channels}, frames]")
+        batch_size, _, length = latents.shape
+        hidden_states = latents.reshape(batch_size * 2, self.config.channel_latent_channels, length)
+        waveform = self.decoder(self.dec_in_proj(hidden_states))
+        return waveform.reshape(batch_size, 2, -1)
+
+    forward = decode

src/minimax_music3_latent_refiner/diffusers_patch.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import types
+
+import torch
+
+from .pipeline import MiniMaxMusic3RefinerPipeline, RefinerOutput
+
+
+def attach_latent_refiner(diffusers_pipeline, refiner: MiniMaxMusic3RefinerPipeline):
+    """Attach `refine_audio` to a loaded MiniMax Music 3 Diffusers pipeline instance."""
+    if not hasattr(diffusers_pipeline, "vocoder") or diffusers_pipeline.vocoder is None:
+        raise TypeError("the Diffusers pipeline must expose its MiniMax Music 3 vocoder")
+    if hasattr(diffusers_pipeline, "refine_audio"):
+        raise ValueError("the Diffusers pipeline already has a refine_audio attribute")
+
+    def refine_audio(
+        self,
+        waveform: torch.Tensor,
+        sample_rate: int,
+        *,
+        steps: int = 32,
+        window_seconds: float | None = 30.0,
+        overlap_seconds: float = 2.0,
+    ) -> RefinerOutput:
+        return refiner(
+            waveform,
+            sample_rate,
+            steps=steps,
+            window_seconds=window_seconds,
+            overlap_seconds=overlap_seconds,
+            decoder=self.vocoder,
+        )
+
+    diffusers_pipeline.refine_audio = types.MethodType(refine_audio, diffusers_pipeline)
+    return diffusers_pipeline

src/minimax_music3_latent_refiner/model.py

@@ -0,0 +1,186 @@
+# Copyright 2026 SimpleTuner contributors
+# Licensed under the Apache License, Version 2.0
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from torch import nn
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, head_dim: int, base: float = 10_000.0):
+        super().__init__()
+        inv_freq = base ** (-torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, length: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+        positions = torch.arange(length, device=device, dtype=torch.float32)
+        angles = positions[:, None] * self.inv_freq.to(device)[None, :]
+        return angles.cos(), angles.sin()
+
+
+def apply_rope(states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    first, second = states.chunk(2, dim=-1)
+    cos = cos[None, None]
+    sin = sin[None, None]
+    return torch.cat((first * cos - second * sin, first * sin + second * cos), dim=-1)
+
+
+class RefinerBlock(nn.Module):
+    def __init__(self, d_model: int, heads: int, cond_dim: int):
+        super().__init__()
+        if d_model % heads:
+            raise ValueError("d_model must be divisible by heads")
+        self.heads = heads
+        self.head_dim = d_model // heads
+        self.attn_norm = nn.LayerNorm(d_model, elementwise_affine=False)
+        self.mlp_norm = nn.LayerNorm(d_model, elementwise_affine=False)
+        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.attn_out = nn.Linear(d_model, d_model, bias=False)
+        self.mlp = nn.Sequential(
+            nn.Linear(d_model, 4 * d_model),
+            nn.SiLU(),
+            nn.Linear(4 * d_model, d_model),
+        )
+        self.adaln = nn.Linear(d_model, 6 * d_model)
+        self.layer_cond_proj = nn.Linear(cond_dim, d_model, bias=False)
+
+    def forward(
+        self,
+        states: torch.Tensor,
+        conditioning: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        layer_conditioning: torch.Tensor,
+    ) -> torch.Tensor:
+        states = states + self.layer_cond_proj(layer_conditioning)
+        shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = self.adaln(conditioning).chunk(6, dim=-1)
+        batch, length, _ = states.shape
+        normed = self.attn_norm(states) * (1 + scale_a[:, None]) + shift_a[:, None]
+        query, key, value = self.qkv(normed).view(batch, length, 3, self.heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        query = apply_rope(query, cos, sin)
+        key = apply_rope(key, cos, sin)
+        attended = F.scaled_dot_product_attention(query, key, value)
+        attended = attended.transpose(1, 2).reshape(batch, length, -1)
+        states = states + gate_a[:, None] * self.attn_out(attended)
+        normed = self.mlp_norm(states) * (1 + scale_m[:, None]) + shift_m[:, None]
+        return states + gate_m[:, None] * self.mlp(normed)
+
+
+class MiniMaxMusic3LatentRefiner(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        latent_dim: int = 128,
+        cond_dim: int = 768,
+        d_model: int = 768,
+        depth: int = 12,
+        heads: int = 12,
+        mert_layer_count: int = 13,
+        mert_input_layer: int = 7,
+        style_dim: int = 512,
+        task_count: int = 3,
+    ):
+        super().__init__()
+        if not 0 <= mert_input_layer < mert_layer_count:
+            raise ValueError("mert_input_layer must select an available MERT layer")
+        self.proj_in = nn.Linear(latent_dim + cond_dim, d_model)
+        self.time_dim = d_model
+        self.time_embed = nn.Sequential(nn.Linear(256, d_model), nn.SiLU(), nn.Linear(d_model, d_model))
+        self.style_proj = nn.Sequential(nn.Linear(style_dim, d_model), nn.SiLU(), nn.Linear(d_model, d_model))
+        self.style_null = nn.Parameter(torch.zeros(style_dim))
+        self.task_embed = nn.Embedding(task_count + 1, d_model)
+        self.context_embed = nn.Parameter(torch.zeros(d_model))
+        self.degraded_in_proj = nn.Linear(latent_dim, d_model, bias=False)
+        self.degraded_null = nn.Parameter(torch.zeros(latent_dim))
+        self.mert_null = nn.Parameter(torch.zeros(mert_layer_count, cond_dim))
+        self.rope = RotaryEmbedding(d_model // heads)
+        self.blocks = nn.ModuleList(RefinerBlock(d_model, heads, cond_dim) for _ in range(depth))
+        self.layer_map = [min(1 + index, mert_layer_count - 2) for index in range(depth)]
+        self.out_norm = nn.LayerNorm(d_model, elementwise_affine=False)
+        self.proj_out = nn.Linear(d_model, latent_dim)
+
+    @staticmethod
+    def timestep_features(timestep: torch.Tensor) -> torch.Tensor:
+        half = 128
+        frequencies = torch.exp(
+            -math.log(10_000.0) * torch.arange(half, device=timestep.device, dtype=torch.float32) / half
+        )
+        angles = timestep[:, None].float() * frequencies[None, :] * 1_000.0
+        return torch.cat((angles.sin(), angles.cos()), dim=-1)
+
+    def forward(
+        self,
+        noisy_latents: torch.Tensor,
+        conditioning: torch.Tensor,
+        timestep: torch.Tensor,
+        layer_conditioning: torch.Tensor,
+        style: torch.Tensor,
+        degraded_latents: torch.Tensor,
+        context_latents: torch.Tensor,
+        task: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        target_length = noisy_latents.shape[1]
+        if conditioning.shape[1] != target_length or degraded_latents.shape[1] != target_length:
+            raise ValueError("conditioning, degraded latents, and noisy latents must share a frame count")
+        if context_latents.shape[1] != target_length:
+            raise ValueError("context latents must share the target frame count")
+
+        states = self.proj_in(torch.cat((noisy_latents, conditioning), dim=-1))
+        states = states + self.degraded_in_proj(degraded_latents)
+        context_states = self.proj_in(torch.cat((context_latents, conditioning), dim=-1))
+        context_states = context_states + self.context_embed[None, None]
+        states = torch.cat((context_states, states), dim=1)
+
+        time_conditioning = self.time_embed(self.timestep_features(timestep))
+        if task is None:
+            task = torch.full(
+                (timestep.shape[0],),
+                self.task_embed.num_embeddings - 1,
+                dtype=torch.long,
+                device=timestep.device,
+            )
+        time_conditioning = time_conditioning + self.task_embed(task)
+        time_conditioning = time_conditioning + self.style_proj(style)
+
+        cos, sin = self.rope(target_length, states.device)
+        cos = torch.cat((cos, cos), dim=0)
+        sin = torch.cat((sin, sin), dim=0)
+        for index, block in enumerate(self.blocks):
+            block_layers = layer_conditioning[:, self.layer_map[index]]
+            block_layers = torch.cat((block_layers, block_layers), dim=1)
+            states = block(states, time_conditioning, cos, sin, block_layers)
+        return self.proj_out(self.out_norm(states[:, target_length:]))
+
+
+@torch.inference_mode()
+def bridge_sample(
+    model: MiniMaxMusic3LatentRefiner,
+    conditioning: torch.Tensor,
+    layer_conditioning: torch.Tensor,
+    degraded_latents: torch.Tensor,
+    style: torch.Tensor,
+    steps: int = 32,
+) -> torch.Tensor:
+    if steps < 1:
+        raise ValueError("steps must be positive")
+    latents = degraded_latents.clone()
+    schedule = torch.linspace(1.0, 0.0, steps + 1, device=latents.device)
+    for index in range(steps):
+        timestep = schedule[index].expand(latents.shape[0])
+        velocity = model(
+            latents,
+            conditioning,
+            timestep,
+            layer_conditioning,
+            style,
+            degraded_latents,
+            degraded_latents,
+        )
+        latents = latents - (schedule[index] - schedule[index + 1]) * velocity
+    return latents

src/minimax_music3_latent_refiner/pipeline.py

@@ -0,0 +1,361 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Protocol
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from transformers import AutoModel, ClapModel, ClapProcessor, Wav2Vec2FeatureExtractor
+
+from .dav import MiniMaxMusic3DAV
+from .model import MiniMaxMusic3LatentRefiner, bridge_sample
+
+SAMPLE_RATE = 44_100
+MERT_SAMPLE_RATE = 24_000
+DAV_HOP = 512
+DEFAULT_MODEL_ID = "terminusresearch/minimax-music3-latent-refiner-v0.10"
+DEFAULT_AUDIO_VAE_ID = "SimpleTuner/MiniMax-Music-3-Encoder"
+DEFAULT_MERT_ID = "m-a-p/MERT-v1-95M"
+DEFAULT_CLAP_ID = "laion/larger_clap_music"
+
+
+class AudioVAEEncoder(Protocol):
+    def encode(self, waveform: torch.Tensor) -> torch.Tensor: ...
+
+
+@dataclass
+class RefinerOutput:
+    latents: torch.Tensor
+    audio: torch.Tensor | None
+    sample_rate: int
+    original_samples: int
+
+
+def _module_device(module: torch.nn.Module) -> torch.device:
+    return next(module.parameters()).device
+
+
+def _module_dtype(module: torch.nn.Module) -> torch.dtype:
+    return next(module.parameters()).dtype
+
+
+def prepare_waveform(waveform: torch.Tensor, sample_rate: int) -> torch.Tensor:
+    if waveform.ndim == 1:
+        waveform = waveform[None, None]
+    elif waveform.ndim == 2:
+        waveform = waveform[None]
+    if waveform.ndim != 3:
+        raise ValueError("waveform must have shape [samples], [channels, samples], or [batch, channels, samples]")
+    if waveform.shape[1] == 1:
+        waveform = waveform.repeat(1, 2, 1)
+    elif waveform.shape[1] != 2:
+        raise ValueError("waveform must be mono or stereo")
+    waveform = waveform.float().cpu()
+    if sample_rate != SAMPLE_RATE:
+        waveform = torchaudio.functional.resample(waveform, sample_rate, SAMPLE_RATE)
+    return waveform
+
+
+def window_starts(total_samples: int, window_samples: int, overlap_samples: int) -> list[int]:
+    if window_samples <= 0:
+        raise ValueError("window_samples must be positive")
+    if not 0 <= overlap_samples < window_samples:
+        raise ValueError("overlap_samples must be non-negative and smaller than window_samples")
+    if total_samples <= window_samples:
+        return [0]
+    stride = window_samples - overlap_samples
+    starts = list(range(0, total_samples - window_samples + 1, stride))
+    final_start = total_samples - window_samples
+    final_start -= final_start % DAV_HOP
+    if final_start > starts[-1]:
+        starts.append(final_start)
+    return starts
+
+
+class MiniMaxMusic3RefinerPipeline:
+    def __init__(
+        self,
+        refiner: MiniMaxMusic3LatentRefiner,
+        audio_vae: MiniMaxMusic3DAV,
+        mert: torch.nn.Module,
+        mert_processor: Wav2Vec2FeatureExtractor,
+        clap: ClapModel,
+        clap_processor: ClapProcessor,
+        normalization: dict[str, torch.Tensor],
+    ):
+        required = {"latent_mean", "latent_std", "layer_mean", "layer_std"}
+        missing = required.difference(normalization)
+        if missing:
+            raise ValueError(f"normalization is missing tensors: {sorted(missing)}")
+        self.refiner = refiner.eval()
+        self.audio_vae = audio_vae.eval()
+        self.mert = mert.eval()
+        self.mert_processor = mert_processor
+        self.clap = clap.eval()
+        self.clap_processor = clap_processor
+        self.normalization = normalization
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id: str | Path = DEFAULT_MODEL_ID,
+        *,
+        audio_vae_id: str = DEFAULT_AUDIO_VAE_ID,
+        mert_id: str = DEFAULT_MERT_ID,
+        clap_id: str = DEFAULT_CLAP_ID,
+        device: str | torch.device = "cpu",
+        dtype: torch.dtype = torch.float32,
+        cache_dir: str | Path | None = None,
+        local_files_only: bool = False,
+    ) -> MiniMaxMusic3RefinerPipeline:
+        device = torch.device(device)
+        model_id_string = str(model_id)
+        refiner = MiniMaxMusic3LatentRefiner.from_pretrained(
+            model_id_string,
+            torch_dtype=dtype,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        ).to(device)
+        if Path(model_id_string).is_dir():
+            normalization_path = Path(model_id_string) / "normalization.safetensors"
+        else:
+            normalization_path = Path(
+                hf_hub_download(
+                    model_id_string,
+                    "normalization.safetensors",
+                    cache_dir=cache_dir,
+                    local_files_only=local_files_only,
+                )
+            )
+        if not normalization_path.is_file():
+            raise FileNotFoundError(f"required normalization file not found: {normalization_path}")
+        normalization = load_file(normalization_path, device=str(device))
+        audio_vae = MiniMaxMusic3DAV.from_pretrained(
+            audio_vae_id,
+            subfolder="audio_vae",
+            torch_dtype=dtype,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        ).to(device)
+        mert = AutoModel.from_pretrained(
+            mert_id,
+            trust_remote_code=True,
+            dtype=dtype,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        ).to(device)
+        mert_processor = Wav2Vec2FeatureExtractor.from_pretrained(
+            mert_id,
+            trust_remote_code=True,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        clap = ClapModel.from_pretrained(
+            clap_id,
+            dtype=dtype,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        ).to(device)
+        clap_processor = ClapProcessor.from_pretrained(
+            clap_id,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        refiner.requires_grad_(False)
+        audio_vae.requires_grad_(False)
+        mert.requires_grad_(False)
+        clap.requires_grad_(False)
+        return cls(refiner, audio_vae, mert, mert_processor, clap, clap_processor, normalization)
+
+    @property
+    def device(self) -> torch.device:
+        return _module_device(self.refiner)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return _module_dtype(self.refiner)
+
+    def _normalization_tensors(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        device = self.device
+        return (
+            self.normalization["latent_mean"].to(device)[None, None],
+            self.normalization["latent_std"].to(device)[None, None],
+            self.normalization["layer_mean"].to(device)[None, :, None],
+            self.normalization["layer_std"].to(device)[None, :, None],
+        )
+
+    @torch.inference_mode()
+    def _mert_layers(self, waveform: torch.Tensor, frame_count: int) -> torch.Tensor:
+        mono = waveform.mean(dim=1)
+        mono_24k = torchaudio.functional.resample(mono, SAMPLE_RATE, MERT_SAMPLE_RATE)
+        inputs = self.mert_processor(
+            [row.numpy() for row in mono_24k],
+            sampling_rate=MERT_SAMPLE_RATE,
+            return_tensors="pt",
+        )
+        outputs = self.mert(inputs["input_values"].to(self.device), output_hidden_states=True)
+        layers = torch.stack(outputs.hidden_states, dim=1).float()
+        batch, layer_count, mert_frames, dim = layers.shape
+        return (
+            F.interpolate(
+                layers.reshape(batch * layer_count, mert_frames, dim).transpose(1, 2),
+                size=frame_count,
+                mode="linear",
+                align_corners=True,
+            )
+            .transpose(1, 2)
+            .reshape(batch, layer_count, frame_count, dim)
+        )
+
+    @torch.inference_mode()
+    def _clap_style(self, waveform: torch.Tensor) -> torch.Tensor:
+        mono = waveform.mean(dim=1)
+        center = mono.shape[-1] // 2
+        half = min(5 * SAMPLE_RATE, center)
+        clip = mono[..., center - half : center + half]
+        clip = torchaudio.functional.resample(clip, SAMPLE_RATE, 48_000)
+        inputs = self.clap_processor(
+            audio=[row.numpy() for row in clip],
+            sampling_rate=48_000,
+            return_tensors="pt",
+        )
+        features = self.clap.get_audio_features(input_features=inputs["input_features"].to(self.device))
+        if not torch.is_tensor(features):
+            features = features.pooler_output
+        if features.shape[-1] != self.clap.config.projection_dim:
+            features = self.clap.audio_projection(features)
+        return features.to(device=self.device, dtype=self.dtype)
+
+    @torch.inference_mode()
+    def _refine_chunk(
+        self,
+        waveform: torch.Tensor,
+        steps: int,
+        audio_encoder: AudioVAEEncoder,
+    ) -> torch.Tensor:
+        encoder_device = _module_device(audio_encoder) if isinstance(audio_encoder, torch.nn.Module) else self.device
+        encoder_dtype = _module_dtype(audio_encoder) if isinstance(audio_encoder, torch.nn.Module) else self.dtype
+        degraded = audio_encoder.encode(waveform.to(device=encoder_device, dtype=encoder_dtype))
+        if degraded.ndim != 3 or degraded.shape[1] != self.refiner.config.latent_dim:
+            raise ValueError("AUDIO_VAE_ENCODER.encode must return [batch, 128, frames]")
+        degraded = degraded.transpose(1, 2).to(device=self.device, dtype=self.dtype)
+        layers = self._mert_layers(waveform, degraded.shape[1])
+        style = self._clap_style(waveform)
+        latent_mean, latent_std, layer_mean, layer_std = self._normalization_tensors()
+        degraded = (degraded - latent_mean) / latent_std
+        layers = ((layers - layer_mean) / layer_std).to(self.dtype)
+        generated = bridge_sample(
+            self.refiner,
+            layers[:, self.refiner.config.mert_input_layer],
+            layers,
+            degraded,
+            style,
+            steps,
+        )
+        return (generated * latent_std + latent_mean).transpose(1, 2)
+
+    @torch.inference_mode()
+    def refine_latents(
+        self,
+        waveform: torch.Tensor,
+        sample_rate: int,
+        *,
+        steps: int = 32,
+        window_seconds: float | None = 30.0,
+        overlap_seconds: float = 2.0,
+        audio_encoder: AudioVAEEncoder | None = None,
+    ) -> tuple[torch.Tensor, int]:
+        waveform = prepare_waveform(waveform, sample_rate)
+        original_samples = waveform.shape[-1]
+        encoder = audio_encoder if audio_encoder is not None else self.audio_vae
+        if not hasattr(encoder, "encode"):
+            raise TypeError("audio_encoder must provide encode(waveform) -> [batch, 128, frames]")
+        if window_seconds is None:
+            return self._refine_chunk(waveform, steps, encoder), original_samples
+        if waveform.shape[0] != 1 and waveform.shape[-1] > round(window_seconds * SAMPLE_RATE):
+            raise ValueError("windowed refinement currently requires batch size 1")
+
+        window_samples = max(DAV_HOP, round(window_seconds * SAMPLE_RATE / DAV_HOP) * DAV_HOP)
+        overlap_samples = max(0, round(overlap_seconds * SAMPLE_RATE / DAV_HOP) * DAV_HOP)
+        starts = window_starts(original_samples, window_samples, overlap_samples)
+        if len(starts) == 1:
+            return self._refine_chunk(waveform, steps, encoder), original_samples
+
+        total_frames = math.ceil(original_samples / DAV_HOP)
+        accumulated = torch.zeros(1, 128, total_frames, device=self.device, dtype=self.dtype)
+        weights = torch.zeros(1, 1, total_frames, device=self.device, dtype=self.dtype)
+        chunk_frames = math.ceil(window_samples / DAV_HOP)
+        for index, start_sample in enumerate(starts):
+            chunk = waveform[..., start_sample : start_sample + window_samples]
+            generated = self._refine_chunk(chunk, steps, encoder)
+            start_frame = start_sample // DAV_HOP
+            available = min(generated.shape[-1], total_frames - start_frame)
+            blend = torch.ones(available, device=self.device, dtype=self.dtype)
+            if index > 0:
+                previous_end = starts[index - 1] // DAV_HOP + chunk_frames
+                left = min(available, max(0, previous_end - start_frame))
+                if left:
+                    blend[:left] = torch.linspace(0.0, 1.0, left + 2, device=self.device, dtype=self.dtype)[1:-1]
+            if index + 1 < len(starts):
+                next_start = starts[index + 1] // DAV_HOP
+                right = min(available, max(0, start_frame + available - next_start))
+                if right:
+                    fade = torch.linspace(1.0, 0.0, right + 2, device=self.device, dtype=self.dtype)[1:-1]
+                    blend[-right:] = torch.minimum(blend[-right:], fade)
+            accumulated[..., start_frame : start_frame + available] += generated[..., :available] * blend
+            weights[..., start_frame : start_frame + available] += blend
+        if torch.any(weights == 0):
+            raise RuntimeError("window overlap left uncovered latent frames")
+        return accumulated / weights, original_samples
+
+    @torch.inference_mode()
+    def decode_latents(
+        self,
+        latents: torch.Tensor,
+        original_samples: int,
+        decoder: torch.nn.Module | None = None,
+    ) -> torch.Tensor:
+        decoder = decoder if decoder is not None else self.audio_vae
+        if isinstance(decoder, torch.nn.Module):
+            decoder_input = latents.to(device=_module_device(decoder), dtype=_module_dtype(decoder))
+        else:
+            decoder_input = latents.cpu()
+        if hasattr(decoder, "decode"):
+            audio = decoder.decode(decoder_input)
+        else:
+            audio = decoder(decoder_input)
+        if audio.ndim != 3:
+            raise ValueError("audio decoder must return a rank-3 waveform tensor")
+        if audio.shape[1] != 2 and audio.shape[-1] == 2:
+            audio = audio.movedim(-1, 1)
+        if audio.shape[1] != 2:
+            raise ValueError("audio decoder must return stereo [batch, 2, samples]")
+        return audio[..., :original_samples].float().cpu()
+
+    def __call__(
+        self,
+        waveform: torch.Tensor,
+        sample_rate: int,
+        *,
+        steps: int = 32,
+        window_seconds: float | None = 30.0,
+        overlap_seconds: float = 2.0,
+        audio_encoder: AudioVAEEncoder | None = None,
+        decoder: torch.nn.Module | None = None,
+        decode: bool = True,
+    ) -> RefinerOutput:
+        latents, original_samples = self.refine_latents(
+            waveform,
+            sample_rate,
+            steps=steps,
+            window_seconds=window_seconds,
+            overlap_seconds=overlap_seconds,
+            audio_encoder=audio_encoder,
+        )
+        audio = self.decode_latents(latents, original_samples, decoder) if decode else None
+        return RefinerOutput(latents, audio, SAMPLE_RATE, original_samples)