| """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] |
| |
| |
| 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, |
| 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 |
|
|
| |
| 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] |
| |
| |
| 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": |
| 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: |
| 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() |
|
|
| |
| 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: |
| 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 |
| 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) |
|
|
