magenta-rt

--- magenta_rt-2.0.2/magenta_rt/__init__.py (excerpt) ---
# 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.

"""Magenta RealTime 2 — real-time streaming audio generation."""

# Activate vendored sequence_layers if it is not installed as a package.
from magenta_rt._vendor import _vendor_hook
_vendor_hook.install()
del _vendor_hook

__version__ = "2.0.2"
__all__ = ["MagentaRT2Jax", "MagentaRT2Mlx", "MagentaRT2Mlxfn"]


def __getattr__(name):
  if name == "MagentaRT2Jax":
    try:
      from magenta_rt.jax.system import MagentaRT2System
    except ImportError as e:
      raise ImportError(
          "MagentaRT2Jax requires JAX dependencies. "
          "Install them with: pip install magenta-rt[jax]"
      ) from e
    return MagentaRT2System
  if name == "MagentaRT2Mlx":
    try:
      from magenta_rt.mlx.system import MagentaRT2System
    except ImportError as e:
      raise ImportError(
          "MagentaRT2Mlx requires MLX dependencies. "
          "Install them with: pip install magenta-rt[mlx]"
      ) from e
    re

--- magenta_rt-2.0.2/magenta_rt/paths.py (excerpt) ---
# 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."""
    g

--- magenta_rt-2.0.2/magenta_rt/mlx/export.py (excerpt) ---
# 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.

import os
import os.path
from pathlib import Path

import mlx.core as mx
import numpy as np

import sequence_layers.mlx as sl
from sequence_layers.mlx import export

from . import model, system
from . import spectrostream
from .load_weights import load_weights, convert_to_bf16
from magenta_rt import paths


def _flatten_state(state):
  """Flatten a nested pytree state into a list of mx.array.

  Handles tuples, lists, and mx.array leaves. Empty tuples contribute
  zero arrays.

  Args:
    state: Nested tuple/list of mx.array.

  Returns:
    (flat_arrays, structure) where structure encodes the nesting.
  """
  flat = []

  def _record(node):
    if isinstance(node, mx.array):
      flat.append(node)
      return 'array'
    elif isinstance(node, tuple):
      children = [_record(child) for child in node]
      return ('tuple', children)
    elif isinstance(node, list):
      children = [_record(child) for child