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)

```bash
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

```python
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

```bash
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 |