Deploy Magenta RealTime 2 HARP endpoint

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+.DS_Store
+.venv/
+outputs/
+checkpoints/

LICENSE

@@ -0,0 +1,202 @@
+
+                                 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,15 +1,22 @@
 ---
-title: Magenta Rt
-emoji: 📊
-colorFrom: red
-colorTo: yellow
+title: Magenta RealTime 2
+emoji: 🎹
+colorFrom: indigo
+colorTo: green
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 5.28.0
+python_version: "3.12"
 app_file: app.py
 pinned: false
 license: apache-2.0
 short_description: Text-guided short music generation with Magenta RealTime 2
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Generate short, steerable music clips with the Magenta RealTime 2 small model.
+
+This Space exposes a HARP-compatible endpoint for use in audio production
+workflows. Model weights are downloaded from the official Google repository at
+runtime and are not stored in this Space.
+
+See [SOURCES.md](SOURCES.md) for source code, model, and license information.

SOURCES.md

@@ -0,0 +1,19 @@
+# Sources
+
+## Magenta RealTime 2
+
+- Model: https://huggingface.co/google/magenta-realtime-2
+- Source: https://github.com/magenta/magenta-realtime
+- Documentation: https://magenta.github.io/magenta-realtime/
+- Code license: Apache License 2.0
+- Model weights license: Creative Commons Attribution 4.0
+
+## PyTorch ZeroGPU Runtime
+
+- Reference Space: https://huggingface.co/spaces/magenta-community/magenta-rt-jam
+- Reference revision: dfd4eb65e8ebf7eaa9be263d9937ddbee1666262
+- AOTI artifacts: https://huggingface.co/magenta-torch/magenta-rt-aoti-small
+
+The `magenta_rt` package in this Space is the PyTorch runtime distributed with
+the reference Space. Its source files retain the original Google copyright and
+Apache License 2.0 notices.

app.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import os
+import uuid
+from pathlib import Path
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import spaces
+import torch
+from huggingface_hub import hf_hub_download
+from pyharp import ModelCard, build_endpoint
+
+from magenta_rt import paths
+from magenta_rt.torch import MagentaRT2
+from magenta_rt.torch.musiccoca import MusicCoCa
+
+
+MODEL_REPO = "google/magenta-realtime-2"
+MODEL_NAME = "mrt2_small"
+CHECKPOINT = f"{MODEL_NAME}.safetensors"
+AOTI_REPO = "magenta-torch/magenta-rt-aoti-small"
+SAMPLE_RATE = 48_000
+FRAMES_PER_SECOND = 25
+OUTPUT_DIR = Path("/tmp/magenta_rt_outputs")
+
+model_root = Path("/data" if Path("/data").is_dir() else "/tmp/magenta")
+magenta_home = model_root / "magenta-rt-v2"
+magenta_home.mkdir(parents=True, exist_ok=True)
+OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+hf_hub_download(
+    repo_id=MODEL_REPO,
+    filename=f"checkpoints/{CHECKPOINT}",
+    local_dir=magenta_home,
+)
+paths.set_magenta_home(magenta_home)
+
+style_model = MusicCoCa(device="cpu")
+model = MagentaRT2(
+    size=MODEL_NAME,
+    device="cuda",
+    dtype=torch.bfloat16,
+    style_model=style_model,
+)
+
+try:
+    model.load_compiled(repo_id=AOTI_REPO)
+except Exception as exc:
+    print(f"AOTI loading failed; using eager inference: {exc}", flush=True)
+
+
+model_card = ModelCard(
+    name="Magenta RealTime 2",
+    description=(
+        "Generate short instrumental music clips from text prompts using "
+        "Google's open-weights Magenta RealTime 2 small model."
+    ),
+    author="Google DeepMind",
+    tags=[
+        "music-generation",
+        "text-to-music",
+        "instrument-synthesis",
+        "real-time-music",
+    ],
+)
+
+
+@spaces.GPU(duration=45)
+def process_fn(
+    prompt: str,
+    duration: str,
+    temperature: float,
+    top_k: float,
+    seed: float,
+) -> str:
+    prompt = (prompt or "").strip()
+    if not prompt:
+        raise gr.Error("Please describe the music you want to generate.")
+    if len(prompt) > 300:
+        raise gr.Error("The prompt must be 300 characters or fewer.")
+
+    duration_seconds = int(duration)
+    expected_samples = duration_seconds * SAMPLE_RATE
+    frames = duration_seconds * FRAMES_PER_SECOND + 1
+
+    try:
+        if style_model.device != "cuda":
+            style_model.to("cuda")
+        audio, _ = model.generate(
+            style=prompt,
+            temperature=float(temperature),
+            top_k=int(top_k),
+            frames=frames,
+            seed=int(seed),
+            flush=True,
+        )
+    except Exception as exc:
+        raise gr.Error(f"Magenta RealTime 2 inference failed: {exc}") from exc
+
+    audio = np.asarray(audio, dtype=np.float32)
+    if audio.ndim != 2 or audio.shape[1] != 2:
+        raise gr.Error("The model returned an unexpected audio shape.")
+    if len(audio) < expected_samples:
+        raise gr.Error("The model returned less audio than requested.")
+
+    output_path = OUTPUT_DIR / f"{uuid.uuid4().hex}.wav"
+    sf.write(
+        output_path,
+        audio[:expected_samples],
+        SAMPLE_RATE,
+        subtype="PCM_16",
+    )
+    return str(output_path)
+
+
+with gr.Blocks(title="Magenta RealTime 2") as demo:
+    input_components = [
+        gr.Textbox(
+            value="warm analog synthesizer with a gentle rhythmic pulse",
+            label="Music Prompt",
+            info="Describe the instruments, texture, style, or mood.",
+            lines=2,
+            max_lines=4,
+        ),
+        gr.Dropdown(
+            choices=["2", "4", "8"],
+            value="4",
+            label="Duration (seconds)",
+            info="Length of the generated clip.",
+        ),
+        gr.Slider(
+            minimum=0.1,
+            maximum=2.0,
+            step=0.1,
+            value=1.1,
+            label="Temperature",
+            info="Higher values produce more variation.",
+        ),
+        gr.Slider(
+            minimum=10,
+            maximum=100,
+            step=5,
+            value=50,
+            label="Top-k",
+            info="Limits each sampling step to the most likely tokens.",
+        ),
+        gr.Number(
+            value=0,
+            minimum=0,
+            maximum=2_147_483_647,
+            precision=0,
+            label="Seed",
+            info="Use the same seed and controls to reproduce a result.",
+        ),
+    ]
+    output_components = [
+        gr.Audio(
+            type="filepath",
+            label="Generated Music",
+        ).set_info("A 48 kHz stereo WAV file."),
+    ]
+    build_endpoint(
+        model_card=model_card,
+        input_components=input_components,
+        output_components=output_components,
+        process_fn=process_fn,
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=1).launch(
+        show_error=True,
+        pwa=True,
+    )

magenta_rt/__init__.py

@@ -0,0 +1,2 @@
+"""Magenta RealTime 2 (PyTorch) — Spaces package."""
+__version__="2.0.3"

magenta_rt/paths.py

@@ -0,0 +1,120 @@
+# Copyright 2026 Google LLC
+#
+# 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.
+
+"""Centralized path resolution for Magenta RT.
+
+All paths resolve under MAGENTA_HOME/magenta-rt-v2 (where MAGENTA_HOME defaults to ~/Documents/Magenta).
+Override with the MAGENTA_HOME environment variable.
+"""
+
+import os
+import pathlib
+from typing import Union
+
+
+# Configurable root for all downloaded assets and models.
+_MAGENTA_BASE = pathlib.Path(
+    os.environ.get("MAGENTA_HOME", pathlib.Path.home() / "Documents" / "Magenta")
+)
+_MAGENTA_HOME = _MAGENTA_BASE / "magenta-rt-v2"
+
+# Default model directory name (under ~/Documents/Magenta/magenta-rt-v2/models/).
+DEFAULT_MODEL_NAME = "mrt2_base"
+DEFAULT_CHECKPOINT = "mrt2_base.safetensors"
+
+
+def magenta_home() -> pathlib.Path:
+    """Returns the magenta home directory (default: ~/Documents/Magenta/magenta-rt-v2)."""
+    return _MAGENTA_HOME
+
+
+def set_magenta_home(path: Union[pathlib.Path, str]) -> None:
+    """Override the magenta home directory at runtime."""
+    global _MAGENTA_HOME
+    if isinstance(path, str):
+        path = pathlib.Path(path)
+    _MAGENTA_HOME = path
+
+
+# ---------------------------------------------------------------------------
+# Resource directories
+# ---------------------------------------------------------------------------
+
+
+def resources_dir() -> pathlib.Path:
+    """~/Documents/Magenta/magenta-rt-v2/resources — shared resource files (musiccoca, spectrostream)."""
+    return _MAGENTA_HOME / "resources"
+
+
+def musiccoca_dir() -> pathlib.Path:
+    """~/Documents/Magenta/magenta-rt-v2/resources/musiccoca — MusicCoCa TFLite models."""
+    return resources_dir() / "musiccoca"
+
+
+def spectrostream_dir() -> pathlib.Path:
+    """~/Documents/Magenta/magenta-rt-v2/resources/spectrostream — SpectroStream weights."""
+    return resources_dir() / "spectrostream"
+
+
+def models_dir() -> pathlib.Path:
+    """~/Documents/Magenta/magenta-rt-v2/models — exported .mlxfn model directories."""
+    return _MAGENTA_HOME / "models"
+
+
+def default_model_dir() -> pathlib.Path:
+    """~/Documents/Magenta/magenta-rt-v2/models/<DEFAULT_MODEL_NAME> — the default model to load."""
+    return models_dir() / DEFAULT_MODEL_NAME
+
+
+def outputs_dir() -> pathlib.Path:
+    """~/Documents/Magenta/magenta-rt-v2/outputs — generation and export outputs."""
+    d = _MAGENTA_HOME / "outputs"
+    d.mkdir(parents=True, exist_ok=True)
+    return d
+
+
+def checkpoints_dir() -> pathlib.Path:
+    """~/Documents/Magenta/magenta-rt-v2/checkpoints — full safetensors from Linen models."""
+    d = _MAGENTA_HOME / "checkpoints"
+    d.mkdir(parents=True, exist_ok=True)
+    return d
+
+
+def resolve_checkpoint(filename: str) -> pathlib.Path:
+    """Resolve a checkpoint file path.
+
+    First check if literal filepath exists; fallback to ~/Documents/Magenta/magenta-rt-v2/checkpoints/<filename>.
+
+    Args:
+        filename: Checkpoint filename ending in `.safetensors`
+
+    Returns:
+        Path to the checkpoint file (may not exist yet).
+    """
+    if os.path.isfile(filename):
+        return filename
+    return checkpoints_dir() / filename
+
+# ---------------------------------------------------------------------------
+# Path resolution without fallbacks — everything in ~/Documents/Magenta/magenta-rt-v2)
+# ---------------------------------------------------------------------------
+
+def resolve_encoder_weights() -> pathlib.Path:
+    """Returns ~/Documents/Magenta/magenta-rt-v2/resources/spectrostream/encoder.safetensors."""
+    return spectrostream_dir() / "encoder.safetensors"
+
+
+def resolve_decoder_weights() -> pathlib.Path:
+    """Returns ~/Documents/Magenta/magenta-rt-v2/resources/spectrostream/decoder.safetensors."""
+    return spectrostream_dir() / "decoder.safetensors"

magenta_rt/torch/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2026 Google LLC
+#
+# 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.
+
+"""Pure-PyTorch port of Magenta RealTime 2."""
+from .system import MagentaRT2  # noqa: F401

magenta_rt/torch/aoti.py

@@ -0,0 +1,134 @@
+# Copyright 2026 Google LLC
+#
+# 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.
+
+"""AOTI export specs for the functional per-frame step.
+
+The temporal and depth steps are export-clean (no inline constant tensors,
+mask-free incremental KV) and exported with dynamic KV-length dims, so a single
+compiled graph serves every frame. Artifacts are GPU-arch specific — compile on
+the same hardware (e.g. the ZeroGPU Blackwell) you run on.
+"""
+import torch
+import torch.nn as nn
+from torch.export import Dim
+
+
+class TemporalStepModule(nn.Module):
+    """forward(prev_frame, self_kv, cross_kv, source) -> (out, new_self, new_cross)."""
+
+    def __init__(self, decoder):
+        super().__init__()
+        self.d = decoder
+
+    def forward(self, prev_frame, self_kv, cross_kv, source):
+        return self.d.temporal_step_fn(prev_frame, self_kv, cross_kv, source)
+
+
+class DepthStepModule(nn.Module):
+    """forward(depth_input, depth_kv) -> (logits, new_kv)."""
+
+    def __init__(self, decoder):
+        super().__init__()
+        self.d = decoder
+
+    def forward(self, depth_input, depth_kv):
+        return self.d.depth_step_fn(depth_input, depth_kv)
+
+
+def _params_ref(decoder):
+    p = next(decoder.parameters())
+    return p.device, p.dtype
+
+
+def temporal_export_inputs(decoder, kv_len=20):
+    """(args, dynamic_shapes) for exporting the temporal step."""
+    dev, dt = _params_ref(decoder)
+    c = decoder.cfg
+    L, nh, uph = c.temporal.num_layers, c.temporal.num_heads, c.temporal.dim_per_head
+    prev = torch.randint(c.num_reserved_tokens, c.num_reserved_tokens + c.codebook_size,
+                         (1, 1, c.num_codebooks), device=dev)
+    mk = lambda: [(torch.randn(1, kv_len, nh, uph, device=dev, dtype=dt),
+                   torch.randn(1, kv_len, nh, uph, device=dev, dtype=dt)) for _ in range(L)]
+    self_kv, cross_kv = mk(), mk()
+    source = torch.randn(1, 1, c.encoder_model_dims, device=dev, dtype=dt)
+    T = Dim("T", min=0, max=c.temporal_max_past + 1)
+    ds = (None, [({1: T}, {1: T}) for _ in range(L)], [({1: T}, {1: T}) for _ in range(L)], None)
+    return (prev, self_kv, cross_kv, source), ds
+
+
+def depth_export_inputs(decoder, kv_len=6):
+    """(args, dynamic_shapes) for exporting the depth step."""
+    dev, dt = _params_ref(decoder)
+    c = decoder.cfg
+    L, nh, uph = c.depth.num_layers, c.depth.num_heads, c.depth.dim_per_head
+    depth_input = torch.randn(1, 1, c.temporal.model_dims, device=dev, dtype=dt)
+    depth_kv = [(torch.randn(1, kv_len, nh, uph, device=dev, dtype=dt),
+                 torch.randn(1, kv_len, nh, uph, device=dev, dtype=dt)) for _ in range(L)]
+    Td = Dim("Td", min=0, max=c.num_codebooks)
+    ds = (None, [({1: Td}, {1: Td}) for _ in range(L)])
+    return (depth_input, depth_kv), ds
+
+
+def export_temporal(decoder):
+    args, ds = temporal_export_inputs(decoder)
+    return torch.export.export(TemporalStepModule(decoder), args, dynamic_shapes=ds)
+
+
+def export_depth(decoder):
+    args, ds = depth_export_inputs(decoder)
+    return torch.export.export(DepthStepModule(decoder), args, dynamic_shapes=ds)
+
+
+def temporal_state_dict(decoder):
+    return TemporalStepModule(decoder).state_dict()
+
+
+def depth_state_dict(decoder):
+    return DepthStepModule(decoder).state_dict()
+
+
+# --- ZeroGPU-native compile / save / load (uses the `spaces` library) --------
+# Artifacts are GPU-arch specific: compile and run on the same hardware.
+
+def compile_step_archives(decoder, inductor_configs=None):
+    """Export + AOTI-compile the temporal and depth steps with spaces.aoti_compile.
+    Returns {'temporal': bytes, 'depth': bytes} — weight-less graph .pt2 blobs.
+    Must run on a GPU (inside @spaces.GPU on ZeroGPU)."""
+    import importlib
+    spaces = importlib.import_module("spaces")  # optional (ZeroGPU); kept off the import graph
+    out = {}
+    for name, ep in (("temporal", export_temporal(decoder)),
+                     ("depth", export_depth(decoder))):
+        cm = spaces.aoti_compile(ep, inductor_configs)
+        out[name] = bytes(cm.archive_file.getbuffer())
+    return out
+
+
+def load_compiled_steps(decoder, repo_id=None, local_dir=None, filenames=("temporal.pt2", "depth.pt2")):
+    """Load weight-less step graphs and bind them to `decoder`'s own weights.
+    Returns (temporal_callable, depth_callable) for step_f(temporal_step=, depth_step=)."""
+    import importlib, os
+    LazyAOTIModel = importlib.import_module("spaces.zero.torch.aoti").LazyAOTIModel
+    paths = {}
+    keys = ("temporal", "depth")
+    if local_dir is not None:
+        for k, fn in zip(keys, filenames):
+            paths[k] = os.path.join(local_dir, fn)
+    else:
+        from huggingface_hub import hf_hub_download
+        for k, fn in zip(keys, filenames):
+            paths[k] = hf_hub_download(repo_id, fn)
+    t = LazyAOTIModel(paths["temporal"]).with_weights(temporal_state_dict(decoder))
+    d = LazyAOTIModel(paths["depth"]).with_weights(depth_state_dict(decoder))
+    return t, d

magenta_rt/torch/depthformer.py

@@ -0,0 +1,403 @@
+# Copyright 2026 Google LLC
+#
+# 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.
+
+"""PyTorch port of the Magenta RealTime 2 Depthformer (encoder + multivariate decoder)."""
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+from . import layers as L
+
+
+@dataclass
+class SpecDims:
+    num_layers: int
+    model_dims: int
+    hidden_dims: int
+    num_heads: int
+    dim_per_head: int
+
+
+@dataclass
+class DepthformerConfig:
+    # encoder (conditioning embedder) dims
+    encoder_model_dims: int
+    musiccoca_rvq: int            # 12
+    musiccoca_per_rvq_vocab: int  # 1031
+    musiccoca_embed_dim: int      # 768
+    regular_num_embeddings_per_channel: list  # per regular channel vocab sizes
+    regular_num_channels: int     # 132
+    # decoder
+    temporal: SpecDims
+    depth: SpecDims
+    temporal_max_past: int        # 41 (small) / 25 (base)
+    depth_max_past: int           # 12
+    num_sinks: int                # 1
+    num_codebooks: int            # 12
+    codebook_size: int            # 1024
+    num_reserved_tokens: int      # 6
+    vocab_size: int               # 12294
+    soft_cap_logits: float = 30.0
+
+
+def _mean_f32(x, axis):
+    return x.float().mean(axis).to(x.dtype)
+
+
+class TransformerStack(nn.Module):
+    """Stack of primer-hybrid transformer layers; optional streaming cross-attn."""
+
+    def __init__(self, spec: SpecDims, max_past, num_sinks, use_cross, source_dim=None):
+        super().__init__()
+        self.use_cross = use_cross
+        self.layers = nn.ModuleList()
+        for _ in range(spec.num_layers):
+            blk = nn.ModuleDict()
+            blk["self_attention"] = L.SelfAttention(
+                spec.model_dims, spec.num_heads, spec.dim_per_head, max_past,
+                num_sinks=num_sinks)
+            if use_cross:
+                blk["cross_attention"] = L.CrossAttention(
+                    spec.model_dims, source_dim, spec.num_heads, spec.dim_per_head,
+                    max_past, num_sinks=num_sinks)
+            blk["ffn"] = L.FFN(spec.model_dims, spec.hidden_dims)
+            self.layers.append(blk)
+
+    def forward(self, x, source=None):
+        for blk in self.layers:
+            x = blk["self_attention"](x)
+            if self.use_cross:
+                x = blk["cross_attention"](x, source)
+            x = blk["ffn"](x)
+        return x
+
+    # ---- streaming ----
+    def init_state(self, batch, device, dtype, source=None):
+        st = []
+        for blk in self.layers:
+            s = {"self": blk["self_attention"].init_state(batch, device, dtype)}
+            if self.use_cross:
+                s["source_kv"] = (None, None)
+            st.append(s)
+        return st
+
+    def step(self, x, state, source_step=None):
+        # source_step: newly-encoded source frame [b,1,source_dim] to append to KV.
+        for blk, s in zip(self.layers, state):
+            x = blk["self_attention"].step(x, s["self"])
+            if self.use_cross:
+                ca = blk["cross_attention"]
+                k, v = ca._kv(source_step)
+                pk, pv = s["source_kv"]
+                k = k if pk is None else torch.cat([pk, k], dim=1)
+                v = v if pv is None else torch.cat([pv, v], dim=1)
+                keep = ca.max_past_horizon + 1
+                if k.shape[1] > keep:
+                    k = k[:, -keep:]
+                    v = v[:, -keep:]
+                s["source_kv"] = (k, v)
+                x = ca.step(x, (k, v))
+            x = blk["ffn"](x)
+        return x
+
+    def step_fn(self, x, self_kv, cross_kv, source_frame):
+        """Functional, export-clean per-frame step over all layers.
+
+        self_kv / cross_kv: lists of (k, v) per layer ([b,T,nh,uph], T may be 0).
+        Returns (out[b,1,d], new_self_kv, new_cross_kv) with untrimmed KV."""
+        new_self = []
+        new_cross = []
+        for i, blk in enumerate(self.layers):
+            x, k, v = blk["self_attention"].step_fn(x, self_kv[i][0], self_kv[i][1])
+            new_self.append((k, v))
+            if self.use_cross:
+                ca = blk["cross_attention"]
+                sk, sv = ca._kv(source_frame)
+                sk = torch.cat([cross_kv[i][0], sk], dim=1)
+                sv = torch.cat([cross_kv[i][1], sv], dim=1)
+                new_cross.append((sk, sv))
+                x = ca.attend_fn(x, sk, sv)
+            x = blk["ffn"](x)
+        return x, new_self, new_cross
+
+
+class EncoderEmbedding(nn.Module):
+    """Embeds the conditioning block [b,t,num_channels] -> source [b,t,enc_dim]."""
+
+    def __init__(self, cfg: DepthformerConfig):
+        super().__init__()
+        self.cfg = cfg
+        m = cfg.musiccoca_rvq
+        self.m = m
+        # mulan branch
+        self.register_buffer("mulan_offset", torch.arange(m) * cfg.musiccoca_per_rvq_vocab)
+        self.mulan_dequantizer = nn.Parameter(
+            torch.zeros(m * cfg.musiccoca_per_rvq_vocab, cfg.musiccoca_embed_dim))
+        self.mulan_adapter = L.JaxLinear(cfg.musiccoca_embed_dim, cfg.encoder_model_dims, use_bias=False)
+        # regular branch (MultiChannelEmbedding)
+        per = cfg.regular_num_embeddings_per_channel
+        total = sum(per)
+        total = (total + 127) // 128 * 128
+        self.regular_embedding = nn.Parameter(torch.zeros(total, cfg.encoder_model_dims))
+        offs = [0]
+        for p in per[:-1]:
+            offs.append(offs[-1] + p)
+        self.register_buffer("regular_offsets", torch.tensor(offs, dtype=torch.long))
+        self.encoder_ln = L.LayerNorm(cfg.encoder_model_dims)
+
+    def forward(self, x):
+        # x: [b,t,num_channels] int
+        m = self.m
+        mulan = x[..., :m]                  # [b,t,m]
+        regular = x[..., m:]                # [b,t,132]
+        off = self.mulan_offset.to(x.device)
+        idx = mulan + off                   # [b,t,m]
+        emb = self.mulan_dequantizer[idx]   # [b,t,m,768]
+        emb = emb.sum(dim=-2)               # [b,t,768]
+        mulan_out = self.mulan_adapter(emb)  # [b,t,enc_dim]
+        # regular
+        ridx = regular + self.regular_offsets.to(x.device)
+        remb = self.regular_embedding[ridx]  # [b,t,132,enc_dim]
+        regular_out = _mean_f32(remb, axis=-2)  # mean over channels in fp32
+        # branch combine via MEAN
+        src = (mulan_out + regular_out) / 2.0
+        return self.encoder_ln(src)
+
+
+class MultivariateDecoder(nn.Module):
+    def __init__(self, cfg: DepthformerConfig):
+        super().__init__()
+        self.cfg = cfg
+        td, dd = cfg.temporal, cfg.depth
+        self.embedding = nn.Parameter(torch.zeros(cfg.vocab_size, td.model_dims))
+        self.embed_scale = math.sqrt(td.model_dims)
+        self.temporal_body = TransformerStack(
+            td, cfg.temporal_max_past, cfg.num_sinks, use_cross=True,
+            source_dim=cfg.encoder_model_dims)
+        # depth input adapter (Dense temporal->depth, no bias) or identity
+        if td.model_dims != dd.model_dims:
+            self.depth_input_adapter = L.JaxLinear(td.model_dims, dd.model_dims, use_bias=False)
+        else:
+            self.depth_input_adapter = None
+        self.depth_body = TransformerStack(
+            dd, cfg.depth_max_past, num_sinks=0, use_cross=False)
+        self.final_ln = L.LayerNorm(dd.model_dims)
+        self.to_logits = L.JaxLinear(dd.model_dims, cfg.vocab_size, use_bias=True)
+
+    def embed(self, tokens):
+        # tokens: [...,] int -> [..., td_dim]
+        return self.embedding[tokens] * self.embed_scale
+
+    def _depth_forward(self, depth_inputs):
+        # depth_inputs: [N, Q, td_dim]
+        h = depth_inputs
+        if self.depth_input_adapter is not None:
+            h = self.depth_input_adapter(h)
+        h = self.depth_body(h)
+        h = self.final_ln(h)
+        logits = self.to_logits(h)
+        return logits
+
+    def forward(self, target, source):
+        """Teacher-forced. target: [b,T,Q] int; source: [b,Tc,enc_dim]. -> logits [b,T,Q,vocab]."""
+        cfg = self.cfg
+        b, T, Q = target.shape
+        sos = target.new_zeros((b, 1, Q))  # sos_id=0
+        x = torch.cat([sos, target], dim=1)             # [b,T+1,Q]
+        embedded = self.embed(x)                         # [b,T+1,Q,D]
+        temporal_inputs = _mean_f32(embedded, axis=-2)[:, :-1]  # [b,T,D]
+        temporal_outputs = self.temporal_body(temporal_inputs, source)  # [b,T,D]
+        depth_inputs = torch.cat(
+            [temporal_outputs[..., None, :], embedded[:, 1:, :-1]], dim=-2)  # [b,T,Q,D]
+        N = b * T
+        logits = self._depth_forward(depth_inputs.reshape(N, Q, -1)).reshape(b, T, Q, -1)
+        if cfg.soft_cap_logits is not None:
+            c = cfg.soft_cap_logits
+            logits = torch.tanh(logits / c) * c
+        return logits
+
+    # ---- functional (AOTI-compilable) streaming -------------------------
+    def temporal_step_fn(self, prev_frame, self_kv, cross_kv, source_frame):
+        """Functional temporal step: prev_frame[b,1,Q] -> (temporal_out, kv...)."""
+        embedded = self.embed(prev_frame)
+        ti = _mean_f32(embedded, axis=-2)
+        return self.temporal_body.step_fn(ti, self_kv, cross_kv, source_frame)
+
+    def depth_step_fn(self, depth_input, depth_kv):
+        """Functional depth step: depth_input[b,1,Dt] + kv -> (logits, new_kv)."""
+        h = depth_input
+        if self.depth_input_adapter is not None:
+            h = self.depth_input_adapter(h)
+        h, nk, _ = self.depth_body.step_fn(h, depth_kv, [], None)
+        h = self.final_ln(h)
+        return self._soft_cap(self.to_logits(h)), nk
+
+    def init_streaming_f(self, batch, device, dtype=torch.float32):
+        td = self.cfg.temporal
+        z = torch.zeros(batch, 0, td.num_heads, td.dim_per_head, device=device, dtype=dtype)
+        kv = [(z, z) for _ in range(td.num_layers)]
+        return {
+            "self": [(k, v) for k, v in kv],
+            "cross": [(k, v) for k, v in kv],
+            "prev": torch.zeros((batch, 1, self.cfg.num_codebooks), dtype=torch.long, device=device),
+        }
+
+    def step_f(self, state, source_frame, sampler=None, forced=None,
+               temporal_step=None, depth_step=None):
+        """One functional frame. temporal_step/depth_step override the eager fns
+        (e.g. with AOTI-compiled callables). Updates state in place; returns [b,1,Q]."""
+        cfg = self.cfg
+        tstep = temporal_step or self.temporal_step_fn
+        dstep = depth_step or self.depth_step_fn
+        to, new_self, new_cross = tstep(state["prev"], state["self"], state["cross"], source_frame)
+        keep = cfg.temporal_max_past + 1
+        state["self"] = [(k[:, -keep:], v[:, -keep:]) for k, v in new_self]
+        state["cross"] = [(k[:, -keep:], v[:, -keep:]) for k, v in new_cross]
+        dd = cfg.depth
+        z = torch.zeros(to.shape[0], 0, dd.num_heads, dd.dim_per_head, device=to.device, dtype=to.dtype)
+        depth_kv = [(z, z) for _ in range(dd.num_layers)]
+        depth_input = to
+        samples = []
+        for q in range(cfg.num_codebooks):
+            logits, depth_kv = dstep(depth_input, depth_kv)
+            lo = cfg.num_reserved_tokens + q * cfg.codebook_size
+            hi = lo + cfg.codebook_size
+            tok = forced[..., q] if forced is not None else sampler(logits.float(), q, lo, hi)
+            samples.append(tok)
+            depth_input = self.embed(tok)
+        frame = torch.stack(samples, dim=-1)
+        state["prev"] = frame
+        return frame
+
+    # ---- streaming generation -------------------------------------------
+    def init_streaming(self, batch, device, dtype=torch.float32):
+        return {
+            "temporal": self.temporal_body.init_state(batch, device, dtype, source=True),
+            "prev": torch.zeros((batch, 1, self.cfg.num_codebooks), dtype=torch.long, device=device),
+        }
+
+    def _soft_cap(self, logits):
+        c = self.cfg.soft_cap_logits
+        return torch.tanh(logits / c) * c if c is not None else logits
+
+    def _depth_step_logits(self, depth_inputs, depth_state):
+        h = depth_inputs
+        if self.depth_input_adapter is not None:
+            h = self.depth_input_adapter(h)
+        h = self.depth_body.step(h, depth_state)
+        h = self.final_ln(h)
+        return self.to_logits(h)
+
+    def step(self, state, source_step, sampler=None, forced_frame=None):
+        """One streaming frame.
+
+        state: from init_streaming (updated in place).
+        source_step: [b,1,enc_dim] encoded conditioning for this frame.
+        sampler: fn(logits[b,1,vocab], rvq_index, valid_lo, valid_hi) -> token[b,1] long.
+        forced_frame: [b,1,Q] tokens to force (teacher forcing) instead of sampling.
+        Returns sampled frame [b,1,Q] long.
+        """
+        cfg = self.cfg
+        prev = state["prev"]
+        embedded = self.embed(prev)                      # [b,1,Q,D]
+        temporal_inputs = _mean_f32(embedded, axis=-2)   # [b,1,D]
+        temporal_out = self.temporal_body.step(temporal_inputs, state["temporal"], source_step)
+        depth_state = self.depth_body.init_state(prev.shape[0], prev.device, temporal_out.dtype)
+        depth_inputs = temporal_out
+        samples = []
+        for q in range(cfg.num_codebooks):
+            logits = self._soft_cap(self._depth_step_logits(depth_inputs, depth_state)).float()
+            lo = cfg.num_reserved_tokens + q * cfg.codebook_size
+            hi = lo + cfg.codebook_size
+            if forced_frame is not None:
+                tok = forced_frame[..., q]               # [b,1]
+            else:
+                tok = sampler(logits, q, lo, hi)         # [b,1]
+            samples.append(tok)
+            depth_inputs = self.embed(tok.unsqueeze(-1)).squeeze(-2) if tok.dim() == 2 else self.embed(tok)
+        frame = torch.stack(samples, dim=-1)             # [b,1,Q]
+        state["prev"] = frame
+        return frame, logits  # last logits returned for debugging
+
+    def streaming_logits(self, target, source):
+        """Re-derive per-(t,q) logits via the streaming step path with forced tokens.
+        Used to validate the KV-cache step path against teacher forcing."""
+        b, T, Q = target.shape
+        state = self.init_streaming(b, target.device)
+        all_logits = []
+        for t in range(T):
+            frame_logits = []
+            prev = state["prev"]
+            embedded = self.embed(prev)
+            temporal_inputs = _mean_f32(embedded, axis=-2)
+            temporal_out = self.temporal_body.step(temporal_inputs, state["temporal"], source[:, t:t+1])
+            depth_state = self.depth_body.init_state(b, target.device, temporal_out.dtype)
+            depth_inputs = temporal_out
+            for q in range(Q):
+                logits = self._soft_cap(self._depth_step_logits(depth_inputs, depth_state))
+                frame_logits.append(logits)
+                tok = target[:, t:t+1, q]
+                depth_inputs = self.embed(tok)
+            state["prev"] = target[:, t:t+1, :]
+            all_logits.append(torch.stack(frame_logits, dim=2))  # [b,1,Q,vocab]
+        return torch.cat(all_logits, dim=1)
+
+
+class Depthformer(nn.Module):
+    def __init__(self, cfg: DepthformerConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.encoder = EncoderEmbedding(cfg)
+        self.decoder = MultivariateDecoder(cfg)
+
+    def encode(self, cond):
+        return self.encoder(cond)
+
+    def forward(self, cond, target):
+        source = self.encode(cond)
+        return self.decoder(target, source)
+
+
+# --- config builders ---------------------------------------------------------
+
+_REGULAR_PER_CHANNEL = [11] * 128 + [9] * 1 + [47] * 2 + [15] * 1  # 132 channels
+
+_COMMON = dict(
+    musiccoca_rvq=12, musiccoca_per_rvq_vocab=1031, musiccoca_embed_dim=768,
+    regular_num_embeddings_per_channel=_REGULAR_PER_CHANNEL, regular_num_channels=132,
+    depth_max_past=12, num_sinks=1, num_codebooks=12, codebook_size=1024,
+    num_reserved_tokens=6, vocab_size=12294, soft_cap_logits=30.0,
+)
+
+
+def config_for(size: str) -> DepthformerConfig:
+    if size == "mrt2_small":
+        return DepthformerConfig(
+            encoder_model_dims=256,
+            temporal=SpecDims(12, 1024, 4096, 8, 128),
+            depth=SpecDims(2, 768, 3072, 6, 128),
+            temporal_max_past=41,
+            **_COMMON,
+        )
+    if size == "mrt2_base":
+        return DepthformerConfig(
+            encoder_model_dims=1024,
+            temporal=SpecDims(20, 3072, 8192, 24, 128),
+            depth=SpecDims(6, 1024, 4096, 8, 128),
+            temporal_max_past=25,
+            **_COMMON,
+        )
+    raise ValueError(f"unknown size {size}")

magenta_rt/torch/layers.py

@@ -0,0 +1,344 @@
+# Copyright 2026 Google LLC
+#
+# 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.
+
+"""PyTorch port of the sequence_layers primitives used by Magenta RealTime 2.
+
+Layouts mirror the JAX/Linen checkpoint exactly so weight loading is a direct
+copy (no transposes for kernels stored as [in, ...]). Reductions for norms and
+softmax run in fp32 to match the reference.
+"""
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+_R_SOFTPLUS_0 = 1.442695041  # 1 / softplus(0); from sequence_layers attention.
+
+
+def gelu_approx(x: torch.Tensor) -> torch.Tensor:
+    """tanh-approximation GELU, matching mlx.nn.gelu_approx / jax gelu(approximate=True)."""
+    return F.gelu(x, approximate="tanh")
+
+
+class JaxLinear(nn.Module):
+    """Linen Dense: y = x @ kernel + bias, kernel stored as [in, out]."""
+
+    def __init__(self, in_features, out_features, use_bias=True, activation=None):
+        super().__init__()
+        self.kernel = nn.Parameter(torch.zeros(in_features, out_features))
+        self.bias = nn.Parameter(torch.zeros(out_features)) if use_bias else None
+        self.activation = activation
+
+    def forward(self, x):
+        y = torch.matmul(x, self.kernel.to(x.dtype))
+        if self.bias is not None:
+            y = y + self.bias.to(x.dtype)
+        if self.activation is not None:
+            y = self.activation(y)
+        return y
+
+
+class RMSNorm(nn.Module):
+    """RMS norm with learned scale; reduction in fp32 (eps 1e-6)."""
+
+    def __init__(self, dim, eps=1e-6, use_scale=True):
+        super().__init__()
+        self.eps = eps
+        self.scale = nn.Parameter(torch.ones(dim)) if use_scale else None
+
+    def forward(self, x):
+        dt = x.dtype
+        v = x.float()
+        v = v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + self.eps)
+        v = v.to(dt)
+        if self.scale is not None:
+            v = v * self.scale.to(dt)
+        return v
+
+
+class LayerNorm(nn.Module):
+    """LayerNorm with scale+bias; reduction in fp32 (eps 1e-6)."""
+
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.scale = nn.Parameter(torch.ones(dim))
+        self.bias = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x):
+        dt = x.dtype
+        v = x.float()
+        mean = v.mean(-1, keepdim=True)
+        var = (v - mean).pow(2).mean(-1, keepdim=True)
+        v = (v - mean) * torch.rsqrt(var + self.eps)
+        v = v.to(dt)
+        return v * self.scale.to(dt) + self.bias.to(dt)
+
+
+def _query_scale_vector(per_dim_scale, units_per_head, dtype):
+    qscale = 1.0 / math.sqrt(units_per_head)
+    if per_dim_scale is not None:
+        scale = _R_SOFTPLUS_0 * qscale
+        softplus = F.softplus(per_dim_scale.to(dtype))
+        return scale * softplus
+    return torch.tensor(qscale, dtype=dtype, device=per_dim_scale.device if per_dim_scale is not None else None)
+
+
+def dot_product_attention(q, k, v, per_dim_scale, sink_k, sink_v, mask):
+    """q,k,v: [b, t, nh, uph]; mask: [b, 1, tq, tkv] bool (True=attend) or None.
+
+    sink_k/sink_v: [num_sink, nh, uph] or None. Sink logits use *unscaled* queries.
+    Returns context [b, tq, nh, uph].
+    """
+    qh = q.transpose(1, 2)  # [b, nh, tq, uph]
+    kh = k.transpose(1, 2)
+    vh = v.transpose(1, 2)
+
+    scale_vec = _query_scale_vector(per_dim_scale, q.shape[-1], qh.dtype)  # [uph]
+
+    if sink_k is not None:
+        # [b, nh, tq, num_sink] using unscaled queries.
+        sink_logits = torch.einsum("bhqd,shd->bhqs", qh, sink_k.to(qh.dtype))
+
+    qs = qh * scale_vec
+    logits = torch.matmul(qs, kh.transpose(-1, -2))  # [b, nh, tq, tkv]
+
+    if sink_k is not None:
+        logits = torch.cat([sink_logits, logits], dim=-1)
+
+    if mask is not None:
+        # Export-clean: scalar masked_fill (no -1e9 constant tensor). Sink
+        # columns (first ns) are always valid, so mask only the kv columns
+        # (avoids building a constant `ones` sink mask).
+        ns = sink_k.shape[0] if sink_k is not None else 0
+        if ns:
+            kv = logits[..., ns:].masked_fill(~mask, -1e9)
+            logits = torch.cat([logits[..., :ns], kv], dim=-1)
+        else:
+            logits = logits.masked_fill(~mask, -1e9)
+
+    weights = torch.softmax(logits.float(), dim=-1).to(vh.dtype)
+
+    if sink_v is not None:
+        b = vh.shape[0]
+        sink_vb = sink_v.to(vh.dtype).permute(1, 0, 2).unsqueeze(0).expand(b, -1, -1, -1)  # [b,nh,num_sink,uph]
+        vh = torch.cat([sink_vb, vh], dim=2)
+
+    ctx = torch.matmul(weights, vh)  # [b, nh, tq, uph]
+    return ctx.transpose(1, 2)  # [b, tq, nh, uph]
+
+
+class AttnProjection(nn.Module):
+    """q/k/v/out projections stored as [in, nh, uph] (Linen attention kernels)."""
+
+    def __init__(self, in_dim, num_heads, units_per_head, has_sinks=False, has_per_dim_scale=True):
+        super().__init__()
+        nh, uph = num_heads, units_per_head
+        self.num_heads, self.units_per_head = nh, uph
+        self.query_projection_kernel = nn.Parameter(torch.zeros(in_dim, nh, uph))
+        self.key_projection_kernel = nn.Parameter(torch.zeros(in_dim, nh, uph))
+        self.value_projection_kernel = nn.Parameter(torch.zeros(in_dim, nh, uph))
+        self.per_dim_scale = nn.Parameter(torch.zeros(uph)) if has_per_dim_scale else None
+        if has_sinks:
+            self.sink_key_embeddings = nn.Parameter(torch.zeros(1, nh, uph))
+            self.sink_value_embeddings = nn.Parameter(torch.zeros(1, nh, uph))
+        else:
+            self.sink_key_embeddings = None
+            self.sink_value_embeddings = None
+
+    def project(self, x, kernel):
+        return torch.einsum("btd,dnh->btnh", x, kernel.to(x.dtype))
+
+
+def banded_causal_mask(tq, tkv, past, future, device):
+    """[1,1,tq,tkv] bool. Query i (global pos offset+i) attends key j with
+    j in [i-past, i+future]. Here tq aligns to the last tq positions of tkv."""
+    offset = tkv - tq
+    row = torch.arange(tq, device=device)[:, None] + offset
+    col = torch.arange(tkv, device=device)[None, :]
+    m = (col <= row + future) & (col >= row - past)
+    return m.view(1, 1, tq, tkv)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, model_dim, num_heads, units_per_head, max_past_horizon,
+                 num_sinks=0, eps=1e-6):
+        super().__init__()
+        self.pre_norm = RMSNorm(model_dim, eps)
+        self.post_norm = RMSNorm(model_dim, eps)
+        self.attention = AttnProjection(model_dim, num_heads, units_per_head,
+                                        has_sinks=num_sinks > 0)
+        self.output_projection_kernel = nn.Parameter(torch.zeros(model_dim, num_heads, units_per_head))
+        self.max_past_horizon = max_past_horizon
+        self.num_heads = num_heads
+        self.units_per_head = units_per_head
+
+    def _branch(self, x):
+        h = self.pre_norm(x)
+        a = self.attention
+        q = a.project(h, a.query_projection_kernel)
+        k = a.project(h, a.key_projection_kernel)
+        v = a.project(h, a.value_projection_kernel)
+        t = x.shape[1]
+        mask = banded_causal_mask(t, t, self.max_past_horizon, 0, x.device)
+        ctx = dot_product_attention(q, k, v, a.per_dim_scale,
+                                    a.sink_key_embeddings, a.sink_value_embeddings, mask)
+        out = torch.einsum("btnh,dnh->btd", ctx, self.output_projection_kernel.to(ctx.dtype))
+        return self.post_norm(out)
+
+    def forward(self, x):
+        return x + self._branch(x)
+
+    # ---- streaming step with KV cache ----
+    def init_state(self, batch, device, dtype):
+        return {"k": None, "v": None}  # lazily grown buffers [b, t, nh, uph]
+
+    def step(self, x, state):
+        # x: [b, 1, d]
+        h = self.pre_norm(x)
+        a = self.attention
+        q = a.project(h, a.query_projection_kernel)
+        k = a.project(h, a.key_projection_kernel)
+        v = a.project(h, a.value_projection_kernel)
+        if state["k"] is None:
+            kk, vv = k, v
+        else:
+            kk = torch.cat([state["k"], k], dim=1)
+            vv = torch.cat([state["v"], v], dim=1)
+        # keep only the window we can attend to (past horizon + current).
+        keep = self.max_past_horizon + 1
+        if kk.shape[1] > keep:
+            kk = kk[:, -keep:]
+            vv = vv[:, -keep:]
+        state["k"], state["v"] = kk, vv
+        # The cache holds only the last `keep` keys, all valid past keys within
+        # the window for the single newest query -> no mask needed.
+        ctx = dot_product_attention(q, kk, vv, a.per_dim_scale,
+                                    a.sink_key_embeddings, a.sink_value_embeddings, None)
+        out = torch.einsum("btnh,dnh->btd", ctx, self.output_projection_kernel.to(ctx.dtype))
+        return x + self.post_norm(out)
+
+    def step_fn(self, x, k_prev, v_prev):
+        """Functional, export-clean single-frame step. k_prev/v_prev: [b,T,nh,uph]
+        (T may be 0). Returns (out[b,1,d], new_k, new_v) — full, untrimmed KV."""
+        h = self.pre_norm(x)
+        a = self.attention
+        q = a.project(h, a.query_projection_kernel)
+        k = torch.cat([k_prev, a.project(h, a.key_projection_kernel)], dim=1)
+        v = torch.cat([v_prev, a.project(h, a.value_projection_kernel)], dim=1)
+        ctx = dot_product_attention(q, k, v, a.per_dim_scale,
+                                    a.sink_key_embeddings, a.sink_value_embeddings, None)
+        out = torch.einsum("btnh,dnh->btd", ctx, self.output_projection_kernel.to(ctx.dtype))
+        return x + self.post_norm(out), k, v
+
+
+class CrossAttention(nn.Module):
+    """Streaming cross-attention to an encoded source sequence."""
+
+    def __init__(self, model_dim, source_dim, num_heads, units_per_head,
+                 max_past_horizon, num_sinks=0, eps=1e-6):
+        super().__init__()
+        self.pre_norm = RMSNorm(model_dim, eps)
+        self.post_norm = RMSNorm(model_dim, eps)
+        # query from decoder (model_dim), key/value from source (source_dim).
+        self.attention = _CrossProj(model_dim, source_dim, num_heads, units_per_head,
+                                    has_sinks=num_sinks > 0)
+        self.output_projection_kernel = nn.Parameter(torch.zeros(model_dim, num_heads, units_per_head))
+        self.max_past_horizon = max_past_horizon
+        self.num_heads = num_heads
+        self.units_per_head = units_per_head
+
+    def _kv(self, source):
+        a = self.attention
+        k = torch.einsum("btd,dnh->btnh", source, a.key_projection_kernel.to(source.dtype))
+        v = torch.einsum("btd,dnh->btnh", source, a.value_projection_kernel.to(source.dtype))
+        return k, v
+
+    def _branch(self, x, source):
+        h = self.pre_norm(x)
+        a = self.attention
+        q = torch.einsum("btd,dnh->btnh", h, a.query_projection_kernel.to(h.dtype))
+        k, v = self._kv(source)
+        tq, tkv = x.shape[1], source.shape[1]
+        # query at decoder time i attends source positions within past horizon, causal.
+        mask = banded_causal_mask(tq, tkv, self.max_past_horizon, 0, x.device)
+        ctx = dot_product_attention(q, k, v, a.per_dim_scale,
+                                    a.sink_key_embeddings, a.sink_value_embeddings, mask)
+        out = torch.einsum("btnh,dnh->btd", ctx, self.output_projection_kernel.to(ctx.dtype))
+        return self.post_norm(out)
+
+    def forward(self, x, source):
+        return x + self._branch(x, source)
+
+    def attend_fn(self, x, k, v):
+        """Functional cross-attention given precomputed source KV [b,T,nh,uph]."""
+        h = self.pre_norm(x)
+        a = self.attention
+        q = torch.einsum("btd,dnh->btnh", h, a.query_projection_kernel.to(h.dtype))
+        ctx = dot_product_attention(q, k, v, a.per_dim_scale,
+                                    a.sink_key_embeddings, a.sink_value_embeddings, None)
+        out = torch.einsum("btnh,dnh->btd", ctx, self.output_projection_kernel.to(ctx.dtype))
+        return x + self.post_norm(out)
+
+    def step(self, x, source_kv):
+        # x: [b,1,d]; source_kv: (k,v) accumulated [b, tkv, nh, uph]
+        h = self.pre_norm(x)
+        a = self.attention
+        q = torch.einsum("btd,dnh->btnh", h, a.query_projection_kernel.to(h.dtype))
+        k, v = source_kv
+        tkv = k.shape[1]
+        keep = self.max_past_horizon + 1
+        if tkv > keep:
+            k = k[:, -keep:]
+            v = v[:, -keep:]
+        ctx = dot_product_attention(q, k, v, a.per_dim_scale,
+                                    a.sink_key_embeddings, a.sink_value_embeddings, None)
+        out = torch.einsum("btnh,dnh->btd", ctx, self.output_projection_kernel.to(ctx.dtype))
+        return x + self.post_norm(out)
+
+
+class _CrossProj(nn.Module):
+    def __init__(self, q_dim, kv_dim, num_heads, units_per_head, has_sinks=False):
+        super().__init__()
+        nh, uph = num_heads, units_per_head
+        self.query_projection_kernel = nn.Parameter(torch.zeros(q_dim, nh, uph))
+        self.key_projection_kernel = nn.Parameter(torch.zeros(kv_dim, nh, uph))
+        self.value_projection_kernel = nn.Parameter(torch.zeros(kv_dim, nh, uph))
+        self.per_dim_scale = nn.Parameter(torch.zeros(uph))
+        if has_sinks:
+            self.sink_key_embeddings = nn.Parameter(torch.zeros(1, nh, uph))
+            self.sink_value_embeddings = nn.Parameter(torch.zeros(1, nh, uph))
+        else:
+            self.sink_key_embeddings = None
+            self.sink_value_embeddings = None
+
+
+class FFN(nn.Module):
+    def __init__(self, model_dim, hidden_dim, eps=1e-6):
+        super().__init__()
+        self.pre_norm = RMSNorm(model_dim, eps)
+        self.post_norm = RMSNorm(model_dim, eps)
+        self.ffn_layer1 = JaxLinear(model_dim, hidden_dim, use_bias=True, activation=gelu_approx)
+        self.ffn_layer2 = JaxLinear(hidden_dim, model_dim, use_bias=True)
+
+    def _branch(self, x):
+        h = self.pre_norm(x)
+        h = self.ffn_layer1(h)
+        h = self.ffn_layer2(h)
+        return self.post_norm(h)
+
+    def forward(self, x):
+        return x + self._branch(x)

magenta_rt/torch/modeling_magenta_rt2.py

@@ -0,0 +1,356 @@
+# Copyright 2026 Google LLC
+#
+# 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.
+
+"""Transformers-compatible Magenta RealTime 2 (PyTorch).
+
+`MagentaRT2ForConditionalGeneration` is a `PreTrainedModel` wrapping the
+Depthformer LLM + SpectroStream codec decoder. Generation is per-frame RVQ
+autoregression with a depth transformer + streaming codec decode, exposed as
+custom `generate` / `stream` methods (it does not fit `GenerationMixin`, whose
+loop is a single token stream). MusicCoCa style encoding is a separate
+`MusicCoCaProcessor`.
+"""
+
+import json
+import os
+
+import numpy as np
+import torch
+from transformers import PreTrainedModel
+
+from .configuration_magenta_rt2 import MagentaRT2Config
+from .depthformer import Depthformer, DepthformerConfig, SpecDims
+from .spectrostream import SpectroStreamDecoder, codes_to_embeddings
+
+# Force `trust_remote_code` to bundle every dependency module. transformers only
+# traces `from .X import ...` (not `from . import X`), so we name them explicitly.
+from .layers import JaxLinear as _ensure_layers  # noqa: F401
+from .musiccoca import MusicCoCa as _ensure_musiccoca  # noqa: F401
+from .processing_musiccoca import MusicCoCaProcessor as _ensure_processor  # noqa: F401
+from .aoti import load_compiled_steps as _ensure_aoti  # noqa: F401
+
+SR = 48000
+FRAME_SAMPLES = 1920
+STREAM_DECODE_CONTEXT = 16
+STREAM_DECODE_MARGIN = 2
+
+
+def discretize_cfg(value, step, max_bin):
+    clamped = max(-1.0, min(7.0, value))
+    return max(0, min(max_bin, int(round((clamped - (-1.0)) / step))))
+
+
+def convert_from_unique_codes(tokens, codebook_size, num_reserved=6):
+    return (tokens - num_reserved) % codebook_size
+
+
+def _float_to_int16(samples, gain=0.5):
+    samples = np.clip(gain * samples, -1, 1)
+    samples = np.round((np.iinfo(np.int16).max + 0.5) * samples - 0.5)
+    return samples.astype(np.int16)
+
+
+def make_sampler(temperature, top_k, generator):
+    """jax-parity: valid-range mask, top-k, gumbel-max. NEG is finite (avoids NaN)."""
+    NEG = -1e9
+
+    def sampler(logits, rvq_index, lo, hi):
+        logits = logits.float()
+        v = logits.shape[-1]
+        idx = torch.arange(v, device=logits.device)
+        valid = (idx >= lo) & (idx < hi)
+        logits = torch.where(valid, logits, torch.full_like(logits, NEG))
+        if top_k is not None and int(top_k) < v:
+            k = min(max(int(top_k), 1), v)
+            kth = torch.topk(logits, k, dim=-1).values[..., -1:]
+            logits = torch.where(logits >= kth, logits, torch.full_like(logits, NEG))
+        if temperature and float(temperature) > 0:
+            u = torch.rand(logits.shape, generator=generator, device=logits.device)
+            gumbel = -torch.log(-torch.log(u.clamp(1e-10, 1 - 1e-7)))
+            logits = logits + gumbel * float(temperature)
+        return logits.argmax(dim=-1)
+
+    return sampler
+
+
+def _depthformer_config(c: MagentaRT2Config) -> DepthformerConfig:
+    return DepthformerConfig(
+        encoder_model_dims=c.encoder_model_dims,
+        musiccoca_rvq=c.musiccoca_rvq,
+        musiccoca_per_rvq_vocab=c.musiccoca_per_rvq_vocab,
+        musiccoca_embed_dim=c.musiccoca_embed_dim,
+        regular_num_embeddings_per_channel=list(c.regular_num_embeddings_per_channel),
+        regular_num_channels=c.regular_num_channels,
+        temporal=SpecDims(*c.temporal),
+        depth=SpecDims(*c.depth),
+        temporal_max_past=c.temporal_max_past,
+        depth_max_past=c.depth_max_past,
+        num_sinks=c.num_sinks,
+        num_codebooks=c.num_codebooks,
+        codebook_size=c.codebook_size,
+        num_reserved_tokens=c.num_reserved_tokens,
+        vocab_size=c.vocab_size,
+        soft_cap_logits=c.soft_cap_logits,
+    )
+
+
+def _empty_codec(shapes):
+    return SpectroStreamDecoder({k: torch.zeros(v) for k, v in shapes.items()})
+
+
+class MagentaRT2PreTrainedModel(PreTrainedModel):
+    config_class = MagentaRT2Config
+    base_model_prefix = "magenta_rt2"
+    _no_split_modules = ["TransformerStack", "SpectroStreamDecoder"]
+    main_input_name = "style_tokens"
+
+    def _init_weights(self, module):
+        pass  # weights come from the checkpoint; no random init needed
+
+
+class MagentaRT2ForConditionalGeneration(MagentaRT2PreTrainedModel):
+    """Depthformer LLM + SpectroStream codec. Custom streaming generation."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.depthformer = Depthformer(_depthformer_config(config))
+        shapes = getattr(config, "codec_param_shapes", None)
+        if not shapes:  # dev/package fallback: file next to the module
+            with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "codec_shapes.json")) as f:
+                shapes = json.load(f)
+        self.codec = _empty_codec(shapes)
+        self.register_buffer("quant", torch.zeros(64, config.codebook_size, 256))
+        self.num_musiccoca = config.musiccoca_rvq
+        self.num_notes = config.num_notes
+        self.num_drums = config.num_drums
+        self.codebook_size = config.codebook_size
+        self.num_reserved_tokens = config.num_reserved_tokens
+        self.sample_rate = config.sample_rate
+        self._temporal_step = None
+        self._depth_step = None
+        self.processor = None
+        self.post_init()
+
+    # ---- helpers ----
+    @property
+    def _dt(self):
+        return next(self.depthformer.parameters()).dtype
+
+    @property
+    def _dev(self):
+        return next(self.depthformer.parameters()).device
+
+    def set_processor(self, processor):
+        """Attach a MusicCoCaProcessor so `generate(style="text"|audio)` works."""
+        self.processor = processor
+        return self
+
+    def load_processor(self, repo_id="magenta-torch/magenta-rt-musiccoca-torch", device=None):
+        """Load + attach the MusicCoCa style processor (text/audio -> RVQ tokens)."""
+        from .processing_musiccoca import MusicCoCaProcessor
+        self.processor = MusicCoCaProcessor.from_pretrained(repo_id, device=device or str(self._dev))
+        return self
+
+    # ---- speedups ----
+    def compile_steps(self, dynamic=True, **kwargs):
+        """`torch.compile` the hot per-frame step paths (dynamic shapes for the
+        growing KV cache; one-time warmup). Portable — works on any CUDA GPU,
+        unlike the prebuilt AOTI artifacts, which are GPU-arch-specific."""
+        dec = self.depthformer.decoder
+        dec.temporal_body.step = torch.compile(dec.temporal_body.step, dynamic=dynamic, **kwargs)
+        dec._depth_step_logits = torch.compile(dec._depth_step_logits, dynamic=dynamic, **kwargs)
+        return self
+
+    # ---- AOTI: export your own ahead-of-time graphs (skip runtime compile) ----
+    def export_aoti(self, out_dir):
+        """AOTInductor-compile the per-frame step to `out_dir` (temporal.pt2 +
+        depth.pt2). Run once on your target GPU; the graphs are architecture-
+        specific. Reload with `load_aoti(out_dir)` to generate with no compile-time."""
+        import os
+        from . import aoti
+        os.makedirs(out_dir, exist_ok=True)
+        dec = self.depthformer.decoder
+        torch._inductor.aoti_compile_and_package(
+            aoti.export_temporal(dec), package_path=os.path.join(out_dir, "temporal.pt2"))
+        torch._inductor.aoti_compile_and_package(
+            aoti.export_depth(dec), package_path=os.path.join(out_dir, "depth.pt2"))
+        return out_dir
+
+    def load_aoti(self, out_dir):
+        """Load AOTI step graphs produced by `export_aoti` and use them for generation."""
+        import os
+        t = torch._inductor.aoti_load_package(os.path.join(out_dir, "temporal.pt2"))
+        d = torch._inductor.aoti_load_package(os.path.join(out_dir, "depth.pt2"))
+        return self.apply_compiled(t, d)
+
+    def apply_compiled(self, temporal_step=None, depth_step=None):
+        if temporal_step is not None:
+            self._temporal_step = temporal_step
+        if depth_step is not None:
+            self._depth_step = depth_step
+        return self
+
+    def load_compiled(self, repo_id=None, local_dir=None):
+        from . import aoti
+        t, d = aoti.load_compiled_steps(self.depthformer.decoder, repo_id=repo_id, local_dir=local_dir)
+        return self.apply_compiled(t, d)
+
+    # ---- conditioning ----
+    def _tokenize_style(self, style):
+        if self.processor is None:
+            raise ValueError("No MusicCoCaProcessor attached; pass `style` as a list of "
+                             f"{self.num_musiccoca} RVQ token ids, or call set_processor().")
+        return np.asarray(self.processor.tokenize(self.processor.embed(style))).tolist()
+
+    def _conditioning(self, style_tokens, notes, drums, cfgs):
+        offset = self.num_reserved_tokens + 1
+        vals = list(style_tokens) + list(notes) + list(drums) + list(cfgs)
+        arr = np.array(vals, dtype=np.int64) + offset
+        return torch.from_numpy(arr).view(1, 1, -1).to(self._dev)
+
+    def _resolve_conditioning(self, style, notes, drums, cfg_musiccoca, cfg_notes, cfg_drums):
+        c = self.config
+        if style is None:
+            style_tokens = [-1] * self.num_musiccoca
+        elif isinstance(style, (list, np.ndarray)) and np.asarray(style).ndim == 1 \
+                and np.asarray(style).dtype.kind in "iu" and len(style) == self.num_musiccoca:
+            style_tokens = list(style)
+        else:
+            style_tokens = self._tokenize_style(style)
+        style_tokens = (list(style_tokens) + [-1] * self.num_musiccoca)[:self.num_musiccoca]
+        notes = notes if notes is not None else [-1] * self.num_notes
+        drums = drums if drums is not None else [-1] * self.num_drums
+        cfgs = [
+            discretize_cfg(c.cfg_musiccoca if cfg_musiccoca is None else cfg_musiccoca, 0.2, 40),
+            discretize_cfg(c.cfg_notes if cfg_notes is None else cfg_notes, 0.2, 40),
+            discretize_cfg(c.cfg_drums if cfg_drums is None else cfg_drums, 1.0, 8),
+        ]
+        return self._conditioning(style_tokens, notes, drums, cfgs)
+
+    # ---- codec ----
+    def _decode_stream(self, history, emitted, context=STREAM_DECODE_CONTEXT,
+                       margin=STREAM_DECODE_MARGIN, flush=False):
+        m = 0 if flush else margin
+        Ttot = history.shape[1]
+        emittable = (Ttot - 1) - m
+        avail = emittable - emitted
+        if avail <= 0:
+            return history.new_zeros((1, 0, 2), dtype=self._dt), emitted
+        w0 = max(0, Ttot - (avail + m + context + 1))
+        window = history[:, w0:]
+        codes = convert_from_unique_codes(window, self.codebook_size, self.num_reserved_tokens)
+        emb = codes_to_embeddings(codes, self.quant)
+        wav = self.codec(emb.to(self._dt))
+        end = wav.shape[1] - m * FRAME_SAMPLES
+        new = wav[:, end - avail * FRAME_SAMPLES: end]
+        return new, emitted + avail
+
+    def init_decode_state(self):
+        """Fresh state dict for streaming decode (decode_stream)."""
+        return {}
+
+    def decode_stream(self, new_codes, state):
+        """Incremental codec decode of new token frames [b, t_new, Q] -> audio [b, N, 2].
+        FLOP-optimal stateful streaming (no overlap-save re-decode); bf16-equivalent to
+        _decode_stream/forward, with a 1-frame (40ms) decoder latency. `state` starts as {}."""
+        codes = convert_from_unique_codes(new_codes, self.codebook_size, self.num_reserved_tokens)
+        emb = codes_to_embeddings(codes, self.quant)
+        return self.codec.decode_streaming(emb.to(self._dt), state)
+
+    # ---- forward: one teacher-forced pass (logits), for parity / training hooks ----
+    def forward(self, style_tokens=None, target=None, source=None, **kwargs):
+        """If `source` is given, returns per-frame logits for `target` [b,T,Q].
+        This is the teacher-forced path; for sampling use `generate`/`stream`."""
+        if source is None:
+            cond = self._resolve_conditioning(style_tokens, None, None, None, None, None)
+            source = self.depthformer.encode(cond).to(self._dt)
+        logits = self.depthformer(target, source)
+        return {"logits": logits, "source": source}
+
+    # ---- generation (custom; not GenerationMixin) ----
+    @torch.no_grad()
+    def generate(self, style=None, notes=None, drums=None, cfg_musiccoca=None,
+                 cfg_notes=None, cfg_drums=None, temperature=None, top_k=None,
+                 frames=25, seed=0, state=None, flush=False, return_int16=False):
+        c = self.config
+        temperature = c.temperature if temperature is None else temperature
+        top_k = c.top_k if top_k is None else top_k
+        cond = self._resolve_conditioning(style, notes, drums, cfg_musiccoca, cfg_notes, cfg_drums)
+        source = self.depthformer.encode(cond).to(self._dt)
+        if state is None:
+            dstate = self.depthformer.decoder.init_streaming_f(1, self._dev, self._dt)
+            gen = torch.Generator(device=self._dev).manual_seed(seed)
+            history = torch.zeros((1, 0, c.num_codebooks), dtype=torch.long, device=self._dev)
+            emitted = 0
+        else:
+            dstate, gen, history, emitted = state["dstate"], state["gen"], state["history"], state["emitted"]
+        sampler = make_sampler(temperature, top_k, gen)
+        toks = [self.depthformer.decoder.step_f(
+            dstate, source, sampler=sampler,
+            temporal_step=self._temporal_step, depth_step=self._depth_step) for _ in range(frames)]
+        history = torch.cat([history] + toks, dim=1)
+        audio, emitted = self._decode_stream(history, emitted, flush=flush)
+        new_state = {"dstate": dstate, "gen": gen, "history": history, "emitted": emitted}
+        wav = audio[0].float().cpu().numpy()
+        i16 = _float_to_int16(wav)
+        out = i16 if return_int16 else i16.astype(np.float32) / 32768.0
+        return out, new_state
+
+    @torch.no_grad()
+    def stream(self, control, chunk_frames=10, max_seconds=55.0, seed=0,
+               time_fn=None, sleep_fn=None, notes=None, drums=None):
+        """Continuous generation. `control()` returns {style_tokens, temperature,
+        top_k, cfg_*} read every chunk for mid-stream steering. Yields int16 [N,2]."""
+        import time as _time
+        time_fn = time_fn or _time.time
+        sleep_fn = sleep_fn or _time.sleep
+        c = self.config
+        dev, dt = self._dev, self._dt
+        notes = notes if notes is not None else [-1] * self.num_notes
+        drums = drums if drums is not None else [-1] * self.num_drums
+        dstate = self.depthformer.decoder.init_streaming_f(1, dev, dt)
+        gen = torch.Generator(device=dev).manual_seed(seed)
+        history = torch.zeros((1, 0, c.num_codebooks), dtype=torch.long, device=dev)
+        emitted = 0
+        cur_tokens = None
+        source = None
+        t0 = time_fn()
+        while time_fn() - t0 < max_seconds:
+            ctl = control()
+            if ctl is None:
+                sleep_fn(0.02)
+                continue
+            tokens = ctl["style_tokens"]
+            if tokens != cur_tokens:
+                cur_tokens = tokens
+                cfgs = [discretize_cfg(ctl.get("cfg_musiccoca", c.cfg_musiccoca), 0.2, 40),
+                        discretize_cfg(ctl.get("cfg_notes", c.cfg_notes), 0.2, 40),
+                        discretize_cfg(ctl.get("cfg_drums", c.cfg_drums), 1.0, 8)]
+                cond = self._conditioning((list(tokens) + [-1] * self.num_musiccoca)[:self.num_musiccoca],
+                                          notes, drums, cfgs)
+                source = self.depthformer.encode(cond).to(dt)
+            sampler = make_sampler(ctl.get("temperature", c.temperature), ctl.get("top_k", c.top_k), gen)
+            toks = [self.depthformer.decoder.step_f(
+                dstate, source, sampler=sampler,
+                temporal_step=self._temporal_step, depth_step=self._depth_step) for _ in range(chunk_frames)]
+            history = torch.cat([history] + toks, dim=1)
+            audio, emitted = self._decode_stream(history, emitted)
+            if audio.shape[1] > 0:
+                yield _float_to_int16(audio[0].float().cpu().numpy())
+            ahead = (emitted * FRAME_SAMPLES / SR) - (time_fn() - t0)
+            if ahead > 1.0:
+                sleep_fn(min(ahead - 1.0, 0.5))
+
+
+__all__ = ["MagentaRT2ForConditionalGeneration", "MagentaRT2PreTrainedModel"]

magenta_rt/torch/musiccoca.py

@@ -0,0 +1,160 @@
+# Copyright 2026 Google LLC
+#
+# 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.
+
+"""Pure-PyTorch MusicCoCa style encoder (text path).
+
+The upstream MusicCoCa ships only as TFLite. The text tower and RVQ quantizer
+were converted TFLite -> ONNX -> torch and traced to TorchScript (token-exact
+vs the TFLite reference). Runtime deps: torch + sentencepiece only (no
+ai_edge_litert / tflite). The audio-prompt tower is not yet ported.
+"""
+import os
+
+import numpy as np
+import torch
+
+MAX_TEXT_LEN = 128
+TARGET_SOS_ID = 1
+EMBEDDING_DIM = 768
+RVQ_DEPTH = 12
+DEFAULT_REPO = "magenta-torch/magenta-rt-musiccoca-torch"
+AUDIO_SR = 16000
+CLIP_SAMPLES = 160000  # 10s @ 16kHz
+_MEL_FL, _MEL_HOP, _MEL_NFFT, _PREEMPH = 400, 160, 2048, 0.97
+
+
+class MusicCoCa:
+    """Text -> 768-d style embedding -> 12 RVQ style tokens, all in torch."""
+
+    def __init__(self, repo_id=DEFAULT_REPO, resource_dir=None, device="cpu"):
+        import sentencepiece
+        if resource_dir is not None:
+            te = os.path.join(resource_dir, "text_encoder.pt")
+            q = os.path.join(resource_dir, "quantizer.pt")
+            spm = os.path.join(resource_dir, "spm.model")
+        else:
+            from huggingface_hub import hf_hub_download
+            te = hf_hub_download(repo_id, "text_encoder.pt")
+            q = hf_hub_download(repo_id, "quantizer.pt")
+            spm = hf_hub_download(repo_id, "spm.model")
+        self.device = device
+        self._te = torch.jit.load(te, map_location=device).eval()
+        self._q = torch.jit.load(q, map_location=device).eval()
+        self._sp = sentencepiece.SentencePieceProcessor()
+        self._sp.Load(spm)
+        self.embedding_dim = EMBEDDING_DIM
+        self.rvq_depth = RVQ_DEPTH
+        # Audio tower (mel preprocessor + music_encoder ViT). Lazy.
+        self._me = None
+        self._mel = None
+        self._hann = None
+        self._resource_dir = resource_dir
+        self._repo_id = repo_id
+
+    def _ensure_audio(self):
+        if self._me is not None:
+            return
+        if self._resource_dir is not None:
+            mep = os.path.join(self._resource_dir, "music_encoder.pt")
+            melp = os.path.join(self._resource_dir, "mel_params.npz")
+        else:
+            from huggingface_hub import hf_hub_download
+            mep = hf_hub_download(self._repo_id, "music_encoder.pt")
+            melp = hf_hub_download(self._repo_id, "mel_params.npz")
+        self._me = torch.jit.load(mep, map_location=self.device).eval()
+        d = np.load(melp)
+        self._mel = torch.from_numpy(d["mel"]).float().to(self.device)
+        self._hann = torch.from_numpy(d["hann"]).float().to(self.device)
+
+    def to(self, device):
+        self.device = device
+        self._te = self._te.to(device)
+        self._q = self._q.to(device)
+        if self._me is not None:
+            self._me = self._me.to(device)
+            self._mel = self._mel.to(device)
+            self._hann = self._hann.to(device)
+        return self
+
+    def _log_mel(self, wav):
+        """wav [S] (16kHz mono float) -> log-mel [992,128] (bit-exact vs TFLite)."""
+        x = wav.to(self.device).float()
+        y = x.clone()
+        y[1:] = x[1:] - _PREEMPH * x[:-1]
+        xp = torch.nn.functional.pad(y, (0, _MEL_FL))
+        nf = (xp.shape[0] - _MEL_FL) // _MEL_HOP + 1
+        idx = (torch.arange(_MEL_FL, device=self.device)[None, :]
+               + torch.arange(nf, device=self.device)[:, None] * _MEL_HOP)
+        fr = xp[idx] * self._hann
+        power = torch.fft.rfft(fr, n=_MEL_NFFT, dim=1).abs() ** 2
+        return torch.log(power[:, 1:1025] @ self._mel + 0.001)[:992]
+
+    @torch.no_grad()
+    def embed_audio(self, samples, sample_rate):
+        """samples [S] or [S,C] float; -> [768] style embedding (clips mean-pooled)."""
+        self._ensure_audio()
+        x = np.asarray(samples, np.float32)
+        if x.ndim == 2:
+            x = x.mean(1)
+        if sample_rate != AUDIO_SR:
+            import importlib
+            resampy = importlib.import_module("resampy")  # optional; off the import graph
+            x = resampy.resample(x, sample_rate, AUDIO_SR)
+        # split into 10s clips (pad last), embed each, mean-pool
+        embs = []
+        for s in range(0, max(len(x), 1), CLIP_SAMPLES):
+            clip = x[s:s + CLIP_SAMPLES]
+            if len(clip) < CLIP_SAMPLES:
+                clip = np.pad(clip, (0, CLIP_SAMPLES - len(clip)))
+            mel = self._log_mel(torch.from_numpy(clip))
+            embs.append(self._me(mel[None]).reshape(-1))
+        return torch.stack(embs).mean(0)
+
+    def _tokenize_text(self, text):
+        labels = self._sp.EncodeAsIds(text.lower())[: MAX_TEXT_LEN - 1]
+        ids = [TARGET_SOS_ID] + labels
+        n = len(ids)
+        ids = ids + [0] * (MAX_TEXT_LEN - len(ids))
+        ids_t = torch.tensor([ids], dtype=torch.int32, device=self.device)
+        pad = torch.ones(1, MAX_TEXT_LEN, device=self.device)
+        pad[0, :n] = 0.0
+        return ids_t, pad
+
+    @torch.no_grad()
+    def embed(self, text_or_audio, *args, **kwargs):
+        """Text string OR audio (Waveform / (samples, sr) / np array @ 16kHz)
+        -> [768] style embedding (torch tensor on self.device)."""
+        if isinstance(text_or_audio, str):
+            ids, pad = self._tokenize_text(text_or_audio)
+            return self._te(ids, pad).reshape(-1)
+        # audio: accept a Waveform-like (has .samples/.sample_rate), (samples,sr), or np
+        obj = text_or_audio
+        if hasattr(obj, "samples") and hasattr(obj, "sample_rate"):
+            return self.embed_audio(obj.samples, obj.sample_rate)
+        if isinstance(obj, tuple) and len(obj) == 2:
+            return self.embed_audio(obj[0], obj[1])
+        return self.embed_audio(obj, AUDIO_SR)
+
+    @torch.no_grad()
+    def tokenize(self, embedding):
+        """[768] embedding -> [12] int RVQ tokens (np.int64). Accepts np or torch.
+        For layering, pass a (weighted) mean of several embeddings."""
+        if isinstance(embedding, np.ndarray):
+            embedding = torch.from_numpy(embedding)
+        embedding = embedding.to(self.device).reshape(1, EMBEDDING_DIM).float()
+        return self._q(embedding).reshape(-1).cpu().numpy().astype(np.int64)
+
+    def embed_tokens(self, text):
+        """Convenience: text -> 12 style tokens (list[int])."""
+        return self.tokenize(self.embed(text)).tolist()

magenta_rt/torch/spectrostream.py

@@ -0,0 +1,359 @@
+# Copyright 2026 Google LLC
+#
+# 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.
+
+"""PyTorch port of the SpectroStream decoder (codes -> embeddings -> waveform).
+
+Feature maps are carried as torch NCHW tensors [b, C, T(time), F(freq)] so
+F.conv2d applies directly. JAX conv kernels are stored [kh, kw, cin, cout]
+(HWIO) and permuted to OIHW. Padding replicates sequence_layers semicausal
+(time) / symmetric (freq) for Conv2D, and causal(time)/same(freq) transpose
+conv via input dilation + explicit pad + valid conv.
+"""
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors import safe_open
+
+# Architecture constants for the 40ms 48kHz stereo config.
+RATIOS = ((1, 2), (1, 2), (1, 3), (1, 2), (1, 2), (2, 2), (2, 1))
+CHANNEL_SPLITS = 2
+INPUT_BINS = 5
+INPUT_CHANNELS = 512
+FRAME_LENGTH = 960
+FRAME_STEP = 480
+FFT_LENGTH = 960
+NUM_BINS = 480
+TOTAL_TIME_STRIDE = 4
+DECODER_LOOKAHEAD = 1
+
+
+def _semicausal_pad(k, s, d=1):
+    eff = (k - 1) * d + 1
+    left = max(eff - s, 0)
+    return left, (eff - 1) - left
+
+
+def _sym_freq_pad(kw, sw, dw=1):
+    pad = max((kw - 1) * dw + 1 - sw, 0)
+    return pad // 2, pad - pad // 2
+
+
+def _transpose_pad(k, s, mode):
+    eff = k
+    if mode == "causal":
+        amt = eff + s - 2
+        left = eff - 1
+        return left, amt - left
+    elif mode == "same":
+        amt = eff + s - 2
+        if s > eff - 1:
+            left = eff - 1
+        else:
+            left = int(np.ceil(amt / 2))
+        return left, amt - left
+    raise ValueError(mode)
+
+
+def _hann_window(n):
+    # periodic raised cosine, a=b=0.5 (matches signal.hann_window)
+    even = 1 - n % 2
+    denom = n + even - 1  # periodic -> n
+    count = np.arange(n)
+    return (0.5 - 0.5 * np.cos(2 * np.pi * count / denom)).astype(np.float32)
+
+
+def _inverse_stft_window(frame_length, frame_step):
+    fwd = _hann_window(frame_length)
+    denom = fwd ** 2
+    overlaps = -(-frame_length // frame_step)
+    denom = np.pad(denom, (0, overlaps * frame_step - frame_length))
+    denom = denom.reshape(overlaps, frame_step).sum(0, keepdims=True)
+    denom = np.tile(denom, (overlaps, 1)).reshape(overlaps * frame_step)[:frame_length]
+    return np.where(denom == 0.0, 0.0, fwd / denom).astype(np.float32)
+
+
+def _overlap_and_add(frames, frame_step):
+    """frames: [..., n_frames, frame_length] -> [..., output] (naive, exact)."""
+    *outer, n, fl = frames.shape
+    out_len = (n - 1) * frame_step + fl
+    out = frames.new_zeros(*outer, out_len)
+    for i in range(n):
+        out[..., i * frame_step: i * frame_step + fl] += frames[..., i, :]
+    return out
+
+
+def _dilate2d(x, strides):
+    sh, sw = strides
+    b, c, h, w = x.shape
+    if sh > 1:
+        y = x.new_zeros(b, c, (h - 1) * sh + 1, w)
+        y[:, :, ::sh, :] = x
+        x = y
+        b, c, h, w = x.shape
+    if sw > 1:
+        y = x.new_zeros(b, c, h, (w - 1) * sw + 1)
+        y[:, :, :, ::sw] = x
+        x = y
+    return x
+
+
+def elu(x):
+    return F.elu(x, alpha=1.0)
+
+
+class SpectroStreamDecoder(nn.Module):
+    """Functional decoder driven by a dict of checkpoint tensors."""
+
+    def __init__(self, weights: dict):
+        super().__init__()
+        # weights: name -> torch tensor (kernels in HWIO; conv biases 1d).
+        self.w = {k: nn.Parameter(v, requires_grad=False) for k, v in weights.items()}
+        self.w = nn.ParameterDict({k.replace("/", "__"): v for k, v in self.w.items()})
+        self.register_buffer("inv_window", torch.from_numpy(
+            _inverse_stft_window(FRAME_LENGTH, FRAME_STEP)))
+
+    def _g(self, name):
+        return self.w[name.replace("/", "__")]
+
+    # ---- conv primitives ----
+    def _conv1x1(self, x, prefix):
+        w = self._g(prefix + "/conv/kernel")  # [1,1,cin,cout]
+        b = self._g(prefix + "/conv/bias")
+        wk = w.permute(3, 2, 0, 1).to(x.dtype)
+        return F.conv2d(x, wk, bias=b.to(x.dtype))
+
+    def _conv2d(self, x, prefix, kh, kw, strides=(1, 1), dil=(1, 1)):
+        w = self._g(prefix + "/conv/kernel")
+        b = self._g(prefix + "/conv/bias")
+        pt = _semicausal_pad(kh, strides[0], dil[0])
+        pf = _sym_freq_pad(kw, strides[1], dil[1])
+        x = F.pad(x, (pf[0], pf[1], pt[0], pt[1]))
+        wk = w.permute(3, 2, 0, 1).to(x.dtype)
+        return F.conv2d(x, wk, bias=b.to(x.dtype), stride=strides, dilation=dil)
+
+    def _conv_transpose(self, x, prefix, kh, kw, strides):
+        w = self._g(prefix + "/conv/kernel")  # [kh,kw,cin,cout]
+        b = self._g(prefix + "/conv/bias")
+        x = _dilate2d(x, strides)
+        pt = _transpose_pad(kh, strides[0], "causal")
+        pf = _transpose_pad(kw, strides[1], "same")
+        x = F.pad(x, (pf[0], pf[1], pt[0], pt[1]))
+        wk = w.permute(3, 2, 0, 1).to(x.dtype)
+        return F.conv2d(x, wk, bias=b.to(x.dtype), stride=1)
+
+    def _upsample(self, x, strides):
+        if strides[0] > 1:
+            x = x.repeat_interleave(strides[0], dim=2)
+        if strides[1] > 1:
+            x = x.repeat_interleave(strides[1], dim=3)
+        return x
+
+    def _residual_unit(self, x, prefix, strides, transposed_resample, kt):
+        """act->[convT or conv3x3_a]->act->conv3x3 + shortcut."""
+        inp = x
+        y = elu(x)
+        if transposed_resample:
+            kh, kw = kt
+            y = self._conv_transpose(y, prefix + "/conv2dtranspose_%dx%d" % (kh, kw), kh, kw, strides)
+        else:
+            y = self._conv2d(y, prefix + "/conv2d_3x3_a", 3, 3)
+        y = elu(y)
+        y = self._conv2d(y, prefix + "/conv2d_3x3", 3, 3)
+        # shortcut
+        sc = inp
+        has_conv = (prefix + "/shortcut_layer/conv1x1/conv/kernel").replace("/", "__") in self.w
+        if has_conv:
+            sc = self._conv1x1(sc, prefix + "/shortcut_layer/conv1x1")
+        if strides != (1, 1):
+            sc = self._upsample(sc, strides)
+        return y + sc
+
+    def decode_embeddings(self, emb):
+        """emb: [b,t,256] -> spectrogram feature map [b,4,T,480] (NCHW)."""
+        b, t, _ = emb.shape
+        x = emb.permute(0, 2, 1).unsqueeze(-1)  # [b,256,t,1]
+        # input_layer residual
+        main = self._conv1x1(x, "input_layer/conv1x1_first")
+        sc = self._conv1x1(x, "input_layer/shortcut_layer/conv1x1_b1")
+        sc = elu(sc)
+        sc = self._conv1x1(sc, "input_layer/shortcut_layer/conv1x1_b2")
+        x = main + sc  # [b,2560,t,1]
+        # reshape (1,2560)->(5,512): [b,2560,t,1]->[b,5,512,t]->[b,512,t,5]
+        x = x.squeeze(-1).view(b, INPUT_BINS, INPUT_CHANNELS, t).permute(0, 2, 3, 1)
+        # input_layers_residual_unit (stride1)
+        x = self._residual_unit(x, "input_layers_residual_unit", (1, 1), False, None)
+        # decoder_0
+        rev = RATIOS[::-1]
+        kt0 = (max(3, 2 * rev[0][0]), max(3, 2 * rev[0][1]))
+        x = self._residual_unit(x, "decoder_0", rev[0], True, kt0)
+        # ParallelChannels(2): split channels, shared decoder_1..6 + output, concat
+        groups = torch.chunk(x, CHANNEL_SPLITS, dim=1)
+        outs = []
+        for g in groups:
+            h = g
+            for i in range(1, len(RATIOS)):
+                s = rev[i]
+                kt = (max(3, 2 * s[0]), max(3, 2 * s[1]))
+                h = self._residual_unit(h, f"decoder_{i}", s, True, kt)
+            # output_layer: act -> conv7x7 (->2)
+            h = elu(h)
+            h = self._conv2d(h, "output_layer/base_conv_last", 7, 7)
+            outs.append(h)
+        x = torch.cat(outs, dim=1)  # [b,4,T,480]
+        # lookahead trim
+        trim = DECODER_LOOKAHEAD * TOTAL_TIME_STRIDE
+        if trim:
+            x = x[:, :, trim:, :]
+        return x
+
+    def forward(self, emb):
+        x = self.decode_embeddings(emb)
+        return self._istft(x)
+
+    # ---- streaming decode (per-frame, stateful) — bit-exact-in-bf16 vs forward,
+    #      FLOP-optimal (no overlap-save re-decode). state = mutable dict of caches. ----
+    def _s_conv2d(self, x, prefix, kh, kw, st, key, strides=(1, 1), dil=(1, 1)):
+        pt = _semicausal_pad(kh, strides[0], dil[0])
+        pf = _sym_freq_pad(kw, strides[1], dil[1])
+        c = st.get(key)
+        if c is None:
+            c = x.new_zeros(x.shape[0], x.shape[1], pt[0], x.shape[3])
+        xc = torch.cat([c, x], dim=2)
+        st[key] = xc[:, :, xc.shape[2] - pt[0]:, :] if pt[0] > 0 else c
+        xp = F.pad(xc, (pf[0], pf[1], 0, pt[1]))
+        w = self._g(prefix + "/conv/kernel"); b = self._g(prefix + "/conv/bias")
+        return F.conv2d(xp, w.permute(3, 2, 0, 1).to(x.dtype), bias=b.to(x.dtype),
+                        stride=strides, dilation=dil)
+
+    def _s_conv_transpose(self, x, prefix, kh, kw, strides, st, key):
+        sh, sw = strides
+        pt = _transpose_pad(kh, sh, "causal"); pf = _transpose_pad(kw, sw, "same")
+        ctx = (pt[0] + sh - 1) // sh + 1
+        c = st.get(key)
+        if c is None:
+            c = x.new_zeros(x.shape[0], x.shape[1], ctx, x.shape[3])
+        C = x.shape[2]
+        xc = torch.cat([c, x], dim=2)
+        st[key] = xc[:, :, xc.shape[2] - ctx:, :]
+        xp = F.pad(_dilate2d(xc, strides), (pf[0], pf[1], pt[0], pt[1]))
+        w = self._g(prefix + "/conv/kernel"); b = self._g(prefix + "/conv/bias")
+        out = F.conv2d(xp, w.permute(3, 2, 0, 1).to(x.dtype), bias=b.to(x.dtype), stride=1)
+        return out[:, :, out.shape[2] - C * sh:, :]
+
+    def _s_resunit(self, x, prefix, strides, transposed, kt, st, key):
+        inp = x; y = elu(x)
+        if transposed:
+            kh, kw = kt
+            y = self._s_conv_transpose(y, prefix + "/conv2dtranspose_%dx%d" % (kh, kw), kh, kw, strides, st, key + "/ct")
+        else:
+            y = self._s_conv2d(y, prefix + "/conv2d_3x3_a", 3, 3, st, key + "/a")
+        y = elu(y)
+        y = self._s_conv2d(y, prefix + "/conv2d_3x3", 3, 3, st, key + "/b")
+        sc = inp
+        if (prefix + "/shortcut_layer/conv1x1/conv/kernel").replace("/", "__") in self.w:
+            sc = self._conv1x1(sc, prefix + "/shortcut_layer/conv1x1")
+        if strides != (1, 1):
+            sc = self._upsample(sc, strides)
+        return y + sc
+
+    def _s_decode_emb(self, emb_new, st):
+        b, t, _ = emb_new.shape
+        x = emb_new.permute(0, 2, 1).unsqueeze(-1)
+        main = self._conv1x1(x, "input_layer/conv1x1_first")
+        sc = self._conv1x1(x, "input_layer/shortcut_layer/conv1x1_b1"); sc = elu(sc)
+        sc = self._conv1x1(sc, "input_layer/shortcut_layer/conv1x1_b2")
+        x = (main + sc).squeeze(-1).view(b, INPUT_BINS, INPUT_CHANNELS, t).permute(0, 2, 3, 1)
+        x = self._s_resunit(x, "input_layers_residual_unit", (1, 1), False, None, st, "ilru")
+        rev = RATIOS[::-1]
+        x = self._s_resunit(x, "decoder_0", rev[0], True, (max(3, 2 * rev[0][0]), max(3, 2 * rev[0][1])), st, "d0")
+        outs = []
+        for gi, g in enumerate(torch.chunk(x, CHANNEL_SPLITS, dim=1)):
+            h = g
+            for i in range(1, len(RATIOS)):
+                s = rev[i]
+                h = self._s_resunit(h, f"decoder_{i}", s, True, (max(3, 2 * s[0]), max(3, 2 * s[1])), st, f"g{gi}/d{i}")
+            h = elu(h)
+            h = self._s_conv2d(h, "output_layer/base_conv_last", 7, 7, st, f"g{gi}/out")
+            outs.append(h)
+        return torch.cat(outs, dim=1)
+
+    def _s_istft(self, xnew, st):
+        v = xnew.permute(0, 2, 3, 1).contiguous(); b, T, nb, nc = v.shape
+        if T == 0:
+            return xnew.new_zeros(b, 0, 2)
+        v = F.pad(v, (0, 0, 0, 1)).float()
+        comp = torch.view_as_complex(v.view(b, T, 481, nc // 2, 2).contiguous())
+        frames = torch.fft.irfft(comp, n=FFT_LENGTH, dim=2) * self.inv_window.view(1, 1, FRAME_LENGTH, 1)
+        fr = frames.permute(0, 3, 1, 2)
+        tail = st.get("_tail")
+        if tail is None:
+            tail = fr.new_zeros(b, 2, FRAME_STEP)
+        emits = []
+        for i in range(T):
+            f = fr[:, :, i, :]; emits.append(tail + f[:, :, :FRAME_STEP]); tail = f[:, :, FRAME_STEP:]
+        st["_tail"] = tail
+        return torch.cat(emits, dim=2).permute(0, 2, 1)
+
+    def decode_streaming(self, emb_new, state):
+        """Incremental decode. `state` is a mutable dict (start with {}). Returns the
+        newly-available audio [b, N, 2] for `emb_new` [b, t_new, 256], carrying overlap
+        + per-layer conv state across calls. Output == forward(full_emb), 1 frame latency."""
+        x = self._s_decode_emb(emb_new, state)
+        wm = state.get("_warm", DECODER_LOOKAHEAD * TOTAL_TIME_STRIDE)
+        if wm > 0:
+            d = min(wm, x.shape[2]); x = x[:, :, d:, :]; state["_warm"] = wm - d
+        return self._s_istft(x, state)
+
+    def _istft(self, x):
+        v = x.permute(0, 2, 3, 1).contiguous()      # [b,T,480,4]
+        b, T, nb, nc = v.shape
+        # pad freq (axis=2) 480 -> 481 (keep_dc: pad right)
+        v = F.pad(v, (0, 0, 0, 1))                   # pad dim=2 right by 1 -> [b,T,481,4]
+        v = v.float()
+        comp = torch.view_as_complex(v.view(b, T, 481, nc // 2, 2).contiguous())  # [b,T,481,2]
+        frames = torch.fft.irfft(comp, n=FFT_LENGTH, dim=2)  # [b,T,960,2]
+        frames = frames * self.inv_window.view(1, 1, FRAME_LENGTH, 1)
+        # overlap-add over (T,960); move channel out: [b,2,T,960]
+        fr = frames.permute(0, 3, 1, 2)              # [b,2,T,960]
+        wav = _overlap_and_add(fr, FRAME_STEP)       # [b,2,samples]
+        trim = max(FRAME_LENGTH - FRAME_STEP, 0)
+        if trim:
+            wav = wav[..., :-trim]
+        return wav.permute(0, 2, 1)                  # [b,samples,2]
+
+
+def codes_to_embeddings(codes, quantizer_embedding):
+    """codes: [b,t,Q] long (0..1023); quantizer_embedding: [64,1024,256]. Sum over levels."""
+    Q = codes.shape[-1]
+    out = None
+    for i in range(Q):
+        e = quantizer_embedding[i][codes[..., i]]  # [b,t,256]
+        out = e if out is None else out + e
+    return out
+
+
+def load_spectrostream_decoder(checkpoint_path, dtype=torch.float32, prefix="params/soundstream"):
+    """Load decoder + quantizer tensors from a safetensors file."""
+    dec_weights = {}
+    quant = None
+    with safe_open(str(checkpoint_path), "numpy") as f:
+        for k in f.keys():
+            if k.startswith(prefix + "/decoder/"):
+                name = k[len(prefix + "/decoder/"):]
+                dec_weights[name] = torch.from_numpy(np.asarray(f.get_tensor(k))).to(dtype)
+            elif k.startswith(prefix + "/quantizer/embedding"):
+                quant = torch.from_numpy(np.asarray(f.get_tensor(k))).to(dtype)
+    dec = SpectroStreamDecoder(dec_weights)
+    return dec, quant

magenta_rt/torch/system.py

@@ -0,0 +1,298 @@
+# Copyright 2026 Google LLC
+#
+# 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.
+
+"""Pure-PyTorch Magenta RealTime 2 streaming system.
+
+Generation uses the torch Depthformer (verified vs JAX) for autoregressive
+token sampling and the torch SpectroStream decoder (verified vs JAX) for audio.
+MusicCoCa style embedding/tokenization is reused from the framework-agnostic
+TFLite component (same as the upstream JAX/MLX paths).
+"""
+import numpy as np
+import torch
+
+from .. import paths
+from .depthformer import Depthformer, config_for
+from .weights import load_depthformer
+from .spectrostream import load_spectrostream_decoder, codes_to_embeddings
+
+NUM_RESERVED_TOKENS = 6
+SR = 48000                    # output sample rate
+FRAME_SAMPLES = 1920          # 48kHz samples per 40ms codec frame
+STREAM_DECODE_CONTEXT = 16    # token frames of left context for overlap-save decode
+STREAM_DECODE_MARGIN = 2      # token frames held back (right context: lookahead + ISTFT overlap)
+_CHECKPOINTS = {"mrt2_small": "mrt2_small.safetensors", "mrt2_base": "mrt2_base.safetensors"}
+
+
+def discretize_cfg(value, step, max_bin):
+    clamped = max(-1.0, min(7.0, value))
+    return max(0, min(max_bin, int(round((clamped - (-1.0)) / step))))
+
+
+def convert_from_unique_codes(tokens, codebook_size=1024):
+    return (tokens - NUM_RESERVED_TOKENS) % codebook_size
+
+
+def _float_to_int16(samples, gain=0.5):
+    samples = np.clip(gain * samples, -1, 1)
+    samples = np.round((np.iinfo(np.int16).max + 0.5) * samples - 0.5)
+    return samples.astype(np.int16)
+
+
+def make_sampler(temperature, top_k, generator):
+    # Matches jax _sample_categorical_with_temperature: valid-range mask, top-k,
+    # then gumbel-max (logits + gumbel * temperature, argmax). Uses a finite
+    # large-negative for masking (a la get_large_negative_number).
+    NEG = -1e9
+
+    def sampler(logits, rvq_index, lo, hi):
+        logits = logits.float()
+        v = logits.shape[-1]
+        idx = torch.arange(v, device=logits.device)
+        valid = (idx >= lo) & (idx < hi)
+        logits = torch.where(valid, logits, torch.full_like(logits, NEG))
+        if top_k is not None and int(top_k) < v:
+            k = min(max(int(top_k), 1), v)
+            kth = torch.topk(logits, k, dim=-1).values[..., -1:]
+            logits = torch.where(logits >= kth, logits, torch.full_like(logits, NEG))
+        if temperature and float(temperature) > 0:
+            u = torch.rand(logits.shape, generator=generator, device=logits.device)
+            gumbel = -torch.log(-torch.log(u.clamp(1e-10, 1 - 1e-7)))
+            logits = logits + gumbel * float(temperature)
+        return logits.argmax(dim=-1)  # [b,1]
+    return sampler
+
+
+class MagentaRT2:
+    def __init__(self, size="mrt2_small", device="cpu", dtype=torch.float32,
+                 temperature=1.3, top_k=40, cfg_musiccoca=3.0, cfg_notes=1.0,
+                 cfg_drums=1.0, style_model=None, compile=False):
+        self.size = size
+        self.device = device
+        self.dtype = dtype
+        self.temperature = temperature
+        self.top_k = top_k
+        self.cfg_musiccoca = cfg_musiccoca
+        self.cfg_notes = cfg_notes
+        self.cfg_drums = cfg_drums
+        self._style_model = style_model
+
+        cfg = config_for(size)
+        self.cfg = cfg
+        self.model = Depthformer(cfg).eval().to(device)
+        ckpt = paths.checkpoints_dir() / _CHECKPOINTS[size]
+        load_depthformer(self.model, ckpt, dtype=dtype)
+        self.model.to(dtype)
+        self.dec, quant = load_spectrostream_decoder(ckpt, dtype=dtype)
+        self.dec = self.dec.eval().to(device)
+        self.quant = quant.to(device)
+
+        # AOTI-compiled per-frame step callables (None = eager). Set via
+        # load_compiled() or apply_compiled().
+        self._temporal_step = None
+        self._depth_step = None
+
+        if compile:
+            # Compile the two hot per-frame step paths (dynamic shapes for the
+            # growing KV caches). One-time warmup; ~1.8x and past real-time.
+            dec = self.model.decoder
+            dec.temporal_body.step = torch.compile(dec.temporal_body.step, dynamic=True)
+            dec._depth_step_logits = torch.compile(dec._depth_step_logits, dynamic=True)
+
+        self.num_musiccoca = cfg.musiccoca_rvq         # 12
+        self.num_notes = 128
+        self.num_drums = 1
+        self.num_cfg = 3
+        self.num_channels = self.num_musiccoca + self.num_notes + self.num_drums + self.num_cfg
+        self.sample_rate = 48000
+        self.codebook_size = cfg.codebook_size
+
+    def apply_compiled(self, temporal_step=None, depth_step=None):
+        """Wire AOTI-compiled per-frame step callables into generation."""
+        if temporal_step is not None:
+            self._temporal_step = temporal_step
+        if depth_step is not None:
+            self._depth_step = depth_step
+
+    def load_compiled(self, repo_id=None, local_dir=None):
+        """Load AOTI artifacts (weight-less, compiled on matching GPU arch) and
+        bind them to this model's weights via aokit."""
+        from . import aoti
+        t, d = aoti.load_compiled_steps(self.model.decoder, repo_id=repo_id, local_dir=local_dir)
+        self.apply_compiled(t, d)
+        return self
+
+    # ---- style ----
+    @property
+    def style_model(self):
+        if self._style_model is None:
+            from .. import musiccoca
+            self._style_model = musiccoca.MusicCoCa()
+        return self._style_model
+
+    def embed_style(self, text_or_audio, **kw):
+        return self.style_model.embed(text_or_audio, **kw)
+
+    def tokenize_style(self, embedding):
+        return self.style_model.tokenize(embedding)
+
+    # ---- conditioning ----
+    def _conditioning(self, style_tokens, notes, drums, cfgs):
+        offset = NUM_RESERVED_TOKENS + 1
+        vals = list(style_tokens) + list(notes) + list(drums) + list(cfgs)
+        arr = np.array(vals, dtype=np.int64) + offset
+        return torch.from_numpy(arr).view(1, 1, -1).to(self.device)
+
+    def _resolve_conditioning(self, style, notes, drums, cfg_musiccoca, cfg_notes, cfg_drums):
+        if style is None:
+            style_tokens = [-1] * self.num_musiccoca
+        elif isinstance(style, (list, np.ndarray)) and len(np.asarray(style).shape) == 1 and np.asarray(style).dtype.kind in "iu" and len(style) == self.num_musiccoca:
+            style_tokens = list(style)
+        else:
+            style_tokens = self.tokenize_style(style).tolist()
+        style_tokens = (style_tokens + [-1] * self.num_musiccoca)[:self.num_musiccoca]
+        notes = notes if notes is not None else [-1] * self.num_notes
+        drums = drums if drums is not None else [-1] * self.num_drums
+        cfgs = [
+            discretize_cfg(self.cfg_musiccoca if cfg_musiccoca is None else cfg_musiccoca, 0.2, 40),
+            discretize_cfg(self.cfg_notes if cfg_notes is None else cfg_notes, 0.2, 40),
+            discretize_cfg(self.cfg_drums if cfg_drums is None else cfg_drums, 1.0, 8),
+        ]
+        return self._conditioning(style_tokens, notes, drums, cfgs)
+
+    def _decode_stream(self, history, emitted, context=STREAM_DECODE_CONTEXT,
+                       margin=STREAM_DECODE_MARGIN, flush=False):
+        """Overlap-save: decode recent token context, emit only the safe new tail.
+
+        The SpectroStream decoder has a small left receptive field and ~1 frame
+        of lookahead + ISTFT overlap on the right. Decoding `[left-context ...
+        new ... right-margin]` and emitting the interior gives output identical
+        to a fully stateful streaming codec. `margin` frames are held back until
+        their future context exists (flush=True emits them at stream end).
+        Returns (new_samples [1,N,2], new_emitted)."""
+        m = 0 if flush else margin
+        Ttot = history.shape[1]
+        emittable = (Ttot - 1) - m
+        avail = emittable - emitted
+        if avail <= 0:
+            return history.new_zeros((1, 0, 2), dtype=self.dtype), emitted
+        w0 = max(0, Ttot - (avail + m + context + 1))
+        window = history[:, w0:]
+        codes = convert_from_unique_codes(window, self.codebook_size)
+        emb = codes_to_embeddings(codes, self.quant)
+        wav = self.dec(emb.to(self.dtype))            # [1, (len(window)-1)*1920, 2]
+        end = wav.shape[1] - m * FRAME_SAMPLES
+        new = wav[:, end - avail * FRAME_SAMPLES: end]
+        return new, emitted + avail
+
+    def init_decode_state(self):
+        """Fresh state dict for streaming decode (decode_stream)."""
+        return {}
+
+    @torch.no_grad()
+    def decode_stream(self, new_codes, state):
+        """Incremental codec decode of new token frames [b, t_new, Q] -> audio [b, N, 2].
+        FLOP-optimal stateful streaming (no overlap-save re-decode); bf16-equivalent to
+        _decode_stream, with a 1-frame (40ms) decoder latency. `state` starts as {}."""
+        codes = convert_from_unique_codes(new_codes, self.codebook_size)
+        emb = codes_to_embeddings(codes, self.quant)
+        return self.dec.decode_streaming(emb.to(self.dtype), state)
+
+    @torch.no_grad()
+    def stream_session(self, control, chunk_frames=10, max_seconds=55.0,
+                       seed=0, time_fn=None, sleep_fn=None, notes=None, drums=None):
+        """Continuous generation for an interactive session. `control()` returns a
+        dict {style_tokens, temperature, top_k, cfg_*} read every chunk, so the
+        prompt can change mid-stream. Yields int16 [N,2] audio chunks. Keeps LLM
+        state across chunks; re-encodes the conditioning `source` when style changes.
+        Paces to ~real-time so steering stays responsive."""
+        import time as _time
+        time_fn = time_fn or _time.time
+        sleep_fn = sleep_fn or _time.sleep
+        dev, dt = self.device, self.dtype
+        notes = notes if notes is not None else [-1] * self.num_notes
+        drums = drums if drums is not None else [-1] * self.num_drums
+        dstate = self.model.decoder.init_streaming_f(1, dev, dt)
+        gen = torch.Generator(device=dev).manual_seed(seed)
+        history = torch.zeros((1, 0, self.cfg.num_codebooks), dtype=torch.long, device=dev)
+        emitted = 0
+        cur_tokens = None
+        source = None
+        t0 = time_fn()
+        while time_fn() - t0 < max_seconds:
+            c = control()
+            if c is None:
+                sleep_fn(0.02)
+                continue
+            tokens = c["style_tokens"]
+            if tokens != cur_tokens:
+                cur_tokens = tokens
+                cfgs = [discretize_cfg(c.get("cfg_musiccoca", self.cfg_musiccoca), 0.2, 40),
+                        discretize_cfg(c.get("cfg_notes", self.cfg_notes), 0.2, 40),
+                        discretize_cfg(c.get("cfg_drums", self.cfg_drums), 1.0, 8)]
+                cond = self._conditioning((list(tokens) + [-1] * self.num_musiccoca)[:self.num_musiccoca],
+                                          notes, drums, cfgs)
+                source = self.model.encode(cond).to(dt)
+            sampler = make_sampler(c.get("temperature", self.temperature), c.get("top_k", self.top_k), gen)
+            toks = []
+            for _ in range(chunk_frames):
+                toks.append(self.model.decoder.step_f(
+                    dstate, source, sampler=sampler,
+                    temporal_step=self._temporal_step, depth_step=self._depth_step))
+            history = torch.cat([history] + toks, dim=1)
+            audio, emitted = self._decode_stream(history, emitted)
+            if audio.shape[1] > 0:
+                yield _float_to_int16(audio[0].float().cpu().numpy())
+            # pace: keep generated audio ~1s ahead of wall-clock
+            ahead = (emitted * FRAME_SAMPLES / SR) - (time_fn() - t0)
+            if ahead > 1.0:
+                sleep_fn(min(ahead - 1.0, 0.5))
+
+    @torch.no_grad()
+    def generate(self, style=None, notes=None, drums=None, cfg_musiccoca=None,
+                 cfg_notes=None, cfg_drums=None, temperature=None, top_k=None,
+                 frames=25, seed=0, state=None, flush=False, return_int16=False):
+        """Generate `frames` of audio. Pass the returned `state` back in to
+        continue seamlessly (continuous/live generation); conditioning args may
+        change between calls to steer the stream. Audio is emitted incrementally
+        (only the newly-available chunk is returned each call). Set flush=True on
+        the final call to emit the held-back tail frames."""
+        temperature = self.temperature if temperature is None else temperature
+        top_k = self.top_k if top_k is None else top_k
+        cond = self._resolve_conditioning(style, notes, drums, cfg_musiccoca, cfg_notes, cfg_drums)
+        source = self.model.encode(cond).to(self.dtype)  # constant per frame this call
+
+        if state is None:
+            dstate = self.model.decoder.init_streaming_f(1, self.device, self.dtype)
+            gen = torch.Generator(device=self.device).manual_seed(seed)
+            history = torch.zeros((1, 0, self.cfg.num_codebooks), dtype=torch.long, device=self.device)
+            emitted = 0
+        else:
+            dstate, gen, history, emitted = state["dstate"], state["gen"], state["history"], state["emitted"]
+
+        sampler = make_sampler(temperature, top_k, gen)
+        toks = []
+        for _ in range(frames):
+            frame = self.model.decoder.step_f(
+                dstate, source, sampler=sampler,
+                temporal_step=self._temporal_step, depth_step=self._depth_step)
+            toks.append(frame)
+        history = torch.cat([history] + toks, dim=1)
+        audio, emitted = self._decode_stream(history, emitted, flush=flush)
+
+        new_state = {"dstate": dstate, "gen": gen, "history": history, "emitted": emitted}
+        wav = audio[0].float().cpu().numpy()
+        i16 = _float_to_int16(wav)  # 0.5 gain + clip, pointwise (chunk-safe)
+        out = i16 if return_int16 else i16.astype(np.float32) / 32768.0
+        return out, new_state

magenta_rt/torch/weights.py

@@ -0,0 +1,111 @@
+# Copyright 2026 Google LLC
+#
+# 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.
+
+"""Map JAX/Linen safetensors checkpoint keys onto the torch Depthformer."""
+import re
+
+import numpy as np
+import torch
+from safetensors import safe_open
+
+
+def _jax_key_to_torch(k: str):
+    """Translate a 'params/depthformer/...' key to a torch named_parameter path,
+    or return None to skip (e.g. soundstream)."""
+    if not k.startswith("params/depthformer/"):
+        return None
+    s = k[len("params/depthformer/"):]
+
+    # Encoder (conditioning embedders).
+    if s.startswith("encoder/"):
+        if s == "encoder/body/encoder_ln/scale":
+            return "encoder.encoder_ln.scale"
+        if s == "encoder/body/encoder_ln/bias":
+            return "encoder.encoder_ln.bias"
+        if s.endswith("mulan_dequantizer/embedding"):
+            return "encoder.mulan_dequantizer"
+        if s.endswith("mulan_embedder/depth_input_adapter/kernel"):
+            return "encoder.mulan_adapter.kernel"
+        if s.endswith("regular_embedder/embedding"):
+            return "encoder.regular_embedding"
+        return None
+
+    # Decoder.
+    if s == "decoder/decoder_embedding/embedding/embedding":
+        return "decoder.embedding"
+    if s == "decoder/depth_body/depth_input_adapter/kernel":
+        return "decoder.depth_input_adapter.kernel"
+    if s.startswith("decoder/depth_body/final_ln/"):
+        return "decoder.final_ln." + s.split("/")[-1]
+    if s.startswith("decoder/depth_body/to_logits/"):
+        return "decoder.to_logits." + s.split("/")[-1]
+
+    m = re.match(r"decoder/(temporal_body|depth_body)/transformer/x_layers_(\d+)/(.*)", s)
+    if m:
+        body, i, rest = m.group(1), int(m.group(2)), m.group(3)
+        prefix = f"decoder.{body}.layers.{i}."
+        return prefix + _layer_subkey(rest)
+
+    return None
+
+
+def _layer_subkey(rest: str):
+    # rest like 'self_attention/attention/query_projection/kernel'
+    parts = rest.split("/")
+    sub = parts[0]  # self_attention | cross_attention | ffn
+    tail = parts[1:]
+    if sub in ("self_attention", "cross_attention"):
+        if tail[0] == "attention":
+            name = tail[1]
+            if name.endswith("_projection"):  # query/key/value_projection/kernel
+                return f"{sub}.attention.{name}_kernel"
+            # per_dim_scale, sink_key_embeddings, sink_value_embeddings
+            return f"{sub}.attention.{name}"
+        if tail[0] == "output_projection":  # output_projection/kernel
+            return f"{sub}.output_projection_kernel"
+        if tail[0] in ("pre_norm", "post_norm"):
+            return f"{sub}.{tail[0]}.scale"
+    if sub == "ffn":
+        if tail[0] in ("ffn_layer1", "ffn_layer2"):
+            return f"ffn.{tail[0]}.{tail[1]}"  # kernel | bias
+        if tail[0] in ("pre_norm", "post_norm"):
+            return f"ffn.{tail[0]}.scale"
+    raise KeyError(f"unhandled layer subkey: {rest}")
+
+
+def load_depthformer(model, checkpoint_path, dtype=torch.float32, strict=True, verbose=False):
+    """Load checkpoint into a torch Depthformer (module with .encoder/.decoder)."""
+    params = dict(model.named_parameters())
+    seen = set()
+    with safe_open(str(checkpoint_path), "numpy") as f:
+        for k in f.keys():
+            tname = _jax_key_to_torch(k)
+            if tname is None:
+                continue
+            if tname not in params:
+                raise KeyError(f"{k} -> {tname} not found in model")
+            arr = f.get_tensor(k)
+            t = torch.from_numpy(np.asarray(arr)).to(dtype)
+            p = params[tname]
+            if tuple(p.shape) != tuple(t.shape):
+                raise ValueError(f"shape mismatch {tname}: model {tuple(p.shape)} ckpt {tuple(t.shape)} ({k})")
+            with torch.no_grad():
+                p.copy_(t)
+            seen.add(tname)
+    missing = [n for n in params if n not in seen]
+    if verbose:
+        print(f"loaded {len(seen)} params, {len(missing)} missing")
+    if strict and missing:
+        raise RuntimeError(f"missing params not loaded: {missing[:20]} ... ({len(missing)} total)")
+    return model

model.json

@@ -0,0 +1,16 @@
+{
+  "name": "Magenta RealTime 2",
+  "model_size": "mrt2_small",
+  "parameters": 230000000,
+  "source_repository": "https://github.com/magenta/magenta-realtime",
+  "model_repository": "google/magenta-realtime-2",
+  "checkpoint": "checkpoints/mrt2_small.safetensors",
+  "model_license": "CC BY 4.0",
+  "runtime": {
+    "backend": "PyTorch",
+    "source": "magenta-community/magenta-rt-jam",
+    "revision": "dfd4eb65e8ebf7eaa9be263d9937ddbee1666262",
+    "license": "Apache-2.0",
+    "aoti_repository": "magenta-torch/magenta-rt-aoti-small"
+  }
+}

requirements.txt

@@ -0,0 +1,9 @@
+gradio==5.28.0
+torch==2.11.0
+git+https://github.com/TEAMuP-dev/pyharp.git@develop
+huggingface-hub>=0.30
+numpy>=1.26
+resampy>=0.4
+safetensors>=0.4
+sentencepiece>=0.2
+soundfile>=0.12