File size: 6,587 Bytes
cfbcb75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | # 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}")
|