magenta-community/magenta-rt-jam
Space, but without the ZeroGPU / spaces infrastructure.This is a thin local driver on top of the official PyTorch port. It does not modify any of Google's model code — everything lives in a singleapp_local.py. Original work © Google LLC, Apache-2.0. See Credits.
app.py assumes:import spaces + @spaces.GPU (ZeroGPU grant) — the browser hangs on "Obtaining
ZeroGPU…" with no cloud backend;/tmp, /data) and re-downloads weights on every boot;app_local.py resolves all of these for local use.The model'sdecode_stream,generate, andstreammethods carry@torch.no_grad, butstep_fandencodedo not. The local stream callsstep_f/encodedirectly (not throughgenerate), so autograd stays on and builds a gradient graph every frame whose activations pile up in VRAM. The official Space never noticed: on ZeroGPU the 80 GB of headroom and the ~55 s grant-reset hide it.
torch.set_grad_enabled(False) # main thread (warm-up) AND the worker thread (thread-local)| Mode | Per-frame | VRAM over 1500 frames | Verdict |
|---|---|---|---|
| eager, autograd on (original) | 70–80 ms, climbing | 1 GB → 26 GB → OOM | unusable |
eager, no_grad | ~70 ms | flat ~1 GB | stable, not real-time |
torch.compile + no_grad | ~24 ms (~60%) | flat ~1 GB | smooth, real-time |
If you maintain the upstream port: adding@torch.no_grad()tostep_f(andencode) would fix this for everyone and likely cut VRAM on ZeroGPU too.This has been proposed upstream: magenta-rt-jam · discussions/1
app_local.py changes vs the upstream app.pyspaces.GPU; the stream runs on your local GPU.torch.set_grad_enabled(False) in the main (warm-up) and worker threads — the fix above.torch.compile (dynamic) on the per-frame step fns + a full warm-up at start
(covers all KV-cache lengths up to temporal_max_past, so no recompiles mid-stream).
The upstream fast path (AOTI) is unavailable: it's A10G-only, and AOTInductor is
hard-blocked on Windows in PyTorch.MRT_LOAD=small|base|both, default small) — 12 GB can't hold both.google/magenta-realtime-21# 1. Get the official PyTorch port (the model code lives here)
2git clone https://huggingface.co/spaces/magenta-community/magenta-rt-jam jam
3# (uses GIT_LFS_SKIP_SMUDGE=1 if you only want the code, not the LFS demo assets)
4
5# 2. Python env
6py -3.11 -m venv .venv
7.venv\Scripts\python -m pip install -r requirements-local.txt
8.venv\Scripts\python -m pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124
9
10# 3. Download weights (~1.1 GB small, ~9.2 GB base) — needs your HF token
11set HF_TOKEN=hf_xxx
12.venv\Scripts\python download_weights.py small # or: base / both
13
14# 4. Drop our driver in next to the model package and run
15copy app_local.py jam\app_local.py
16cd jam
17..\.venv\Scripts\python app_local.py| Var | Default | Meaning |
|---|---|---|
MAGENTA_HOME | ./magenta-home | where weights live (<home>/magenta-rt-v2/checkpoints) |
MRT_LOAD | small | small | base | both (12 GB → keep small) |
MRT_COMPILE | 1 | 1 = torch.compile (smooth); 0 = eager (slower, no warm-up wait) |
torch.compile warms every KV-cache shape before serving.mrt2_base (2.4B) does not fit alongside small on 12 GB; load it alone (MRT_LOAD=base).torch.compile needs MSVC — stay on 2.6.magenta-community/magenta-rt-jam
Space and the google/magenta-realtime-2 weights.app_local.py (the local driver) and this documentation. It is an
independent, unofficial integration and is not affiliated with or endorsed by Google.NOTICE for attribution details.