code/PYTHON.md

tt_diffrhythm — Python reference

DiffRhythm 1.2 Full on one Tenstorrent Blackhole p100a. The CFM transformer and the stable-audio VAE decoder run on the card; everything else stays on the CPU in DiffRhythm's own environment.

Two processes

DiffRhythm's environment and the tt-metal runtime need different torch and numpy builds, so they run as two processes talking over a Unix socket in a directory bind-mounted into both. Tensors cross as raw little-endian bytes with a dtype and a shape (wire.py).

Host side (inside the DiffRhythm environment)

TTWorker(checkpoint, image=None, log=None, reply_timeout_s=120.0, model="diffrhythm", vae=None)

checkpoint is cfm_model.pt and vae is vae_model.pt; Hugging Face cache links are resolved. image is the tt-metal / ttnn image the port was built against; it has no portable default, so set it here or through MUSIC_TT_IMAGE. .start() starts the worker container on the card, .wait_ready(timeout_s) blocks until the weights are on the device, .close() stops the container and waits for it to exit, so the card is free only after that.

RemoteDiffRhythmDiT(worker, transformer)

A drop-in for cfm.transformer in CFM.sample (batch 1): same call, same output. Assign it and the upstream sampler runs unchanged.

RemoteVae(worker, upsample=2048)

[B, 64, T] latents to [B, 2, T * upsample] float32 audio on the CPU. It picks the largest VAE window from shapes.VAE_WINDOWS that fits the song.

Worker side (inside the tt-metal image)

python -m tt_diffrhythm.worker --socket /ipc/worker.sock \
  --checkpoint /checkpoint/cfm_model.pt --vae /vae/vae_model.pt

Protocol (plain dicts, tensors via wire):

-> {"op": "ready", ...}	once the weights are on the device
<- {"op": "prepare", "text_embed": {drop_text: [T, 512]}, "cond": [T, 64]}	-> {"op": "prepared", ...}
<- {"op": "forward", "x": [T, 64], "drop_text": bool, "style": [512], "c": [512]}	-> {"op": "velocity", "v": [T, 64]}

Direct use on the card

from tt_diffrhythm import DiffRhythmDiT, DRPrecision, OobleckTT, open_device, close_device
from tt_diffrhythm.diffrhythm_host import Checkpoint

dev = open_device()
dit = DiffRhythmDiT(dev, Checkpoint("cfm_model.pt"))     # DRPrecision() by default
vae = OobleckTT(dev, "vae_model.pt")
...
close_device(dev)

DRPrecision

Extends the shared Precision with the choices CFG forces. pred + 4 (pred − null) multiplies the difference between the two branches by four, so the activations, the residual and the output stay in fp32 and matmuls default to HiFi4. HiFi2 is about 14% faster and measurably worse (final latent PCC 0.9931 against 0.9954).

Kernel cache

python -m tt_diffrhythm.warm_cache

Compiles every sequence bucket of a 96–285 s song, both CFG branches, and the VAE windows in shapes.VAE_WINDOWS. The first decode at a VAE shape costs 33.9 s against 0.33 s warm, so this is worth running after any change to the code or the image.

Modules

Module	Role
remote.py	host side: TTWorker, RemoteDiffRhythmDiT, RemoteVae, the worker container
worker.py	worker side: loads the CFM transformer and the VAE and answers one host pipeline
wire.py	the tensor transport between the two torch builds
shapes.py	shapes both sides need, with no torch or ttnn import
diffrhythm_dit.py	the CFM transformer on the card, and DRPrecision
diffrhythm_host.py	checkpoint access, the torch reference and the sequence buckets
oobleck_vae.py	the audio VAE decoder on the card
common.py	device and host helpers shared with the tt-acestep port
warm_cache.py	ahead-of-time kernel compilation for every bucket