Add files using upload-large-folder tool

README.md

@@ -0,0 +1,232 @@
+---
+tags:
+- blackhole
+- p100a
+- tenstorrent
+- ttnn
+- tt-metal
+- tt-model-catalog
+- tt-nn
+- text-to-audio
+- music-generation
+- diffrhythm
+- tt-dit-server
+- tt-model-cache
+- tt-model-container
+pipeline_tag: text-to-audio
+license: apache-2.0
+license_link: https://huggingface.co/ASLP-lab/DiffRhythm-1_2-full
+base_model:
+- ASLP-lab/DiffRhythm-1_2-full
+---
+
+# diffrhythm-1.2-p100a
+
+DiffRhythm 1.2 Full, ported to one Tenstorrent Blackhole p100a.
+Weights: [ASLP-lab/DiffRhythm-1_2-full](https://huggingface.co/ASLP-lab/DiffRhythm-1_2-full) · Upstream code: [ASLP-lab/DiffRhythm](https://github.com/ASLP-lab/DiffRhythm) · Port: [SunandBean/tt-diffrhythm](https://github.com/SunandBean/tt-diffrhythm)
+
+Runs on **p100a** (mesh 1×1). The **CFM transformer and the stable-audio VAE decoder run on the card**; the style encoder, the sampler and the post-processing stay on the CPU.
+
+**What this port adds** — the CFM transformer and the stable-audio-tools Oobleck decoder on the card, sharing the worker, transport and host helpers with the sibling [tt-acestep](https://github.com/SunandBean/tt-acestep) port.
+**What it builds on** — DiffRhythm's own pipeline (Apache-2.0), which runs unchanged on the CPU; only those two stages move.
+
+**A 96-second song in 10.0 s** end to end — the CFM transformer 7.7 s across 62 calls, the audio VAE decode 2.0 s. The same song takes 9.7–10.1 s on an RTX 5070 Ti.
+
+## Listen
+
+Exactly what the container returned on run 1 of the measurement below — `POST /predict` with
+`{"prompt": "warm lo-fi hip hop, mellow rhodes piano, soft vinyl crackle, relaxed drums",
+"instrumental": true, "duration": 96, "seed": 42}`, generated in **10.05 s** on the card. The file is
+the response's WAV bytes, unmodified: 96.0 s stereo 44.1 kHz 16-bit, sha256 `7e89a573e6d94e4f`.
+
+<audio controls src="https://huggingface.co/sunandbean/diffrhythm-1.2-p100a/resolve/main/media/diffrhythm_96s_instrumental.wav"></audio>
+
+[`media/diffrhythm_96s_instrumental.wav`](media/diffrhythm_96s_instrumental.wav) (16 MB)
+
+## How it is wired
+
+Two stages move to the card and nothing else changes: `remote.py` replaces `cfm.transformer` with a `RemoteDiffRhythmDiT` and decodes through a `RemoteVae`, so the upstream `CFM.sample` loop runs unmodified.
+
+Both shims reach the card through exactly one method, `worker.call(msg)`, and there are two things that answer it.
+
+**Two processes** (`TTWorker`), when DiffRhythm and tt-metal cannot share an interpreter — which is the deployment this port came from, where DiffRhythm resolved `torch 2.10.0+cu128` and `numpy>=2` against ttnn's `torch 2.11.0+cpu` and `numpy 1.x`:
+
+```
+  DiffRhythm venv (CPU)                    TT image (p100a)
+  ─────────────────────                    ────────────────
+  style encoder                ──socket──► worker.py
+  CFM sampler loop                         DiffRhythmDiT  (on the card)
+  post-processing                          OobleckTT      (on the card)
+```
+
+Tensors cross as raw little-endian bytes with a dtype and a shape (`wire.py`), because the two sides cannot share a pickle of a torch tensor.
+
+**One process** (`LocalWorker`), when one environment satisfies both — which is what the container below is. `LocalWorker` answers the same `call(msg)` in-process, so the shims, their window arithmetic and their drop-in signatures are reused unchanged rather than reimplemented.
+
+Messages still carry `wire`-encoded tensors in that case even though both ends are one process. That is deliberate: it is what lets the shims run against either worker, and the encoding is a float32 copy of a few MB against card work measured in seconds.
+
+## Quickstart
+
+Prerequisite: a tt-metal / ttnn environment for the worker, and DiffRhythm installed in its own environment for the host side.
+
+```bash
+hf download sunandbean/diffrhythm-1.2-p100a --include "code/*" --local-dir diffrhythm-1.2-p100a && cd diffrhythm-1.2-p100a
+pip install -e code/
+```
+
+Host side, inside the DiffRhythm environment:
+
+```python
+from tt_diffrhythm import TTWorker, RemoteDiffRhythmDiT, RemoteVae
+
+worker = TTWorker(cfm_model_pt, model="diffrhythm", vae=vae_model_pt).start()
+worker.wait_ready()
+cfm.transformer = RemoteDiffRhythmDiT(worker, cfm.transformer)
+...                                          # upstream CFM.sample runs unchanged
+audio = RemoteVae(worker, upsample=2048)(latents)
+worker.close()
+```
+
+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
+python -m tt_diffrhythm.warm_cache     # compile every sequence bucket ahead of time
+```
+
+| | |
+|---|---|
+| **Duration** | 96–285 seconds |
+| **Lyrics** | Plain lyrics with estimated line timing, or a timestamped LRC |
+| **Guidance** | CFG: `pred + 4 (pred − null)` — every step runs two DiT calls |
+| **Shapes** | Song lengths fall into sequence buckets (`diffrhythm_host.seq_buckets`), both CFG branches, plus the VAE windows. `warm_cache` compiles all of them. |
+
+## Performance
+
+Measured on an idle p100a. A 96-second song is 2067 latent frames, padded to a sequence of 2560.
+
+| | |
+|---|---:|
+| **DiT, whole 96-second song** | **8.9 s** (7.7 s at HiFi2) |
+| One DiT call at sequence 2560 | 121 ms |
+| Per call during the device check | 0.143–0.186 s |
+
+Where a 121 ms call goes (`dit_profile.py`, 156.8 ms profiled):
+
+| Op | Time | Calls |
+|---|---:|---:|
+| matmul | 66.6 ms | 66 |
+| SDPA | 21.9 ms | 16 |
+| add | 12.0 ms | 47 |
+| RMS norm | 7.7 ms | 32 |
+| slice | 7.3 ms | 62 |
+| rotary embedding | 5.2 ms | 32 |
+
+The audio VAE decoder, against a CPU decode of the same latents:
+
+| Latent frames | CPU | p100a (first / warm) |
+|---:|---:|---:|
+| 512 | 10.67 s | 33.94 / **0.33** s |
+| 1024 | — | 0.81 s |
+
+The first decode at a shape pays the kernel compilation; afterwards it is 30× faster than the CPU.
+
+### From the published container
+
+The same image anyone can `tt-model pull`, reached over its own HTTP server on an idle card
+(2026-10-09). Three 96-second instrumentals, seed 42, from a cold start:
+
+| | run 1 | run 2 | run 3 |
+|---|---:|---:|---:|
+| **Generation** | **10.05 s** | 10.06 s | 10.04 s |
+| CFM transformer, 62 calls | 7.68 s | 7.68 s | 7.69 s |
+| VAE decode | 2.03 s | 2.05 s | 2.01 s |
+
+All three returned the same 96.0 s stereo 44.1 kHz WAV. Load to ready was 50 s, of which the VAE
+warm-up is 3.1 s; device DRAM after load is 2.9 GB of 29.9 — the card is nowhere near full, and
+the port holds that headroom deliberately so the sequence buckets can grow to 285 seconds.
+
+The whole request takes about 30 s of wall time, because the response carries the song: 17 MB of
+WAV, base64-encoded. Generation is the 10 s above; the rest is serialising and shipping it.
+
+### Against an RTX 5070 Ti
+
+DiffRhythm's own GPU path with the deployment's stage-by-stage placement (one model in VRAM at a
+time, which is what keeps it inside 16 GB), same prompt and duration, a fresh seed per run:
+
+| | p100a (container) | RTX 5070 Ti |
+|---|---:|---:|
+| 96-second instrumental | 10.04–10.06 s | **9.73–10.05 s** |
+| First run | 10.05 s | 12.35 s |
+| Peak device memory | 2.9 GB of 29.9 | 2.5 GB of 15.5 |
+
+The same, within noise. The card reaches it with the CFM transformer and the audio VAE decoder in
+TTNN and everything else on the CPU; the GPU reaches it by moving three models in and out of VRAM.
+
+### Accuracy
+
+**The torch reference in this port against upstream DiffRhythm** (float32, per CFM call): PCC **0.9999997–0.9999999**. For scale, upstream float32 against the GPU fp16 run is 0.999995–0.999999 — so the reimplementation is closer to upstream than upstream is to its own GPU run.
+
+**The card against this port's reference**: minimum PCC **0.99994** across the whole 96-second song.
+
+**Final latents against the GPU fp16 run**, which is what actually matters after CFG amplification:
+
+| | Final latent PCC | Relative RMS |
+|---|---:|---:|
+| CPU float32 | 0.99997 | 0.008 |
+| **p100a (HiFi4, the default)** | **0.9954** | 0.096 |
+| p100a (HiFi2) | 0.9931 | 0.118 |
+| CPU bf16 | 0.99129 | 0.132 |
+
+**The card beats a CPU bf16 run.** Decoded audio against the GPU: PCC 0.9721 and 2.06 dB spectral distance, against 0.9691 and 3.07 dB for CPU bf16 (CPU float32 is 0.9996 / 0.51 dB). The audio VAE decoder itself is PCC 0.99963 with 0.93 dB spectral distance against the CPU.
+
+CFG is why the margin matters: `pred + 4 (pred − null)` multiplies the difference between the two branches by four, so a small per-call error is amplified before it reaches the audio. `DRPrecision` therefore keeps the activations, the residual and the output in fp32 and defaults to HiFi4 matmuls — HiFi2 is about 14% faster and measurably worse.
+
+## Caveats
+
+- **Single card only.** Mesh 1×1.
+- **Seeds do not reproduce across backends.** A song made on the p100a cannot be reproduced bit-for-bit on a GPU with the same seed.
+- **Two environments, or one.** Where DiffRhythm's environment and the tt-metal runtime cannot be merged, the host side runs inside DiffRhythm's own environment and the worker inside a tt-metal image. Where they can — which the container below is — `LocalWorker` runs both in one process. Which one you get is a property of the environment, not of the port.
+- **CFG doubles the work.** Every sampler step is two DiT calls, one conditional and one null.
+- The first song of each shape pays the kernel compilation, and so does the first VAE decode at a shape (33.9 s against 0.33 s warm). Run `warm_cache` after changing the code or the image.
+- A render holds the card for the length of the song. On a shared card that blocks everything else.
+
+## Licensing
+
+- Weights: [ASLP-lab/DiffRhythm-1_2-full](https://huggingface.co/ASLP-lab/DiffRhythm-1_2-full), Apache-2.0. Not redistributed here.
+- Port code (`code/tt_diffrhythm/`): Apache-2.0, from [SunandBean/tt-diffrhythm](https://github.com/SunandBean/tt-diffrhythm).
+- `code/{model,infer,g2p,config,dataset}/` are **DiffRhythm itself** (Apache-2.0), at [`28ad63c`](https://github.com/ASLP-lab/DiffRhythm/commit/28ad63c0f096fe2ee258bcabbcf081d5d9366afd), unmodified. They ship here because the container runs them: upstream is a research checkout with no packaging, so it is vendored as source rather than resolved by pip. `code/models/` is [tt-metal](https://github.com/tenstorrent/tt-metal)'s `models/common` and `models/tt_dit` (Apache-2.0), for the same reason.
+- `code/tt_diffrhythm/common.py` and the worker, transport and audio VAE modules are shared with the sibling [tt-acestep](https://github.com/SunandBean/tt-acestep) port (Apache-2.0); `oobleck_vae.py` implements both DiffRhythm's stable-audio-tools decoder and ACE-Step's diffusers Oobleck decoder.
+
+## Run it as a container
+
+Everything below is generated by `tt-model package` from the manifest this image was built from.
+
+```bash
+tt-model pull  sunandbean/diffrhythm-1.2-p100a --with-weights
+tt-model serve sunandbean/diffrhythm-1.2-p100a
+```
+
+`pull --with-weights` downloads the Docker image and the [`ASLP-lab/DiffRhythm-1_2-full`](https://huggingface.co/ASLP-lab/DiffRhythm-1_2-full) weights (into your HF cache; they are not in the image). `serve` starts the model's own HTTP server on port 20000 (or the next free port, if that one is busy); the first start compiles kernels for your device, which takes several minutes, and the server is ready when it logs `Application startup complete`.
+
+### Request shape
+
+`POST /predict` takes `{"prompt": str, "lyrics": str, "instrumental": bool,
+"duration": int, "seed": int}` and returns a base64 WAV with `timing_s` and the `lrc` the
+song was actually sung to.
+
+`prompt` is the style description. Vocals need `lyrics`: real LRC is used as given, and
+plain lines get approximate starts derived from their word counts -- a convenience, not
+learned alignment. `duration` is 96-285 seconds; use the ACE-Step port for longer songs.
+The schedule is fixed at **32 steps**; any other `inference_steps` is a 422.
+
+### Provenance
+
+The exact sources the image was built from — the model code under `code/` is byte-identical to what is inside the image (`pyproject.toml`, `PYTHON.md` and `examples/` are packaging files this repo adds, and are not in the image):
+
+| component | built from |
+| --- | --- |
+| tt-metal | [`58648d44443223041148125843ca7c7d4780463e`](https://github.com/tenstorrent/tt-metal/commit/58648d44443223041148125843ca7c7d4780463e) |
+| `code/` digest | `6fe96af2d7e7a1b8` (sha256, first 16 hex digits) |
+| built | 2026-10-09T14:54:37+00:00 by tt-model 0.1.0 |

code/PYTHON.md

@@ -0,0 +1,89 @@
+# 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 |

code/config/accelerate_config.yaml

@@ -0,0 +1,18 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+rdzv_backend: static
+same_network: true
+num_processes: 1
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
+main_process_port: 12342
+

code/config/default.ini

@@ -0,0 +1,85 @@
+
+[DEFAULTS]
+
+#name of the run
+exp_name = diffrhythm-test
+
+# the batch size
+batch_size = 8 
+
+# the chunk size
+max_frames = 3000 
+min_frames = 10
+
+# number of CPU workers for the DataLoader
+num_workers = 4
+
+# the random seed
+seed = 42
+
+# Batches for gradient accumulation
+accum_batches = 1
+
+# Number of steps between checkpoints
+checkpoint_every = 10000                              
+                     
+# trainer checkpoint file to restart training from
+ckpt_path = ''
+
+# model checkpoint file to start a new training run from
+pretrained_ckpt_path = ''
+
+# Checkpoint path for the pretransform model if needed
+pretransform_ckpt_path = ''
+
+# configuration model specifying model hyperparameters
+model_config = ''
+
+# configuration for datasets
+dataset_config = ''
+
+# directory to save the checkpoints in
+save_dir = ''
+
+# grad norm
+max_grad_norm = 1.0
+
+# grad accu
+grad_accumulation_steps = 1
+
+# lr 1e-5
+learning_rate = 7.5e-5
+
+# epoch
+epochs = 110 
+
+# warmup steps
+num_warmup_updates = 20
+
+# save checkpoint per steps
+save_per_updates = 5000
+
+# save last checkpoint per steps
+last_per_steps = 5000
+
+file_path = "dataset/train.scp"
+
+audio_drop_prob = 0.3
+
+cond_drop_prob = 0.0
+
+style_drop_prob = 0.1
+
+lrc_drop_prob = 0.1
+
+use_style_prompt = 1
+
+reset_lr = 0
+
+resumable_with_seed = 666
+
+downsample_rate = 2048
+
+grad_ckpt = 0
+
+sampling_rate = 44100

code/config/diffrhythm-1b.json

@@ -0,0 +1,13 @@
+{
+    "model_type": "diffrhythm",
+    "model": {
+        "dim": 2048,
+        "depth": 16, 
+        "heads": 32,
+        "ff_mult": 4,
+        "text_dim": 512,
+        "conv_layers": 4,
+        "mel_dim": 64,
+        "text_num_embeds": 363
+    }
+}

code/dataset/dataset.py

@@ -0,0 +1,171 @@
+# Copyright (c) 2025 ASLP-LAB
+#               2025 Huakang Chen  (huakang@mail.nwpu.edu.cn)
+#
+# 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 torch
+import random
+
+class DiffusionDataset(torch.utils.data.Dataset):
+    def __init__(self, file_path, max_frames=2048, min_frames=512, sampling_rate=44100, downsample_rate=2048, precision='fp16'):
+
+        self.max_frames = max_frames
+        self.min_frames = min_frames
+        self.sampling_rate = sampling_rate
+        self.downsample_rate = 2048
+        self.max_secs = max_frames / (sampling_rate / downsample_rate)
+        
+        self.file_path = file_path
+        with open(file_path, 'r') as f:
+            self.file_lst = [line.strip() for line in f.readlines()]
+
+        self.pad_token_id = 0
+        self.comma_token_id = 1
+        self.period_token_id = 2
+        self.start_token_id = 355
+
+        if precision == 'fp16':
+            self.feature_dtype = torch.float16
+        elif precision == 'bf16':
+            self.feature_dtype = torch.bfloat16
+        elif precision == 'fp32':
+            self.feature_dtype = torch.float32
+
+        random.seed(42)
+        random.shuffle(self.file_lst)
+
+    def load_item(self, item, field):
+        try:
+            item, reader_idx = item[field]
+            item = self.lance_connections[reader_idx].get_datas_by_rowids([item._rowid])[0]
+        except Exception as e:
+            return None
+        return item
+    
+    def get_triple(self, item):
+        utt, lrc_path, latent_path, style_path = item.split("|")
+
+        time_lrc = torch.load(lrc_path, map_location='cpu')
+        input_times = time_lrc['time']
+        input_lrcs = time_lrc['lrc']
+        lrc_with_time = list(zip(input_times, input_lrcs))
+        
+        latent = torch.load(latent_path, map_location='cpu') # [b, d, t]
+        latent = latent.squeeze(0)
+        
+        prompt = torch.load(style_path, map_location='cpu') # [b, d]
+        prompt = prompt.squeeze(0)
+        
+        max_start_frame = max(0, latent.shape[-1] - self.max_frames)
+        start_frame = random.randint(0, max_start_frame)
+        start_time = start_frame * self.downsample_rate / self.sampling_rate
+        normalized_start_time = start_frame / latent.shape[-1]
+        latent = latent[:, start_frame:]
+    
+        lrc_with_time = [(time_start - start_time, line) for (time_start, line) in lrc_with_time if (time_start - start_time) >= 0] # empty for pure music
+        lrc_with_time = [(time_start, line) for (time_start, line) in lrc_with_time if time_start < self.max_secs] # drop time longer than max_secs
+
+        if len(lrc_with_time) >= 1:
+            latent_end_time = lrc_with_time[-1][0]
+        else:
+            raise
+        
+        if self.max_frames == 2048:
+            lrc_with_time = lrc_with_time[:-1] if len(lrc_with_time) >= 1 else lrc_with_time # drop last, can be empty
+
+        lrc = torch.zeros((self.max_frames,), dtype=torch.long)
+        
+        tokens_count = 0
+        last_end_pos = 0
+        for time_start, line in lrc_with_time:
+            tokens = [token if token != self.period_token_id else self.comma_token_id for token in line] + [self.period_token_id]
+            tokens = torch.tensor(tokens, dtype=torch.long)
+            num_tokens = tokens.shape[0]
+
+            gt_frame_start = int(time_start * self.sampling_rate / self.downsample_rate)
+            
+            frame_shift = 0
+
+            frame_start = max(gt_frame_start - frame_shift, last_end_pos)
+            frame_len = min(num_tokens, self.max_frames - frame_start)
+
+            lrc[frame_start:frame_start + frame_len] = tokens[:frame_len]
+
+            tokens_count += num_tokens
+            last_end_pos = frame_start + frame_len        
+
+        latent = latent[:, :int(latent_end_time * self.sampling_rate / self.downsample_rate)]
+
+        latent = latent.to(self.feature_dtype)
+        prompt = prompt.to(self.feature_dtype)
+
+        return prompt, lrc, latent, normalized_start_time
+
+    def __getitem__(self, index):
+        idx = index
+        while True:
+            try:
+                prompt, lrc, latent, start_time = self.get_triple(self.file_lst[idx])
+                if latent.shape[-1] < self.min_frames: # Too short
+                    raise
+                item = {'prompt': prompt, "lrc": lrc, "latent": latent, "start_time": start_time}
+                return item
+            except Exception as e:
+                idx = random.randint(0, self.__len__() - 1)
+                continue
+
+    def __len__(self):
+        return len(self.file_lst)
+
+    def custom_collate_fn(self, batch):
+        latent_list = [item['latent'] for item in batch]
+        prompt_list = [item['prompt'] for item in batch]
+        lrc_list = [item['lrc'] for item in batch]
+        start_time_list = [item['start_time'] for item in batch]
+
+        latent_lengths = torch.LongTensor([latent.shape[-1] for latent in latent_list])
+        prompt_lengths = torch.LongTensor([prompt.shape[-1] for prompt in prompt_list])
+        lrc_lengths = torch.LongTensor([lrc.shape[-1] for lrc in lrc_list])
+
+        max_prompt_length = prompt_lengths.amax()
+
+        padded_prompt_list = []
+        for prompt in prompt_list:
+            padded_prompt = torch.nn.functional.pad(prompt, (0, max_prompt_length - prompt.shape[-1]))
+            padded_prompt_list.append(padded_prompt)
+
+        padded_latent_list = []
+        for latent in latent_list:
+            padded_latent = torch.nn.functional.pad(latent, (0, self.max_frames - latent.shape[-1]))
+            padded_latent_list.append(padded_latent)
+
+        padded_start_time_list = []
+        for start_time in start_time_list:
+            padded_start_time = start_time
+            padded_start_time_list.append(padded_start_time)
+
+        prompt_tensor = torch.stack(padded_prompt_list)
+        lrc_tensor = torch.stack(lrc_list)
+        latent_tensor = torch.stack(padded_latent_list)
+        start_time_tensor = torch.tensor(padded_start_time_list)
+
+        return {'prompt': prompt_tensor, 'lrc': lrc_tensor, 'latent': latent_tensor, \
+                "prompt_lengths": prompt_lengths, "lrc_lengths": lrc_lengths, "latent_lengths": latent_lengths, \
+                "start_time": start_time_tensor}
+
+
+if __name__ == "__main__":
+    dd = DiffusionDataset("train.scp", 2048, 512)
+    x = dd[0]
+    import pdb; pdb.set_trace()
+    print(x)

code/dataset/train.scp

@@ -0,0 +1,4 @@
+2626046476|dataset/lrc/2626046476.pt|dataset/latent/2626046476.pt|dataset/style/2626046476.pt
+549162241|dataset/lrc/549162241.pt|dataset/latent/549162241.pt|dataset/style/549162241.pt
+547512549|dataset/lrc/547512549.pt|dataset/latent/547512549.pt|dataset/style/547512549.pt
+28528706|dataset/lrc/28528706.pt|dataset/latent/28528706.pt|dataset/style/28528706.pt

code/g2p/g2p_generation.py

@@ -0,0 +1,133 @@
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import sys
+
+from g2p.g2p import PhonemeBpeTokenizer
+from g2p.utils.g2p import phonemizer_g2p
+import tqdm
+from typing import List
+import json
+import os
+import re
+
+
+def ph_g2p(text, language):
+
+    return phonemizer_g2p(text=text, language=language)
+
+
+def g2p(text, sentence, language):
+
+    return text_tokenizer.tokenize(text=text, sentence=sentence, language=language)
+
+
+def is_chinese(char):
+    if char >= "\u4e00" and char <= "\u9fa5":
+        return True
+    else:
+        return False
+
+
+def is_alphabet(char):
+    if (char >= "\u0041" and char <= "\u005a") or (
+        char >= "\u0061" and char <= "\u007a"
+    ):
+        return True
+    else:
+        return False
+
+
+def is_other(char):
+    if not (is_chinese(char) or is_alphabet(char)):
+        return True
+    else:
+        return False
+
+
+def get_segment(text: str) -> List[str]:
+    # sentence --> [ch_part, en_part, ch_part, ...]
+    segments = []
+    types = []
+    flag = 0
+    temp_seg = ""
+    temp_lang = ""
+
+    # Determine the type of each character. type: blank, chinese, alphabet, number, unk and point.
+    for i, ch in enumerate(text):
+        if is_chinese(ch):
+            types.append("zh")
+        elif is_alphabet(ch):
+            types.append("en")
+        else:
+            types.append("other")
+
+    assert len(types) == len(text)
+
+    for i in range(len(types)):
+        # find the first char of the seg
+        if flag == 0:
+            temp_seg += text[i]
+            temp_lang = types[i]
+            flag = 1
+        else:
+            if temp_lang == "other":
+                if types[i] == temp_lang:
+                    temp_seg += text[i]
+                else:
+                    temp_seg += text[i]
+                    temp_lang = types[i]
+            else:
+                if types[i] == temp_lang:
+                    temp_seg += text[i]
+                elif types[i] == "other":
+                    temp_seg += text[i]
+                else:
+                    segments.append((temp_seg, temp_lang))
+                    temp_seg = text[i]
+                    temp_lang = types[i]
+                    flag = 1
+
+    segments.append((temp_seg, temp_lang))
+    return segments
+
+
+def chn_eng_g2p(text: str):
+    # now only en and ch
+    segments = get_segment(text)
+    all_phoneme = ""
+    all_tokens = []
+
+    for index in range(len(segments)):
+        seg = segments[index]
+        phoneme, token = g2p(seg[0], text, seg[1])
+        all_phoneme += phoneme + "|"
+        all_tokens += token
+
+        if seg[1] == "en" and index == len(segments) - 1 and all_phoneme[-2] == "_":
+            all_phoneme = all_phoneme[:-2]
+            all_tokens = all_tokens[:-1]
+    return all_phoneme, all_tokens
+
+
+text_tokenizer = PhonemeBpeTokenizer()
+with open("./g2p/g2p/vocab.json", "r", encoding='utf-8') as f:
+    json_data = f.read()
+data = json.loads(json_data)
+vocab = data["vocab"]
+
+if __name__ == '__main__':
+    phone, token = chn_eng_g2p("你好，hello world")
+    phone, token = chn_eng_g2p("你好，hello world, Bonjour, 테스트 해 보겠습니다, 五月雨緑")
+    print(phone)
+    print(token)
+
+    #phone, token = text_tokenizer.tokenize("你好，hello world, Bonjour, 테스트 해 보겠습니다, 五月雨緑", "", "auto")
+    phone, token = text_tokenizer.tokenize("緑", "", "auto")
+    #phone, token = text_tokenizer.tokenize("आइए इसका परीक्षण करें", "", "auto")
+    #phone, token = text_tokenizer.tokenize("आइए इसका परीक्षण करें", "", "other")
+    print(phone)
+    print(token)

code/g2p/sources/bpmf_2_pinyin.txt

@@ -0,0 +1,41 @@
+b	ㄅ
+p	ㄆ
+m	ㄇ
+f	ㄈ
+d	ㄉ
+t	ㄊ
+n	ㄋ
+l	ㄌ
+g	ㄍ
+k	ㄎ
+h	ㄏ
+j	ㄐ
+q	ㄑ
+x	ㄒ
+zh	ㄓ
+ch	ㄔ
+sh	ㄕ
+r	ㄖ
+z	ㄗ
+c	ㄘ
+s	ㄙ
+i	ㄧ
+u	ㄨ
+v	ㄩ
+a	ㄚ
+o	ㄛ
+e	ㄜ
+e	ㄝ
+ai	ㄞ
+ei	ㄟ
+ao	ㄠ
+ou	ㄡ
+an	ㄢ
+en	ㄣ
+ang	ㄤ
+eng	ㄥ
+er	ㄦ
+2	ˊ
+3	ˇ
+4	ˋ
+0	˙

code/g2p/sources/pinyin_2_bpmf.txt

@@ -0,0 +1,429 @@
+a	ㄚ
+ai	ㄞ
+an	ㄢ
+ang	ㄤ
+ao	ㄠ
+ba	ㄅㄚ
+bai	ㄅㄞ
+ban	ㄅㄢ
+bang	ㄅㄤ
+bao	ㄅㄠ
+bei	ㄅㄟ
+ben	ㄅㄣ
+beng	ㄅㄥ
+bi	ㄅㄧ
+bian	ㄅㄧㄢ
+biang	ㄅㄧㄤ
+biao	ㄅㄧㄠ
+bie	ㄅㄧㄝ
+bin	ㄅㄧㄣ
+bing	ㄅㄧㄥ
+bo	ㄅㄛ
+bu	ㄅㄨ
+ca	ㄘㄚ
+cai	ㄘㄞ
+can	ㄘㄢ
+cang	ㄘㄤ
+cao	ㄘㄠ
+ce	ㄘㄜ
+cen	ㄘㄣ
+ceng	ㄘㄥ
+cha	ㄔㄚ
+chai	ㄔㄞ
+chan	ㄔㄢ
+chang	ㄔㄤ
+chao	ㄔㄠ
+che	ㄔㄜ
+chen	ㄔㄣ
+cheng	ㄔㄥ
+chi	ㄔ
+chong	ㄔㄨㄥ
+chou	ㄔㄡ
+chu	ㄔㄨ
+chua	ㄔㄨㄚ
+chuai	ㄔㄨㄞ
+chuan	ㄔㄨㄢ
+chuang	ㄔㄨㄤ
+chui	ㄔㄨㄟ
+chun	ㄔㄨㄣ
+chuo	ㄔㄨㄛ
+ci	ㄘ
+cong	ㄘㄨㄥ
+cou	ㄘㄡ
+cu	ㄘㄨ
+cuan	ㄘㄨㄢ
+cui	ㄘㄨㄟ
+cun	ㄘㄨㄣ
+cuo	ㄘㄨㄛ
+da	ㄉㄚ
+dai	ㄉㄞ
+dan	ㄉㄢ
+dang	ㄉㄤ
+dao	ㄉㄠ
+de	ㄉㄜ
+dei	ㄉㄟ
+den	ㄉㄣ
+deng	ㄉㄥ
+di	ㄉㄧ
+dia	ㄉㄧㄚ
+dian	ㄉㄧㄢ
+diao	ㄉㄧㄠ
+die	ㄉㄧㄝ
+din	ㄉㄧㄣ
+ding	ㄉㄧㄥ
+diu	ㄉㄧㄡ
+dong	ㄉㄨㄥ
+dou	ㄉㄡ
+du	ㄉㄨ
+duan	ㄉㄨㄢ
+dui	ㄉㄨㄟ
+dun	ㄉㄨㄣ
+duo	ㄉㄨㄛ
+e	ㄜ
+ei	ㄟ
+en	ㄣ
+eng	ㄥ
+er	ㄦ
+fa	ㄈㄚ
+fan	ㄈㄢ
+fang	ㄈㄤ
+fei	ㄈㄟ
+fen	ㄈㄣ
+feng	ㄈㄥ
+fo	ㄈㄛ
+fou	ㄈㄡ
+fu	ㄈㄨ
+ga	ㄍㄚ
+gai	ㄍㄞ
+gan	ㄍㄢ
+gang	ㄍㄤ
+gao	ㄍㄠ
+ge	ㄍㄜ
+gei	ㄍㄟ
+gen	ㄍㄣ
+geng	ㄍㄥ
+gong	ㄍㄨㄥ
+gou	ㄍㄡ
+gu	ㄍㄨ
+gua	ㄍㄨㄚ
+guai	ㄍㄨㄞ
+guan	ㄍㄨㄢ
+guang	ㄍㄨㄤ
+gui	ㄍㄨㄟ
+gun	ㄍㄨㄣ
+guo	ㄍㄨㄛ
+ha	ㄏㄚ
+hai	ㄏㄞ
+han	ㄏㄢ
+hang	ㄏㄤ
+hao	ㄏㄠ
+he	ㄏㄜ
+hei	ㄏㄟ
+hen	ㄏㄣ
+heng	ㄏㄥ
+hm	ㄏㄇ
+hong	ㄏㄨㄥ
+hou	ㄏㄡ
+hu	ㄏㄨ
+hua	ㄏㄨㄚ
+huai	ㄏㄨㄞ
+huan	ㄏㄨㄢ
+huang	ㄏㄨㄤ
+hui	ㄏㄨㄟ
+hun	ㄏㄨㄣ
+huo	ㄏㄨㄛ
+ji	ㄐㄧ
+jia	ㄐㄧㄚ
+jian	ㄐㄧㄢ
+jiang	ㄐㄧㄤ
+jiao	ㄐㄧㄠ
+jie	ㄐㄧㄝ
+jin	ㄐㄧㄣ
+jing	ㄐㄧㄥ
+jiong	ㄐㄩㄥ
+jiu	ㄐㄧㄡ
+ju	ㄐㄩ
+jv	ㄐㄩ
+juan	ㄐㄩㄢ
+jvan	ㄐㄩㄢ
+jue	ㄐㄩㄝ
+jve	ㄐㄩㄝ
+jun	ㄐㄩㄣ
+ka	ㄎㄚ
+kai	ㄎㄞ
+kan	ㄎㄢ
+kang	ㄎㄤ
+kao	ㄎㄠ
+ke	ㄎㄜ
+kei	ㄎㄟ
+ken	ㄎㄣ
+keng	ㄎㄥ
+kong	ㄎㄨㄥ
+kou	ㄎㄡ
+ku	ㄎㄨ
+kua	ㄎㄨㄚ
+kuai	ㄎㄨㄞ
+kuan	ㄎㄨㄢ
+kuang	ㄎㄨㄤ
+kui	ㄎㄨㄟ
+kun	ㄎㄨㄣ
+kuo	ㄎㄨㄛ
+la	ㄌㄚ
+lai	ㄌㄞ
+lan	ㄌㄢ
+lang	ㄌㄤ
+lao	ㄌㄠ
+le	ㄌㄜ
+lei	ㄌㄟ
+leng	ㄌㄥ
+li	ㄌㄧ
+lia	ㄌㄧㄚ
+lian	ㄌㄧㄢ
+liang	ㄌㄧㄤ
+liao	ㄌㄧㄠ
+lie	ㄌㄧㄝ
+lin	ㄌㄧㄣ
+ling	ㄌㄧㄥ
+liu	ㄌㄧㄡ
+lo	ㄌㄛ
+long	ㄌㄨㄥ
+lou	ㄌㄡ
+lu	ㄌㄨ
+luan	ㄌㄨㄢ
+lue	ㄌㄩㄝ
+lun	ㄌㄨㄣ
+luo	ㄌㄨㄛ
+lv	ㄌㄩ
+lve	ㄌㄩㄝ
+m	ㄇㄨ
+ma	ㄇㄚ
+mai	ㄇㄞ
+man	ㄇㄢ
+mang	ㄇㄤ
+mao	ㄇㄠ
+me	ㄇㄜ
+mei	ㄇㄟ
+men	ㄇㄣ
+meng	ㄇㄥ
+mi	ㄇㄧ
+mian	ㄇㄧㄢ
+miao	ㄇㄧㄠ
+mie	ㄇㄧㄝ
+min	ㄇㄧㄣ
+ming	ㄇㄧㄥ
+miu	ㄇㄧㄡ
+mo	ㄇㄛ
+mou	ㄇㄡ
+mu	ㄇㄨ
+n	ㄣ
+na	ㄋㄚ
+nai	ㄋㄞ
+nan	ㄋㄢ
+nang	ㄋㄤ
+nao	ㄋㄠ
+ne	ㄋㄜ
+nei	ㄋㄟ
+nen	ㄋㄣ
+neng	ㄋㄥ
+ng	ㄣ
+ni	ㄋㄧ
+nian	ㄋㄧㄢ
+niang	ㄋㄧㄤ
+niao	ㄋㄧㄠ
+nie	ㄋㄧㄝ
+nin	ㄋㄧㄣ
+ning	ㄋㄧㄥ
+niu	ㄋㄧㄡ
+nong	ㄋㄨㄥ
+nou	ㄋㄡ
+nu	ㄋㄨ
+nuan	ㄋㄨㄢ
+nue	ㄋㄩㄝ
+nun	ㄋㄨㄣ
+nuo	ㄋㄨㄛ
+nv	ㄋㄩ
+nve	ㄋㄩㄝ
+o	ㄛ
+ou	ㄡ
+pa	ㄆㄚ
+pai	ㄆㄞ
+pan	ㄆㄢ
+pang	ㄆㄤ
+pao	ㄆㄠ
+pei	ㄆㄟ
+pen	ㄆㄣ
+peng	ㄆㄥ
+pi	ㄆㄧ
+pian	ㄆㄧㄢ
+piao	ㄆㄧㄠ
+pie	ㄆㄧㄝ
+pin	ㄆㄧㄣ
+ping	ㄆㄧㄥ
+po	ㄆㄛ
+pou	ㄆㄡ
+pu	ㄆㄨ
+qi	ㄑㄧ
+qia	ㄑㄧㄚ
+qian	ㄑㄧㄢ
+qiang	ㄑㄧㄤ
+qiao	ㄑㄧㄠ
+qie	ㄑㄧㄝ
+qin	ㄑㄧㄣ
+qing	ㄑㄧㄥ
+qiong	ㄑㄩㄥ
+qiu	ㄑㄧㄡ
+qu	ㄑㄩ
+quan	ㄑㄩㄢ
+qvan	ㄑㄩㄢ
+que	ㄑㄩㄝ
+qun	ㄑㄩㄣ
+ran	ㄖㄢ
+rang	ㄖㄤ
+rao	ㄖㄠ
+re	ㄖㄜ
+ren	ㄖㄣ
+reng	ㄖㄥ
+ri	ㄖ
+rong	ㄖㄨㄥ
+rou	ㄖㄡ
+ru	ㄖㄨ
+rua	ㄖㄨㄚ
+ruan	ㄖㄨㄢ
+rui	ㄖㄨㄟ
+run	ㄖㄨㄣ
+ruo	ㄖㄨㄛ
+sa	ㄙㄚ
+sai	ㄙㄞ
+san	ㄙㄢ
+sang	ㄙㄤ
+sao	ㄙㄠ
+se	ㄙㄜ
+sen	ㄙㄣ
+seng	ㄙㄥ
+sha	ㄕㄚ
+shai	ㄕㄞ
+shan	ㄕㄢ
+shang	ㄕㄤ
+shao	ㄕㄠ
+she	ㄕㄜ
+shei	ㄕㄟ
+shen	ㄕㄣ
+sheng	ㄕㄥ
+shi	ㄕ
+shou	ㄕㄡ
+shu	ㄕㄨ
+shua	ㄕㄨㄚ
+shuai	ㄕㄨㄞ
+shuan	ㄕㄨㄢ
+shuang	ㄕㄨㄤ
+shui	ㄕㄨㄟ
+shun	ㄕㄨㄣ
+shuo	ㄕㄨㄛ
+si	ㄙ
+song	ㄙㄨㄥ
+sou	ㄙㄡ
+su	ㄙㄨ
+suan	ㄙㄨㄢ
+sui	ㄙㄨㄟ
+sun	ㄙㄨㄣ
+suo	ㄙㄨㄛ
+ta	ㄊㄚ
+tai	ㄊㄞ
+tan	ㄊㄢ
+tang	ㄊㄤ
+tao	ㄊㄠ
+te	ㄊㄜ
+tei	ㄊㄟ
+teng	ㄊㄥ
+ti	ㄊㄧ
+tian	ㄊㄧㄢ
+tiao	ㄊㄧㄠ
+tie	ㄊㄧㄝ
+ting	ㄊㄧㄥ
+tong	ㄊㄨㄥ
+tou	ㄊㄡ
+tsuo	ㄘㄨㄛ
+tu	ㄊㄨ
+tuan	ㄊㄨㄢ
+tui	ㄊㄨㄟ
+tun	ㄊㄨㄣ
+tuo	ㄊㄨㄛ
+tzan	ㄗㄢ
+wa	ㄨㄚ
+wai	ㄨㄞ
+wan	ㄨㄢ
+wang	ㄨㄤ
+wei	ㄨㄟ
+wen	ㄨㄣ
+weng	ㄨㄥ
+wo	ㄨㄛ
+wong	ㄨㄥ
+wu	ㄨ
+xi	ㄒㄧ
+xia	ㄒㄧㄚ
+xian	ㄒㄧㄢ
+xiang	ㄒㄧㄤ
+xiao	ㄒㄧㄠ
+xie	ㄒㄧㄝ
+xin	ㄒㄧㄣ
+xing	ㄒㄧㄥ
+xiong	ㄒㄩㄥ
+xiu	ㄒㄧㄡ
+xu	ㄒㄩ
+xuan	ㄒㄩㄢ
+xue	ㄒㄩㄝ
+xun	ㄒㄩㄣ
+ya	ㄧㄚ
+yai	ㄧㄞ
+yan	ㄧㄢ
+yang	ㄧㄤ
+yao	ㄧㄠ
+ye	ㄧㄝ
+yi	ㄧ
+yin	ㄧㄣ
+ying	ㄧㄥ
+yo	ㄧㄛ
+yong	ㄩㄥ
+you	ㄧㄡ
+yu	ㄩ
+yuan	ㄩㄢ
+yue	ㄩㄝ
+yve	ㄩㄝ
+yun	ㄩㄣ
+za	ㄗㄚ
+zai	ㄗㄞ
+zan	ㄗㄢ
+zang	ㄗㄤ
+zao	ㄗㄠ
+ze	ㄗㄜ
+zei	ㄗㄟ
+zen	ㄗㄣ
+zeng	ㄗㄥ
+zha	ㄓㄚ
+zhai	ㄓㄞ
+zhan	ㄓㄢ
+zhang	ㄓㄤ
+zhao	ㄓㄠ
+zhe	ㄓㄜ
+zhei	ㄓㄟ
+zhen	ㄓㄣ
+zheng	ㄓㄥ
+zhi	ㄓ
+zhong	ㄓㄨㄥ
+zhou	ㄓㄡ
+zhu	ㄓㄨ
+zhua	ㄓㄨㄚ
+zhuai	ㄓㄨㄞ
+zhuan	ㄓㄨㄢ
+zhuang	ㄓㄨㄤ
+zhui	ㄓㄨㄟ
+zhun	ㄓㄨㄣ
+zhuo	ㄓㄨㄛ
+zi	ㄗ
+zong	ㄗㄨㄥ
+zou	ㄗㄡ
+zu	ㄗㄨ
+zuan	ㄗㄨㄢ
+zui	ㄗㄨㄟ
+zun	ㄗㄨㄣ
+zuo	ㄗㄨㄛ

code/infer/infer.py

@@ -0,0 +1,228 @@
+# Copyright (c) 2025 ASLP-LAB
+#               2025 Huakang Chen  (huakang@mail.nwpu.edu.cn)
+#               2025 Guobin Ma     (guobin.ma@gmail.com)
+#
+# 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 argparse
+import os
+import time
+import random
+
+import torch
+import torchaudio
+from einops import rearrange
+
+print("Current working directory:", os.getcwd())
+
+from infer_utils import (
+    decode_audio,
+    get_lrc_token,
+    get_negative_style_prompt,
+    get_reference_latent,
+    get_style_prompt,
+    prepare_model,
+)
+
+
+def inference(
+    cfm_model,
+    vae_model,
+    cond,
+    text,
+    duration,
+    style_prompt,
+    negative_style_prompt,
+    start_time,
+    pred_frames,
+    batch_infer_num,
+    song_duration,
+    chunked=False,
+):
+    with torch.inference_mode():
+        latents, _ = cfm_model.sample(
+            cond=cond,
+            text=text,
+            duration=duration,
+            style_prompt=style_prompt,
+            max_duration=duration,
+            song_duration=song_duration, 
+            negative_style_prompt=negative_style_prompt,
+            steps=32,
+            cfg_strength=4.0,
+            start_time=start_time,
+            latent_pred_segments=pred_frames,
+            batch_infer_num=batch_infer_num
+        )
+
+        outputs = []
+        for latent in latents:
+            latent = latent.to(torch.float32)
+            latent = latent.transpose(1, 2)  # [b d t]
+
+            output = decode_audio(latent, vae_model, chunked=chunked)
+
+            # Rearrange audio batch to a single sequence
+            output = rearrange(output, "b d n -> d (b n)")
+            # Peak normalize, clip, convert to int16, and save to file
+            output = (
+                output.to(torch.float32)
+                .div(torch.max(torch.abs(output)))
+                .clamp(-1, 1)
+                .mul(32767)
+                .to(torch.int16)
+                .cpu()
+            )
+            outputs.append(output)
+
+        return outputs
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--lrc-path",
+        type=str,
+        help="lyrics of target song",
+    )  # lyrics of target song
+    parser.add_argument(
+        "--ref-prompt",
+        type=str,
+        help="reference prompt as style prompt for target song",
+        required=False,
+    )  # reference prompt as style prompt for target song
+    parser.add_argument(
+        "--ref-audio-path",
+        type=str,
+        help="reference audio as style prompt for target song",
+        required=False,
+    )  # reference audio as style prompt for target song
+    parser.add_argument(
+        "--chunked",
+        action="store_true",
+        help="whether to use chunked decoding",
+    )  # whether to use chunked decoding
+    parser.add_argument(
+        "--audio-length",
+        type=int,
+        default=95,
+        # choices=[95, 285],
+        help="length of generated song, upported values are exactly 95 or any value between 96 and 285 (inclusive).",
+    )  # length of target song
+    # parser.add_argument(
+    #     "--repo-id", type=str, default="ASLP-lab/DiffRhythm-base", help="target model"
+    # )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="infer/example/output",
+        help="output directory fo generated song",
+    )  # output directory of target song
+    parser.add_argument(
+        "--edit",
+        action="store_true",
+        help="whether to open edit mode",
+    )  # edit flag
+    parser.add_argument(
+        "--ref-song",
+        type=str,
+        required=False,
+        help="reference prompt as latent prompt for editing",
+    )  # reference prompt as latent prompt for editing
+    parser.add_argument(
+        "--edit-segments",
+        type=str,
+        required=False,
+        help="Time segments to edit (in seconds). Format: `[[start1,end1],...]`. "
+             "Use `-1` for audio start/end (e.g., `[[-1,25], [50.0,-1]]`)."
+    )  # edit segments of target song
+    parser.add_argument(
+        "--batch-infer-num",
+        type=int,
+        default=1,
+        required=False,
+        help="number of songs per batch",
+    )  # number of songs per batch
+    args = parser.parse_args()
+
+    assert (
+        args.ref_prompt or args.ref_audio_path
+    ), "either ref_prompt or ref_audio_path should be provided"
+    assert not (
+        args.ref_prompt and args.ref_audio_path
+    ), "only one of them should be provided"
+    if args.edit:
+        assert (
+            args.ref_song and args.edit_segments
+        ), "reference song and edit segments should be provided for editing"
+
+    device = "cpu"
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.mps.is_available():
+        device = "mps"
+
+    audio_length = args.audio_length
+    if audio_length == 95:
+        max_frames = 2048
+    elif 95 < audio_length <= 285:
+        max_frames = 6144
+    else:
+        raise ValueError(
+            f"Invalid audio_length: {audio_length}. "
+            "Supported values are exactly 95 or any value between 96 and 285 (inclusive)."
+        )
+
+    cfm, tokenizer, muq, vae = prepare_model(max_frames, device)
+
+    if args.lrc_path:
+        with open(args.lrc_path, "r", encoding='utf-8') as f:
+            lrc = f.read()
+    else:
+        lrc = ""
+    lrc_prompt, start_time, end_frame, song_duration = get_lrc_token(max_frames, lrc, tokenizer, audio_length, device)
+
+    if args.ref_audio_path:
+        style_prompt = get_style_prompt(muq, args.ref_audio_path)
+    else:
+        style_prompt = get_style_prompt(muq, prompt=args.ref_prompt)
+
+    negative_style_prompt = get_negative_style_prompt(device)
+
+    latent_prompt, pred_frames = get_reference_latent(device, max_frames, args.edit, args.edit_segments, args.ref_song, vae)
+
+    s_t = time.time()
+    generated_songs = inference(
+        cfm_model=cfm,
+        vae_model=vae,
+        cond=latent_prompt,
+        text=lrc_prompt,
+        duration=end_frame,
+        style_prompt=style_prompt,
+        negative_style_prompt=negative_style_prompt,
+        start_time=start_time,
+        pred_frames=pred_frames,
+        chunked=args.chunked,
+        batch_infer_num=args.batch_infer_num,
+        song_duration=song_duration
+    )
+    e_t = time.time() - s_t
+    print(f"inference cost {e_t:.2f} seconds")
+    
+    generated_song = random.sample(generated_songs, 1)[0]
+
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    output_path = os.path.join(output_dir, "output.wav")
+    torchaudio.save(output_path, generated_song, sample_rate=44100)

code/infer/infer_utils.py

@@ -0,0 +1,459 @@
+# Copyright (c) 2025 ASLP-LAB
+#               2025 Huakang Chen  (huakang@mail.nwpu.edu.cn)
+#               2025 Guobin Ma     (guobin.ma@gmail.com)
+#
+# 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 torch
+import librosa
+import torchaudio
+import random
+import json
+from muq import MuQMuLan
+from mutagen.mp3 import MP3
+import os
+import numpy as np
+from huggingface_hub import hf_hub_download
+
+from sys import path
+path.append(os.getcwd())
+
+from model import DiT, CFM
+
+def vae_sample(mean, scale):
+    stdev = torch.nn.functional.softplus(scale) + 1e-4
+    var = stdev * stdev
+    logvar = torch.log(var)
+    latents = torch.randn_like(mean) * stdev + mean
+
+    kl = (mean * mean + var - logvar - 1).sum(1).mean()
+
+    return latents, kl
+
+def normalize_audio(y, target_dbfs=0):
+    max_amplitude = torch.max(torch.abs(y))
+
+    target_amplitude = 10.0**(target_dbfs / 20.0)
+    scale_factor = target_amplitude / max_amplitude
+
+    normalized_audio = y * scale_factor
+
+    return normalized_audio
+
+def set_audio_channels(audio, target_channels):
+    if target_channels == 1:
+        # Convert to mono
+        audio = audio.mean(1, keepdim=True)
+    elif target_channels == 2:
+        # Convert to stereo
+        if audio.shape[1] == 1:
+            audio = audio.repeat(1, 2, 1)
+        elif audio.shape[1] > 2:
+            audio = audio[:, :2, :]
+    return audio
+
+class PadCrop(torch.nn.Module):
+    def __init__(self, n_samples, randomize=True):
+        super().__init__()
+        self.n_samples = n_samples
+        self.randomize = randomize
+
+    def __call__(self, signal):
+        n, s = signal.shape
+        start = 0 if (not self.randomize) else torch.randint(0, max(0, s - self.n_samples) + 1, []).item()
+        end = start + self.n_samples
+        output = signal.new_zeros([n, self.n_samples])
+        output[:, :min(s, self.n_samples)] = signal[:, start:end]
+        return output
+
+def prepare_audio(audio, in_sr, target_sr, target_length, target_channels, device):
+    
+    audio = audio.to(device)
+
+    if in_sr != target_sr:
+        resample_tf = torchaudio.functional.Resample(in_sr, target_sr).to(device)
+        audio = resample_tf(audio)
+    if target_length is None:
+        target_length = audio.shape[-1]
+    audio = PadCrop(target_length, randomize=False)(audio)
+
+    # Add batch dimension
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0).unsqueeze(0)
+    elif audio.dim() == 2:
+        audio = audio.unsqueeze(0)
+
+    audio = set_audio_channels(audio, target_channels)
+
+    return audio
+
+def decode_audio(latents, vae_model, chunked=False, overlap=32, chunk_size=128):
+    downsampling_ratio = 2048
+    io_channels = 2
+    if not chunked:
+        return vae_model.decode_export(latents)
+    else:
+        # chunked decoding
+        hop_size = chunk_size - overlap
+        total_size = latents.shape[2]
+        batch_size = latents.shape[0]
+        chunks = []
+        i = 0
+        for i in range(0, total_size - chunk_size + 1, hop_size):
+            chunk = latents[:, :, i : i + chunk_size]
+            chunks.append(chunk)
+        if i + chunk_size != total_size:
+            # Final chunk
+            chunk = latents[:, :, -chunk_size:]
+            chunks.append(chunk)
+        chunks = torch.stack(chunks)
+        num_chunks = chunks.shape[0]
+        # samples_per_latent is just the downsampling ratio
+        samples_per_latent = downsampling_ratio
+        # Create an empty waveform, we will populate it with chunks as decode them
+        y_size = total_size * samples_per_latent
+        y_final = torch.zeros((batch_size, io_channels, y_size)).to(latents.device)
+        for i in range(num_chunks):
+            x_chunk = chunks[i, :]
+            # decode the chunk
+            y_chunk = vae_model.decode_export(x_chunk)
+            # figure out where to put the audio along the time domain
+            if i == num_chunks - 1:
+                # final chunk always goes at the end
+                t_end = y_size
+                t_start = t_end - y_chunk.shape[2]
+            else:
+                t_start = i * hop_size * samples_per_latent
+                t_end = t_start + chunk_size * samples_per_latent
+            #  remove the edges of the overlaps
+            ol = (overlap // 2) * samples_per_latent
+            chunk_start = 0
+            chunk_end = y_chunk.shape[2]
+            if i > 0:
+                # no overlap for the start of the first chunk
+                t_start += ol
+                chunk_start += ol
+            if i < num_chunks - 1:
+                # no overlap for the end of the last chunk
+                t_end -= ol
+                chunk_end -= ol
+            # paste the chunked audio into our y_final output audio
+            y_final[:, :, t_start:t_end] = y_chunk[:, :, chunk_start:chunk_end]
+        return y_final
+
+def encode_audio(audio, vae_model, chunked=False, overlap=32, chunk_size=128):
+    downsampling_ratio = 2048
+    latent_dim = 128
+    if not chunked:
+        # default behavior. Encode the entire audio in parallel
+        return vae_model.encode_export(audio)
+    else:
+        # CHUNKED ENCODING
+        # samples_per_latent is just the downsampling ratio (which is also the upsampling ratio)
+        samples_per_latent = downsampling_ratio
+        total_size = audio.shape[2] # in samples
+        batch_size = audio.shape[0]
+        chunk_size *= samples_per_latent # converting metric in latents to samples
+        overlap *= samples_per_latent # converting metric in latents to samples
+        hop_size = chunk_size - overlap
+        chunks = []
+        for i in range(0, total_size - chunk_size + 1, hop_size):
+            chunk = audio[:,:,i:i+chunk_size]
+            chunks.append(chunk)
+        if i+chunk_size != total_size:
+            # Final chunk
+            chunk = audio[:,:,-chunk_size:]
+            chunks.append(chunk)
+        chunks = torch.stack(chunks)
+        num_chunks = chunks.shape[0]
+        # Note: y_size might be a different value from the latent length used in diffusion training
+        # because we can encode audio of varying lengths
+        # However, the audio should've been padded to a multiple of samples_per_latent by now.
+        y_size = total_size // samples_per_latent
+        # Create an empty latent, we will populate it with chunks as we encode them
+        y_final = torch.zeros((batch_size,latent_dim,y_size)).to(audio.device)
+        for i in range(num_chunks):
+            x_chunk = chunks[i,:]
+            # encode the chunk
+            y_chunk = vae_model.encode_export(x_chunk)
+            # figure out where to put the audio along the time domain
+            if i == num_chunks-1:
+                # final chunk always goes at the end
+                t_end = y_size
+                t_start = t_end - y_chunk.shape[2]
+            else:
+                t_start = i * hop_size // samples_per_latent
+                t_end = t_start + chunk_size // samples_per_latent
+            #  remove the edges of the overlaps
+            ol = overlap//samples_per_latent//2
+            chunk_start = 0
+            chunk_end = y_chunk.shape[2]
+            if i > 0:
+                # no overlap for the start of the first chunk
+                t_start += ol
+                chunk_start += ol
+            if i < num_chunks-1:
+                # no overlap for the end of the last chunk
+                t_end -= ol
+                chunk_end -= ol
+            # paste the chunked audio into our y_final output audio
+            y_final[:,:,t_start:t_end] = y_chunk[:,:,chunk_start:chunk_end]
+        return y_final
+
+def prepare_model(max_frames, device, repo_id="ASLP-lab/DiffRhythm-1_2"):
+    # prepare cfm model
+    if max_frames == 2048:
+        repo_id = "ASLP-lab/DiffRhythm-1_2"
+    else:
+        repo_id = "ASLP-lab/DiffRhythm-1_2-full"
+        
+    dit_ckpt_path = hf_hub_download(
+        repo_id=repo_id, filename="cfm_model.pt", cache_dir="./pretrained"
+    )
+        
+    dit_config_path = "./config/diffrhythm-1b.json"
+    with open(dit_config_path) as f:
+        model_config = json.load(f)
+    dit_model_cls = DiT
+    cfm = CFM(
+        transformer=dit_model_cls(**model_config["model"], max_frames=max_frames),
+        num_channels=model_config["model"]["mel_dim"],
+        max_frames=max_frames
+    )
+    cfm = cfm.to(device)
+    cfm = load_checkpoint(cfm, dit_ckpt_path, device=device, use_ema=False)
+
+    # prepare tokenizer
+    tokenizer = CNENTokenizer()
+
+    # prepare muq
+    muq = MuQMuLan.from_pretrained("OpenMuQ/MuQ-MuLan-large", cache_dir="./pretrained")
+    muq = muq.to(device).eval()
+
+    # prepare vae
+    vae_ckpt_path = hf_hub_download(
+        repo_id="ASLP-lab/DiffRhythm-vae",
+        filename="vae_model.pt",
+        cache_dir="./pretrained",
+    )
+    vae = torch.jit.load(vae_ckpt_path, map_location="cpu").to(device)
+
+    return cfm, tokenizer, muq, vae
+
+
+# for song edit, will be added in the future
+def get_reference_latent(device, max_frames, edit, pred_segments, ref_song, vae_model):
+    sampling_rate = 44100
+    downsample_rate = 2048
+    io_channels = 2
+    if edit:
+        input_audio, in_sr = torchaudio.load(ref_song)
+        input_audio = prepare_audio(input_audio, in_sr=in_sr, target_sr=sampling_rate, target_length=None, target_channels=io_channels, device=device)
+        input_audio = normalize_audio(input_audio, -6)
+        
+        with torch.no_grad():
+            latent = encode_audio(input_audio, vae_model, chunked=True) # [b d t]
+            mean, scale = latent.chunk(2, dim=1)
+            prompt, _ = vae_sample(mean, scale)
+            prompt = prompt.transpose(1, 2) # [b t d]
+        
+        pred_segments = json.loads(pred_segments)
+        
+        pred_frames = []
+        for st, et in pred_segments:
+            sf = 0 if st == -1 else int(st * sampling_rate / downsample_rate)
+            ef = max_frames if et == -1 else int(et * sampling_rate / downsample_rate)
+            pred_frames.append((sf, ef))
+
+        return prompt, pred_frames
+    else:
+        prompt = torch.zeros(1, max_frames, 64).to(device)
+        pred_frames = [(0, max_frames)]
+        return prompt, pred_frames
+
+
+def get_negative_style_prompt(device):
+    file_path = "infer/example/vocal.npy"
+    vocal_stlye = np.load(file_path)
+
+    vocal_stlye = torch.from_numpy(vocal_stlye).to(device)  # [1, 512]
+    vocal_stlye = vocal_stlye.half()
+
+    return vocal_stlye
+
+
+@torch.no_grad()
+def get_style_prompt(model, wav_path=None, prompt=None):
+    mulan = model
+
+    if prompt is not None:
+        return mulan(texts=prompt).half()
+
+    ext = os.path.splitext(wav_path)[-1].lower()
+    if ext == ".mp3":
+        meta = MP3(wav_path)
+        audio_len = meta.info.length
+    elif ext in [".wav", ".flac"]:
+        audio_len = librosa.get_duration(path=wav_path)
+    else:
+        raise ValueError("Unsupported file format: {}".format(ext))
+
+    if audio_len < 10:
+        print(
+            f"Warning: The audio file {wav_path} is too short ({audio_len:.2f} seconds). Expected at least 10 seconds."
+        )
+
+    assert audio_len >= 10
+
+    mid_time = audio_len // 2
+    start_time = mid_time - 5
+    wav, _ = librosa.load(wav_path, sr=24000, offset=start_time, duration=10)
+
+    wav = torch.tensor(wav).unsqueeze(0).to(model.device)
+
+    with torch.no_grad():
+        audio_emb = mulan(wavs=wav)  # [1, 512]
+
+    audio_emb = audio_emb
+    audio_emb = audio_emb.half()
+
+    return audio_emb
+
+
+def parse_lyrics(lyrics: str):
+    lyrics_with_time = []
+    lyrics = lyrics.strip()
+    for line in lyrics.split("\n"):
+        try:
+            time, lyric = line[1:9], line[10:]
+            lyric = lyric.strip()
+            mins, secs = time.split(":")
+            secs = int(mins) * 60 + float(secs)
+            lyrics_with_time.append((secs, lyric))
+        except:
+            continue
+    return lyrics_with_time
+
+
+class CNENTokenizer:
+    def __init__(self):
+        with open("./g2p/g2p/vocab.json", "r", encoding='utf-8') as file:
+            self.phone2id: dict = json.load(file)["vocab"]
+        self.id2phone = {v: k for (k, v) in self.phone2id.items()}
+        from g2p.g2p_generation import chn_eng_g2p
+
+        self.tokenizer = chn_eng_g2p
+
+    def encode(self, text):
+        phone, token = self.tokenizer(text)
+        token = [x + 1 for x in token]
+        return token
+
+    def decode(self, token):
+        return "|".join([self.id2phone[x - 1] for x in token])
+
+
+def get_lrc_token(max_frames, text, tokenizer, max_secs, device):
+
+    lyrics_shift = 0
+    sampling_rate = 44100
+    downsample_rate = 2048
+    # max_secs = max_frames / (sampling_rate / downsample_rate)
+
+    comma_token_id = 1
+    period_token_id = 2
+
+    lrc_with_time = parse_lyrics(text)
+
+    modified_lrc_with_time = []
+    for i in range(len(lrc_with_time)):
+        time, line = lrc_with_time[i]
+        line_token = tokenizer.encode(line)
+        modified_lrc_with_time.append((time, line_token))
+    lrc_with_time = modified_lrc_with_time
+
+    lrc_with_time = [
+        (time_start, line)
+        for (time_start, line) in lrc_with_time
+        if time_start < max_secs
+    ]
+    if max_frames == 2048:
+        lrc_with_time = lrc_with_time[:-1] if len(lrc_with_time) >= 1 else lrc_with_time
+        
+    end_frame = max_frames if max_frames == 2048 else int(max_secs * (sampling_rate / downsample_rate))
+    end_frame = min(end_frame, max_frames) 
+
+    normalized_start_time = 0.0
+    
+    normalized_duration = end_frame / max_frames
+
+    lrc = torch.zeros((end_frame,), dtype=torch.long)
+
+    tokens_count = 0
+    last_end_pos = 0
+    for time_start, line in lrc_with_time:
+        tokens = [
+            token if token != period_token_id else comma_token_id for token in line
+        ] + [period_token_id]
+        tokens = torch.tensor(tokens, dtype=torch.long)
+        num_tokens = tokens.shape[0]
+
+        gt_frame_start = int(time_start * sampling_rate / downsample_rate)
+
+        frame_shift = random.randint(int(-lyrics_shift), int(lyrics_shift))
+
+        frame_start = max(gt_frame_start - frame_shift, last_end_pos)
+        frame_len = min(num_tokens, end_frame - frame_start)
+
+        lrc[frame_start : frame_start + frame_len] = tokens[:frame_len]
+
+        tokens_count += num_tokens
+        last_end_pos = frame_start + frame_len
+
+    lrc_emb = lrc.unsqueeze(0).to(device)
+
+    normalized_start_time = torch.tensor(normalized_start_time).unsqueeze(0).to(device)
+    normalized_start_time = normalized_start_time.half()
+    
+    normalized_duration = torch.tensor(normalized_duration).unsqueeze(0).to(device)
+    normalized_duration = normalized_duration.half()
+
+    return lrc_emb, normalized_start_time, end_frame, normalized_duration
+
+
+def load_checkpoint(model, ckpt_path, device, use_ema=True):
+    model = model.half()
+
+    ckpt_type = ckpt_path.split(".")[-1]
+    if ckpt_type == "safetensors":
+        from safetensors.torch import load_file
+
+        checkpoint = load_file(ckpt_path)
+    else:
+        checkpoint = torch.load(ckpt_path, weights_only=True)
+
+    if use_ema:
+        if ckpt_type == "safetensors":
+            checkpoint = {"ema_model_state_dict": checkpoint}
+        checkpoint["model_state_dict"] = {
+            k.replace("ema_model.", ""): v
+            for k, v in checkpoint["ema_model_state_dict"].items()
+            if k not in ["initted", "step"]
+        }
+        model.load_state_dict(checkpoint["model_state_dict"], strict=False)
+    else:
+        if ckpt_type == "safetensors":
+            checkpoint = {"model_state_dict": checkpoint}
+        model.load_state_dict(checkpoint["model_state_dict"], strict=False)
+
+    return model.to(device)

code/model/__init__.py

@@ -0,0 +1,6 @@
+from model.cfm import CFM
+from model.dit import DiT
+from model.trainer import Trainer
+
+
+__all__ = ["CFM"]

code/model/cfm.py

@@ -0,0 +1,318 @@
+# Copyright (c) 2025 ASLP-LAB
+#               2025 Ziqian Ning   (ningziqian@mail.nwpu.edu.cn)
+#               2025 Huakang Chen  (huakang@mail.nwpu.edu.cn)
+#               2025 Guobin Ma     (guobin.ma@mail.nwpu.edu.cn)
+#
+# 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.
+
+""" This implementation is adapted from github repo:
+    https://github.com/SWivid/F5-TTS.
+"""
+
+from __future__ import annotations
+from typing import Callable
+from random import random
+
+import torch
+from torch import nn
+import torch
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+
+from torchdiffeq import odeint
+
+from model.utils import (
+    exists,
+    list_str_to_idx,
+    list_str_to_tensor,
+    lens_to_mask,
+    mask_from_frac_lengths,
+)
+
+def custom_mask_from_start_end_indices(
+    seq_len: int["b"],  # noqa: F821
+    latent_pred_segments,
+    device,
+    max_seq_len
+):
+    max_seq_len = max_seq_len
+    seq = torch.arange(max_seq_len, device=device).long()
+
+    res_mask = torch.zeros(max_seq_len, device=device, dtype=torch.bool)
+    
+    for start, end in latent_pred_segments:
+        start = start.unsqueeze(0)
+        end = end.unsqueeze(0)
+        start_mask = seq[None, :] >= start[:, None]
+        end_mask = seq[None, :] < end[:, None]
+        res_mask = res_mask | (start_mask & end_mask)
+    
+    return res_mask
+
+class CFM(nn.Module):
+    def __init__(
+        self,
+        transformer: nn.Module,
+        sigma=0.0,
+        odeint_kwargs: dict = dict(
+            method="euler"
+        ),
+        odeint_options: dict = dict(
+            min_step=0.05
+        ),
+        audio_drop_prob=0.3,
+        cond_drop_prob=0.2,
+        style_drop_prob=0.1,
+        lrc_drop_prob=0.1,
+        num_channels=None,
+        frac_lengths_mask: tuple[float, float] = (0.7, 1.0),
+        vocab_char_map: dict[str:int] | None = None,
+        max_frames=2048
+    ):
+        super().__init__()
+
+        self.frac_lengths_mask = frac_lengths_mask
+
+        self.num_channels = num_channels
+
+        # classifier-free guidance
+        self.audio_drop_prob = audio_drop_prob
+        self.cond_drop_prob = cond_drop_prob
+        self.style_drop_prob = style_drop_prob
+        self.lrc_drop_prob = lrc_drop_prob
+
+        # transformer
+        self.transformer = transformer
+        dim = transformer.dim
+        self.dim = dim
+
+        # conditional flow related
+        self.sigma = sigma
+
+        # sampling related
+        self.odeint_kwargs = odeint_kwargs
+        
+        self.odeint_options = odeint_options
+
+        # vocab map for tokenization
+        self.vocab_char_map = vocab_char_map
+        
+        self.max_frames = max_frames
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @torch.no_grad()
+    def sample(
+        self,
+        cond: float["b n d"] | float["b nw"],  # noqa: F722
+        text: int["b nt"] | list[str],  # noqa: F722
+        duration: int | int["b"],  # noqa: F821
+        *,
+        style_prompt = None,
+        style_prompt_lens = None,
+        negative_style_prompt = None,
+        lens: int["b"] | None = None,  # noqa: F821
+        steps=32,
+        cfg_strength=4.0,
+        sway_sampling_coef=None,
+        seed: int | None = None,
+        max_duration=6144,
+        vocoder: Callable[[float["b d n"]], float["b nw"]] | None = None,  # noqa: F722
+        no_ref_audio=False,
+        duplicate_test=False,
+        t_inter=0.1,
+        edit_mask=None,
+        start_time=None,
+        latent_pred_segments=None,
+        song_duration=None,
+        batch_infer_num=1
+    ):
+        self.eval()
+
+        if next(self.parameters()).dtype == torch.float16:
+            cond = cond.half()
+
+        # raw wave
+        if cond.shape[1] > duration:
+            cond = cond[:, :duration, :]
+
+        if cond.ndim == 2:
+            cond = self.mel_spec(cond)
+            cond = cond.permute(0, 2, 1)
+            assert cond.shape[-1] == self.num_channels
+
+        batch, cond_seq_len, device = *cond.shape[:2], cond.device
+        if not exists(lens):
+            lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)
+
+        # text
+        if isinstance(text, list):
+            if exists(self.vocab_char_map):
+                text = list_str_to_idx(text, self.vocab_char_map).to(device)
+            else:
+                text = list_str_to_tensor(text).to(device)
+            assert text.shape[0] == batch
+
+        # duration
+        cond_mask = lens_to_mask(lens)
+        if edit_mask is not None:
+            cond_mask = cond_mask & edit_mask
+
+        latent_pred_segments = torch.tensor(latent_pred_segments).to(cond.device)
+        fixed_span_mask = custom_mask_from_start_end_indices(cond_seq_len, latent_pred_segments, device=cond.device, max_seq_len=duration)
+        fixed_span_mask = fixed_span_mask.unsqueeze(-1)
+        step_cond = torch.where(fixed_span_mask, torch.zeros_like(cond), cond)
+
+        if isinstance(duration, int):
+            duration = torch.full((batch_infer_num,), duration, device=device, dtype=torch.long)
+
+        duration = duration.clamp(max=max_duration)
+        max_duration = duration.amax()
+
+        # duplicate test corner for inner time step oberservation
+        if duplicate_test:
+            test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2 * cond_seq_len), value=0.0)
+
+        if batch > 1:
+            mask = lens_to_mask(duration)
+        else:  # save memory and speed up, as single inference need no mask currently
+            mask = None
+
+        # test for no ref audio
+        if no_ref_audio:
+            cond = torch.zeros_like(cond)
+
+        cond = cond.repeat(batch_infer_num, 1, 1)
+        step_cond = step_cond.repeat(batch_infer_num, 1, 1)
+        text = text.repeat(batch_infer_num, 1)
+        style_prompt = style_prompt.repeat(batch_infer_num, 1)
+        negative_style_prompt = negative_style_prompt.repeat(batch_infer_num, 1)
+        start_time = start_time.repeat(batch_infer_num)
+        fixed_span_mask = fixed_span_mask.repeat(batch_infer_num, 1, 1)
+        song_duration = song_duration.repeat(batch_infer_num)
+
+        def fn(t, x):
+            # predict flow
+            pred = self.transformer(
+                x=x, cond=step_cond, text=text, time=t, drop_audio_cond=False, drop_text=False, drop_prompt=False,
+                style_prompt=style_prompt, start_time=start_time, duration=song_duration
+            )
+            if cfg_strength < 1e-5:
+                return pred
+
+            null_pred = self.transformer(
+                x=x, cond=step_cond, text=text, time=t, drop_audio_cond=True, drop_text=True, drop_prompt=False,
+                style_prompt=negative_style_prompt, start_time=start_time, duration=song_duration
+            )
+            return pred + (pred - null_pred) * cfg_strength
+
+        # noise input
+        # to make sure batch inference result is same with different batch size, and for sure single inference
+        # still some difference maybe due to convolutional layers
+        y0 = []
+        for dur in duration:
+            if exists(seed):
+                torch.manual_seed(seed)
+            y0.append(torch.randn(dur, self.num_channels, device=self.device, dtype=step_cond.dtype))
+        y0 = pad_sequence(y0, padding_value=0, batch_first=True)
+
+        t_start = 0
+
+        # duplicate test corner for inner time step oberservation
+        if duplicate_test:
+            t_start = t_inter
+            y0 = (1 - t_start) * y0 + t_start * test_cond
+            steps = int(steps * (1 - t_start))
+        
+        t = torch.linspace(t_start, 1, steps, device=self.device, dtype=step_cond.dtype)
+        if sway_sampling_coef is not None:
+            t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
+
+        trajectory = odeint(fn, y0, t, **self.odeint_kwargs)
+
+        sampled = trajectory[-1]
+        out = sampled
+        out = torch.where(fixed_span_mask, out, cond)
+
+        if exists(vocoder):
+            out = out.permute(0, 2, 1)
+            out = vocoder(out)
+
+        out = torch.chunk(out, batch_infer_num, dim=0)
+        return out, trajectory
+
+    def forward(
+        self,
+        inp: float["b n d"] | float["b nw"],  # mel or raw wave  # noqa: F722
+        text: int["b nt"] | list[str],  # noqa: F722
+        style_prompt = None,
+        style_prompt_lens = None,
+        lens: int["b"] | None = None,  # noqa: F821
+        noise_scheduler: str | None = None,
+        grad_ckpt = False,
+        start_time = None,
+    ):
+
+        batch, seq_len, dtype, device, _σ1 = *inp.shape[:2], inp.dtype, self.device, self.sigma
+
+        # lens and mask
+        if not exists(lens):
+            lens = torch.full((batch,), seq_len, device=device)
+
+        mask = lens_to_mask(lens, length=seq_len)  # useless here, as collate_fn will pad to max length in batch
+
+        # get a random span to mask out for training conditionally
+        frac_lengths = torch.zeros((batch,), device=self.device).float().uniform_(*self.frac_lengths_mask)
+        rand_span_mask = mask_from_frac_lengths(lens, frac_lengths, self.max_frames)
+
+        if exists(mask):
+            rand_span_mask = mask
+
+        # mel is x1
+        x1 = inp
+
+        # x0 is gaussian noise
+        x0 = torch.randn_like(x1)
+
+        # time step
+        time = torch.normal(mean=0, std=1, size=(batch,), device=self.device)
+        time = torch.nn.functional.sigmoid(time)
+        # TODO. noise_scheduler
+
+        # sample xt (��_t(x) in the paper)
+        t = time.unsqueeze(-1).unsqueeze(-1)
+        φ = (1 - t) * x0 + t * x1
+        flow = x1 - x0
+
+        # only predict what is within the random mask span for infilling
+        cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1)
+
+        # transformer and cfg training with a drop rate
+        drop_audio_cond = random() < self.audio_drop_prob  # p_drop in voicebox paper
+        drop_text = random() < self.lrc_drop_prob
+        drop_prompt = random() < self.style_drop_prob
+
+        # if want rigourously mask out padding, record in collate_fn in dataset.py, and pass in here
+        # adding mask will use more memory, thus also need to adjust batchsampler with scaled down threshold for long sequences
+        pred = self.transformer(
+            x=φ, cond=cond, text=text, time=time, drop_audio_cond=drop_audio_cond, drop_text=drop_text, drop_prompt=drop_prompt,
+            style_prompt=style_prompt, start_time=start_time
+        )
+
+        # flow matching loss
+        loss = F.mse_loss(pred, flow, reduction="none")
+        loss = loss[rand_span_mask]
+
+        return loss.mean(), cond, pred

code/model/dit.py

@@ -0,0 +1,224 @@
+# Copyright (c) 2025 ASLP-LAB
+#               2025 Ziqian Ning   (ningziqian@mail.nwpu.edu.cn)
+#               2025 Huakang Chen  (huakang@mail.nwpu.edu.cn)
+#               2025 Yuepeng Jiang (Jiangyp@mail.nwpu.edu.cn)
+#
+# 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.
+
+""" This implementation is adapted from github repo:
+    https://github.com/SWivid/F5-TTS.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch
+
+from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRotaryEmbedding
+from transformers.models.llama import LlamaConfig
+
+from model.modules import (
+    TimestepEmbedding,
+    ConvNeXtV2Block,
+    ConvPositionEmbedding,
+    AdaLayerNormZero_Final,
+    precompute_freqs_cis,
+    get_pos_embed_indices,
+    _prepare_decoder_attention_mask,
+)
+
+# Text embedding
+class TextEmbedding(nn.Module):
+    def __init__(self, text_num_embeds, text_dim, max_pos, conv_layers=0, conv_mult=2):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
+
+        if conv_layers > 0:
+            self.extra_modeling = True
+            self.precompute_max_pos = max_pos  # ~44s of 24khz audio
+            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)
+            self.text_blocks = nn.Sequential(
+                *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]
+            )
+        else:
+            self.extra_modeling = False
+
+    def forward(self, text: int["b nt"], seq_len, drop_text=False):  # noqa: F722
+        batch, text_len = text.shape[0], text.shape[1]
+
+        if drop_text:  # cfg for text
+            text = torch.zeros_like(text)
+
+        text = self.text_embed(text)  # b n -> b n d
+
+        # possible extra modeling
+        if self.extra_modeling:
+            # sinus pos emb
+            batch_start = torch.zeros((batch,), dtype=torch.long)
+            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
+            text_pos_embed = self.freqs_cis[pos_idx]
+            text = text + text_pos_embed
+
+            # convnextv2 blocks
+            text = self.text_blocks(text)
+
+        return text
+
+
+# noised input audio and context mixing embedding
+class InputEmbedding(nn.Module):
+    def __init__(self, mel_dim, text_dim, out_dim, cond_dim):
+        super().__init__()
+        self.proj = nn.Linear(mel_dim * 2 + text_dim + cond_dim * 2, out_dim)
+        self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
+
+    def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], style_emb, time_emb, drop_audio_cond=False):  # noqa: F722
+        if drop_audio_cond:  # cfg for cond audio
+            cond = torch.zeros_like(cond)
+        style_emb = style_emb.unsqueeze(1).repeat(1, x.shape[1], 1)
+        time_emb = time_emb.unsqueeze(1).repeat(1, x.shape[1], 1)
+        x = self.proj(torch.cat((x, cond, text_embed, style_emb, time_emb), dim=-1))
+        x = self.conv_pos_embed(x) + x
+        return x
+
+
+# Transformer backbone using Llama blocks
+class DiT(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth=8,
+        heads=8,
+        dim_head=64,
+        dropout=0.1,
+        ff_mult=4,
+        mel_dim=100,
+        text_num_embeds=256,
+        text_dim=None,
+        conv_layers=0,
+        long_skip_connection=False,
+        max_frames=2048
+    ):
+        super().__init__()
+        
+        self.max_frames = max_frames
+
+        cond_dim = 512
+        self.time_embed = TimestepEmbedding(cond_dim)
+        self.start_time_embed = TimestepEmbedding(cond_dim)
+        self.duration_time_embed = TimestepEmbedding(cond_dim) if self.max_frames == 6144 else None
+        if text_dim is None:
+            text_dim = mel_dim
+        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers, max_pos=self.max_frames)
+        self.input_embed = InputEmbedding(mel_dim, text_dim, dim, cond_dim=cond_dim)
+
+        self.dim = dim
+        self.depth = depth
+
+        llama_config = LlamaConfig(hidden_size=dim, intermediate_size=dim * ff_mult, hidden_act='silu', max_position_embeddings=self.max_frames)
+        llama_config._attn_implementation = 'sdpa'
+        self.transformer_blocks = nn.ModuleList(
+            [LlamaDecoderLayer(llama_config, layer_idx=i) for i in range(depth)]
+        )
+        self.rotary_emb = LlamaRotaryEmbedding(config=llama_config)
+        self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None
+
+        self.text_fusion_linears = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Linear(cond_dim, dim),
+                    nn.SiLU()
+                ) for i in range(depth // 2)
+            ]
+        )
+        for layer in self.text_fusion_linears:
+            for p in layer.parameters():
+                p.detach().zero_()
+
+        self.norm_out = AdaLayerNormZero_Final(dim, cond_dim)  # final modulation
+        self.proj_out = nn.Linear(dim, mel_dim)
+
+    def forward_timestep_invariant(self, text, seq_len, drop_text, start_time):
+        s_t = self.start_time_embed(start_time)
+        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
+        text_residuals = []
+        for layer in self.text_fusion_linears:
+            text_residual = layer(text_embed)
+            text_residuals.append(text_residual)
+        return s_t, text_embed, text_residuals
+
+
+    def forward(
+        self,
+        x: float["b n d"],  # nosied input audio  # noqa: F722
+        cond: float["b n d"],  # masked cond audio  # noqa: F722
+        text: int["b nt"],  # text  # noqa: F722
+        time: float["b"] | float[""],  # time step  # noqa: F821 F722
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,  # cfg for text
+        drop_prompt=False,
+        style_prompt=None, # [b d t]
+        start_time=None,
+        duration=None
+    ):
+
+        batch, seq_len = x.shape[0], x.shape[1]
+        if time.ndim == 0:
+            time = time.repeat(batch)
+
+        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio
+        t = self.time_embed(time)
+        s_t = self.start_time_embed(start_time)
+        d_t = self.duration_time_embed(duration) if self.max_frames == 6144 else torch.zeros_like(s_t)
+        c = t + s_t + d_t
+        text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
+
+        if drop_prompt:
+            style_prompt = torch.zeros_like(style_prompt)
+        
+        style_embed = style_prompt # [b, 512]
+
+        x = self.input_embed(x, cond, text_embed, style_embed, c, drop_audio_cond=drop_audio_cond)
+
+        if self.long_skip_connection is not None:
+            residual = x
+
+        pos_ids = torch.arange(x.shape[1], device=x.device)
+        pos_ids = pos_ids.unsqueeze(0).repeat(x.shape[0], 1)
+        rotary_embed = self.rotary_emb(x, pos_ids)
+        
+        attention_mask = torch.ones(
+            (batch, seq_len),
+            dtype=torch.bool,
+            device=x.device,
+        )
+        attention_mask = _prepare_decoder_attention_mask(
+            attention_mask,
+            (batch, seq_len),
+            x,
+        )
+
+        for i, block in enumerate(self.transformer_blocks):
+            x, *_ = block(x, attention_mask=attention_mask, position_embeddings=rotary_embed)
+            if i < self.depth // 2:
+                x = x + self.text_fusion_linears[i](text_embed)
+
+        if self.long_skip_connection is not None:
+            x = self.long_skip_connection(torch.cat((x, residual), dim=-1))
+
+        x = self.norm_out(x, c)
+        output = self.proj_out(x)
+
+        return output

code/model/modules.py

@@ -0,0 +1,660 @@
+# Copyright (c) 2025 ASLP-LAB
+#
+# 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.
+
+""" This implementation is adapted from github repo:
+    https://github.com/SWivid/F5-TTS.
+"""
+
+from __future__ import annotations
+from typing import Optional
+import math
+
+import torch
+from torch import nn
+import torch
+import torch.nn.functional as F
+import torchaudio
+
+from x_transformers.x_transformers import apply_rotary_pos_emb
+
+
+
+class FiLMLayer(nn.Module):
+    """
+    Feature-wise Linear Modulation (FiLM) layer
+    Reference: https://arxiv.org/abs/1709.07871
+    """
+    def __init__(self, in_channels, cond_channels):
+
+        super(FiLMLayer, self).__init__()
+        self.in_channels = in_channels
+        self.film = nn.Conv1d(cond_channels, in_channels * 2, 1)
+
+    def forward(self, x, c):
+        gamma, beta = torch.chunk(self.film(c.unsqueeze(2)), chunks=2, dim=1)
+        gamma = gamma.transpose(1, 2)
+        beta = beta.transpose(1, 2)
+        return gamma * x + beta
+
+# raw wav to mel spec
+
+
+class MelSpec(nn.Module):
+    def __init__(
+        self,
+        filter_length=1024,
+        hop_length=256,
+        win_length=1024,
+        n_mel_channels=100,
+        target_sample_rate=24_000,
+        normalize=False,
+        power=1,
+        norm=None,
+        center=True,
+    ):
+        super().__init__()
+        self.n_mel_channels = n_mel_channels
+
+        self.mel_stft = torchaudio.transforms.MelSpectrogram(
+            sample_rate=target_sample_rate,
+            n_fft=filter_length,
+            win_length=win_length,
+            hop_length=hop_length,
+            n_mels=n_mel_channels,
+            power=power,
+            center=center,
+            normalized=normalize,
+            norm=norm,
+        )
+
+        self.register_buffer("dummy", torch.tensor(0), persistent=False)
+
+    def forward(self, inp):
+        if len(inp.shape) == 3:
+            inp = inp.squeeze(1)  # 'b 1 nw -> b nw'
+
+        assert len(inp.shape) == 2
+
+        if self.dummy.device != inp.device:
+            self.to(inp.device)
+
+        mel = self.mel_stft(inp)
+        mel = mel.clamp(min=1e-5).log()
+        return mel
+
+
+# sinusoidal position embedding
+
+
+class SinusPositionEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x, scale=1000):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+# convolutional position embedding
+
+
+class ConvPositionEmbedding(nn.Module):
+    def __init__(self, dim, kernel_size=31, groups=16):
+        super().__init__()
+        assert kernel_size % 2 != 0
+        self.conv1d = nn.Sequential(
+            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
+            nn.Mish(),
+            nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),
+            nn.Mish(),
+        )
+
+    def forward(self, x: float["b n d"], mask: bool["b n"] | None = None):  # noqa: F722
+        if mask is not None:
+            mask = mask[..., None]
+            x = x.masked_fill(~mask, 0.0)
+
+        x = x.permute(0, 2, 1)
+        x = self.conv1d(x)
+        out = x.permute(0, 2, 1)
+
+        if mask is not None:
+            out = out.masked_fill(~mask, 0.0)
+
+        return out
+
+
+# rotary positional embedding related
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):
+    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+    # has some connection to NTK literature
+    # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+    # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
+    theta *= theta_rescale_factor ** (dim / (dim - 2))
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t, freqs).float()  # type: ignore
+    freqs_cos = torch.cos(freqs)  # real part
+    freqs_sin = torch.sin(freqs)  # imaginary part
+    return torch.cat([freqs_cos, freqs_sin], dim=-1)
+
+
+def get_pos_embed_indices(start, length, max_pos, scale=1.0):
+    # length = length if isinstance(length, int) else length.max()
+    scale = scale * torch.ones_like(start, dtype=torch.float32)  # in case scale is a scalar
+    pos = (
+        start.unsqueeze(1)
+        + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()
+    )
+    # avoid extra long error.
+    pos = torch.where(pos < max_pos, pos, max_pos - 1)
+    return pos
+
+
+# Global Response Normalization layer (Instance Normalization ?)
+
+
+class GRN(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=1, keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py
+# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108
+
+
+class ConvNeXtV2Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        dilation: int = 1,
+    ):
+        super().__init__()
+        padding = (dilation * (7 - 1)) // 2
+        self.dwconv = nn.Conv1d(
+            dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation
+        )  # depthwise conv
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.grn = GRN(intermediate_dim)
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = x.transpose(1, 2)  # b n d -> b d n
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # b d n -> b n d
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        return residual + x
+
+
+# AdaLayerNormZero
+# return with modulated x for attn input, and params for later mlp modulation
+
+
+class AdaLayerNormZero(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 6)
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb=None):
+        emb = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+# AdaLayerNormZero for final layer
+# return only with modulated x for attn input, cuz no more mlp modulation
+
+
+class AdaLayerNormZero_Final(nn.Module):
+    def __init__(self, dim, cond_dim):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(cond_dim, dim * 2)
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb):
+        emb = self.linear(self.silu(emb))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+# FeedForward
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        activation = nn.GELU(approximate=approximate)
+        #activation = nn.SiLU()
+        project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
+        self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
+
+    def forward(self, x):
+        return self.ff(x)
+
+
+# Attention with possible joint part
+# modified from diffusers/src/diffusers/models/attention_processor.py
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        processor: JointAttnProcessor | AttnProcessor,
+        dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        context_dim: Optional[int] = None,  # if not None -> joint attention
+        context_pre_only=None,
+    ):
+        super().__init__()
+
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+        self.processor = processor
+
+        self.dim = dim
+        self.heads = heads
+        self.inner_dim = dim_head * heads
+        self.dropout = dropout
+
+        self.context_dim = context_dim
+        self.context_pre_only = context_pre_only
+
+        self.to_q = nn.Linear(dim, self.inner_dim)
+        self.to_k = nn.Linear(dim, self.inner_dim)
+        self.to_v = nn.Linear(dim, self.inner_dim)
+
+        if self.context_dim is not None:
+            self.to_k_c = nn.Linear(context_dim, self.inner_dim)
+            self.to_v_c = nn.Linear(context_dim, self.inner_dim)
+            if self.context_pre_only is not None:
+                self.to_q_c = nn.Linear(context_dim, self.inner_dim)
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, dim))
+        self.to_out.append(nn.Dropout(dropout))
+
+        if self.context_pre_only is not None and not self.context_pre_only:
+            self.to_out_c = nn.Linear(self.inner_dim, dim)
+
+    def forward(
+        self,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        c: float["b n d"] = None,  # context c  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding for x
+        c_rope=None,  # rotary position embedding for c
+    ) -> torch.Tensor:
+        if c is not None:
+            return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope)
+        else:
+            return self.processor(self, x, mask=mask, rope=rope)
+
+
+# Attention processor
+
+
+class AttnProcessor:
+    def __init__(self):
+        pass
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding
+    ) -> torch.FloatTensor:
+        batch_size = x.shape[0]
+
+        # `sample` projections.
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # apply rotary position embedding
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+
+            query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+            key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        # attention
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # mask. e.g. inference got a batch with different target durations, mask out the padding
+        if mask is not None:
+            attn_mask = mask
+            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # 'b n -> b 1 1 n'
+            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+
+        return x
+
+
+# Joint Attention processor for MM-DiT
+# modified from diffusers/src/diffusers/models/attention_processor.py
+
+
+class JointAttnProcessor:
+    def __init__(self):
+        pass
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b n d"],  # noised input x  # noqa: F722
+        c: float["b nt d"] = None,  # context c, here text # noqa: F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        rope=None,  # rotary position embedding for x
+        c_rope=None,  # rotary position embedding for c
+    ) -> torch.FloatTensor:
+        residual = x
+
+        batch_size = c.shape[0]
+
+        # `sample` projections.
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # `context` projections.
+        c_query = attn.to_q_c(c)
+        c_key = attn.to_k_c(c)
+        c_value = attn.to_v_c(c)
+
+        # apply rope for context and noised input independently
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+            query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+            key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+        if c_rope is not None:
+            freqs, xpos_scale = c_rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+            c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)
+            c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)
+
+        # attention
+        query = torch.cat([query, c_query], dim=1)
+        key = torch.cat([key, c_key], dim=1)
+        value = torch.cat([value, c_value], dim=1)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # mask. e.g. inference got a batch with different target durations, mask out the padding
+        if mask is not None:
+            attn_mask = F.pad(mask, (0, c.shape[1]), value=True)  # no mask for c (text)
+            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # 'b n -> b 1 1 n'
+            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # Split the attention outputs.
+        x, c = (
+            x[:, : residual.shape[1]],
+            x[:, residual.shape[1] :],
+        )
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+        if not attn.context_pre_only:
+            c = attn.to_out_c(c)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+            # c = c.masked_fill(~mask, 0.)  # no mask for c (text)
+
+        return x, c
+
+
+# DiT Block
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, use_style_prompt=False):
+        super().__init__()
+
+        self.attn_norm = AdaLayerNormZero(dim)
+        self.attn = Attention(
+            processor=AttnProcessor(),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+        )
+
+        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+
+        self.use_style_prompt = use_style_prompt
+        if use_style_prompt:
+            #self.film = FiLMLayer(dim, dim)
+            self.prompt_norm = AdaLayerNormZero_Final(dim)
+
+    def forward(self, x, t, c=None, mask=None, rope=None):  # x: noised input, t: time embedding
+        if c is not None and self.use_style_prompt:
+            #x = self.film(x, c)
+            x = self.prompt_norm(x, c)
+
+        # pre-norm & modulation for attention input
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
+
+        # attention
+        attn_output = self.attn(x=norm, mask=mask, rope=rope)
+
+        # process attention output for input x
+        x = x + gate_msa.unsqueeze(1) * attn_output
+
+        norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm)
+        x = x + gate_mlp.unsqueeze(1) * ff_output
+
+        return x
+
+
+# MMDiT Block https://arxiv.org/abs/2403.03206
+
+
+class MMDiTBlock(nn.Module):
+    r"""
+    modified from diffusers/src/diffusers/models/attention.py
+
+    notes.
+    _c: context related. text, cond, etc. (left part in sd3 fig2.b)
+    _x: noised input related. (right part)
+    context_pre_only: last layer only do prenorm + modulation cuz no more ffn
+    """
+
+    def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, context_pre_only=False):
+        super().__init__()
+
+        self.context_pre_only = context_pre_only
+
+        self.attn_norm_c = AdaLayerNormZero_Final(dim) if context_pre_only else AdaLayerNormZero(dim)
+        self.attn_norm_x = AdaLayerNormZero(dim)
+        self.attn = Attention(
+            processor=JointAttnProcessor(),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+            context_dim=dim,
+            context_pre_only=context_pre_only,
+        )
+
+        if not context_pre_only:
+            self.ff_norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+            self.ff_c = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+        else:
+            self.ff_norm_c = None
+            self.ff_c = None
+        self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
+
+    def forward(self, x, c, t, mask=None, rope=None, c_rope=None):  # x: noised input, c: context, t: time embedding
+        # pre-norm & modulation for attention input
+        if self.context_pre_only:
+            norm_c = self.attn_norm_c(c, t)
+        else:
+            norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t)
+        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t)
+
+        # attention
+        x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope)
+
+        # process attention output for context c
+        if self.context_pre_only:
+            c = None
+        else:  # if not last layer
+            c = c + c_gate_msa.unsqueeze(1) * c_attn_output
+
+            norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+            c_ff_output = self.ff_c(norm_c)
+            c = c + c_gate_mlp.unsqueeze(1) * c_ff_output
+
+        # process attention output for input x
+        x = x + x_gate_msa.unsqueeze(1) * x_attn_output
+
+        norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None]
+        x_ff_output = self.ff_x(norm_x)
+        x = x + x_gate_mlp.unsqueeze(1) * x_ff_output
+
+        return c, x
+
+
+# time step conditioning embedding
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, dim, freq_embed_dim=256):
+        super().__init__()
+        self.time_embed = SinusPositionEmbedding(freq_embed_dim)
+        self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+
+    def forward(self, timestep: float["b"]):  # noqa: F821
+        time_hidden = self.time_embed(timestep)
+        time_hidden = time_hidden.to(timestep.dtype)
+        time = self.time_mlp(time_hidden)  # b d
+        return time
+
+
+# attention mask realated
+
+
+def _prepare_decoder_attention_mask(attention_mask, input_shape, inputs_embeds):
+    # create noncausal mask
+    # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+    combined_attention_mask = None
+
+    def _expand_mask(
+        mask: torch.Tensor, dtype: torch.dtype, tgt_len: int = None
+    ):
+        """
+        Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+        """
+        bsz, src_len = mask.size()
+        tgt_len = tgt_len if tgt_len is not None else src_len
+
+        expanded_mask = (
+            mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+        )
+
+        inverted_mask = 1.0 - expanded_mask
+
+        return inverted_mask.masked_fill(
+            inverted_mask.to(torch.bool), torch.finfo(dtype).min
+        )
+
+    if attention_mask is not None:
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        expanded_attn_mask = _expand_mask(
+            attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+        ).to(inputs_embeds.device)
+        combined_attention_mask = (
+            expanded_attn_mask
+            if combined_attention_mask is None
+            else expanded_attn_mask + combined_attention_mask
+        )
+
+    return combined_attention_mask

code/model/trainer.py

@@ -0,0 +1,334 @@
+# Copyright (c) 2025 ASLP-LAB
+#               2025 Ziqian Ning   (ningziqian@mail.nwpu.edu.cn)
+#               2025 Huakang Chen  (huakang@mail.nwpu.edu.cn)
+#
+# 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.
+
+""" This implementation is adapted from github repo:
+    https://github.com/SWivid/F5-TTS.
+"""
+
+from __future__ import annotations
+
+import os
+import gc
+from tqdm import tqdm
+import wandb
+
+import torch
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import LinearLR, SequentialLR, ConstantLR
+
+from accelerate import Accelerator
+from accelerate.utils import DistributedDataParallelKwargs
+from dataset.dataset import DiffusionDataset
+
+from torch.utils.data import DataLoader
+
+from ema_pytorch import EMA
+
+from model import CFM
+from model.utils import exists, default
+
+class Trainer:
+    def __init__(
+        self,
+        model: CFM,
+        args,
+        epochs,
+        learning_rate,
+        num_warmup_updates=20000,
+        save_per_updates=1000,
+        checkpoint_path=None,
+        batch_size=32,
+        batch_size_type: str = "sample",
+        max_samples=32,
+        grad_accumulation_steps=1,
+        max_grad_norm=1.0,
+        noise_scheduler: str | None = None,
+        duration_predictor: torch.nn.Module | None = None,
+        wandb_project="test_e2-tts",
+        wandb_run_name="test_run",
+        wandb_resume_id: str = None,
+        last_per_steps=None,
+        accelerate_kwargs: dict = dict(),
+        ema_kwargs: dict = dict(),
+        bnb_optimizer: bool = False,
+        reset_lr: bool = False,
+        use_style_prompt: bool = False,
+        grad_ckpt: bool = False
+    ):
+        self.args = args
+
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False, )
+
+        logger = "wandb" if wandb.api.api_key else None
+        
+        self.accelerator = Accelerator(
+            log_with=logger,
+            kwargs_handlers=[ddp_kwargs],
+            gradient_accumulation_steps=grad_accumulation_steps,
+            **accelerate_kwargs,
+        )
+
+        if logger == "wandb":
+            if exists(wandb_resume_id):
+                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name, "id": wandb_resume_id}}
+            else:
+                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name}}
+            self.accelerator.init_trackers(
+                project_name=wandb_project,
+                init_kwargs=init_kwargs,
+                config={
+                    "epochs": epochs,
+                    "learning_rate": learning_rate,
+                    "num_warmup_updates": num_warmup_updates,
+                    "batch_size": batch_size,
+                    "batch_size_type": batch_size_type,
+                    "max_samples": max_samples,
+                    "grad_accumulation_steps": grad_accumulation_steps,
+                    "max_grad_norm": max_grad_norm,
+                    "gpus": self.accelerator.num_processes,
+                    "noise_scheduler": noise_scheduler,
+                },
+            )
+
+        self.precision = self.accelerator.state.mixed_precision
+        self.precision = self.precision.replace("no", "fp32")
+
+        self.model = model
+
+        if self.is_main:
+            self.ema_model = EMA(model, include_online_model=False, **ema_kwargs)
+
+            self.ema_model.to(self.accelerator.device)
+            if self.accelerator.state.distributed_type in ["DEEPSPEED", "FSDP"]:
+                self.ema_model.half()
+
+        self.epochs = epochs
+        self.num_warmup_updates = num_warmup_updates
+        self.save_per_updates = save_per_updates
+        self.last_per_steps = default(last_per_steps, save_per_updates * grad_accumulation_steps)
+        self.checkpoint_path = default(checkpoint_path, "ckpts/test_e2-tts")
+
+        self.max_samples = max_samples
+        self.grad_accumulation_steps = grad_accumulation_steps
+        self.max_grad_norm = max_grad_norm
+
+        self.noise_scheduler = noise_scheduler
+
+        self.duration_predictor = duration_predictor
+
+        self.reset_lr = reset_lr
+
+        self.use_style_prompt = use_style_prompt
+        
+        self.grad_ckpt = grad_ckpt
+
+        if bnb_optimizer:
+            import bitsandbytes as bnb
+
+            self.optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=learning_rate)
+        else:
+            self.optimizer = AdamW(model.parameters(), lr=learning_rate)
+
+        if self.accelerator.state.distributed_type == "DEEPSPEED":
+            self.accelerator.state.deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu'] = batch_size
+        
+        self.get_dataloader()
+        self.get_scheduler()
+
+        self.model, self.optimizer, self.scheduler, self.train_dataloader = self.accelerator.prepare(self.model, self.optimizer, self.scheduler, self.train_dataloader)
+
+    def get_scheduler(self):
+        warmup_steps = (
+            self.num_warmup_updates * self.accelerator.num_processes
+        )  # consider a fixed warmup steps while using accelerate multi-gpu ddp
+        total_steps = len(self.train_dataloader) * self.epochs / self.grad_accumulation_steps
+        decay_steps = total_steps - warmup_steps
+        warmup_scheduler = LinearLR(self.optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps)
+        decay_scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps)
+        self.scheduler = SequentialLR(
+            self.optimizer, schedulers=[warmup_scheduler, decay_scheduler], milestones=[warmup_steps]
+        )
+
+    def get_constant_scheduler(self):
+        total_steps = len(self.train_dataloader) * self.epochs / self.grad_accumulation_steps
+        self.scheduler = ConstantLR(self.optimizer, factor=1, total_iters=total_steps)
+
+    def get_dataloader(self):
+        print(self.args)
+        dd = DiffusionDataset(self.args.file_path, self.args.max_frames, self.args.min_frames, self.args.sampling_rate, self.args.downsample_rate, self.precision)
+        self.train_dataloader = DataLoader(
+            dataset=dd,
+            batch_size=self.args.batch_size,
+            shuffle=True,
+            num_workers=4,
+            pin_memory=True,
+            collate_fn=dd.custom_collate_fn,
+            persistent_workers=True
+        )
+
+
+    @property
+    def is_main(self):
+        return self.accelerator.is_main_process
+
+    def save_checkpoint(self, step, last=False):
+        self.accelerator.wait_for_everyone()
+        if self.is_main:
+            checkpoint = dict(
+                model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
+                optimizer_state_dict=self.accelerator.unwrap_model(self.optimizer).state_dict(),
+                ema_model_state_dict=self.ema_model.state_dict(),
+                scheduler_state_dict=self.scheduler.state_dict(),
+                step=step,
+            )
+            if not os.path.exists(self.checkpoint_path):
+                os.makedirs(self.checkpoint_path)
+            if last:
+                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
+                print(f"Saved last checkpoint at step {step}")
+            else:
+                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")
+
+    def load_checkpoint(self):
+        if (
+            not exists(self.checkpoint_path)
+            or not os.path.exists(self.checkpoint_path)
+            or not os.listdir(self.checkpoint_path)
+        ):
+            return 0
+
+        self.accelerator.wait_for_everyone()
+        if "model_last.pt" in os.listdir(self.checkpoint_path):
+            latest_checkpoint = "model_last.pt"
+        else:
+            latest_checkpoint = sorted(
+                [f for f in os.listdir(self.checkpoint_path) if f.endswith(".pt")],
+                key=lambda x: int("".join(filter(str.isdigit, x))),
+            )[-1]
+        
+        checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", map_location="cpu")
+
+        if self.is_main:
+            ema_dict = self.ema_model.state_dict()
+            ema_checkpoint_dict = checkpoint["ema_model_state_dict"]
+
+            filtered_ema_dict = {
+                k: v for k, v in ema_checkpoint_dict.items()
+                if k in ema_dict and ema_dict[k].shape == v.shape 
+            }
+
+            self.ema_model.load_state_dict(filtered_ema_dict, strict=False)
+
+        model_dict = self.accelerator.unwrap_model(self.model).state_dict()
+        checkpoint_model_dict = checkpoint["model_state_dict"]
+
+        filtered_model_dict = {
+            k: v for k, v in checkpoint_model_dict.items()
+            if k in model_dict and model_dict[k].shape == v.shape  
+        }
+
+        self.accelerator.unwrap_model(self.model).load_state_dict(filtered_model_dict, strict=False)
+        
+        if "step" in checkpoint:
+            if self.scheduler and not self.reset_lr:
+                self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
+            step = checkpoint["step"]
+        else:
+            step = 0
+
+        del checkpoint
+        gc.collect()
+        print("Checkpoint loaded at step", step)
+        return step
+
+    def train(self, resumable_with_seed: int = None):
+        train_dataloader = self.train_dataloader
+
+        start_step = self.load_checkpoint()
+        global_step = start_step
+
+        if resumable_with_seed > 0:
+            orig_epoch_step = len(train_dataloader)
+            skipped_epoch = int(start_step // orig_epoch_step)
+            skipped_batch = start_step % orig_epoch_step
+            skipped_dataloader = self.accelerator.skip_first_batches(train_dataloader, num_batches=skipped_batch)
+        else:
+            skipped_epoch = 0
+
+        for epoch in range(skipped_epoch, self.epochs):
+            self.model.train()
+            if resumable_with_seed > 0 and epoch == skipped_epoch:
+                progress_bar = tqdm(
+                    skipped_dataloader,
+                    desc=f"Epoch {epoch+1}/{self.epochs}",
+                    unit="step",
+                    disable=not self.accelerator.is_local_main_process,
+                    initial=skipped_batch,
+                    total=orig_epoch_step,
+                    smoothing=0.15
+                )
+            else:
+                progress_bar = tqdm(
+                    train_dataloader,
+                    desc=f"Epoch {epoch+1}/{self.epochs}",
+                    unit="step",
+                    disable=not self.accelerator.is_local_main_process,
+                    smoothing=0.15
+                )
+
+            for batch in progress_bar:
+                with self.accelerator.accumulate(self.model):
+                    text_inputs = batch["lrc"]
+                    mel_spec = batch["latent"].permute(0, 2, 1)
+                    mel_lengths = batch["latent_lengths"]
+                    style_prompt = batch["prompt"]
+                    style_prompt_lens = batch["prompt_lengths"]
+                    start_time = batch["start_time"]
+
+                    loss, cond, pred = self.model(
+                        mel_spec, text=text_inputs, lens=mel_lengths, noise_scheduler=self.noise_scheduler,
+                        style_prompt=style_prompt,
+                        style_prompt_lens=style_prompt_lens,
+                        grad_ckpt=self.grad_ckpt, start_time=start_time
+                    )
+                    self.accelerator.backward(loss)
+
+                    if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
+                        self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+
+                    self.optimizer.step()
+                    self.scheduler.step()
+                    self.optimizer.zero_grad()
+
+                if self.is_main:
+                    self.ema_model.update()
+
+                global_step += 1
+
+                if self.accelerator.is_local_main_process:
+                    self.accelerator.log({"loss": loss.item(), "lr": self.scheduler.get_last_lr()[0]}, step=global_step)
+
+                progress_bar.set_postfix(step=str(global_step), loss=loss.item())
+
+                if global_step % (self.save_per_updates * self.grad_accumulation_steps) == 0:
+                    self.save_checkpoint(global_step)
+
+                if global_step % self.last_per_steps == 0:
+                    self.save_checkpoint(global_step, last=True)
+
+        self.save_checkpoint(global_step, last=True)
+
+        self.accelerator.end_training()

code/model/utils.py

@@ -0,0 +1,200 @@
+# Copyright (c) 2025 ASLP-LAB
+#
+# 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.
+
+""" This implementation is adapted from github repo:
+    https://github.com/SWivid/F5-TTS.
+"""
+
+from __future__ import annotations
+
+import os
+import random
+from collections import defaultdict
+from importlib.resources import files
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+
+# seed everything
+
+
+def seed_everything(seed=0):
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+# helpers
+
+
+def exists(v):
+    return v is not None
+
+
+def default(v, d):
+    return v if exists(v) else d
+
+
+# tensor helpers
+
+
+def lens_to_mask(t: int["b"], length: int | None = None) -> bool["b n"]:  # noqa: F722 F821
+    if not exists(length):
+        length = t.amax()
+
+    seq = torch.arange(length, device=t.device)
+    return seq[None, :] < t[:, None]
+
+
+def mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"], max_frames):  # noqa: F722 F821
+    max_seq_len = max_frames
+    seq = torch.arange(max_seq_len, device=start.device).long()
+    start_mask = seq[None, :] >= start[:, None]
+    end_mask = seq[None, :] < end[:, None]
+    return start_mask & end_mask
+
+
+def mask_from_frac_lengths(seq_len: int["b"], frac_lengths: float["b"], max_frames):  # noqa: F722 F821
+    lengths = (frac_lengths * seq_len).long()
+    max_start = seq_len - lengths
+
+    rand = torch.rand_like(frac_lengths)
+    start = (max_start * rand).long().clamp(min=0)
+    end = start + lengths
+
+    return mask_from_start_end_indices(seq_len, start, end, max_frames)
+
+
+def maybe_masked_mean(t: float["b n d"], mask: bool["b n"] = None) -> float["b d"]:  # noqa: F722
+    if not exists(mask):
+        return t.mean(dim=1)
+
+    t = torch.where(mask[:, :, None], t, torch.tensor(0.0, device=t.device))
+    num = t.sum(dim=1)
+    den = mask.float().sum(dim=1)
+
+    return num / den.clamp(min=1.0)
+
+
+# simple utf-8 tokenizer, since paper went character based
+def list_str_to_tensor(text: list[str], padding_value=-1) -> int["b nt"]:  # noqa: F722
+    list_tensors = [torch.tensor([*bytes(t, "UTF-8")]) for t in text]  # ByT5 style
+    text = pad_sequence(list_tensors, padding_value=padding_value, batch_first=True)
+    return text
+
+
+# char tokenizer, based on custom dataset's extracted .txt file
+def list_str_to_idx(
+    text: list[str] | list[list[str]],
+    vocab_char_map: dict[str, int],  # {char: idx}
+    padding_value=-1,
+) -> int["b nt"]:  # noqa: F722
+    list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text]  # pinyin or char style
+    text = pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True)
+    return text
+
+
+# Get tokenizer
+
+
+def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
+    """
+    tokenizer   - "pinyin" do g2p for only chinese characters, need .txt vocab_file
+                - "char" for char-wise tokenizer, need .txt vocab_file
+                - "byte" for utf-8 tokenizer
+                - "custom" if you're directly passing in a path to the vocab.txt you want to use
+    vocab_size  - if use "pinyin", all available pinyin types, common alphabets (also those with accent) and symbols
+                - if use "char", derived from unfiltered character & symbol counts of custom dataset
+                - if use "byte", set to 256 (unicode byte range)
+    """
+    if tokenizer in ["pinyin", "char"]:
+        tokenizer_path = os.path.join(files("diffrhythm").joinpath("../../data"), f"{dataset_name}_{tokenizer}/vocab.txt")
+        with open(tokenizer_path, "r", encoding="utf-8") as f:
+            vocab_char_map = {}
+            for i, char in enumerate(f):
+                vocab_char_map[char[:-1]] = i
+        vocab_size = len(vocab_char_map)
+        assert vocab_char_map[" "] == 0, "make sure space is of idx 0 in vocab.txt, cuz 0 is used for unknown char"
+
+    elif tokenizer == "byte":
+        vocab_char_map = None
+        vocab_size = 256
+
+    elif tokenizer == "custom":
+        with open(dataset_name, "r", encoding="utf-8") as f:
+            vocab_char_map = {}
+            for i, char in enumerate(f):
+                vocab_char_map[char[:-1]] = i
+        vocab_size = len(vocab_char_map)
+
+    return vocab_char_map, vocab_size
+
+
+# convert char to pinyin
+
+
+def convert_char_to_pinyin(text_list, polyphone=True):
+    final_text_list = []
+    god_knows_why_en_testset_contains_zh_quote = str.maketrans(
+        {"“": '"', "”": '"', "‘": "'", "’": "'"}
+    )  # in case librispeech (orig no-pc) test-clean
+    custom_trans = str.maketrans({";": ","})  # add custom trans here, to address oov
+    for text in text_list:
+        char_list = []
+        text = text.translate(god_knows_why_en_testset_contains_zh_quote)
+        text = text.translate(custom_trans)
+        for seg in jieba.cut(text):
+            seg_byte_len = len(bytes(seg, "UTF-8"))
+            if seg_byte_len == len(seg):  # if pure alphabets and symbols
+                if char_list and seg_byte_len > 1 and char_list[-1] not in " :'\"":
+                    char_list.append(" ")
+                char_list.extend(seg)
+            elif polyphone and seg_byte_len == 3 * len(seg):  # if pure chinese characters
+                seg = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True)
+                for c in seg:
+                    if c not in "。，、；：？！《》【】—…":
+                        char_list.append(" ")
+                    char_list.append(c)
+            else:  # if mixed chinese characters, alphabets and symbols
+                for c in seg:
+                    if ord(c) < 256:
+                        char_list.extend(c)
+                    else:
+                        if c not in "。，、；：？！《》【】—…":
+                            char_list.append(" ")
+                            char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True))
+                        else:  # if is zh punc
+                            char_list.append(c)
+        final_text_list.append(char_list)
+
+    return final_text_list
+
+
+# filter func for dirty data with many repetitions
+
+
+def repetition_found(text, length=2, tolerance=10):
+    pattern_count = defaultdict(int)
+    for i in range(len(text) - length + 1):
+        pattern = text[i : i + length]
+        pattern_count[pattern] += 1
+    for pattern, count in pattern_count.items():
+        if count > tolerance:
+            return True
+    return False

code/pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "tt-diffrhythm"
+version = "0.1.0"
+description = "DiffRhythm 1.2 Full on one Tenstorrent Blackhole p100a"
+readme = "PYTHON.md"
+requires-python = ">=3.10"
+license = { text = "Apache-2.0" }
+# ttnn / tt-metal is not on PyPI: it comes from the tt-metal environment this runs in.
+dependencies = [
+    "torch",
+    "numpy",
+    "safetensors",
+    "huggingface_hub",
+]
+
+[project.optional-dependencies]
+test = ["pytest"]
+
+[tool.setuptools.packages.find]
+include = ["tt_diffrhythm*"]

code/tt_diffrhythm/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DiffRhythm 1.2 Full on one Tenstorrent Blackhole p100a.
+
+The CFM transformer and the stable-audio VAE decoder run on the card in TTNN. The style
+encoder, the sampler and the post-processing stay on the CPU inside DiffRhythm's own
+environment. `LocalWorker` runs both in one process when one environment satisfies DiffRhythm
+and ttnn at once; `TTWorker` keeps them apart, talking to the card through a worker process over
+a Unix socket, when it does not.
+"""
+from .common import close_device, dram_stats, open_device
+from .diffrhythm_dit import DiffRhythmDiT, DRPrecision
+from .local import LocalWorker
+from .oobleck_vae import OobleckTT
+from .remote import RemoteDiffRhythmDiT, RemoteVae, TTWorker
+
+__all__ = ["DiffRhythmDiT", "DRPrecision", "OobleckTT", "TTWorker", "LocalWorker", "RemoteDiffRhythmDiT",
+           "RemoteVae", "open_device", "close_device", "dram_stats"]
+__version__ = "0.1.0"

code/tt_diffrhythm/common.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Device and host helpers shared with the ACE-Step port.
+
+Extracted verbatim from `tt-acestep` (Apache-2.0), where the same code opens the card, reports
+DRAM, moves tensors, and prepares checkpoint tensors for ttnn: the matmul layout, the
+adjacent-pair head permutation ttnn's rotary kernel expects, and the tiled rotation matrix.
+
+The two ports also share `worker.py`, `remote.py`, `wire.py` and `oobleck_vae.py`; this module
+exists so that this repository stands alone.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import ttnn
+
+MEM = ttnn.DRAM_MEMORY_CONFIG
+
+
+def open_device(l1_small_size: int = 98304):  # convs (VAE) keep their config tensors in L1_SMALL
+    return ttnn.open_mesh_device(mesh_shape=ttnn.MeshShape(1, 1),
+                                 dispatch_core_config=ttnn.DispatchCoreConfig(ttnn.DispatchCoreType.WORKER),
+                                 l1_small_size=l1_small_size)
+
+
+def close_device(dev):
+    ttnn.close_mesh_device(dev)
+
+
+def dram_stats(dev) -> dict:
+    try:
+        s = ttnn.get_memory_view(dev, ttnn.BufferType.DRAM)
+        return {"allocated_bytes": s.total_bytes_allocated_per_bank * s.num_banks,
+                "free_bytes": s.total_bytes_free_per_bank * s.num_banks}
+    except Exception:  # pragma: no cover - diagnostics only
+        return {}
+
+
+@dataclass
+class Precision:
+    weight_dtype: ttnn.DataType = ttnn.bfloat16
+    mm_fidelity: ttnn.MathFidelity = ttnn.MathFidelity.HiFi2
+    sdpa_fidelity: ttnn.MathFidelity = ttnn.MathFidelity.HiFi4
+    sdpa_fp32_acc: bool = True
+
+
+def _dev(dev, t: torch.Tensor, dtype=ttnn.bfloat16):
+    return ttnn.from_torch(t.contiguous(), dtype=dtype, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=MEM)
+
+
+def _row(dev, v: torch.Tensor):
+    """[D] -> [1, 1, 1, D] bf16 row (norm gamma, shift, gate, bias)."""
+    return _dev(dev, v.to(torch.bfloat16).reshape(1, 1, 1, -1))
+
+
+TILE = 32
+SEQ_MULTIPLE = 128  # sequence lengths are padded to multiples of this at least (SDPA chunk sizes divide it)
+
+
+NEG = -1e9  # additive mask value for excluded keys (finite, representable in bf16)
+
+
+def linear_to_mm(w: torch.Tensor) -> torch.Tensor:
+    """nn.Linear weight [out, in] -> matmul weight [in, out]."""
+    return w.t().contiguous()
+
+
+def interleave_pairs_permutation(head_dim: int) -> torch.Tensor:
+    """new[j] = old[p[j]]: rotate-half pairs (i, i + D/2) -> adjacent pairs (2i, 2i + 1).
+
+    Applying it to the rows of q_proj/k_proj (per head) and to q_norm/k_norm leaves every q.k product unchanged and
+    lets tt's rotary_embedding_llama (adjacent pairs) replace upstream's rotate_half RoPE."""
+    half = head_dim // 2
+    p = torch.empty(head_dim, dtype=torch.long)
+    p[0::2] = torch.arange(half)
+    p[1::2] = torch.arange(half) + half
+    return p
+
+
+def permute_heads_rows(w_out_in: torch.Tensor, n_heads: int, head_dim: int, perm: torch.Tensor) -> torch.Tensor:
+    out, inp = w_out_in.shape
+    assert out == n_heads * head_dim
+    return w_out_in.view(n_heads, head_dim, inp)[:, perm, :].reshape(out, inp).contiguous()
+
+
+def rot_transformation_mat(tile: int = TILE) -> torch.Tensor:
+    """[1, 1, 32, 32] T with (x @ T)[2k] = -x[2k+1], (x @ T)[2k+1] = x[2k] (adjacent-pair rotation)."""
+    m = torch.zeros(1, 1, tile, tile)
+    m[..., torch.arange(0, tile, 2), torch.arange(1, tile, 2)] = 1.0
+    m[..., torch.arange(1, tile, 2), torch.arange(0, tile, 2)] = -1.0
+    return m
+
+
+def rms(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
+    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * w
+
+
+def round_up(n: int, m: int = SEQ_MULTIPLE) -> int:
+    return (n + m - 1) // m * m
+
+
+def swiglu_interleave(gate_out_in: torch.Tensor, up_out_in: torch.Tensor, tile: int = TILE) -> torch.Tensor:
+    """[K, 2N] weight for minimal_matmul(fuse_swiglu=True): column tile 2p = gate tile p, 2p+1 = up tile p."""
+    g, u = linear_to_mm(gate_out_in), linear_to_mm(up_out_in)
+    K, N = g.shape
+    assert N % tile == 0, N
+    return torch.stack([g.view(K, N // tile, tile), u.view(K, N // tile, tile)], dim=2).reshape(K, 2 * N).contiguous()
+
+
+def apply_rope_adjacent(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    rot = torch.stack([-x[..., 1::2], x[..., 0::2]], dim=-1).flatten(-2)
+    return x * cos + rot * sin

code/tt_diffrhythm/diffrhythm_dit.py

@@ -0,0 +1,401 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DiffRhythm 1.2 DiT on ONE Blackhole (P100a), TTNN, batch 1.
+
+Mirrors diffrhythm_host.RefDecoder op for op, with the block kernels of acestep_dit.py (fused QKV minimal_matmul,
+adjacent-pair rotary_embedding_llama, fused SwiGLU, fp32-accumulated SDPA).
+
+Per song (prepare): for each CFG branch, the timestep-invariant part of the input projection
+[text_embed @ Wtext + cond @ Wcond] and the 8 text-fusion residuals silu(text_embed @ Wf + bf) are computed on the
+device from the host's text embedding and kept there; plus the RoPE tables and the masks.
+Per call (forward): x @ Wx + that base + the call's row (style/time), zero past the song, two grouped k31 convs with
+Mish (zeroing the padding before each), 16 Llama layers (32 x 64 heads, padded keys excluded), the fusion residuals
+after the first 8, then LayerNorm * (1 + scale) + shift and proj_out.
+
+Speed (dit_profile.py, sdpa_check.py, conv_speed.py, mm_sweep.py): the attention excludes the padding with
+block-diagonal windows [0, frames, S] built on the device instead of reading an [S, S] mask for every head (30 % faster
+at 6144 frames); each position conv is one matmul over the 31 shifted rows (unfold) instead of ttnn's grouped conv,
+which expands the 16 groups densely anyway; the matmuls use blockings measured for these shapes.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import torch
+import ttnn
+
+from . import diffrhythm_host as host
+from .common import MEM, Precision, _dev, _row, close_device, dram_stats, open_device  # noqa: F401
+from .common import rot_transformation_mat
+
+try:
+    from models.tt_dit.utils.matmul import get_matmul_config
+except ImportError:  # pragma: no cover - only on hosts without tt-metal models/
+    get_matmul_config = None
+
+
+@dataclass
+class DRPrecision(Precision):
+    """Precision on top of acestep_dit.Precision. CFG (pred + 4 (pred - null)) amplifies the error of the small
+    difference between the two branches, and 31 steps accumulate it, so a DiT call that matches the GPU at PCC 0.9999
+    still drifts over a song. Closed-loop check (device_check.py --closed-loop, 96 s song, final latents vs the GPU's;
+    CPU float32 0.99997, CPU all-bfloat16 0.991): all bfloat16 0.970; the defaults below 0.994-0.995 for +1.2 s per
+    song. Splitting the weights (bf16 hi + lo) or keeping the branch tensors in float32 gained nothing more
+    (experiments/tt-diffrhythm/results/precision_sweep_d96.json)."""
+
+    mm_fidelity: ttnn.MathFidelity = ttnn.MathFidelity.HiFi4
+    x_fp32: bool = True  # upload the noisy latents in float32 (the input projection reads them)
+    residual_fp32: bool = True  # keep the residual stream in float32 (branch outputs are added into it)
+    out_fp32: bool = True  # the final projection writes float32
+    rows_fp32: bool = True  # the per-call rows (style/time input row, output scale and shift) in float32: their
+    # rounding is the same for every frame and step, so it adds up instead of averaging out
+    consts_fp32: bool = True  # the same for the constant per-channel rows: RMSNorm weights and proj_out's bias
+    windowed_attention: bool = True  # padded keys excluded by on-device windows (False: an [S, S] additive mask)
+    unfold_conv: bool = True  # position convs as one matmul over the unfolded rows (False: ttnn.conv1d, groups=16)
+
+
+# minimal_matmul blocks (M, K, N, (subblock h, w)) per (K, N) on the P100a's 11x10 grid (mm_sweep.py: 5-15 % faster
+# than tt-metal's 8x8x8 default); a list of (largest M, blocks) where the best depends on the length
+MM_BLOCKS = {
+    (2048, 6144): [(None, (4, 8, 4, (4, 1)))],  # qkv
+    (2048, 2048): [(None, (4, 4, 4, (4, 1)))],  # o
+    (2048, 16384): [(3072, (8, 8, 8, (2, 2))), (None, (4, 8, 8, (4, 1)))],  # gate/up, fused SwiGLU
+    (8192, 2048): [(None, (4, 8, 4, (4, 1)))],  # down
+}
+
+
+def conv_unfold_weight(w: torch.Tensor, groups: int) -> torch.Tensor:
+    """Grouped conv1d weight [C, C / groups, k] -> matmul weight [k * C, C] for the unfolded input (row t holds the
+    input rows t - k // 2 ... t + k // 2 side by side): block-diagonal per tap, row j * C + (g * cg + i), column
+    g * cg + o."""
+    C, cg, k = w.shape
+    out = torch.zeros(k, C, C)
+    wg = w.view(groups, cg, cg, k)  # g, o, i, tap
+    for g in range(groups):
+        out[:, g * cg:(g + 1) * cg, g * cg:(g + 1) * cg] = wg[g].permute(2, 1, 0)  # tap, i, o
+    return out.reshape(k * C, C)
+
+
+class Layer:
+    def __init__(self, dev, w: host.LayerWeights, prec: Precision):
+        wd = prec.weight_dtype
+        self.wqkv = _dev(dev, w.wqkv.to(torch.bfloat16), wd)
+        self.wo = _dev(dev, w.wo.to(torch.bfloat16), wd)
+        self.w_gateup = _dev(dev, w.w_gateup.to(torch.bfloat16), wd)
+        self.w_down = _dev(dev, w.w_down.to(torch.bfloat16), wd)
+        if getattr(prec, "consts_fp32", False):
+            self.ln1, self.ln2 = (_dev(dev, v.float().reshape(1, 1, 1, -1), ttnn.float32) for v in (w.ln1, w.ln2))
+        else:
+            self.ln1, self.ln2 = _row(dev, w.ln1), _row(dev, w.ln2)
+
+
+@dataclass
+class Song:
+    """Device state of one generation: per-branch base and fusion residuals, tables and masks."""
+
+    geo: host.Geometry
+    branches: Dict[bool, dict]  # drop_text -> {"base": tensor, "fusion": [8 tensors]}
+    cos: ttnn.Tensor
+    sin: ttnn.Tensor
+    key_mask: Optional[ttnn.Tensor]  # [S, S] additive mask, or None with windowed attention
+    windows: Optional[ttnn.Tensor]  # cu_window_seqlens [0, frames, S] (or [0, S]) for windowed attention
+    row_mask: ttnn.Tensor
+
+    def tensors(self) -> List[ttnn.Tensor]:
+        out = [t for t in (self.cos, self.sin, self.key_mask, self.windows, self.row_mask) if t is not None]
+        for b in self.branches.values():
+            out += [b["base"], *b["fusion"]]
+        return out
+
+
+class DiffRhythmDiT:
+    def __init__(self, dev, ckpt: host.Checkpoint, prec: Optional[Precision] = None, keep_ref: bool = False):
+        self.dev, self.cfg = dev, ckpt.cfg
+        self.prec = prec or DRPrecision()
+        self.res_dtype = ttnn.float32 if getattr(self.prec, "residual_fp32", False) else ttnn.bfloat16
+        self.ref = host.RefDecoder(ckpt)  # host weights: layouts and the per-call rows
+        self.grid = dev.compute_with_storage_grid_size()
+        r = self.ref
+        self.layers = [Layer(dev, w, self.prec) for w in r.layers]
+        bf = lambda t: _dev(dev, t.to(torch.bfloat16))
+        self.wx, self.wcond, self.wtext = bf(r.iw.wx), bf(r.iw.wcond), bf(r.iw.wtext)
+        self.fusion = [(bf(w), _row(dev, b)) for w, b in r.fusion]
+        self.unfold = getattr(self.prec, "unfold_conv", False)
+        if self.unfold:
+            self.conv = [(bf(conv_unfold_weight(w, self.cfg.conv_groups)), _row(dev, b)) for w, b in r.conv]
+        else:
+            self.conv = [(ttnn.from_torch(w.unsqueeze(2).contiguous(), dtype=ttnn.float32),
+                          ttnn.from_torch(b.reshape(1, 1, 1, -1), dtype=ttnn.float32)) for w, b in r.conv]
+        self.w_out, self.b_out = bf(r.w_out), _row(dev, r.b_out)
+        self.b_out32 = (_dev(dev, r.b_out.float().reshape(1, 1, 1, -1), ttnn.float32)
+                        if getattr(self.prec, "consts_fp32", False) else None)
+        self.trans_mat = _dev(dev, rot_transformation_mat())
+        arch = dev.arch()
+        ck = lambda fid, fp32: ttnn.init_device_compute_kernel_config(
+            arch, math_fidelity=fid, math_approx_mode=False, fp32_dest_acc_en=fp32, packer_l1_acc=False)
+        self.ck_mm = ck(self.prec.mm_fidelity, True)
+        self.ck_norm = ck(ttnn.MathFidelity.HiFi4, True)
+        self.ck_sdpa = ck(self.prec.sdpa_fidelity, self.prec.sdpa_fp32_acc)
+        self.ck_rope = ck(ttnn.MathFidelity.HiFi4, True)
+        self.ck_conv = ck(ttnn.MathFidelity.HiFi4, True)
+        self._mm_cfg_cache: Dict[tuple, object] = {}
+        self.windowed = getattr(self.prec, "windowed_attention", False)
+        if not self.unfold:
+            # ttnn's conv moves its weights onto the device (in its own layout) on the first call. Done now, before
+            # any song, so they do not land behind a song's tensors (the VAE's later allocations, and so its compiled
+            # kernels, would then depend on the song length). The prepared weights serve every sequence bucket
+            # (bit-identical to weights prepared at each length).
+            S = host.seq_buckets()[0]
+            z = _dev(dev, torch.zeros(1, 1, S, self.cfg.dim, dtype=torch.bfloat16))
+            for i in range(2):
+                ttnn.deallocate(self._conv(z, i, S))
+            ttnn.deallocate(z)
+        if not keep_ref:  # the per-call rows need only the small input/output weights on the host
+            r.layers = None
+
+    # ------------------------------------------------------------------ song
+    def prepare(self, text_embed: Dict[bool, torch.Tensor], cond: torch.Tensor) -> Song:
+        """text_embed {drop_text: [S, 512]} from upstream's TextEmbedding, cond [S, 64] -> device state."""
+        c = self.cfg
+        geo = host.Geometry(frames=cond.shape[0])
+        S = geo.seq_pad
+        branches = {}
+        cond_d = _dev(self.dev, host.pad_rows(cond.float(), S).to(torch.bfloat16).reshape(1, 1, S, -1))
+        for null, te in text_embed.items():
+            t = _dev(self.dev, host.pad_rows(te.float(), S).to(torch.bfloat16).reshape(1, 1, S, -1))
+            bd = ttnn.bfloat16
+            base = ttnn.linear(t, self.wtext, compute_kernel_config=self.ck_mm, memory_config=MEM, dtype=bd)
+            if not null:  # the CFG branch drops the audio condition (zeros)
+                bc = ttnn.linear(cond_d, self.wcond, compute_kernel_config=self.ck_mm, memory_config=MEM, dtype=bd)
+                base = self._add_free(base, bc)
+            fusion = []
+            for w, b in self.fusion:
+                f = ttnn.linear(t, w, bias=b, compute_kernel_config=self.ck_mm, memory_config=MEM, dtype=bd)
+                fusion.append(ttnn.silu(f, memory_config=MEM))
+                ttnn.deallocate(f)
+            ttnn.deallocate(t)
+            branches[null] = {"base": base, "fusion": fusion}
+        ttnn.deallocate(cond_d)
+        cos, sin = host.rope_tables(S, c)
+        f = lambda m: _dev(self.dev, m.to(torch.bfloat16)[None, None])
+        rows = host.row_mask(geo).expand(S, c.dim)
+        if self.windowed:
+            bounds = [0, geo.frames, S] if geo.frames < S else [0, S]
+            windows = ttnn.from_torch(torch.tensor(bounds, dtype=torch.int32), dtype=ttnn.int32,
+                                      layout=ttnn.ROW_MAJOR_LAYOUT, device=self.dev)
+            key_mask = None
+        else:
+            windows, key_mask = None, f(host.key_mask(geo))
+        return Song(geo, branches, f(cos), f(sin), key_mask, windows, f(rows))
+
+    def release(self, song: Song):
+        for t in song.tensors():
+            ttnn.deallocate(t)
+
+    # ------------------------------------------------------------------ ops
+    def _mm_config(self, M, K, N):
+        for max_m, (mb, kb, nb, (sh, sw)) in MM_BLOCKS.get((K, N), []):
+            if max_m is None or M <= max_m:
+                return ttnn.MinimalMatmulConfig(M_block_size=mb, K_block_size=kb, N_block_size=nb, subblock_h=sh,
+                                                subblock_w=sw, compute_with_storage_grid_size=self.grid)
+        return get_matmul_config(M, K, N, self.grid)
+
+    def _mm(self, x, w, M, K, N, fuse_swiglu=False, dtype=ttnn.bfloat16):
+        key = (M, K, N)
+        if key not in self._mm_cfg_cache:
+            self._mm_cfg_cache[key] = self._mm_config(M, K, N)
+        out = ttnn.experimental.minimal_matmul(x, w, config=self._mm_cfg_cache[key], compute_kernel_config=self.ck_mm,
+                                               dtype=dtype, memory_config=MEM, fuse_swiglu=fuse_swiglu)
+        if len(out.shape) != 4:
+            out = ttnn.reshape(out, [1, 1, M, out.shape[-1]])
+        return out
+
+    @staticmethod
+    def _add_free(a, b, dtype=None):
+        out = ttnn.add(a, b, memory_config=MEM, dtype=dtype) if dtype is not None else ttnn.add(a, b, memory_config=MEM)
+        ttnn.deallocate(a)
+        ttnn.deallocate(b)
+        return out
+
+    def _residual(self, h, o):
+        """h + o in the residual stream's dtype; frees both."""
+        if self.res_dtype == ttnn.bfloat16:
+            return self._add_free(h, o)
+        if o.dtype != ttnn.float32:
+            o32 = ttnn.typecast(o, ttnn.float32)
+            ttnn.deallocate(o)
+            o = o32
+        return self._add_free(h, o)
+
+    def _bf16(self, x):
+        """x as bfloat16 for a matmul; frees x when it had to be converted."""
+        if x.dtype == ttnn.bfloat16:
+            return x
+        out = ttnn.typecast(x, ttnn.bfloat16)
+        ttnn.deallocate(x)
+        return out
+
+    def _rms(self, x, w):
+        n = ttnn.rms_norm(x, epsilon=self.cfg.eps, weight=w, compute_kernel_config=self.ck_norm, memory_config=MEM)
+        return self._bf16(n)
+
+    def _mish(self, x):
+        if hasattr(ttnn, "mish"):
+            return ttnn.mish(x, memory_config=MEM)
+        sp = ttnn.softplus(x, memory_config=MEM)
+        th = ttnn.tanh(sp, memory_config=MEM)
+        ttnn.deallocate(sp)
+        out = ttnn.multiply(x, th, memory_config=MEM)
+        ttnn.deallocate(th)
+        return out
+
+    def _conv(self, x, i, S):
+        if self.unfold:
+            return self._conv_unfold(x, i, S)
+        c = self.cfg
+        w, b = self.conv[i]
+        cfg = ttnn.Conv2dConfig(weights_dtype=ttnn.bfloat16, config_tensors_in_dram=True)
+        out, out_len, (w, b) = ttnn.conv1d(
+            input_tensor=x, weight_tensor=w, bias_tensor=b, device=self.dev, in_channels=c.dim, out_channels=c.dim,
+            batch_size=1, input_length=S, kernel_size=c.conv_kernel, stride=1, padding=c.conv_kernel // 2, dilation=1,
+            groups=c.conv_groups, dtype=ttnn.bfloat16, compute_config=self.ck_conv, conv_config=cfg,
+            return_output_dim=True, return_weights_and_bias=True)
+        self.conv[i] = (w, b)  # prepared on the device by the first call
+        if out.layout != ttnn.TILE_LAYOUT:
+            out = ttnn.to_layout(out, ttnn.TILE_LAYOUT)
+        return ttnn.reshape(out, [1, 1, out_len, c.dim])
+
+    def _conv_unfold(self, x, i, S):
+        """[1, 1, S, C] -> the same conv as one matmul: rows padded by k // 2 zeros at both ends, the k shifted
+        copies side by side ([S, k * C]), times the block-diagonal weight."""
+        c = self.cfg
+        k, pad = c.conv_kernel, c.conv_kernel // 2
+        w, b = self.conv[i]
+        xr = ttnn.reshape(ttnn.to_layout(x, ttnn.ROW_MAJOR_LAYOUT), [S, c.dim])
+        xp = ttnn.pad(xr, [(pad, pad), (0, 0)], 0.0)
+        ttnn.deallocate(xr)
+        parts = [ttnn.slice(xp, [j, 0], [j + S, c.dim]) for j in range(k)]
+        ttnn.deallocate(xp)
+        xu = ttnn.concat(parts, dim=1, memory_config=MEM)
+        for t in parts:
+            ttnn.deallocate(t)
+        xt = ttnn.to_layout(ttnn.reshape(xu, [1, 1, S, k * c.dim]), ttnn.TILE_LAYOUT)
+        ttnn.deallocate(xu)
+        y = self._mm(xt, w, S, k * c.dim, c.dim)
+        ttnn.deallocate(xt)
+        out = ttnn.add(y, b, memory_config=MEM)
+        ttnn.deallocate(y)
+        return out
+
+    def _conv_pos(self, h, song: Song, S):
+        """Two k31 convs with Mish; the rows past the song are zeroed before each (upstream's sequence ends there)."""
+        y = h
+        for i in range(2):
+            m = ttnn.multiply(y, song.row_mask, memory_config=MEM)
+            if y is not h:
+                ttnn.deallocate(y)
+            conv = self._conv(m, i, S)
+            ttnn.deallocate(m)
+            y = self._mish(conv)
+            ttnn.deallocate(conv)
+        return y
+
+    def _sdpa(self, q, k, v, song: Song, S):
+        # sdpa_check.py: 256 x 256 chunks are fastest from 3072 frames on; at 2560 windowed attention prefers 128-row Q
+        # chunks (S is a multiple of SEQ_BUCKET = 256, so both divide it)
+        q_chunk = 128 if self.windowed and S <= 2560 else 256
+        cfg = ttnn.SDPAProgramConfig(compute_with_storage_grid_size=self.grid, q_chunk_size=q_chunk, k_chunk_size=256,
+                                     exp_approx_mode=False)
+        exclude = {"cu_window_seqlens": song.windows} if self.windowed else {"attn_mask": song.key_mask}
+        return ttnn.transformer.scaled_dot_product_attention(
+            q, k, v, is_causal=False, scale=self.cfg.head_dim ** -0.5, program_config=cfg,
+            compute_kernel_config=self.ck_sdpa, memory_config=MEM, **exclude)
+
+    def _attention(self, h, L: Layer, S, song: Song):
+        c = self.cfg
+        qkv = self._mm(h, L.wqkv, S, c.dim, 3 * c.dim)
+        q, k, v = ttnn.experimental.nlp_create_qkv_heads(qkv, num_heads=c.heads, num_kv_heads=c.heads,
+                                                         transpose_k_heads=False, memory_config=MEM)
+        ttnn.deallocate(qkv)
+        roped = []
+        for t in (q, k):
+            roped.append(ttnn.experimental.rotary_embedding_llama(t, song.cos, song.sin, self.trans_mat,
+                                                                  is_decode_mode=False, compute_kernel_config=self.ck_rope))
+            ttnn.deallocate(t)
+        a = self._sdpa(roped[0], roped[1], v, song, S)
+        for t in (*roped, v):
+            ttnn.deallocate(t)
+        a2 = ttnn.transformer.concatenate_heads(a, memory_config=MEM)
+        ttnn.deallocate(a)
+        o = self._mm(a2, L.wo, S, c.dim, c.dim, dtype=self.res_dtype)  # added into the residual stream as is
+        ttnn.deallocate(a2)
+        return o
+
+    # ------------------------------------------------------------------ public
+    def forward(self, x_t: torch.Tensor, drop_text: bool, style: torch.Tensor, c_emb: torch.Tensor, song: Song,
+                taps: Optional[list] = None) -> torch.Tensor:
+        """[S, 64] noisy latents, the branch, style [512] and c = time + start + duration embeddings [512]
+        -> [S, 64] float32 prediction."""
+        c, geo = self.cfg, song.geo
+        S = geo.seq_pad
+        branch = song.branches[drop_text]
+        rows = self.ref.rows(style.float(), c_emb.float())
+        if getattr(self.prec, "rows_fp32", False):
+            row = lambda v: _dev(self.dev, v.float().reshape(1, 1, 1, -1), ttnn.float32)
+        else:
+            row = lambda v: _row(self.dev, v)
+        in_row, gamma, shift = row(rows.input_row), row(1 + rows.out_scale), row(rows.out_shift)
+        p = self.prec
+        xs = host.pad_rows(x_t.float(), S).reshape(1, 1, S, -1)
+        if getattr(p, "x_fp32", False):
+            x = _dev(self.dev, xs, ttnn.float32)
+        else:
+            x = _dev(self.dev, xs.to(torch.bfloat16))
+        xw = ttnn.linear(x, self.wx, compute_kernel_config=self.ck_mm, memory_config=MEM, dtype=self.res_dtype)
+        ttnn.deallocate(x)
+        h0 = ttnn.add(xw, branch["base"], memory_config=MEM, dtype=self.res_dtype)
+        ttnn.deallocate(xw)
+        h1 = ttnn.add(h0, in_row, memory_config=MEM, dtype=self.res_dtype)
+        ttnn.deallocate(h0)
+        h = ttnn.multiply(h1, song.row_mask, memory_config=MEM, dtype=self.res_dtype)
+        ttnn.deallocate(h1)
+        conv_in = ttnn.typecast(h, ttnn.bfloat16) if self.res_dtype != ttnn.bfloat16 else h
+        pos = self._conv_pos(conv_in, song, S)
+        if conv_in is not h:
+            ttnn.deallocate(conv_in)
+        h = self._residual(h, pos)
+        tap = (lambda v: taps.append(ttnn.to_torch(v)[0, 0].float())) if taps is not None else (lambda v: None)
+        for i, L in enumerate(self.layers):
+            a = self._rms(h, L.ln1)
+            o = self._attention(a, L, S, song)
+            ttnn.deallocate(a)
+            h = self._residual(h, o)
+            m = self._rms(h, L.ln2)
+            gu = self._mm(m, L.w_gateup, S, c.dim, 2 * c.intermediate, fuse_swiglu=True)
+            ttnn.deallocate(m)
+            d = self._mm(gu, L.w_down, S, c.intermediate, c.dim, dtype=self.res_dtype)
+            ttnn.deallocate(gu)
+            h = self._residual(h, d)
+            if i < c.fusion_layers:
+                h2 = ttnn.add(h, branch["fusion"][i], memory_config=MEM, dtype=self.res_dtype)
+                ttnn.deallocate(h)
+                h = h2
+            tap(h)
+        n = ttnn.layer_norm(h, epsilon=c.eps, compute_kernel_config=self.ck_norm, memory_config=MEM)
+        ttnn.deallocate(h)
+        g = ttnn.multiply(n, gamma, memory_config=MEM)
+        ttnn.deallocate(n)
+        mod = self._bf16(ttnn.add(g, shift, memory_config=MEM))
+        ttnn.deallocate(g)
+        if self.b_out32 is not None:
+            o = ttnn.linear(mod, self.w_out, compute_kernel_config=self.ck_mm, memory_config=MEM, dtype=ttnn.float32)
+            out = ttnn.add(o, self.b_out32, memory_config=MEM)
+            ttnn.deallocate(o)
+        else:
+            out = ttnn.linear(mod, self.w_out, bias=self.b_out, compute_kernel_config=self.ck_mm, memory_config=MEM,
+                              dtype=ttnn.float32 if getattr(p, "out_fp32", False) else ttnn.bfloat16)
+        ttnn.deallocate(mod)
+        host_out = ttnn.to_torch(out)[0, 0, :geo.frames, : c.mel_dim].float()
+        for t in (out, in_row, gamma, shift):
+            ttnn.deallocate(t)
+        return host_out

code/tt_diffrhythm/diffrhythm_host.py

@@ -0,0 +1,224 @@
+"""Host (torch, CPU) side of the DiffRhythm 1.2 DiT port to one Tenstorrent P100a.
+
+Upstream (DiffRhythm model/dit.py DiT.forward) as used by CFM.sample, per call:
+  c = time_embed(t) + start_time_embed(start) + duration_time_embed(duration)                    [512]
+  text_embed = TextEmbedding(text)  (embedding + sinusoid + 4 ConvNeXtV2 blocks; zeros in for drop_text)  [S, 512]
+  h = proj(cat(x, cond, text_embed, style, c))  -> [S, 2048];  h = ConvPos(h) + h   (two k31 groups-16 convs + Mish)
+  16 x LlamaDecoderLayer (RMSNorm, 32 x 64 heads, RoPE theta 1e4, SwiGLU 8192; no mask, bidirectional);
+       after each of the first 8: h += silu(text_fusion_linears[i](text_embed))
+  out = proj_out(LayerNorm(h) * (1 + scale) + shift),  scale, shift = linear(silu(c)).chunk(2)
+The CFG branch (drop_audio_cond + drop_text) zeroes cond and the text tokens and uses the negative style.
+
+Split: the timestep-invariant text embedding and the small time MLPs stay in the DiffRhythm process on the host
+(upstream modules, exact); everything weight-heavy runs on the card. The input projection is split by input:
+  h = x @ Wx + [cond @ Wcond + text_embed @ Wtext]  (per song and branch)  + [style @ Ws + c @ Wc + b]  (per call row)
+RefDecoder runs this formulation in float32 so it can be checked against upstream without a device.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict
+
+import torch
+
+from .common import (NEG, TILE, apply_rope_adjacent, interleave_pairs_permutation, linear_to_mm,
+                     permute_heads_rows, rms, round_up, swiglu_interleave)
+
+SEQ_BUCKET = 256  # frames padded to multiples of this: 96-285 s songs (2067-6136 frames) -> 16 shapes
+LATENT_RATE = 44100 / 2048  # latent frames per second (get_lrc_token: int(seconds * 44100 / 2048))
+MIN_FRAMES, MAX_FRAMES = int(96 * LATENT_RATE), int(285 * LATENT_RATE)  # the 96-285 s the port was exercised over
+
+
+def seq_buckets():
+    """Every padded length a DiffRhythm song can have (warm_cache.py compiles them ahead of time)."""
+    return sorted({round_up(n, SEQ_BUCKET) for n in range(MIN_FRAMES, MAX_FRAMES + 1)})
+
+
+@dataclass(frozen=True)
+class Config:
+    dim: int = 2048
+    depth: int = 16
+    heads: int = 32
+    head_dim: int = 64
+    intermediate: int = 8192
+    text_dim: int = 512
+    cond_dim: int = 512
+    mel_dim: int = 64
+    fusion_layers: int = 8
+    conv_kernel: int = 31
+    conv_groups: int = 16
+    rope_theta: float = 10000.0
+    eps: float = 1e-6
+
+
+class Checkpoint:
+    """transformer.* of DiffRhythm's cfm_model.pt (fp16 on disk), read as float32."""
+
+    def __init__(self, path):
+        sd = torch.load(str(path), weights_only=True, map_location="cpu")["model_state_dict"]
+        self._sd = {k[len("transformer."):]: v for k, v in sd.items() if k.startswith("transformer.")}
+        self.cfg = Config()
+
+    def get(self, name: str) -> torch.Tensor:
+        return self._sd[name].float()
+
+
+@dataclass
+class LayerWeights:
+    wqkv: torch.Tensor  # [2048, 3 * 2048], q/k rows permuted to adjacent pairs
+    wo: torch.Tensor
+    w_gateup: torch.Tensor  # [2048, 2 * 8192] tile-interleaved
+    w_down: torch.Tensor
+    ln1: torch.Tensor
+    ln2: torch.Tensor
+
+
+def layer_weights(ckpt: Checkpoint, i: int) -> LayerWeights:
+    c = ckpt.cfg
+    g = lambda k: ckpt.get(f"transformer_blocks.{i}.{k}")
+    perm = interleave_pairs_permutation(c.head_dim)
+    wq = permute_heads_rows(g("self_attn.q_proj.weight"), c.heads, c.head_dim, perm)
+    wk = permute_heads_rows(g("self_attn.k_proj.weight"), c.heads, c.head_dim, perm)
+    return LayerWeights(
+        wqkv=torch.cat([linear_to_mm(wq), linear_to_mm(wk), linear_to_mm(g("self_attn.v_proj.weight"))], dim=1),
+        wo=linear_to_mm(g("self_attn.o_proj.weight")),
+        w_gateup=swiglu_interleave(g("mlp.gate_proj.weight"), g("mlp.up_proj.weight")),
+        w_down=linear_to_mm(g("mlp.down_proj.weight")),
+        ln1=g("input_layernorm.weight"), ln2=g("post_attention_layernorm.weight"))
+
+
+@dataclass
+class InputWeights:
+    """input_embed.proj split by input (matmul layout [in, out])."""
+
+    wx: torch.Tensor  # [64, 2048]
+    wcond: torch.Tensor  # [64, 2048]
+    wtext: torch.Tensor  # [512, 2048]
+    wstyle: torch.Tensor  # [512, 2048]
+    wc: torch.Tensor  # [512, 2048]
+    b: torch.Tensor
+
+
+def input_weights(ckpt: Checkpoint) -> InputWeights:
+    c = ckpt.cfg
+    w = linear_to_mm(ckpt.get("input_embed.proj.weight"))  # [1664, 2048]
+    edges = [0, c.mel_dim, 2 * c.mel_dim, 2 * c.mel_dim + c.text_dim, 2 * c.mel_dim + c.text_dim + c.cond_dim]
+    parts = [w[a:b] for a, b in zip(edges, edges[1:] + [w.shape[0]])]
+    return InputWeights(*[p.contiguous() for p in parts], ckpt.get("input_embed.proj.bias"))
+
+
+@dataclass(frozen=True)
+class Geometry:
+    frames: int
+
+    @property
+    def seq_pad(self) -> int:
+        return round_up(self.frames, SEQ_BUCKET)
+
+
+def key_mask(geo: Geometry) -> torch.Tensor:
+    """Additive [S_pad, S_pad] mask excluding padded keys (upstream attends over exactly S frames)."""
+    valid = torch.arange(geo.seq_pad)[None, :] < geo.frames
+    return torch.where(valid, 0.0, NEG).expand(geo.seq_pad, -1).contiguous()
+
+
+def row_mask(geo: Geometry) -> torch.Tensor:
+    """[S_pad, 1]: 1 on real frames, 0 on padding (the position convs must see zeros past the song, like upstream)."""
+    return (torch.arange(geo.seq_pad) < geo.frames).float()[:, None]
+
+
+def rope_tables(seq_pad: int, cfg: Config, dtype=torch.float32):
+    inv = 1.0 / (cfg.rope_theta ** (torch.arange(0, cfg.head_dim, 2, dtype=torch.float64) / cfg.head_dim))
+    ang = torch.outer(torch.arange(seq_pad, dtype=torch.float64), inv).float()
+    return torch.cos(ang).repeat_interleave(2, -1).to(dtype), torch.sin(ang).repeat_interleave(2, -1).to(dtype)
+
+
+def pad_rows(t: torch.Tensor, rows: int) -> torch.Tensor:
+    out = torch.zeros(rows, t.shape[-1], dtype=t.dtype)
+    out[: t.shape[0]] = t
+    return out
+
+
+@dataclass
+class CallRows:
+    """What changes per call: the style/time part of the input projection and the output modulation."""
+
+    input_row: torch.Tensor  # [2048]
+    out_scale: torch.Tensor  # [2048]
+    out_shift: torch.Tensor  # [2048]
+
+
+def call_rows(ckpt_or_iw, norm_linear_w, norm_linear_b, style: torch.Tensor, c: torch.Tensor) -> CallRows:
+    """style [512], c [512] (time + start + duration embeddings, from upstream on the host)."""
+    iw = ckpt_or_iw
+    input_row = style @ iw.wstyle + c @ iw.wc + iw.b
+    emb = torch.nn.functional.silu(c) @ norm_linear_w.t() + norm_linear_b
+    scale, shift = emb.chunk(2, dim=-1)
+    return CallRows(input_row, scale, shift)
+
+
+class RefDecoder:
+    """float32 CPU DiT in the TT formulation."""
+
+    def __init__(self, ckpt: Checkpoint):
+        self.ckpt, self.cfg = ckpt, ckpt.cfg
+        c = self.cfg
+        self.layers = [layer_weights(ckpt, i) for i in range(c.depth)]
+        self.iw = input_weights(ckpt)
+        self.conv = [(ckpt.get(f"input_embed.conv_pos_embed.conv1d.{i}.weight"),
+                      ckpt.get(f"input_embed.conv_pos_embed.conv1d.{i}.bias")) for i in (0, 2)]
+        self.fusion = [(linear_to_mm(ckpt.get(f"text_fusion_linears.{i}.0.weight")),
+                        ckpt.get(f"text_fusion_linears.{i}.0.bias")) for i in range(c.fusion_layers)]
+        self.w_out, self.b_out = linear_to_mm(ckpt.get("proj_out.weight")), ckpt.get("proj_out.bias")
+        self.norm_w, self.norm_b = ckpt.get("norm_out.linear.weight"), ckpt.get("norm_out.linear.bias")
+
+    def rows(self, style: torch.Tensor, c: torch.Tensor) -> CallRows:
+        return call_rows(self.iw, self.norm_w, self.norm_b, style, c)
+
+    def branch(self, text_embed: torch.Tensor, cond, geo: Geometry) -> Dict[str, object]:
+        """Per song and CFG branch: [cond @ Wcond + text @ Wtext] and the 8 text-fusion residuals, padded."""
+        t = pad_rows(text_embed.float(), geo.seq_pad)
+        base = t @ self.iw.wtext
+        if cond is not None:
+            base = base + pad_rows(cond.float(), geo.seq_pad) @ self.iw.wcond
+        fusion = [torch.nn.functional.silu(t @ w + b) for w, b in self.fusion]
+        return {"base": base, "fusion": fusion}
+
+    def conv_pos(self, h: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+        """Two k31 convs with Mish. Upstream's sequence ends at the song, so each conv sees zeros past it: the rows
+        past the song are zeroed before every conv (the first conv's Mish output is not zero there)."""
+        c = self.cfg
+        y = h
+        for w, b in self.conv:
+            if mask is not None:
+                y = y * mask
+            y = torch.nn.functional.mish(torch.nn.functional.conv1d(y.t()[None], w, b, padding=c.conv_kernel // 2,
+                                                                    groups=c.conv_groups))[0].t()
+        return y
+
+    def forward(self, x_t: torch.Tensor, branch: Dict[str, object], rows: CallRows, geo: Geometry, taps: list = None):
+        """[S, 64] noisy latents -> [S, 64] prediction."""
+        c = self.cfg
+        mask = row_mask(geo)
+        h = (pad_rows(x_t.float(), geo.seq_pad) @ self.iw.wx + branch["base"] + rows.input_row) * mask
+        h = self.conv_pos(h, mask) + h
+        cos, sin = rope_tables(geo.seq_pad, c)
+        kmask = key_mask(geo)
+        for i, w in enumerate(self.layers):
+            a = rms(h, w.ln1, c.eps)
+            q, k, v = (a @ w.wqkv).split(c.dim, -1)
+            q = apply_rope_adjacent(q.view(-1, c.heads, c.head_dim).transpose(0, 1), cos, sin)
+            k = apply_rope_adjacent(k.view(-1, c.heads, c.head_dim).transpose(0, 1), cos, sin)
+            v = v.view(-1, c.heads, c.head_dim).transpose(0, 1)
+            scores = q @ k.transpose(-1, -2) * c.head_dim ** -0.5 + kmask
+            att = (torch.softmax(scores, -1) @ v).transpose(0, 1).reshape(geo.seq_pad, -1)
+            h = h + att @ w.wo
+            m = rms(h, w.ln2, c.eps)
+            gu = (m @ w.w_gateup).view(geo.seq_pad, -1, 2, TILE)
+            h = h + (torch.nn.functional.silu(gu[:, :, 0]) * gu[:, :, 1]).reshape(geo.seq_pad, -1) @ w.w_down
+            if i < c.fusion_layers:
+                h = h + branch["fusion"][i]
+            if taps is not None:
+                taps.append(h.clone())
+        h = torch.nn.functional.layer_norm(h, (c.dim,), eps=c.eps) * (1 + rows.out_scale) + rows.out_shift
+        return (h @ self.w_out + self.b_out)[: geo.frames]

code/tt_diffrhythm/local.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The card in THIS process, behind the same `.call(msg)` surface `TTWorker` gives.
+
+`RemoteDiffRhythmDiT` and `RemoteVae` reach the card through exactly one method -- `worker.call(msg)`
+-- so swapping the transport is enough to run everything in one process: the shims, their window
+arithmetic and their drop-in signatures are reused unchanged rather than reimplemented here.
+
+Why there are two transports at all: DiffRhythm resolves `torch 2.10.0+cu128` and `numpy>=2`, while
+ttnn's extension modules are built against tt-metal's own `torch 2.11.0+cpu` and `numpy 1.x`. Two
+environments cannot be one process, so the deployment this port came from ran DiffRhythm on the host
+and the card in a container, joined by a Unix socket (`remote.py` / `worker.py`). A container built
+for this model alone has no such split -- one environment satisfies both -- and then the socket is
+pure cost. Use `LocalWorker` there, `TTWorker` when the two environments must stay apart.
+
+    from tt_diffrhythm import LocalWorker, RemoteDiffRhythmDiT, RemoteVae
+
+    worker = LocalWorker(cfm_model_pt, vae=vae_model_pt).start()
+    cfm.transformer = RemoteDiffRhythmDiT(worker, cfm.transformer)
+    waveform = RemoteVae(worker, upsample=worker.upsample)(latent)
+    ...
+    worker.close()
+"""
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from . import diffrhythm_dit, oobleck_vae, wire
+from .shapes import VAE_WINDOWS
+from .worker import DiffRhythm
+
+MODELS = {"diffrhythm": DiffRhythm}
+
+
+class LocalWorker:
+    """`TTWorker`'s surface, answered in-process. `checkpoint` is cfm_model.pt and `vae` is
+    vae_model.pt; both are the files DiffRhythm's own Hugging Face cache holds."""
+
+    def __init__(self, checkpoint, vae, model: str = "diffrhythm", warm: bool = True):
+        self.checkpoint = Path(checkpoint)
+        self.vae_path = Path(vae)
+        self.model_name = model
+        self.warm = warm
+        self.dev = None
+        self.model = self.vae = None
+        self.song = None
+        self.stats: dict = {}
+        self.failed = False
+        self.upsample: Optional[int] = None
+
+    # ------------------------------------------------------------------ lifecycle
+    def start(self) -> "LocalWorker":
+        started = time.monotonic()
+        self.dev = diffrhythm_dit.open_device()
+        try:
+            self.model = MODELS[self.model_name](self.dev, str(self.checkpoint))
+            self.vae = oobleck_vae.OobleckTT(self.dev, str(self.vae_path))
+            self.upsample = self.vae.upsample
+            if self.warm:
+                # Same reason worker.main() warms here: built now, the VAE's conv kernels compile
+                # in one device state and every later decode reuses the in-process programs. Built
+                # after a song's DiT steps instead, the state -- and so the kernels -- depends on
+                # the song length and recompiles for 20-40 s.
+                t0 = time.monotonic()
+                for window in VAE_WINDOWS:
+                    self.vae.decode(torch.zeros(64, window))
+                self.stats["vae_warm_s"] = round(time.monotonic() - t0, 2)
+            self.stats["load_s"] = round(time.monotonic() - started, 2)
+            self.stats["dram"] = diffrhythm_dit.dram_stats(self.dev)
+        except BaseException:
+            self.close()
+            raise
+        return self
+
+    def wait_ready(self, timeout_s: float = 0.0) -> dict:
+        """`TTWorker.wait_ready`'s shape. Loading already finished in `start`, so this only reports."""
+        if self.model is None:
+            raise RuntimeError("LocalWorker.start() has not run")
+        return {"op": "ready", "model": self.model_name, "upsample": self.upsample, **self.stats}
+
+    def died(self) -> bool:
+        """There is no child process to lose; a failure here raised in the caller's own stack."""
+        return False
+
+    def close(self) -> None:
+        if self.song is not None and self.model is not None:
+            self.model.dit.release(self.song)
+            self.song = None
+        self.model = self.vae = None
+        if self.dev is not None:
+            diffrhythm_dit.close_device(self.dev)
+            self.dev = None
+
+    def __enter__(self):
+        return self.start()
+
+    def __exit__(self, *exc):
+        self.close()
+
+    # ------------------------------------------------------------------ calls
+    def call(self, msg: dict) -> dict:
+        """`worker.py:serve`'s dispatch with the socket taken out and nothing else changed.
+
+        Messages still carry `wire`-encoded tensors even though both ends are now one process.
+        That is deliberate: it is what lets `RemoteDiffRhythmDiT` and `RemoteVae` run against this
+        worker unmodified, and the encoding is a float32 copy of a few MB against card work
+        measured in seconds."""
+        if self.model is None:
+            raise RuntimeError("LocalWorker.start() has not run")
+        op = msg.get("op")
+        started = time.monotonic()
+        try:
+            if op == "prepare":
+                if self.song is not None:
+                    self.model.dit.release(self.song)
+                    self.song = None
+                self.song, geometry = self.model.prepare(msg)
+                return {"op": "prepared", "seconds": time.monotonic() - started, "geometry": geometry}
+            if op == "forward":
+                if self.song is None:
+                    raise RuntimeError("forward before prepare")
+                v = self.model.forward(msg, self.song)
+                return {"op": "velocity", "v": wire.encode(v), "seconds": time.monotonic() - started}
+            if op == "vae":
+                if self.song is not None:
+                    # Sampling is over; freeing the song's tensors first gives the VAE the same
+                    # device allocation state for every song length, so its compiled conv kernels
+                    # are reused instead of recompiled.
+                    self.model.dit.release(self.song)
+                    self.song = None
+                audio = self.vae.decode(wire.decode(msg["z"]))
+                return {"op": "audio", "audio": wire.encode(audio), "seconds": time.monotonic() - started}
+            if op == "close":
+                self.close()
+                return {"op": "closed", "seconds": time.monotonic() - started}
+        except BaseException:
+            self.failed = True  # the card may be unhealthy, same contract as TTWorker
+            raise
+        raise RuntimeError(f"unknown op {op!r}")

code/tt_diffrhythm/oobleck_vae.py

@@ -0,0 +1,296 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Oobleck audio VAE decoders on ONE Blackhole (P100a), TTNN: ACE-Step's (diffusers AutoencoderOobleck, 48 kHz
+stereo, strides 10, 6, 4, 4, 2) and DiffRhythm's (stable-audio-tools, TorchScript, 44.1 kHz stereo, strides
+8, 8, 4, 4, 2). Same architecture, different weight files:
+
+  conv1 (64 -> 2048, k7)
+  5 x block: Snake -> ConvTranspose1d (k = 2s, stride s) -> 3 x residual unit (dilation 1, 3, 9):
+             x + conv_k1(Snake(conv_k7_dilated(Snake(x))))
+  Snake -> conv2 (128 -> 2, k7, no bias)
+Snake(x) = x + 1 / (exp(beta) + 1e-9) * sin(exp(alpha) * x)^2 per channel. Weight norm is folded on the host
+(diffusers) or already folded (the TorchScript export).
+
+Activations are [1, 1, L, C] (channels last), as ttnn.conv1d / conv_transpose2d (height 1) expect.
+"""
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional
+
+import torch
+import ttnn
+
+MEM = ttnn.DRAM_MEMORY_CONFIG
+
+
+def fold_weight_norm(g: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """torch weight_norm(dim=0): w = g * v / ||v|| over every dim but the first."""
+    norm = v.flatten(1).norm(dim=1).view(-1, *([1] * (v.dim() - 1)))
+    return g * v / norm
+
+
+class VaeWeights:
+    """decoder.* of checkpoints/vae, float32, weight norm folded."""
+
+    def __init__(self, vae_dir):
+        from safetensors import safe_open
+
+        self.dir = Path(vae_dir)
+        self.config = json.loads((self.dir / "config.json").read_text())
+        self._f = safe_open(str(self.dir / "diffusion_pytorch_model.safetensors"), framework="pt")
+        self.keys = set(self._f.keys())
+
+    def get(self, name):
+        return self._f.get_tensor(f"decoder.{name}").float()
+
+    def conv(self, name):
+        w = fold_weight_norm(self.get(f"{name}.weight_g"), self.get(f"{name}.weight_v"))
+        b = self.get(f"{name}.bias") if f"decoder.{name}.bias" in self.keys else None
+        return w, b
+
+    def snake(self, name):
+        alpha = torch.exp(self.get(f"{name}.alpha")).flatten()
+        inv_beta = 1.0 / (torch.exp(self.get(f"{name}.beta")).flatten() + 1e-9)
+        return alpha, inv_beta
+
+
+class StableAudioVaeWeights:
+    """decoder.* of DiffRhythm's vae_model.pt (a stable-audio-tools OobleckDecoder exported with TorchScript, weight
+    norm already folded), under VaeWeights' diffusers names."""
+
+    def __init__(self, path):
+        sd = torch.jit.load(str(path), map_location="cpu").state_dict()
+        self._sd = {k[len("decoder."):]: v for k, v in sd.items() if k.startswith("decoder.")}
+        strides = [self._sd[f"layers.{i + 1}.layers.1.weight"].shape[-1] // 2 for i in range(5)]
+        self.config = {"downsampling_ratios": strides}
+
+    @staticmethod
+    def _key(name):
+        parts = name.split(".")
+        if name == "conv1":
+            return "layers.0"
+        if name == "snake1":
+            return "layers.6"
+        if name == "conv2":
+            return "layers.7"
+        block = f"layers.{int(parts[1]) + 1}"
+        if parts[2] == "snake1":
+            return f"{block}.layers.0"
+        if parts[2] == "conv_t1":
+            return f"{block}.layers.1"
+        unit = f"{block}.layers.{int(parts[2][len('res_unit'):]) + 1}"
+        return f"{unit}.layers.{ {'snake1': 0, 'conv1': 1, 'snake2': 2, 'conv2': 3}[parts[3]] }"
+
+    def conv(self, name):
+        k = self._key(name)
+        return self._sd[f"{k}.weight"].float(), (self._sd[f"{k}.bias"].float() if f"{k}.bias" in self._sd else None)
+
+    def snake(self, name):
+        k = self._key(name)
+        alpha = torch.exp(self._sd[f"{k}.alpha"].float()).flatten()
+        inv_beta = 1.0 / (torch.exp(self._sd[f"{k}.beta"].float()).flatten() + 1e-9)
+        return alpha, inv_beta
+
+
+def vae_weights(path):
+    """A diffusers VAE directory (ACE-Step) or a stable-audio TorchScript file (DiffRhythm)."""
+    path = Path(path)
+    if path.is_dir() and (path / "config.json").exists():
+        return VaeWeights(path)
+    return StableAudioVaeWeights(path / "vae_model.pt" if path.is_dir() else path)
+
+
+@dataclass
+class Conv:
+    weight: object  # ttnn host tensor, replaced by the prepared device tensor after the first call
+    bias: object
+    in_ch: int
+    out_ch: int
+    kernel: int
+    pad: int
+    dilation: int = 1
+    stride: int = 1
+    transpose: bool = False
+
+
+@dataclass
+class Snake:
+    alpha: ttnn.Tensor
+    inv_beta: ttnn.Tensor
+
+
+@dataclass
+class ResUnit:
+    snake1: Snake
+    conv1: Conv
+    snake2: Snake
+    conv2: Conv
+
+
+@dataclass
+class Block:
+    snake: Snake
+    up: Conv
+    units: List[ResUnit]
+
+
+@dataclass
+class Precision:
+    fidelity: ttnn.MathFidelity = ttnn.MathFidelity.HiFi4  # HiFi2/LoFi are no faster here: the convs are not math-bound
+    act_dtype: ttnn.DataType = ttnn.bfloat16
+    # vae_conv_sweep.py: activation blocks of 64 rows make the 48 kHz k7 convs 4.9x and the output conv 6x faster
+    # than the default block height, same output; double buffering and bfp8 weights were slower. Applied to the
+    # layers with <= 256 channels (the long, late ones); the wide early layers keep the default (the override fails
+    # or is slower there).
+    act_block_h: int = 64
+    transpose_act_block_h: int = 64
+    act_block_h_max_channels: int = 256
+    config_tensors_in_dram: bool = True
+    fused_snake: bool = True
+
+
+class OobleckTT:
+    def __init__(self, dev, vae_path, prec: Optional[Precision] = None):
+        self.dev = dev
+        self.prec = prec or Precision()
+        w = vae_weights(vae_path)
+        self.upsample = math.prod(w.config["downsampling_ratios"])
+        self.conv1 = self._conv(*w.conv("conv1"), pad=3)
+        self.blocks = []
+        for i in range(len(w.config["downsampling_ratios"])):
+            p = f"block.{i}"
+            wt, bt = w.conv(f"{p}.conv_t1")  # ConvTranspose1d weight [in, out, 2s]
+            stride = wt.shape[-1] // 2
+            units = []
+            for j, dilation in enumerate((1, 3, 9)):
+                u = f"{p}.res_unit{j + 1}"
+                units.append(ResUnit(self._snake(*w.snake(f"{u}.snake1")),
+                                     self._conv(*w.conv(f"{u}.conv1"), pad=3 * dilation, dilation=dilation),
+                                     self._snake(*w.snake(f"{u}.snake2")), self._conv(*w.conv(f"{u}.conv2"), pad=0)))
+            self.blocks.append(Block(self._snake(*w.snake(f"{p}.snake1")),
+                                     self._conv(wt, bt, pad=math.ceil(stride / 2), stride=stride, transpose=True), units))
+        self.snake_out = self._snake(*w.snake("snake1"))
+        self.conv2 = self._conv(*w.conv("conv2"), pad=3)
+        self.ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=self.prec.fidelity,
+                                                         math_approx_mode=False, fp32_dest_acc_en=True,
+                                                         packer_l1_acc=False)
+        self.profile = None  # a list to record (where, op, length, channels, seconds) per op, synchronizing each
+        self._where = ""
+
+    def _timed(self, op, length, channels, fn, *args):
+        if self.profile is None:
+            return fn(*args)
+        import time
+
+        ttnn.synchronize_device(self.dev)
+        started = time.perf_counter()
+        out = fn(*args)
+        ttnn.synchronize_device(self.dev)
+        self.profile.append((self._where, op, length, channels, time.perf_counter() - started))
+        return out
+
+    # ------------------------------------------------------------------ weights
+    def _conv(self, w, b, pad, dilation=1, stride=1, transpose=False):
+        if transpose:  # [in, out, k] -> [in, out, 1, k]
+            in_ch, out_ch, k = w.shape
+        else:  # [out, in, k] -> [out, in, 1, k]
+            out_ch, in_ch, k = w.shape
+        weight = ttnn.from_torch(w.unsqueeze(2).contiguous(), dtype=ttnn.float32)
+        bias = ttnn.from_torch(b.reshape(1, 1, 1, -1), dtype=ttnn.float32) if b is not None else None
+        return Conv(weight, bias, in_ch, out_ch, k, pad, dilation, stride, transpose)
+
+    def _snake(self, alpha, inv_beta):
+        row = lambda v: ttnn.from_torch(v.float().reshape(1, 1, 1, -1), dtype=self.prec.act_dtype,
+                                        layout=ttnn.TILE_LAYOUT, device=self.dev, memory_config=MEM)
+        return Snake(row(alpha), row(inv_beta))
+
+    # ------------------------------------------------------------------ ops
+    def _run_conv(self, x, c: Conv, length: int):
+        block_h = self.prec.transpose_act_block_h if c.transpose else self.prec.act_block_h
+        if max(c.in_ch, c.out_ch) > self.prec.act_block_h_max_channels:
+            block_h = 0
+        # smaller activation blocks need larger per-core config tensors; keep them in DRAM, not the 96 KiB L1_SMALL
+        cfg = ttnn.Conv2dConfig(weights_dtype=ttnn.bfloat16, config_tensors_in_dram=self.prec.config_tensors_in_dram)
+        if block_h:
+            cfg.act_block_h_override = block_h
+        common = dict(device=self.dev, in_channels=c.in_ch, out_channels=c.out_ch, batch_size=1, bias_tensor=c.bias,
+                      dtype=self.prec.act_dtype, compute_config=self.ck, conv_config=cfg,
+                      return_weights_and_bias=True)
+        if c.transpose:
+            out, out_len, (c.weight, c.bias) = ttnn.conv_transpose2d(
+                input_tensor=x, weight_tensor=c.weight, input_height=1, input_width=length, kernel_size=(1, c.kernel),
+                stride=(1, c.stride), padding=(0, c.pad), output_padding=(0, 0), dilation=(1, 1), groups=1,
+                return_output_dim=True, **common)
+            out_len = out_len[1] if isinstance(out_len, (tuple, list)) else out_len
+        else:
+            out, out_len, (c.weight, c.bias) = ttnn.conv1d(
+                input_tensor=x, weight_tensor=c.weight, input_length=length, kernel_size=c.kernel, stride=1,
+                padding=c.pad, dilation=c.dilation, groups=1, return_output_dim=True, **common)
+        if out.layout != ttnn.TILE_LAYOUT:
+            out = ttnn.to_layout(out, ttnn.TILE_LAYOUT)
+        return ttnn.reshape(out, [1, 1, out_len, c.out_ch]), out_len
+
+    def _snake_op(self, x, s: Snake):
+        if self.prec.fused_snake:  # sin(alpha x)^2 in one multiply, then x + inv_beta * that in one addcmul: 2 passes
+            unary = lambda op: ttnn.UnaryWithParam(op)
+            sq = ttnn.multiply(x, s.alpha, activations=[unary(ttnn.UnaryOpType.SIN), unary(ttnn.UnaryOpType.SQUARE)],
+                               memory_config=MEM)
+            out = ttnn.addcmul(x, sq, s.inv_beta, value=1.0, memory_config=MEM)
+            ttnn.deallocate(sq)
+            return out
+        ax = ttnn.multiply(x, s.alpha, memory_config=MEM)
+        sn = ttnn.sin(ax, memory_config=MEM)
+        ttnn.deallocate(ax)
+        sq = ttnn.multiply(sn, sn, memory_config=MEM)
+        ttnn.deallocate(sn)
+        t = ttnn.multiply(sq, s.inv_beta, memory_config=MEM)
+        ttnn.deallocate(sq)
+        out = ttnn.add(x, t, memory_config=MEM)
+        ttnn.deallocate(t)
+        return out
+
+    def _unit(self, x, u: ResUnit, length):
+        c = u.conv1.out_ch
+        h = self._timed("snake", length, c, self._snake_op, x, u.snake1)
+        y, _ = self._timed(f"conv_k7_d{u.conv1.dilation}", length, c, self._run_conv, h, u.conv1, length)
+        ttnn.deallocate(h)
+        h = self._timed("snake", length, c, self._snake_op, y, u.snake2)
+        ttnn.deallocate(y)
+        y, _ = self._timed("conv_k1", length, c, self._run_conv, h, u.conv2, length)
+        ttnn.deallocate(h)
+        out = self._timed("add", length, c, lambda a, b: ttnn.add(a, b, memory_config=MEM), x, y)
+        ttnn.deallocate(y)
+        ttnn.deallocate(x)
+        return out
+
+    # ------------------------------------------------------------------ public
+    def decode(self, latents: torch.Tensor, taps: Optional[list] = None) -> torch.Tensor:
+        """[64, T] latents -> [2, T * upsample] float32 audio (1920 for ACE-Step, 2048 for DiffRhythm)."""
+        length = latents.shape[-1]
+        x = ttnn.from_torch(latents.t().contiguous().reshape(1, 1, length, -1).float(), dtype=self.prec.act_dtype,
+                            layout=ttnn.TILE_LAYOUT, device=self.dev, memory_config=MEM)
+        tap = (lambda t, n: taps.append(ttnn.to_torch(t)[0, 0, :n].float())) if taps is not None else (lambda t, n: None)
+        self._where = "in"
+        h, length = self._timed("conv_in", length, self.conv1.out_ch, self._run_conv, x, self.conv1, length)
+        ttnn.deallocate(x)
+        tap(h, length)
+        for i, b in enumerate(self.blocks):
+            self._where = f"block{i}"
+            s = self._timed("snake", length, b.up.in_ch, self._snake_op, h, b.snake)
+            ttnn.deallocate(h)
+            h, length = self._timed(f"conv_t_s{b.up.stride}", length, b.up.out_ch, self._run_conv, s, b.up, length)
+            ttnn.deallocate(s)
+            for u in b.units:
+                h = self._unit(h, u, length)
+            tap(h, length)
+        self._where = "out"
+        s = self._timed("snake", length, self.snake_out.alpha.shape[-1], self._snake_op, h, self.snake_out)
+        ttnn.deallocate(h)
+        out, length = self._timed("conv_out", length, 2, self._run_conv, s, self.conv2, length)
+        ttnn.deallocate(s)
+        audio = ttnn.to_torch(out)[0, 0, :length, :2].float().t().contiguous()
+        ttnn.deallocate(out)
+        return audio

code/tt_diffrhythm/remote.py

@@ -0,0 +1,264 @@
+"""Host side of the P100a models: start the TT worker container on the card and stand in for the DiT and
+the audio VAE decoder of the host pipeline.
+
+The ACE-Step port is the sibling repo, tt-acestep: it shares this worker and transport.
+
+Used from the DiffRhythm host pipeline: TTWorker(cfm_model.pt, model="diffrhythm", vae=vae_model.pt),
+cfm.transformer = RemoteDiffRhythmDiT(worker, cfm.transformer), decoding with RemoteVae(worker, upsample=2048).
+
+Runs in the model's own environment; needs only the standard library and torch.
+"""
+from __future__ import annotations
+
+from multiprocessing.connection import Client
+import os
+from pathlib import Path
+import secrets
+import shutil
+import subprocess
+import tempfile
+import time
+
+import torch
+
+from . import wire
+from .shapes import VAE_WINDOWS
+
+ROOT = Path(__file__).resolve().parents[1]
+# The tt-metal image the worker runs in. There is no portable default: set MUSIC_TT_IMAGE,
+# or pass image= , to the tt-metal / ttnn image you built the port against.
+DEFAULT_IMAGE = ""
+
+
+class WorkerError(RuntimeError):
+    """The P100a worker failed (died, hung or reported an error)."""
+
+
+class TTWorker:
+    """checkpoint is cfm_model.pt and vae is vae_model.pt (Hugging Face cache links are resolved)."""
+
+    def __init__(self, checkpoint: Path, image: str = None, log=None,
+                 reply_timeout_s: float = 120.0,  # the slowest call (an 8-minute song's DiT step or VAE window) is < 1 s
+                 model: str = "diffrhythm", vae: Path = None):
+        self.checkpoint = Path(checkpoint)
+        self.model = model
+        self.vae = Path(vae) if vae is not None else None
+        self.image = image or os.environ.get("MUSIC_TT_IMAGE") or DEFAULT_IMAGE
+        if not self.image:
+            raise ValueError("set MUSIC_TT_IMAGE (or pass image=) to the tt-metal image the worker runs in")
+        self.log = log
+        self.reply_timeout_s = reply_timeout_s
+        self.proc = self.conn = self.ipc = None
+        self.name = f"music-tt-{os.getpid()}-{secrets.token_hex(3)}"
+        self.stats = {}
+        self.failed = False  # set when the worker died, hung or reported an error: the card may be unhealthy
+
+    # ------------------------------------------------------------------ lifecycle
+    def start(self) -> "TTWorker":
+        try:
+            cache = ROOT / "data/tt-cache"
+            cache.mkdir(parents=True, exist_ok=True)
+            run_dir = ROOT / "data/tt-run"
+            run_dir.mkdir(parents=True, exist_ok=True)
+            self.ipc = Path(tempfile.mkdtemp(prefix="dit-", dir=run_dir))
+            self.authkey = secrets.token_hex(16)
+            mounts = [("/dev/hugepages-1G", "/dev/hugepages-1G", False), (ROOT, "/work", True),
+                      *self._model_mounts(), (cache, "/cache", False), (self.ipc, "/ipc", False)]
+            cmd = ["docker", "run", "--rm", "--name", self.name, "--ipc", "host", "--device", "/dev/tenstorrent"]
+            for src, dst, ro in mounts:
+                cmd += ["--mount", f"type=bind,src={src},dst={dst}" + (",readonly" if ro else "")]
+            for k, v in {"TT_METAL_VISIBLE_DEVICES": "0", "MESH_DEVICE": "P100", "PYTHONUNBUFFERED": "1",
+                         "TT_METAL_OPERATION_TIMEOUT_SECONDS": "90", "MUSIC_TT_AUTHKEY": self.authkey}.items():
+                cmd += ["-e", f"{k}={v}"]
+            cmd += [self.image, "python", "/work/tt_diffrhythm/worker.py", "--socket", "/ipc/worker.sock", "--model", self.model]
+            # own session: a supervisor may stop the host process by signalling its process group, and a signal reaching
+            # `docker run` directly would stop the container while this process still has to wait for it
+            self.proc = subprocess.Popen(cmd, stdout=self.log or subprocess.DEVNULL, stderr=subprocess.STDOUT,
+                                         start_new_session=True)
+            self._started = time.monotonic()
+        except BaseException:
+            self.close()
+            raise
+        return self
+
+    def _model_mounts(self):
+        if self.model == "diffrhythm":  # files, mounted where worker.py expects them
+            return [(self.checkpoint.resolve(), "/checkpoint/cfm_model.pt", True),
+                    (self.vae.resolve(), "/vae/vae_model.pt", True)]
+        raise ValueError(f"unknown model {self.model!r}")
+
+    def died(self) -> bool:
+        """The worker container exited on its own (not through close())."""
+        return self.proc is not None and self.proc.poll() is not None
+
+    def _fail(self, message):
+        self.failed = True
+        raise WorkerError(message)
+
+    def wait_ready(self, timeout_s: float = 300.0):
+        sock = self.ipc / "worker.sock"
+        deadline = time.monotonic() + timeout_s
+        while not sock.exists():
+            self._check_alive()
+            if time.monotonic() > deadline:
+                self._fail("TT worker did not open its socket")
+            time.sleep(0.2)
+        try:
+            self.conn = Client(str(sock), family="AF_UNIX", authkey=self.authkey.encode())
+        except (OSError, EOFError) as exc:  # the worker died while loading
+            self._fail(f"TT worker connection failed: {exc}")
+        try:
+            ready = self._recv(timeout_s)
+        except (OSError, EOFError) as exc:
+            self._fail(f"TT worker connection lost: {exc}")
+        self.stats.update(worker_load_s=round(ready["load_s"], 2), start_to_ready_s=round(time.monotonic() - self._started, 2))
+        return ready
+
+    def close(self):
+        """Stop the worker and wait for the container to exit; only then give the card back."""
+        try:
+            if self.conn is not None:
+                try:
+                    self.conn.send({"op": "close"})
+                except OSError:
+                    pass
+                self.conn.close()
+            if self.proc is not None:
+                try:
+                    self.proc.wait(60)
+                except subprocess.TimeoutExpired:
+                    subprocess.run(["docker", "stop", "-t", "20", self.name], capture_output=True, timeout=60)
+                    self.proc.wait(60)
+        finally:
+            self.conn = None
+            if self.ipc is not None:
+                shutil.rmtree(self.ipc, ignore_errors=True)
+
+    def __enter__(self):
+        return self.start()
+
+    def __exit__(self, *exc):
+        self.close()
+
+    # ------------------------------------------------------------------ calls
+    def _check_alive(self):
+        if self.proc is not None and self.proc.poll() is not None:
+            self._fail(f"TT worker exited ({self.proc.returncode}); see the worker log")
+
+    def _recv(self, timeout_s):
+        deadline = time.monotonic() + timeout_s
+        while not self.conn.poll(1.0):
+            self._check_alive()
+            if time.monotonic() > deadline:
+                self._fail("TT worker did not answer in time")
+        msg = self.conn.recv()
+        if msg.get("op") == "error":
+            self._fail(msg["error"])
+        return msg
+
+    def call(self, msg: dict) -> dict:
+        try:
+            self.conn.send(msg)
+            return self._recv(self.reply_timeout_s)
+        except (OSError, EOFError) as exc:  # the socket broke: the worker is gone
+            self._fail(f"TT worker connection lost: {exc}")
+
+
+class RemoteDiffRhythmDiT(torch.nn.Module):
+    """Drop-in for DiffRhythm's DiT (cfm.transformer) in CFM.sample (batch 1): same call, same output.
+
+    Keeps upstream's small modules on the host, exactly as upstream runs them: the text embedding (computed once per
+    song and CFG branch; upstream recomputes the same thing every call) and the time / start / duration embeddings.
+    Everything weight-heavy (input projection, position convs, 16 Llama layers, output) runs on the card."""
+
+    def __init__(self, worker: TTWorker, dit):
+        super().__init__()
+        self.worker = worker
+        self.text_embed, self.time_embed = dit.text_embed, dit.time_embed
+        self.start_time_embed, self.duration_time_embed = dit.start_time_embed, dit.duration_time_embed
+        self._song = None
+        self.seconds = []
+
+    def reset(self) -> None:
+        """Forget which song is prepared on the card.
+
+        The worker releases the song's tensors when the VAE runs, so that every song length meets
+        the same device allocation state and the VAE's compiled kernels are reused. After that the
+        card holds no song, and this cache -- which exists to skip re-sending an unchanged text and
+        conditioning -- would otherwise send a `forward` for a song that is no longer there.
+
+        A caller that builds a shim per song never notices. One that keeps it across songs, as a
+        server does, has to call this between them.
+        """
+        self._song = None
+
+    def forward(self, x, cond, text, time, drop_audio_cond, drop_text, drop_prompt=False, style_prompt=None,
+                start_time=None, duration=None):
+        if x.shape[0] != 1:
+            raise ValueError("The P100a DiT runs batch 1")
+        if bool(drop_audio_cond) != bool(drop_text):
+            raise ValueError("The P100a DiT runs CFM.sample's two branches (both conditions kept or both dropped)")
+        seq_len = x.shape[1]
+        if self._song is None or not (torch.equal(text, self._song[0]) and torch.equal(cond, self._song[1])):
+            embeds = {null: wire.encode(self.text_embed(text, seq_len, drop_text=null)[0]) for null in (False, True)}
+            self.worker.call({"op": "prepare", "text_embed": embeds, "cond": wire.encode(cond[0])})
+            self._song = (text.clone(), cond.clone())
+        if time.ndim == 0:
+            time = time.repeat(1)
+        dtype = next(self.time_embed.parameters()).dtype  # upstream hands over fp16 style/start/duration tensors
+        c = (self.time_embed(time.to(dtype)) + self.start_time_embed(start_time.to(dtype))
+             + self.duration_time_embed(duration.to(dtype)))
+        style = torch.zeros_like(style_prompt) if drop_prompt else style_prompt
+        reply = self.worker.call({"op": "forward", "x": wire.encode(x[0]), "drop_text": bool(drop_text),
+                                  "style": wire.encode(style[0]), "c": wire.encode(c[0])})
+        self.seconds.append(reply["seconds"])
+        return wire.decode(reply["v"]).to(x.dtype)[None]
+
+
+class RemoteVae:
+    """Stands in for handler.tiled_decode: overlap-discard tiling like upstream's, but every window of a song has the
+    same length (the last one is moved back to end at the song's end), so the P100a compiles only a few shapes.
+
+    The window is the largest of `chunks` that fits the song: 1024 frames (41 s) decode at 0.61 ms per frame with
+    12.5 % overlap, 512 frames at the same rate with 25 % overlap (vae_profile.py); longer windows get slower per
+    frame. Songs of 30-41 s use 512, so no song compiles a shape of its own.
+
+    latents [B, 64, T] -> audio [B, 2, T * upsample] float32 on the CPU (1920 for ACE-Step, 2048 for DiffRhythm)."""
+
+    def __init__(self, worker: TTWorker, chunks=VAE_WINDOWS, overlap: int = 64, upsample: int = 1920):
+        self.worker, self.chunks, self.overlap, self.upsample = worker, tuple(sorted(chunks, reverse=True)), overlap, upsample
+        self.seconds = []
+
+    def chunk_for(self, frames: int) -> int:
+        return next((c for c in self.chunks if c <= frames), self.chunks[-1])
+
+    @staticmethod
+    def windows(frames: int, chunk: int, overlap: int):
+        """(window_start, core_start, core_end) covering [0, frames); windows are exactly `chunk` frames long
+        (or the whole song when it is shorter) and keep >= overlap frames of context around each core."""
+        if frames <= chunk:
+            return [(0, 0, frames)]
+        stride = chunk - 2 * overlap
+        out = []
+        for core_start in range(0, frames, stride):
+            core_end = min(core_start + stride, frames)
+            out.append((min(max(core_start - overlap, 0), frames - chunk), core_start, core_end))
+        return out
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        reply = self.worker.call({"op": "vae", "z": wire.encode(z)})
+        self.seconds.append(reply["seconds"])
+        return wire.decode(reply["audio"])
+
+    def __call__(self, latents, chunk_size=None, overlap=None, offload_wav_to_cpu=None):
+        out = []
+        for z in latents.detach().float().cpu():
+            frames = z.shape[-1]
+            chunk = self.chunk_for(frames)
+            pieces = []
+            for win_start, core_start, core_end in self.windows(frames, chunk, self.overlap):
+                win_end = min(win_start + chunk, frames)
+                audio = self._decode(z[:, win_start:win_end].contiguous())
+                pieces.append(audio[:, (core_start - win_start) * self.upsample:(core_end - win_start) * self.upsample])
+            out.append(torch.cat(pieces, dim=-1))
+        return torch.stack(out)

code/tt_diffrhythm/server.py

@@ -0,0 +1,334 @@
+# SPDX-License-Identifier: Apache-2.0
+"""HTTP server for DiffRhythm 1.2 Full on one Blackhole p100a, everything in one process.
+
+    uvicorn tt_diffrhythm.server:app --host 0.0.0.0 --port 20000
+
+DiffRhythm's own sampling loop runs here on the CPU and the CFM transformer and the stable-audio
+VAE decoder run on the card, joined by `LocalWorker` instead of the Unix socket `remote.py` uses.
+That split existed because DiffRhythm and tt-metal could not share an interpreter in the
+deployment this port came from; an image built for this model alone resolves one environment that
+satisfies both, so the socket is unnecessary there. See `local.py`.
+
+Requires DiffRhythm itself (its `model`, `infer` and `g2p` trees) importable beside this one -- it
+is the model, and this package only moves two of its stages onto the card.
+"""
+from __future__ import annotations
+
+import base64
+import io
+import os
+import random
+import re
+import sys
+import time
+from contextlib import asynccontextmanager
+from pathlib import Path
+from threading import Lock
+from typing import Optional
+
+import torch
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, ConfigDict, Field
+
+from .local import LocalWorker
+from .remote import RemoteDiffRhythmDiT, RemoteVae
+
+LICENSE = "apache-2.0"
+WEIGHTS_REPO = "ASLP-lab/DiffRhythm-1_2-full"
+VAE_REPO = "ASLP-lab/DiffRhythm-vae"
+STYLE_REPO = "OpenMuQ/MuQ-MuLan-large"
+STEPS = 32  # the schedule this port's shapes are compiled for
+CFG_STRENGTH = 4.0
+MAX_FRAMES = 6144  # the "full" checkpoint: 285 seconds at 21.5 frames per second
+MIN_DURATION, MAX_DURATION = 96, 285
+SAMPLE_RATE = 44100
+UPSAMPLE = 2048  # VAE frames -> audio samples
+TURN_WAIT_S = float(os.environ.get("MUSIC_TURN_WAIT_S", "1800"))
+
+# DiffRhythm reads `./config/diffrhythm-1b.json` and `./g2p/g2p/vocab.json` relative to the working
+# directory, so the server runs from the checkout rather than rewriting upstream's path handling.
+DIFFRHYTHM_ROOT = os.environ.get("DIFFRHYTHM_ROOT", "")
+
+TIMESTAMP = re.compile(r"^\[(\d{1,2}):(\d{2}(?:\.\d{1,3})?)\]\s*(.*)$")
+SECTION = re.compile(r"^\[[A-Za-z][^\]]*\]$")
+
+STATE: dict = {"status": "loading", "error": None, "generating": False, "load_s": None, "dram": None}
+WORKER: Optional[LocalWorker] = None
+CFM = TOKENIZER = STYLE = None
+LOCK = Lock()
+
+
+def _root() -> Path:
+    """The DiffRhythm checkout: DIFFRHYTHM_ROOT, or the one shipped beside this package."""
+    if DIFFRHYTHM_ROOT:
+        return Path(DIFFRHYTHM_ROOT).resolve()
+    here = Path(__file__).resolve().parent.parent
+    if (here / "infer" / "infer_utils.py").is_file():
+        return here
+    raise RuntimeError("set DIFFRHYTHM_ROOT to the DiffRhythm checkout (the directory holding infer/)")
+
+
+def _point_phonemizer_at_a_bundled_espeak() -> None:
+    """Let phonemizer find a libespeak-ng that came from a wheel instead of a system package.
+
+    `g2p.utils.g2p` builds its `EspeakBackend`s at import, so espeak has to be resolvable before
+    upstream is imported at all. Everywhere but a container that means `apt install espeak-ng`;
+    the `espeakng-loader` wheel carries the library and its data instead. Both have to be set --
+    the library alone leaves espeak looking for `phontab` under the path it was built at, which
+    does not exist here.
+
+    Does nothing when the wheel is absent, or when PHONEMIZER_ESPEAK_LIBRARY already names a
+    library: a system espeak-ng is the better answer where there is one, and an explicit choice
+    is not ours to override. (Asking `EspeakWrapper` whether one is set is not the way to decide
+    -- `library_path` is a property on the class, so reading it off the class hands back the
+    property object, which is never None.)
+    """
+    if os.environ.get("PHONEMIZER_ESPEAK_LIBRARY"):
+        return
+    try:
+        import espeakng_loader
+        from phonemizer.backend.espeak.wrapper import EspeakWrapper
+    except ImportError:
+        return
+    EspeakWrapper.set_library(str(espeakng_loader.get_library_path()))
+    if hasattr(EspeakWrapper, "set_data_path"):
+        EspeakWrapper.set_data_path(str(espeakng_loader.get_data_path()))
+    else:  # older phonemizer: espeak reads the data path from its own environment variable
+        os.environ.setdefault("ESPEAK_DATA_PATH", str(espeakng_loader.get_data_path()))
+
+
+def _enter_checkout() -> Path:
+    """Make upstream importable and its relative paths resolve, the way its own scripts do."""
+    root = _root()
+    _point_phonemizer_at_a_bundled_espeak()
+    for path in (root, root / "infer"):
+        if str(path) not in sys.path:
+            sys.path.insert(0, str(path))
+    os.chdir(root)
+    return root
+
+
+def _load_cfm(root: Path, device: str = "cpu"):
+    """`infer_utils.prepare_model` without the VAE (the card decodes) and without its
+    CWD-relative download cache, so the weights land in the ordinary HF cache."""
+    import json
+
+    from huggingface_hub import hf_hub_download
+    from infer_utils import CNENTokenizer, load_checkpoint
+    from model import CFM as CFMModel
+    from model import DiT
+
+    with open(root / "config" / "diffrhythm-1b.json") as f:
+        config = json.load(f)
+    cfm = CFMModel(transformer=DiT(**config["model"], max_frames=MAX_FRAMES),
+                   num_channels=config["model"]["mel_dim"], max_frames=MAX_FRAMES).to(device)
+    cfm = load_checkpoint(cfm, hf_hub_download(WEIGHTS_REPO, "cfm_model.pt"), device=device, use_ema=False)
+    return cfm, CNENTokenizer()
+
+
+def _style_model(device: str = "cpu"):
+    from muq import MuQMuLan
+
+    return MuQMuLan.from_pretrained(STYLE_REPO).to(device).eval().float()
+
+
+@asynccontextmanager
+async def lifespan(_app: FastAPI):
+    global WORKER, CFM, TOKENIZER, STYLE
+    started = time.monotonic()
+    try:
+        import gc
+
+        from huggingface_hub import hf_hub_download
+
+        os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+        root = _enter_checkout()
+        WORKER = LocalWorker(hf_hub_download(WEIGHTS_REPO, "cfm_model.pt"),
+                             vae=hf_hub_download(VAE_REPO, "vae_model.pt")).start()
+        cfm, TOKENIZER = _load_cfm(root)
+        # Replacing the transformer drops the host copy of the big weights; the card holds them.
+        cfm.transformer = RemoteDiffRhythmDiT(WORKER, cfm.transformer)
+        CFM = cfm.float().eval()
+        STYLE = _style_model()
+        gc.collect()
+        STATE.update(status="ok", load_s=round(time.monotonic() - started, 2),
+                     dram=WORKER.stats.get("dram"), vae_warm_s=WORKER.stats.get("vae_warm_s"))
+    except Exception as exc:
+        STATE.update(status="error", error=f"{type(exc).__name__}: {exc}")
+    try:
+        yield
+    finally:
+        CFM = TOKENIZER = STYLE = None
+        if WORKER is not None:
+            WORKER.close()
+            WORKER = None
+
+
+app = FastAPI(title="DiffRhythm 1.2 Full on p100a", lifespan=lifespan, docs_url=None, redoc_url=None)
+
+
+class Request(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    prompt: str = Field(min_length=1, max_length=4000, description="the style description")
+    lyrics: str = Field(default="", max_length=20000, description="LRC, or plain lines to be timed")
+    instrumental: bool = False
+    duration: int = 180
+    seed: Optional[int] = Field(default=None, ge=0, le=4294967295)
+    inference_steps: Optional[int] = None
+
+
+def _stamp(seconds: float) -> str:
+    centiseconds = round(seconds * 100)
+    minutes, remainder = divmod(centiseconds, 6000)
+    return f"[{minutes:02d}:{remainder / 100:05.2f}]"
+
+
+def _lrc(lyrics: str, duration: int, instrumental: bool) -> str:
+    """Real LRC as given, or approximate starts derived for untimed lines.
+
+    The model needs a timestamp on every line. Deriving them by word count is a convenience,
+    not learned lyric alignment -- pass LRC when the timing matters."""
+    if instrumental:
+        return ""
+    lines = [line.strip() for line in lyrics.splitlines() if line.strip()]
+    lines = [line for line in lines if not SECTION.fullmatch(line)]
+    if not lines:
+        raise ValueError("vocals need lyrics; set instrumental=true for music without them")
+    if len(lines) > 400:
+        raise ValueError("at most 400 lyric lines")
+    matches = [TIMESTAMP.match(line) for line in lines]
+    if any(matches):
+        if not all(matches):
+            raise ValueError("use either plain lyrics or a timestamp on every line (LRC)")
+        out, previous = [], -1.0
+        for match in matches:
+            minutes, seconds, line = match.groups()
+            start = int(minutes) * 60 + float(seconds)
+            if float(seconds) >= 60 or not previous <= start < duration:
+                raise ValueError("LRC timestamps must be ordered and shorter than the song")
+            previous = start
+            if line.strip():  # an empty timed line marks a break; the tokenizer cannot encode it
+                out.append(f"{_stamp(start)}{line.strip()}")
+        if not out:
+            raise ValueError("the LRC contains no singable lyrics")
+        return "\n".join(out)
+    if any(re.match(r"^\[\d", line) for line in lines):
+        raise ValueError("invalid LRC timestamp; use [mm:ss.xx] before each line")
+    intro, outro = min(12.0, duration * 0.06), min(10.0, duration * 0.05)
+    weights = [max(3, len(line.split())) for line in lines]
+    total, elapsed, out = sum(weights), 0, []
+    for line, weight in zip(lines, weights):
+        out.append(f"{_stamp(intro + (duration - intro - outro) * elapsed / total)}{line}")
+        elapsed += weight
+    return "\n".join(out)
+
+
+def _wav(waveform: torch.Tensor) -> bytes:
+    """[channels, samples] float -> peak-normalized 16-bit 44.1 kHz WAV."""
+    import soundfile as sf
+
+    waveform = waveform.float().cpu()
+    if not torch.isfinite(waveform).all():
+        raise RuntimeError("non-finite audio; retry with a different seed")
+    peak = waveform.abs().max().item()
+    if peak < 1e-7:
+        raise RuntimeError("silent audio; retry with a different seed")
+    buf = io.BytesIO()
+    sf.write(buf, (waveform / peak).clamp(-1, 1).T.numpy(), SAMPLE_RATE, subtype="PCM_16", format="WAV")
+    return buf.getvalue()
+
+
+@app.get("/health")
+def health():
+    return dict(STATE)
+
+
+@app.get("/info")
+def info():
+    return dict(STATE) | {
+        "model": WEIGHTS_REPO,
+        "license": LICENSE,
+        "device": "Tenstorrent Blackhole p100a",
+        "on_card": ["CFM transformer", "stable-audio VAE decoder"],
+        "on_cpu": ["MuQ-MuLan style encoder", "lyric tokenizer", "sampler", "post-processing"],
+        "inference_steps": STEPS,
+        "cfg_strength": CFG_STRENGTH,
+        "task_modes": ["text-to-music"],
+        "duration_limits": {"min": MIN_DURATION, "max": MAX_DURATION},
+    }
+
+
+@app.post("/predict")
+def predict(request: Request):
+    if not request.prompt.strip():
+        raise HTTPException(400, "Empty prompt")
+    try:
+        if request.inference_steps is not None and request.inference_steps != STEPS:
+            raise ValueError(f"this port runs {STEPS} steps (the schedule its shapes are compiled for)")
+        if not MIN_DURATION <= request.duration <= MAX_DURATION:
+            raise ValueError(f"duration must be between {MIN_DURATION} and {MAX_DURATION} seconds "
+                             "(use ACE-Step for longer songs)")
+        lrc = _lrc(request.lyrics, request.duration, request.instrumental)
+    except ValueError as exc:
+        raise HTTPException(422, str(exc)) from exc
+    if STATE["status"] != "ok" or CFM is None:
+        raise HTTPException(503, STATE["error"] or "Model is not ready")
+    if not LOCK.acquire(timeout=TURN_WAIT_S):
+        raise HTTPException(409, "Another song is in progress")
+    try:
+        import numpy as np
+        from infer_utils import (get_lrc_token, get_negative_style_prompt, get_reference_latent,
+                                 get_style_prompt)
+
+        STATE["generating"] = True
+        t0 = time.perf_counter()
+        seed = request.seed if request.seed is not None else random.SystemRandom().randrange(2**32)
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        with torch.inference_mode():
+            style = get_style_prompt(STYLE, prompt=request.prompt).float()
+            text, start_time, end_frame, song_duration = get_lrc_token(
+                MAX_FRAMES, lrc, TOKENIZER, request.duration, "cpu")
+            negative = get_negative_style_prompt("cpu").float()
+            cond, segments = get_reference_latent("cpu", MAX_FRAMES, False, None, None, None)
+            latents, _ = CFM.sample(
+                cond=cond, text=text, duration=end_frame, max_duration=end_frame,
+                song_duration=song_duration, style_prompt=style, negative_style_prompt=negative,
+                steps=STEPS, cfg_strength=CFG_STRENGTH, start_time=start_time,
+                latent_pred_segments=segments, batch_infer_num=1, seed=seed)
+            latent = latents[0].float().transpose(1, 2).contiguous()
+            vae = RemoteVae(WORKER, upsample=UPSAMPLE)
+            waveform = vae(latent)[0]
+        wav = _wav(waveform)
+        timing = {
+            "dit_s": round(sum(CFM.transformer.seconds), 2),
+            "dit_calls": len(CFM.transformer.seconds),
+            "vae_s": round(sum(vae.seconds), 2),
+            "total_s": round(time.perf_counter() - t0, 2),
+        }
+        CFM.transformer.seconds.clear()
+        # The card holds no song once the VAE has run, so the shim's "already prepared" cache has
+        # to go with it -- the next request's text and conditioning differ and would otherwise be
+        # skipped. A per-song shim would not need this; a server keeps one across songs.
+        CFM.transformer.reset()
+        return {
+            "audio": base64.b64encode(wav).decode(),
+            "format": "wav",
+            "model": WEIGHTS_REPO,
+            "license": LICENSE,
+            "seed": seed,
+            "duration": request.duration,
+            "lrc": lrc,
+            "inference_steps": STEPS,
+            "timing_s": timing,
+        }
+    except Exception as exc:
+        if CFM is not None:
+            CFM.transformer.reset()
+        if WORKER is not None and WORKER.failed:
+            STATE.update(status="error", error=str(exc))
+        raise HTTPException(503, str(exc)) from exc
+    finally:
+        STATE["generating"] = False
+        LOCK.release()

code/tt_diffrhythm/shapes.py

@@ -0,0 +1,5 @@
+"""Shapes shared by the host pipeline and the TT worker (no torch or ttnn imports)."""
+
+# VAE windows in latent frames (remote.py RemoteVae picks the largest that fits the song); the worker builds
+# these programs at start and warm_cache.py compiles them ahead of time
+VAE_WINDOWS = (1024, 512)

code/tt_diffrhythm/warm_cache.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+"""Compile every P100a kernel shape this port can hit, ahead of the first songs.
+
+The card compiles kernels per shape and keeps them in TT_METAL_CACHE (data/tt-cache). Without this, the first song
+of each new length bucket pays that compilation (tens of seconds). Shapes, per model (--model):
+           ENC_BUCKET), and both VAE windows (remote.py RemoteVae)
+  diffrhythm: every DiT sequence bucket of a 96-285 s song (diffrhythm_host.seq_buckets), both CFG branches, and
+           both VAE windows
+Random inputs: only the shapes matter. Run it again after changing this package or the TT image.
+"""
+import argparse
+import json
+import time
+
+import torch
+
+from . import diffrhythm_dit, diffrhythm_host as dhost, oobleck_vae
+from .shapes import VAE_WINDOWS
+
+def warm_vae(vae, report):
+    for window in VAE_WINDOWS:  # exactly as the worker does right after loading (worker.py)
+        t0 = time.monotonic()
+        audio = vae.decode(torch.zeros(64, window))
+        rec = {"vae_window": window, "s": round(time.monotonic() - t0, 2), "finite": bool(torch.isfinite(audio).all())}
+        report["vae"].append(rec)
+        print(json.dumps(rec), flush=True)
+
+
+def warm_diffrhythm(args):
+    g = torch.Generator().manual_seed(0)
+    started = time.monotonic()
+    dev = diffrhythm_dit.open_device()
+    report = {"seq_buckets": dhost.seq_buckets(), "dit": [], "vae": []}
+    try:
+        dit = diffrhythm_dit.DiffRhythmDiT(dev, dhost.Checkpoint(args.checkpoint or "/checkpoint/cfm_model.pt"))
+        vae = oobleck_vae.OobleckTT(dev, args.vae or "/vae/vae_model.pt")
+        report["load_s"] = round(time.monotonic() - started, 1)
+        warm_vae(vae, report)
+        for seq_pad in report["seq_buckets"]:
+            frames = min(seq_pad, dhost.MAX_FRAMES)  # any length in the bucket gives the same shapes
+            t0 = time.monotonic()
+            song = dit.prepare({null: torch.randn(frames, 512, generator=g) for null in (False, True)},
+                               torch.randn(frames, 64, generator=g))
+            assert song.geo.seq_pad == seq_pad
+            finite = True
+            for null in (False, True):
+                v = dit.forward(torch.randn(frames, 64, generator=g), null, torch.randn(512, generator=g),
+                                torch.randn(512, generator=g), song)
+                finite &= bool(torch.isfinite(v).all())
+            dit.release(song)
+            rec = {"seq_pad": seq_pad, "s": round(time.monotonic() - t0, 2), "finite": finite}
+            report["dit"].append(rec)
+            print(json.dumps(rec), flush=True)
+    finally:
+        diffrhythm_dit.close_device(dev)
+    report["total_s"] = round(time.monotonic() - started, 1)
+    print(json.dumps({k: report[k] for k in ("seq_buckets", "load_s", "total_s")}), flush=True)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--vae")
+    args = parser.parse_args()
+    return warm_diffrhythm(args)
+
+
+if __name__ == "__main__":
+    main()

code/tt_diffrhythm/wire.py

@@ -0,0 +1,18 @@
+"""Tensors between the host ACE-Step process and the TT DiT worker.
+
+The two sides run different numpy/torch builds (ACE-Step venv vs the TT image), so tensors cross the socket as raw
+little-endian bytes with dtype and shape, inside plain dicts that pickle identically on both sides."""
+from __future__ import annotations
+
+import torch
+
+_DTYPES = {"float32": torch.float32}
+
+
+def encode(t: torch.Tensor) -> dict:
+    t = t.detach().to("cpu", torch.float32).contiguous()
+    return {"dtype": "float32", "shape": list(t.shape), "data": t.numpy().tobytes()}
+
+
+def decode(d: dict) -> torch.Tensor:
+    return torch.frombuffer(bytearray(d["data"]), dtype=_DTYPES[d["dtype"]]).reshape(d["shape"])

code/tt_diffrhythm/worker.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""TT worker: loads a music model's DiT and its audio VAE decoder onto the P100a and answers one host pipeline.
+
+Runs inside the TT image; the host (remote.py) owns the container for its whole life and talks to
+this process over a Unix socket in a directory bind-mounted into both:
+
+  python /work/tt_diffrhythm/worker.py --socket /ipc/worker.sock [--model diffrhythm] --checkpoint ... --vae ...
+
+Messages are plain dicts (tensors via wire.py):
+  -> {"op": "ready", ...}                                           once the weights are on the device
+  <- {"op": "prepare", ...}                                         -> {"op": "prepared", ...}
+       DiffRhythm: "text_embed": {drop_text: [T, 512]}, "cond": [T, 64]
+  <- {"op": "forward", ...}                                         -> {"op": "velocity", "v": [T, 64], ...}
+       DiffRhythm: "x": [T, 64], "drop_text": bool, "style": [512], "c": [512]
+  <- {"op": "vae", "z": [64, L]}                                    -> {"op": "audio", "audio": [2, L * upsample], ...}
+  <- {"op": "close"}
+Any failure is sent back as {"op": "error", "error": "..."} and ends the worker.
+"""
+import argparse
+from multiprocessing.connection import Listener
+import os
+import signal
+import sys
+import time
+import traceback
+
+import torch
+
+from . import diffrhythm_dit, oobleck_vae, wire
+from .shapes import VAE_WINDOWS
+
+
+class DiffRhythm:
+    def __init__(self, dev, checkpoint):
+        from . import diffrhythm_host
+
+        self.dit = diffrhythm_dit.DiffRhythmDiT(dev, diffrhythm_host.Checkpoint(checkpoint))
+
+    def prepare(self, msg):
+        text_embed = {bool(k): wire.decode(v) for k, v in msg["text_embed"].items()}
+        song = self.dit.prepare(text_embed, wire.decode(msg["cond"]))
+        return song, {"frames": song.geo.frames, "seq_pad": song.geo.seq_pad}
+
+    def forward(self, msg, song):
+        return self.dit.forward(wire.decode(msg["x"]), bool(msg["drop_text"]), wire.decode(msg["style"]),
+                                wire.decode(msg["c"]), song)
+
+
+MODELS = {"diffrhythm": (DiffRhythm, "/checkpoint/cfm_model.pt", "/vae/vae_model.pt")}
+
+
+def serve(conn, model, vae):
+    dit = model.dit
+    song = None
+    while True:
+        msg = conn.recv()
+        op = msg.get("op")
+        started = time.monotonic()
+        if op == "prepare":
+            if song is not None:
+                dit.release(song)
+                song = None
+            song, geometry = model.prepare(msg)
+            conn.send({"op": "prepared", "seconds": time.monotonic() - started, "geometry": geometry})
+        elif op == "forward":
+            if song is None:
+                raise RuntimeError("forward before prepare")
+            v = model.forward(msg, song)
+            conn.send({"op": "velocity", "v": wire.encode(v), "seconds": time.monotonic() - started})
+        elif op == "vae":
+            if song is not None:  # sampling is over; freeing the song's tensors first gives the VAE the same device
+                dit.release(song)  # allocation state for every song length, so its compiled conv kernels are reused
+                song = None
+            audio = vae.decode(wire.decode(msg["z"]))
+            conn.send({"op": "audio", "audio": wire.encode(audio), "seconds": time.monotonic() - started})
+        elif op == "close":
+            if song is not None:
+                dit.release(song)
+            return
+        else:
+            raise RuntimeError(f"unknown op {op!r}")
+
+
+def main():
+    # docker stop / a signal: leave through the finally below so the device is closed, not dropped
+    signal.signal(signal.SIGTERM, lambda *_: sys.exit(143))
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--socket", required=True)
+    parser.add_argument("--model", choices=sorted(MODELS), default="diffrhythm")
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--vae")
+    args = parser.parse_args()
+    model_class, checkpoint, vae_path = MODELS[args.model]
+    authkey = os.environ.get("MUSIC_TT_AUTHKEY", "").encode() or None
+    listener = Listener(args.socket, family="AF_UNIX", authkey=authkey)
+    started = time.monotonic()
+    dev = diffrhythm_dit.open_device()
+    conn = None
+    try:
+        model = model_class(dev, args.checkpoint or checkpoint)
+        vae = oobleck_vae.OobleckTT(dev, args.vae or vae_path)
+        # Build the VAE programs now, in the same device state for every song: their conv kernels then compile
+        # once (TT_METAL_CACHE) and later decodes reuse the in-process programs. Built after a song's DiT steps
+        # instead, the device state (and so the kernels) depends on the song length and recompiles for 20-40 s.
+        warm = time.monotonic()
+        for window in VAE_WINDOWS:
+            vae.decode(torch.zeros(64, window))
+        ready = {"op": "ready", "model": args.model, "load_s": time.monotonic() - started,
+                 "vae_warm_s": time.monotonic() - warm, "upsample": vae.upsample, "dram": diffrhythm_dit.dram_stats(dev)}
+        conn = listener.accept()  # one pipeline per container
+        conn.send(ready)
+        serve(conn, model, vae)
+    except BaseException as exc:
+        traceback.print_exc()
+        if conn is not None:
+            try:
+                conn.send({"op": "error", "error": f"{type(exc).__name__}: {exc}"})
+            except OSError:
+                pass
+        raise
+    finally:
+        diffrhythm_dit.close_device(dev)
+        listener.close()
+
+
+if __name__ == "__main__":
+    main()

image/blobs/sha256/0cc125e439ef6882a01bae57bece67f14ec2ad123aec1fb12e86bc50175a4b44

@@ -0,0 +1,28 @@
+{
+  "schemaVersion": 2,
+  "mediaType": "application/vnd.oci.image.index.v1+json",
+  "manifests": [
+    {
+      "mediaType": "application/vnd.oci.image.manifest.v1+json",
+      "digest": "sha256:81aae6dd123609902a2f56719e50ef0bcf3aa00aff71fc66b02648b3f4aea2f2",
+      "size": 4499,
+      "platform": {
+        "architecture": "amd64",
+        "os": "linux"
+      }
+    },
+    {
+      "mediaType": "application/vnd.oci.image.manifest.v1+json",
+      "digest": "sha256:8f94a918ad7d7da27e2b025adda011abaa34d892741d6b28221015c75a4ac5ba",
+      "size": 837,
+      "annotations": {
+        "vnd.docker.reference.digest": "sha256:81aae6dd123609902a2f56719e50ef0bcf3aa00aff71fc66b02648b3f4aea2f2",
+        "vnd.docker.reference.type": "attestation-manifest"
+      },
+      "platform": {
+        "architecture": "unknown",
+        "os": "unknown"
+      }
+    }
+  ]
+}

image/blobs/sha256/44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a

@@ -0,0 +1 @@
+{}

image/blobs/sha256/81aae6dd123609902a2f56719e50ef0bcf3aa00aff71fc66b02648b3f4aea2f2

@@ -0,0 +1,121 @@
+{
+  "schemaVersion": 2,
+  "mediaType": "application/vnd.oci.image.manifest.v1+json",
+  "config": {
+    "mediaType": "application/vnd.oci.image.config.v1+json",
+    "digest": "sha256:eae56eb3c9610901618ab9c30e2ba8cdcf6981cdd3af7ed5755aba2cc46db6e1",
+    "size": 14328
+  },
+  "layers": [
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:98c4455a98982b35380ec62f90d5cf0edff02eeb4b3bb18595d871fdeb0c0c67",
+      "size": 29751627
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:3375e6602aa7264d6d67ae69ccb99690c8f40af7ab48d0af060dac1a2e9c24c9",
+      "size": 69483284
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:340cf5e2c047e0fc940f2973810428f58250242238231682812dec6f6d54f3ac",
+      "size": 4399
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:1f47921d53a3b8de72172e2fd5f0a0d27f1e8f4b77016c6b8a1a4b591bb83ae5",
+      "size": 9565639
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:4ec2b03c373212fbe7fda97dbf9eaadb1bd42c2b0afc8b8a49806fe1dbc507f7",
+      "size": 138968721
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:b7c2e1ec67bfd9d4578a81ca6987f095b74869db5f518bf2fc9d3ffeb30821e5",
+      "size": 66025096
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:428f99cfe94d9b1e42af2f3d8104598f8ca6e6105925c7687595a18c78d21011",
+      "size": 503666016
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:ee5adbd6eba4dfd9ba7e59692c46379b3d60a88f8445ec6a045d4aeb336f5bb9",
+      "size": 116
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:672d584c094e51e29aa5e46a6dd22cbfb993f4f839609e954d2c299912ea5d77",
+      "size": 139080178
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:2d16f4104fc6e34c556d42d6a35adedd5ffa396e71d59ef6dd7d5e88c01225e0",
+      "size": 39231465
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:556ec3c4c0e12c97f4c9b8dd0cd65ae321238c55438fbf498055f0ad7afc05bb",
+      "size": 39231323
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:b9ddf00cad9d9a3bc91d09c8ee0f4fef42056861bb5ca1209fcc204b68950abf",
+      "size": 66525731
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:0b4ec6da503185d389dcfeefc2c5d03ab13e0d76620b5fb6ce022f79bb64140f",
+      "size": 12761466
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:49353ead98e6a241dd015dda937bf8dcf335e4eed27fe732df8163dd955d042f",
+      "size": 1898379
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:c2a425aad5e9144add8f2a619c6e3d452d30ddbb5d4f4d3e45caab4f7bb3d233",
+      "size": 6878
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:e211ac567cceb6b5be62db5e135a1a5f3e77bcacebdac6a8cf9e0cdc24158c09",
+      "size": 108850224
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:c6bd350dd2ed57d0ea34248e5b27d985a6786330d2d3ec29f5563e03ca9f00da",
+      "size": 1080
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:8f09ce5e10995b3fa6362dae9945d9402a5ccabce3d5bf20df0a53c75d9f94ab",
+      "size": 982
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1",
+      "size": 32
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:e600ba821186d2d540ea813a478dda1e629c539eef283f062f31309cb8c3af47",
+      "size": 674
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:3c7495db974af283bccc2dee306a75da6491b2ca664d55597da926c4fd104dd5",
+      "size": 216662
+    },
+    {
+      "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
+      "digest": "sha256:6acee8cd93aa0889cc3a8b09f1b2564234ee82aa7c31e87ab61591552d15e894",
+      "size": 5608
+    }
+  ]
+}

image/blobs/sha256/8f94a918ad7d7da27e2b025adda011abaa34d892741d6b28221015c75a4ac5ba

@@ -0,0 +1,26 @@
+{
+  "schemaVersion": 2,
+  "mediaType": "application/vnd.oci.image.manifest.v1+json",
+  "artifactType": "application/vnd.docker.attestation.manifest.v1+json",
+  "config": {
+    "mediaType": "application/vnd.oci.empty.v1+json",
+    "digest": "sha256:44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a",
+    "size": 2,
+    "data": "e30="
+  },
+  "layers": [
+    {
+      "mediaType": "application/vnd.in-toto+json",
+      "digest": "sha256:9b3433e39d3d40a2f64d1f1e8cb155387651684c5f5c9473bd733c72056dee58",
+      "size": 1923,
+      "annotations": {
+        "in-toto.io/predicate-type": "https://slsa.dev/provenance/v1"
+      }
+    }
+  ],
+  "subject": {
+    "mediaType": "application/vnd.oci.image.manifest.v1+json",
+    "digest": "sha256:81aae6dd123609902a2f56719e50ef0bcf3aa00aff71fc66b02648b3f4aea2f2",
+    "size": 4499
+  }
+}

image/blobs/sha256/9b3433e39d3d40a2f64d1f1e8cb155387651684c5f5c9473bd733c72056dee58

@@ -0,0 +1 @@
+{"_type":"https://in-toto.io/Statement/v1","predicateType":"https://slsa.dev/provenance/v1","subject":[{"name":"pkg:docker/tt-model/diffrhythm-1.2-p100a@build-58648d444?platform=linux%2Famd64","digest":{"sha256":"81aae6dd123609902a2f56719e50ef0bcf3aa00aff71fc66b02648b3f4aea2f2"}}],"predicate":{"buildDefinition":{"buildType":"https://github.com/moby/buildkit/blob/master/docs/attestations/slsa-definitions.md","resolvedDependencies":[{"uri":"pkg:docker/docker/dockerfile@1.7-labs","digest":{"sha256":"b99fecfe00268a8b556fad7d9c37ee25d716ae08a5d7320e6d51c4dd83246894"}},{"uri":"pkg:docker/ubuntu@22.04?platform=linux%2Famd64","digest":{"sha256":"5ec03bb3441e8b0bf3b4f9cd4629a1ae763010dc3035bb8da3ae6cf026486401"}},{"uri":"pkg:docker/ghcr.io/tenstorrent/tt-metal/tt-metalium/ubuntu-22.04-dev-amd64@latest?platform=linux%2Famd64","digest":{"sha256":"df9d279c7f85c17c6fad982d196802682d669cca1b7ced9cbaad8181339cd5fc"}}],"externalParameters":{"configSource":{"path":"Dockerfile"},"request":{"frontend":"gateway.v0","args":{"cmdline":"docker/dockerfile:1.7-labs","context:metalsrc":"local:metalsrc","frontend.caps":"moby.buildkit.frontend.contexts+forward","sharedkey:localdir:metalsrc":"empty-metalsrc:f5e176f12fb9c045","source":"docker/dockerfile:1.7-labs"},"locals":[{"name":"context"},{"name":"dockerfile"},{"name":"metalsrc"}],"root":{"configSource":{"path":"Dockerfile"},"request":{"args":{"context:metalsrc":"local:metalsrc","frontend.caps":"moby.buildkit.frontend.contexts+forward","sharedkey:localdir:metalsrc":"empty-metalsrc:f5e176f12fb9c045"}}},"compatibilityVersion":30}},"internalParameters":{"builderPlatform":"linux/amd64"}},"runDetails":{"builder":{"id":""},"metadata":{"invocationId":"lt0mpnhkaz55xrvi2wqvfduge","startedOn":"2026-10-09T23:02:08.374050954-04:00","finishedOn":"2026-10-09T23:07:18.826107934-04:00","buildkit_metadata":{},"buildkit_completeness":{"request":false,"resolvedDependencies":false}}}}}

image/blobs/sha256/eae56eb3c9610901618ab9c30e2ba8cdcf6981cdd3af7ed5755aba2cc46db6e1

@@ -0,0 +1 @@
+{"architecture":"amd64","config":{"User":"tt","Env":["PATH=/opt/tt-venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin","VENV=/opt/tt-venv","VIRTUAL_ENV=/opt/tt-venv","TT_METAL_RUNTIME_ROOT=/opt/tt-metal","TT_METAL_HOME=/opt/tt-metal","PYTHONPATH=/opt/tt-metal","LD_LIBRARY_PATH=/opt/tt-metal/build/lib:/opt/openmpi-v5.0.7-ulfm/lib","EXTRA_MODELS_DIR=","TT_VLLM_BUILTIN_MODELS=","TT_MODEL_KIND=tt-dit-server","HF_HOME=/hf","TT_METAL_CACHE=/cache","HOME=/home/tt","USER=tt","LOGNAME=tt"],"Entrypoint":["/usr/local/bin/entrypoint.sh"],"Cmd":["/usr/local/bin/serve-default.sh"],"WorkingDir":"/home/tt/work","Labels":{"org.opencontainers.image.revision":"58648d44443223041148125843ca7c7d4780463e","org.opencontainers.image.version":"22.04","org.tenstorrent.tt-model":"diffrhythm-1.2-p100a","org.tenstorrent.tt-model.arch":"blackhole","org.tenstorrent.tt-model.kind":"tt-dit-server","org.tenstorrent.tt-model.plugin":"","org.tenstorrent.tt-model.profiles":"default","org.tenstorrent.tt-model.repo":"sunandbean/diffrhythm-1.2-p100a","org.tenstorrent.tt-model.tt-metal":"","org.tenstorrent.tt-model.weights":"ASLP-lab/DiffRhythm-1_2-full"},"ArgsEscaped":true},"created":"2026-10-09T23:06:45.672498279-04:00","history":[{"created":"2026-09-24T22:24:04.688908979Z","created_by":"/bin/sh -c #(nop)  ARG RELEASE","empty_layer":true},{"created":"2026-09-24T22:24:04.723527376Z","created_by":"/bin/sh -c #(nop)  ARG LAUNCHPAD_BUILD_ARCH","empty_layer":true},{"created":"2026-09-24T22:24:04.751455651Z","created_by":"/bin/sh -c #(nop)  LABEL org.opencontainers.image.version=22.04","empty_layer":true},{"created":"2026-09-24T22:24:06.837737736Z","created_by":"/bin/sh -c #(nop) ADD file:b0bf3f64519bf10a51e00d4f8ab9c8693620659a0c4a4a5090a9765e754e1d1c in / "},{"created":"2026-09-24T22:24:07.239071539Z","created_by":"/bin/sh -c #(nop)  CMD [\"/bin/bash\"]","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG OMPI_DIR=/opt/openmpi-v5.0.7-ulfm","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG EXTRA_MODELS_DIR=","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG TT_MODEL_KIND","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG MODEL_NAME","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG MODEL_REPO","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG MODEL_WEIGHTS","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG MODEL_ARCH","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG MODEL_PROFILES","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG MODEL_TT_METAL_SHA","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG MODEL_TT_METAL_DESCRIBE","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"ARG MODEL_PLUGIN_SHA","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T09:22:08.388097186-04:00","created_by":"RUN |11 OMPI_DIR=/opt/openmpi-v5.0.7-ulfm EXTRA_MODELS_DIR= TT_MODEL_KIND=tt-dit-server MODEL_NAME=diffrhythm-1.2-p100a MODEL_REPO=sunandbean/diffrhythm-1.2-p100a MODEL_WEIGHTS=ASLP-lab/DiffRhythm-1_2-full MODEL_ARCH=blackhole MODEL_PROFILES=default MODEL_TT_METAL_SHA=58648d44443223041148125843ca7c7d4780463e MODEL_TT_METAL_DESCRIBE= MODEL_PLUGIN_SHA= /bin/sh -c apt-get update \u0026\u0026 apt-get install -y --no-install-recommends     libhwloc15 libnuma1 libatomic1 libudev1 libcap2 zlib1g     libmpc3 libmpfr6 libgmp10 libzstd1     libevent-core-2.1-7 libevent-pthreads-2.1-7     libgl1 libsndfile1 ca-certificates     \u0026\u0026 apt-get clean \u0026\u0026 rm -rf /var/lib/apt/lists/* # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T09:22:08.570265459-04:00","created_by":"RUN |11 OMPI_DIR=/opt/openmpi-v5.0.7-ulfm EXTRA_MODELS_DIR= TT_MODEL_KIND=tt-dit-server MODEL_NAME=diffrhythm-1.2-p100a MODEL_REPO=sunandbean/diffrhythm-1.2-p100a MODEL_WEIGHTS=ASLP-lab/DiffRhythm-1_2-full MODEL_ARCH=blackhole MODEL_PROFILES=default MODEL_TT_METAL_SHA=58648d44443223041148125843ca7c7d4780463e MODEL_TT_METAL_DESCRIBE= MODEL_PLUGIN_SHA= /bin/sh -c existing=\"$(getent passwd 1000 | cut -d: -f1)\"     \u0026\u0026 if [ -n \"$existing\" ]; then          userdel -r \"$existing\" 2\u003e/dev/null || userdel \"$existing\";        fi     \u0026\u0026 useradd --uid 1000 --create-home --home-dir /home/tt --shell /bin/bash tt     \u0026\u0026 mkdir -p /home/tt/work/logs /cache /opt/tt-metal     \u0026\u0026 chown -R tt:tt /home/tt /cache /opt/tt-metal     \u0026\u0026 chmod 1777 /home/tt /home/tt/work /home/tt/work/logs /cache # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T09:26:50.393285226-04:00","created_by":"COPY /opt/openmpi-v5.0.7-ulfm /opt/openmpi-v5.0.7-ulfm # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T09:26:50.967022517-04:00","created_by":"COPY /opt/tenstorrent /opt/tenstorrent # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:49.28099242-04:00","created_by":"COPY /usr/local/share/uv /usr/local/share/uv # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:52.554077344-04:00","created_by":"COPY /opt/tt-venv /opt/tt-venv # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:52.741256841-04:00","created_by":"COPY /opt/vllm /opt/vllm # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:53.245850738-04:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/runtime /opt/tt-metal/runtime # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:53.42384338-04:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/build_Release /opt/tt-metal/build_Release # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:53.584846845-04:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/build /opt/tt-metal/build # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:54.25191112-04:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/tt_metal /opt/tt-metal/tt_metal # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:54.825250753-04:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/ttnn /opt/tt-metal/ttnn # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:54.915803026-04:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/tools /opt/tt-metal/tools # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:54.941505166-04:00","created_by":"COPY --chown=tt:tt /opt/tt-metal/setup.py /opt/tt-metal/pyproject.toml /opt/tt-metal/ # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:55.213574197-04:00","created_by":"COPY --chown=tt:tt code/ /opt/tt-metal/ # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:55.238685903-04:00","created_by":"COPY entrypoint.sh /usr/local/bin/entrypoint.sh # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"COPY --chmod=0755 serve-default.sh /usr/local/bin/serve-default.sh # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV VENV=/opt/tt-venv","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV VIRTUAL_ENV=/opt/tt-venv","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV PATH=/opt/tt-venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV TT_METAL_RUNTIME_ROOT=/opt/tt-metal","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV TT_METAL_HOME=/opt/tt-metal","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV PYTHONPATH=/opt/tt-metal","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV LD_LIBRARY_PATH=/opt/tt-metal/build/lib:/opt/openmpi-v5.0.7-ulfm/lib","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV EXTRA_MODELS_DIR=","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ARG TT_VLLM_BUILTIN_MODELS=","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV TT_VLLM_BUILTIN_MODELS=","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV TT_MODEL_KIND=tt-dit-server","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV HF_HOME=/hf","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV TT_METAL_CACHE=/cache","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV HOME=/home/tt","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV USER=tt","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"ENV LOGNAME=tt","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.256601142-04:00","created_by":"USER tt","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:05:55.273892093-04:00","created_by":"WORKDIR /home/tt/work","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:05:55.28789741-04:00","created_by":"COPY verify.sh /ctx/verify.sh # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:06:45.494946817-04:00","created_by":"RUN |12 OMPI_DIR=/opt/openmpi-v5.0.7-ulfm EXTRA_MODELS_DIR= TT_MODEL_KIND=tt-dit-server MODEL_NAME=diffrhythm-1.2-p100a MODEL_REPO=sunandbean/diffrhythm-1.2-p100a MODEL_WEIGHTS=ASLP-lab/DiffRhythm-1_2-full MODEL_ARCH=blackhole MODEL_PROFILES=default MODEL_TT_METAL_SHA=58648d44443223041148125843ca7c7d4780463e MODEL_TT_METAL_DESCRIBE= MODEL_PLUGIN_SHA= TT_VLLM_BUILTIN_MODELS= /bin/sh -c bash /ctx/verify.sh # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:06:45.494946817-04:00","created_by":"USER root","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:06:45.672498279-04:00","created_by":"RUN |12 OMPI_DIR=/opt/openmpi-v5.0.7-ulfm EXTRA_MODELS_DIR= TT_MODEL_KIND=tt-dit-server MODEL_NAME=diffrhythm-1.2-p100a MODEL_REPO=sunandbean/diffrhythm-1.2-p100a MODEL_WEIGHTS=ASLP-lab/DiffRhythm-1_2-full MODEL_ARCH=blackhole MODEL_PROFILES=default MODEL_TT_METAL_SHA=58648d44443223041148125843ca7c7d4780463e MODEL_TT_METAL_DESCRIBE= MODEL_PLUGIN_SHA= TT_VLLM_BUILTIN_MODELS= /bin/sh -c chmod -R a+rwX /home/tt # buildkit","comment":"buildkit.dockerfile.v0"},{"created":"2026-10-09T23:06:45.672498279-04:00","created_by":"USER tt","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:06:45.672498279-04:00","created_by":"LABEL org.tenstorrent.tt-model=diffrhythm-1.2-p100a org.tenstorrent.tt-model.repo=sunandbean/diffrhythm-1.2-p100a org.tenstorrent.tt-model.weights=ASLP-lab/DiffRhythm-1_2-full org.tenstorrent.tt-model.arch=blackhole org.tenstorrent.tt-model.kind=tt-dit-server org.tenstorrent.tt-model.profiles=default org.opencontainers.image.revision=58648d44443223041148125843ca7c7d4780463e org.tenstorrent.tt-model.tt-metal= org.tenstorrent.tt-model.plugin=","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:06:45.672498279-04:00","created_by":"ENTRYPOINT [\"/usr/local/bin/entrypoint.sh\"]","comment":"buildkit.dockerfile.v0","empty_layer":true},{"created":"2026-10-09T23:06:45.672498279-04:00","created_by":"CMD [\"/usr/local/bin/serve-default.sh\"]","comment":"buildkit.dockerfile.v0","empty_layer":true}],"os":"linux","rootfs":{"type":"layers","diff_ids":["sha256:cbaf391e670933f026c05c7dedec3faf8f194b6c2d83b0d5694ad68323c8077c","sha256:09daac6f134b2dba27b8d4a0ef3121a37537cad3f15daa032b9431f80ff87448","sha256:5287d0203a2b8dcd23409b5e9013a5c575f805d95599f294f57401c395c68522","sha256:542109b0dfa66e7d5c45e42504d6acaf67a10027553a1855ed50d97cf40b5fd4","sha256:85747f2ebe464ef004163a156c985fd7654ef8463f0e1f673d5fc3ebb8ca5036","sha256:0a843c5fa2e20d2fb2d102bd7c52914df44527da4a07f6a90d1b7cbdc99d4139","sha256:3c8124f890f49acc6c287be3d6967d2c7aea4b0613442eecb2d12649e346dd19","sha256:84150c24aea3517574d06f1ccb0fd8cbb705ef7fe1360b737bd73f41fe5fe802","sha256:7c7537d0f65122e0a4fb6e585e744529be140d9f30775d4ac3e7d71678841766","sha256:c55f647fbe1619ed1658092a93e0de5543b510df376925b0077fbf87c2b6a606","sha256:0f111d51bc656a83c00484d1ecaf7c053336f5d87914e3e169fdfb1b8730c0b1","sha256:b2b6c01d259d2ba562d3ee67e8481bb3c12c484fa4817411822773177d83122e","sha256:558e4e46be0c2d5366bf9e13cb1076dffcfa9b386b7a702a2be672df0abbfe26","sha256:2aa75a7513dd7fd2e2df7fb28da3f3de2bddfe115eb51b562b3a9cf7faea0719","sha256:ac5907e681ddb595c56a03bfe14231bfdd907438015951cbeaa6313e4b5072af","sha256:5fc6fd4ae4f4966d156d66670a5f1c1ec3087a656ecb06b8af21f27be8e94870","sha256:89c2380aa2116794479b618d5e7a373ba519c24b3a5e67756a1c09e6ca089fa5","sha256:568be329db8740036ac4055096cd89a8b909b0375fefe0f2885a287c7ffb9855","sha256:5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef","sha256:7b7966ac1b44afb1004b16928394cc98583c8854c533405e3e68d01de8600c5b","sha256:228b2f895919d9b9f799eb573f98d3e01f7e6e45b01a406366b69b1aa333fa91","sha256:369854526b98a8f50c544e3aca82a2cc696396fbe0f96ca020dca91a87502fd5"]}}

image/index.json

@@ -0,0 +1 @@
+{"schemaVersion":2,"mediaType":"application/vnd.oci.image.index.v1+json","manifests":[{"mediaType":"application/vnd.oci.image.index.v1+json","digest":"sha256:0cc125e439ef6882a01bae57bece67f14ec2ad123aec1fb12e86bc50175a4b44","size":856,"annotations":{"io.containerd.image.name":"docker.io/tt-model/diffrhythm-1.2-p100a:0cc125e439ef","org.opencontainers.image.ref.name":"0cc125e439ef"}}]}

image/manifest.json

@@ -0,0 +1 @@
+[{"Config":"blobs/sha256/eae56eb3c9610901618ab9c30e2ba8cdcf6981cdd3af7ed5755aba2cc46db6e1","RepoTags":["tt-model/diffrhythm-1.2-p100a:0cc125e439ef"],"Layers":["blobs/sha256/98c4455a98982b35380ec62f90d5cf0edff02eeb4b3bb18595d871fdeb0c0c67","blobs/sha256/3375e6602aa7264d6d67ae69ccb99690c8f40af7ab48d0af060dac1a2e9c24c9","blobs/sha256/340cf5e2c047e0fc940f2973810428f58250242238231682812dec6f6d54f3ac","blobs/sha256/1f47921d53a3b8de72172e2fd5f0a0d27f1e8f4b77016c6b8a1a4b591bb83ae5","blobs/sha256/4ec2b03c373212fbe7fda97dbf9eaadb1bd42c2b0afc8b8a49806fe1dbc507f7","blobs/sha256/b7c2e1ec67bfd9d4578a81ca6987f095b74869db5f518bf2fc9d3ffeb30821e5","blobs/sha256/428f99cfe94d9b1e42af2f3d8104598f8ca6e6105925c7687595a18c78d21011","blobs/sha256/ee5adbd6eba4dfd9ba7e59692c46379b3d60a88f8445ec6a045d4aeb336f5bb9","blobs/sha256/672d584c094e51e29aa5e46a6dd22cbfb993f4f839609e954d2c299912ea5d77","blobs/sha256/2d16f4104fc6e34c556d42d6a35adedd5ffa396e71d59ef6dd7d5e88c01225e0","blobs/sha256/556ec3c4c0e12c97f4c9b8dd0cd65ae321238c55438fbf498055f0ad7afc05bb","blobs/sha256/b9ddf00cad9d9a3bc91d09c8ee0f4fef42056861bb5ca1209fcc204b68950abf","blobs/sha256/0b4ec6da503185d389dcfeefc2c5d03ab13e0d76620b5fb6ce022f79bb64140f","blobs/sha256/49353ead98e6a241dd015dda937bf8dcf335e4eed27fe732df8163dd955d042f","blobs/sha256/c2a425aad5e9144add8f2a619c6e3d452d30ddbb5d4f4d3e45caab4f7bb3d233","blobs/sha256/e211ac567cceb6b5be62db5e135a1a5f3e77bcacebdac6a8cf9e0cdc24158c09","blobs/sha256/c6bd350dd2ed57d0ea34248e5b27d985a6786330d2d3ec29f5563e03ca9f00da","blobs/sha256/8f09ce5e10995b3fa6362dae9945d9402a5ccabce3d5bf20df0a53c75d9f94ab","blobs/sha256/4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1","blobs/sha256/e600ba821186d2d540ea813a478dda1e629c539eef283f062f31309cb8c3af47","blobs/sha256/3c7495db974af283bccc2dee306a75da6491b2ca664d55597da926c4fd104dd5","blobs/sha256/6acee8cd93aa0889cc3a8b09f1b2564234ee82aa7c31e87ab61591552d15e894"]}]

image/oci-layout

@@ -0,0 +1 @@
+{"imageLayoutVersion":"1.0.0"}

requirements.lock

@@ -0,0 +1,148 @@
+accelerate==1.15.0
+annotated-doc==0.0.5
+annotated-types==0.8.0
+anyio==4.15.1
+attrs==26.1.0
+beartype==0.22.9
+certifi==2026.7.22
+cffi==2.1.1
+cfgv==3.5.0
+charset-normalizer==3.5.2
+click==8.5.0
+cloudpickle==3.1.2
+cn2an==0.5.24
+contourpy==1.3.3
+cycler==0.12.1
+decorator==5.3.1
+Deprecated==3.0.0
+distlib==0.4.3
+distro==1.9.0
+dlinfo==2.0.0
+easydict==1.13
+einops==0.8.2
+einx==0.4.3
+elastic-transport==9.4.2
+elasticsearch==9.5.1
+ema-pytorch==0.8.3
+espeakng-loader==0.2.4
+fastapi==0.143.0
+filelock==4.1.0
+flatbuffers==25.12.19
+fonttools==4.66.1
+frozendict==2.4.7
+fsspec==2026.9.0
+ftfy==6.3.1
+googleapis-common-protos==1.75.5
+graphviz==0.21
+h11==0.16.0
+hf-xet==1.7.0
+httpcore2==2.13.1
+httpx2==2.13.1
+huggingface_hub==2.2.0
+identify==2.6.20
+idna==3.20
+inflect==7.5.0
+jaconv==0.5.0
+jieba==0.42.1
+Jinja2==3.1.6
+joblib==1.6.0
+kiwisolver==1.5.1
+lazy-loader==0.6
+librosa==1.0.0
+linkify-it-py==2.2.0
+llvmlite==0.50.0
+loguru==0.7.3
+markdown-it-py==4.2.0
+MarkupSafe==3.0.4
+matplotlib==3.11.2
+mdit-py-plugins==0.6.1
+mdurl==0.1.2
+ml_dtypes==0.5.4
+more-itertools==11.1.0
+mpmath==1.3.0
+msgpack==1.2.3
+muq==0.1.0
+mutagen==1.48.1
+narwhals==2.26.0
+networkx==3.7
+nnaudio==0.3.4
+nodeenv==1.11.0
+numba==0.68.0
+numpy==2.5.3
+onnxruntime==1.31.0
+opentelemetry-api==1.45.1
+opentelemetry-exporter-http-transport==0.66b1
+opentelemetry-exporter-otlp-common==0.66b1
+opentelemetry-exporter-otlp-proto-common==1.45.1
+opentelemetry-exporter-otlp-proto-http==1.45.1
+opentelemetry-proto==1.45.1
+opentelemetry-sdk==1.45.1
+opentelemetry-semantic-conventions==0.66b1
+packaging==26.3
+pandas==3.0.6
+phonemizer==3.4.0
+pillow==12.3.0
+platformdirs==4.12.4
+pooch==1.9.0
+pre_commit==4.6.2
+proces==0.1.7
+protobuf==7.36.2
+psutil==7.2.2
+py3langid==0.2.2
+pycparser==3.11
+pydantic==2.14.0
+pydantic_core==2.50.0
+Pygments==2.21.0
+pykakasi==2.3.0
+pyluwen==0.10.0
+pyparsing==3.3.3
+pypinyin==0.55.0
+python-dateutil==2.9.0.post0
+python-discovery==1.6.2
+PyYAML==6.0.3
+regex==2026.9.29
+requests==2.34.2
+rich==15.0.0
+safetensors==0.8.0
+scikit-learn==1.9.1
+scipy==1.18.1
+seaborn==0.13.2
+setuptools==80.10.2
+setuptools-scm==8.1.0
+shellingham==1.5.4
+six==1.17.0
+sniffio==1.3.1
+soundfile==0.14.0
+soxr==1.1.0
+starlette==1.7.0
+sympy==1.14.0
+textual==8.2.8
+threadpoolctl==3.7.0
+tokenizers==0.23.3
+tomli==2.5.0
+torch==2.14.1+cpu
+torch-einops-utils==0.1.36
+torchaudio==2.11.0+cpu
+torchdiffeq==0.2.5
+torchvision==0.29.1+cpu
+tqdm==4.70.1
+transformers==5.19.0
+truststore==0.10.4
+tt-smi==6.7.0
+tt-tools-common==1.6.0
+tt-umd==0.9.12
+typeguard==4.6.0
+typer==0.27.3
+typing-inspection==0.4.4
+typing_extensions==4.16.0
+Unidecode==1.4.0
+urllib3==2.8.0
+uvicorn==0.54.0
+virtualenv==21.14.6
+wandb==0.30.0
+wcwidth==0.9.2
+wheel==0.48.0
+wrapt==2.5.0
+x-clip==0.14.4
+x-transformers==2.31.7
+xxhash==4.0.1

tt_kernel_manifest.json

@@ -0,0 +1,130 @@
+{
+  "schema_version": "5.1",
+  "name": "diffrhythm-1.2-p100a",
+  "tt_metal_version": "0.0.0.dev0",
+  "arch": "blackhole",
+  "device_count": 1,
+  "producer": {
+    "tt_kernel_version": "0.1.0",
+    "created_at": "2026-10-10T03:07:37.410607+00:00",
+    "hostname": ""
+  },
+  "weights": {
+    "repo_id": "ASLP-lab/DiffRhythm-1_2-full",
+    "revision": null,
+    "allow_patterns": null,
+    "ignore_patterns": null,
+    "repo_type": "model"
+  },
+  "mesh": null,
+  "entrypoint": null,
+  "resources": null,
+  "capabilities": null,
+  "env": {},
+  "bundled": null,
+  "deps": null,
+  "container": {
+    "image": {
+      "registry": "hf",
+      "repository": "diffrhythm-1.2-p100a",
+      "tag": "tt-model/diffrhythm-1.2-p100a:0cc125e439ef",
+      "digest": "sha256:0cc125e439ef6882a01bae57bece67f14ec2ad123aec1fb12e86bc50175a4b44"
+    },
+    "kind": "tt-dit-server",
+    "runtime": {
+      "app": "tt_diffrhythm.server:app",
+      "mesh_shape_env": "DIFFRHYTHM_MESH_SHAPE",
+      "packages": [
+        "torchaudio==2.11.0+cpu",
+        "torchdiffeq",
+        "x-transformers",
+        "transformers",
+        "librosa",
+        "soundfile",
+        "muq",
+        "mutagen",
+        "ema-pytorch",
+        "accelerate",
+        "einops",
+        "tqdm",
+        "wandb",
+        "scipy",
+        "numpy",
+        "huggingface_hub",
+        "phonemizer",
+        "espeakng-loader",
+        "pykakasi",
+        "jieba",
+        "onnxruntime",
+        "cn2an",
+        "pypinyin",
+        "Unidecode",
+        "inflect",
+        "py3langid==0.2.2"
+      ],
+      "lock": "requirements.lock"
+    },
+    "serve": {
+      "hardware": "p100a",
+      "mesh_device": "P100",
+      "port": 8000,
+      "max_model_len": null,
+      "max_num_seqs": null,
+      "block_size": null,
+      "server_timeout": null,
+      "capabilities": null,
+      "additional_config": {},
+      "args": [],
+      "env": {
+        "ARCH_NAME": "blackhole",
+        "TT_METAL_OPERATION_TIMEOUT_SECONDS": "600.0"
+      }
+    },
+    "serve_profiles": [
+      {
+        "hardware": null,
+        "mesh_device": null,
+        "port": null,
+        "max_model_len": null,
+        "max_num_seqs": null,
+        "block_size": null,
+        "server_timeout": null,
+        "capabilities": null,
+        "additional_config": {},
+        "args": [],
+        "env": {},
+        "name": "default",
+        "description": null
+      }
+    ],
+    "default_profile": null,
+    "code_dir": "code",
+    "verify": [
+      "import tt_diffrhythm.server as s; assert s.app",
+      "from tt_diffrhythm.local import LocalWorker; assert LocalWorker",
+      "import tt_diffrhythm.server as s; s._enter_checkout(); import infer_utils, model; assert model.CFM and infer_utils.load_checkpoint",
+      "from tt_diffrhythm.server import _enter_checkout; _enter_checkout(); import os; assert os.path.isfile('infer/example/vocal.npy'), 'the negative style prompt did not ship'",
+      "from tt_diffrhythm.server import _enter_checkout; _enter_checkout(); import json; assert json.load(open('config/diffrhythm-1b.json'))['model']",
+      "from tt_diffrhythm.server import _enter_checkout; _enter_checkout(); from infer_utils import CNENTokenizer; assert CNENTokenizer().encode('a quiet morning')",
+      "from tt_diffrhythm.server import _point_phonemizer_at_a_bundled_espeak as fix; fix(); from phonemizer.backend import EspeakBackend; assert EspeakBackend('en-us').phonemize(['hello'])[0].strip()"
+    ],
+    "built": {
+      "image": "tt-model/diffrhythm-1.2-p100a:0cc125e439ef",
+      "repo": "sunandbean/diffrhythm-1.2-p100a",
+      "tt_model_version": "0.1.0",
+      "created_at": "2026-10-10T03:02:08+00:00",
+      "tt_metal": {
+        "sha": "58648d44443223041148125843ca7c7d4780463e",
+        "describe": null,
+        "dirty": false,
+        "scm_version": "0.0.0.dev0",
+        "mode": "git",
+        "remote": "https://github.com/tenstorrent/tt-metal",
+        "branch": null,
+        "pushed": null
+      },
+      "code_sha256": "322cdea4755e9a8117adc934aed2edfd1e3a7631fbbe9949bdbdc46c077cf184",
+      "image_digest": "sha256:0cc125e439ef6882a01bae57bece67f14ec2ad123aec1fb12e86bc50175a4b44"
+    }
+  }
+}