SoundGame Dev
f0f0a9b50e
- 共享引擎（app/src/main）：GestureEngine 6 手势识别、AudioEngine 三层音频 （ExoPlayer BGM 交叉淡入 / 对白 ducking / SoundPool 低延迟）、Game/GameEngine 状态机接口、MainActivity 全屏封面 + TalkBack - 《声命线》storyline：末世互动广播剧，3 幕 8 抉择点 4 结局 - 《声东击西》reflex：听声辨位反应，12 波战斗 + BGM 三档 BPM 变速 - 《暗室》escape：3x3 网格密室逃脱，钥匙→铁柜→保险丝→电箱→大门谜题链 - 文档：setup（从零搭建）/ architecture（架构）/ engine-api（引擎 API）/ asset-pipeline（资产流水线）/ games/*（三款游戏设计文档） - 资产管线：Qwen3-TTS 对白 + 程序合成音效/BGM + Mureka BGM 的 gen 脚本与台账， 全部 ogg 资产入库可复现 - env.d 环境脚本：android/tts 入库；sfx/mureka 含密钥不入库，仅提交 .example.sh 模板
1766 lines
97 KiB
Python
1766 lines
97 KiB
Python
"""SA3 text-to-audio via TensorRT — full-pipeline CUDA-graph variant.
|
||
|
||
Same as sa3_trt.py but captures the entire fast-path inference pipeline
|
||
(T5 encode + DiT 8-step loop + decoder + narrow + DtoH) in ONE CUDA graph.
|
||
A single g.replay() runs everything from input_ids to a pinned-host int16
|
||
PCM tensor, eliminating per-stage Python / CUDA submission overhead.
|
||
|
||
Mega-graph path requirements:
|
||
--cfg 1.0 (the default)
|
||
--init-noise-level 1.0 (the default)
|
||
no --inpaint-range
|
||
no --init-audio
|
||
Anything else delegates to sa3_trt_core.main()'s eager pipeline. That handoff happens
|
||
BEFORE this module loads any engine — canon.main() loads its own set, so loading here
|
||
first left two full sets resident and doubled peak VRAM (see _delegate_to_eager).
|
||
|
||
Usage: same CLI as sa3_trt.py. Optional --no-mega-graph forces the eager
|
||
path even for the fast case (useful for ablation).
|
||
|
||
Bit-exact parity with sa3_trt.py: the build() warmup deliberately replays
|
||
canonical's exact pre-capture call sequence (3 warmups with zero T5/DiT/
|
||
decoder inputs, then 3 more 8-step DiT loops to mimic canonical's
|
||
graph_sampler.build) so that the decoder context's internal state at the
|
||
captured-call boundary matches canonical's "real" call. With this match,
|
||
each fresh process produces a WAV byte-identical to sa3_trt.py's.
|
||
|
||
Caveat — multi-inference drift: BOTH this script and sa3_trt.py drift
|
||
after the first inference within a single process, because the SAME-S/L
|
||
decoder engine carries internal state across calls (its TRT workspace
|
||
contents evolve). This is a property of the engine, not the wrapper.
|
||
The canonical CLI invokes one inference per process, so users never see
|
||
the drift in practice.
|
||
"""
|
||
from __future__ import annotations
|
||
import argparse, math, os, random, sys, threading, time, wave
|
||
from pathlib import Path
|
||
import numpy as np
|
||
|
||
# Reuse everything from the canonical script — including its lazy torch/trt
|
||
# imports, engine path resolution, and all the helper classes.
|
||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||
import tensorrt as trt # windowed decode builds its own contexts
|
||
import sa3_trt_core as canon
|
||
from sa3_trt_core import (
|
||
# Module-level state we touch:
|
||
SAMPLE_RATE, SAMPLES_PER_LATENT, IO_CHANNELS, T5_MAX_LEN, COND_DIM,
|
||
DIT_CHOICES, DECODER_PATHS, ENCODER_PATHS,
|
||
DIT_ENGINE_FILES, DECODER_FILES, ENCODER_FILES, SHARED_FILES,
|
||
HF_REPO_ID, ARCH, DECODER_MIN_L,
|
||
# Helpers:
|
||
_import_heavy, _ensure_files, _init_nvml, _stage_vram, _my_vram_bytes,
|
||
_STAGE_VRAMS, _silence_fd, _fmt_mem,
|
||
stage, sub, banner, bold, dim, cyan, yellow, green, magenta, red, rule,
|
||
prompt_user_if_missing, _arrow_pick, _HelpfulParser,
|
||
TRTRunner, DiTRunner, GraphPingpongSampler,
|
||
t5gemma_encode, encoder_encode, decoder_decode,
|
||
build_pingpong_schedule, sample_flow_pingpong,
|
||
save_wav, read_wav,
|
||
# Checked enqueue for capture regions — shared so both graphs behave alike.
|
||
_enqueue_captured as _enqueue,
|
||
)
|
||
|
||
|
||
_WRAPPER_ONLY_FLAGS = ("--mega-graph", "--no-mega-graph")
|
||
|
||
|
||
def _delegate_to_eager(args) -> None:
|
||
"""Hand off to canon.main() for any configuration the mega-graph cannot serve.
|
||
|
||
MUST be called before this module loads a single engine. canon.main() loads its
|
||
own T5 + DiT + decoder, so loading them here first left two full sets resident —
|
||
this frame keeps the first set referenced for as long as canon.main() runs. That
|
||
doubled peak VRAM (measured 26.7-26.8 GB against 14.7 GB for medium + same-l) and
|
||
on a 32 GB card it OOM'd *every* eager configuration: --cfg != 1,
|
||
audio-to-audio, inpaint, --init-noise-level != 1, and --no-mega-graph. It is also
|
||
the "malloc_consolidate(): invalid chunk size" abort that had been recorded
|
||
against medium + same-l + cfg>1 — the trigger was never CFG itself, only that
|
||
asking for CFG routes you here.
|
||
|
||
canon.main() re-parses sys.argv from scratch, so anything this wrapper already
|
||
resolved has to be pinned into argv or it gets re-derived differently. --seed is
|
||
the one that silently corrupts results: prompt_user_if_missing() picks a RANDOM
|
||
seed when the user omits it and the banner has already printed that value, so
|
||
letting canon draw its own meant the render did not match the seed we reported
|
||
and could not be reproduced from it.
|
||
"""
|
||
print()
|
||
print(dim(" Mega-graph path unavailable for this configuration; "
|
||
"delegating to sa3_trt_core.main()..."))
|
||
argv = [a for a in sys.argv if a not in _WRAPPER_ONLY_FLAGS]
|
||
for flag, val in (("--dit", args.dit), ("--decoder", args.decoder),
|
||
("--precision", args.precision), ("--seed", args.seed)):
|
||
if val is not None and flag not in argv:
|
||
argv += [flag, str(val)]
|
||
sys.argv = argv
|
||
canon.main()
|
||
|
||
|
||
# ─── Full-pipeline CUDA-graph runner ─────────────────────────────────────
|
||
class _SecondContext:
|
||
"""A second TRT execution context on an ALREADY-LOADED engine.
|
||
|
||
The eager path (CFG / audio-to-audio / inpaint / σmax≠1) has to call
|
||
set_input_shape + set_tensor_address on every step, at whatever L the request
|
||
asks for. Doing that on the same context the mega-graph captured is unsafe:
|
||
DiTRunner._setup() reallocates its velocity buffer when L changes, and the
|
||
captured graph still launches kernels pointing at the OLD allocation — which
|
||
torch is then free to hand to someone else. The failure is silent (garbage or
|
||
stale audio on the next fast-path replay), so the eager path gets its own
|
||
context instead.
|
||
|
||
Costs activation memory only — the engine's weights are shared, not re-loaded.
|
||
"""
|
||
|
||
def __init__(self, runner: TRTRunner, scratch=None):
|
||
"""scratch: optional shared device buffer (see SA3Inference._share_scratch).
|
||
|
||
A separate CONTEXT is required (the reason is above), but the separate SCRATCH is not: the
|
||
eager path and the captured graph never run concurrently -- the inference lock serialises
|
||
them -- and TRT device memory is pure scratch, re-derived on every execution rather than
|
||
carried between calls. Left STATIC, each eager context re-added the engine's full private
|
||
worst case (~980 MiB on `large`), giving back most of what sharing just saved.
|
||
"""
|
||
torch = canon.torch
|
||
self.engine = runner.engine
|
||
if scratch is not None:
|
||
import tensorrt as _trt
|
||
self.context = runner.engine.create_execution_context(
|
||
_trt.ExecutionContextAllocationStrategy.USER_MANAGED)
|
||
self.context.set_device_memory(scratch.data_ptr(), scratch.numel())
|
||
else:
|
||
self.context = runner.engine.create_execution_context()
|
||
self.stream = torch.cuda.Stream()
|
||
self.in_dtype = dict(runner.in_dtype)
|
||
self.out_dtype = dict(runner.out_dtype)
|
||
|
||
def free(self):
|
||
del self.context
|
||
|
||
|
||
class BatchedDiT:
|
||
"""Batch-N DiT forward on a var-batch engine (one enqueue, not N).
|
||
|
||
The canonical engine bakes batch=1, so CFG has to run cond and uncond as two sequential
|
||
calls. The var-batch engine (traced with axis 0 dynamic, built `--batches 1,2`) carries
|
||
one optimization profile per batch size, so the pair goes through as a single batch-2
|
||
call with tactics tuned for that shape.
|
||
|
||
`attach()` returns None when the engine has no profile covering this batch, so callers
|
||
fall back to the sequential dual pass without inspecting TRT themselves.
|
||
"""
|
||
|
||
def __init__(self, engine, batch, profile_index):
|
||
torch = canon.torch
|
||
self.engine = engine
|
||
self.batch = batch
|
||
self.profile_index = profile_index
|
||
self.context = engine.create_execution_context()
|
||
self.stream = torch.cuda.Stream()
|
||
self.context.set_optimization_profile_async(profile_index, self.stream.cuda_stream)
|
||
self.stream.synchronize()
|
||
self._L = None
|
||
self._vel = None
|
||
|
||
@classmethod
|
||
def attach(cls, runner, batch=2):
|
||
eng = runner.engine
|
||
for p in range(eng.num_optimization_profiles):
|
||
try:
|
||
mn, _opt, mx = eng.get_tensor_profile_shape("x", p)
|
||
except Exception:
|
||
continue
|
||
if mn[0] <= batch <= mx[0]:
|
||
return cls(eng, batch, p)
|
||
return None
|
||
|
||
def _setup(self, L):
|
||
if L == self._L:
|
||
return
|
||
torch = canon.torch
|
||
b, ctx = self.batch, self.context
|
||
ctx.set_input_shape("x", (b, IO_CHANNELS, L))
|
||
ctx.set_input_shape("t", (b,))
|
||
ctx.set_input_shape("t5_hidden", (b, T5_MAX_LEN, COND_DIM))
|
||
ctx.set_input_shape("t5_mask", (b, T5_MAX_LEN))
|
||
ctx.set_input_shape("seconds_total", (b,))
|
||
ctx.set_input_shape("local_add_cond", (b, 257, L))
|
||
self._vel = torch.empty(tuple(ctx.get_tensor_shape("velocity")),
|
||
dtype=torch.float32, device="cuda")
|
||
self._L = L
|
||
|
||
def step(self, x, t, t5_hidden, t5_mask, seconds, local_add_cond):
|
||
"""All inputs batch-N on cuda. Returns velocity (N, 256, L)."""
|
||
torch = canon.torch
|
||
self._setup(x.shape[-1])
|
||
ctx = self.context
|
||
held = [] # keep the contiguous copies alive
|
||
for name, ten in (("x", x), ("t", t), ("t5_hidden", t5_hidden),
|
||
("t5_mask", t5_mask), ("seconds_total", seconds),
|
||
("local_add_cond", local_add_cond)):
|
||
c = ten.float().contiguous()
|
||
held.append(c)
|
||
ctx.set_tensor_address(name, c.data_ptr())
|
||
ctx.set_tensor_address("velocity", self._vel.data_ptr())
|
||
self.stream.wait_stream(torch.cuda.current_stream())
|
||
if not ctx.execute_async_v3(self.stream.cuda_stream):
|
||
raise RuntimeError(f"batched DiT enqueue failed (batch={self.batch}, "
|
||
f"L={x.shape[-1]})")
|
||
self.stream.synchronize()
|
||
return self._vel.clone()
|
||
|
||
|
||
class FullPipelineGraph:
|
||
"""Captures T5 + DiT-loop + decoder + narrow + DtoH in ONE CUDA graph.
|
||
|
||
Layout (all on capture_stream):
|
||
1. T5 engine reads input_ids_buf / attn_mask_buf, writes hidden_buf
|
||
2. t5_mask_buf <- attn_mask_buf.float()
|
||
3. dit._x_buf <- initial_noise_buf
|
||
4. 8-step DiT pingpong loop (same body as GraphPingpongSampler.build):
|
||
for i in range(steps):
|
||
dit._t_buf[0] = sigma_curr_bufs[i]
|
||
dit_engine.execute_async_v3
|
||
... pingpong math reads dit._vel_buf / dit._x_buf / noise_bufs[i] ...
|
||
5. decoder_in_buf <- final_latents_buf
|
||
6. decoder engine reads decoder_in_buf, writes pcm_int32_buf
|
||
(1, T_lat*4096, 2) int32
|
||
7. pcm_int16_buf <- pcm_int32_buf.to(int16) (narrow cast)
|
||
8. pinned_host_pcm.copy_(pcm_int16_buf[:requested_samples], non_blocking)
|
||
|
||
Per-inference (replay path):
|
||
1. Copy host int64 input_ids/attn_mask into pinned staging buffers
|
||
(Python tokenizer output) — those staging buffers were registered
|
||
as persistent device-readable srcs of HtoD copies into
|
||
input_ids_buf / attn_mask_buf inside the captured graph.
|
||
2. Refresh the 7 in-loop noise buffers + initial_noise_buf with
|
||
fresh seeded randn, on the capture stream.
|
||
3. Write seconds_total scalar into dit._sec_buf via copy_().
|
||
4. self._graph.replay()
|
||
5. sync — host pcm is ready in pinned_host_pcm[:requested_samples].
|
||
"""
|
||
|
||
def __init__(self, t5_runner: TRTRunner, dit: DiTRunner, dec_runner: TRTRunner,
|
||
L: int, steps: int, requested_samples: int, dec_chunk_lat: int = 0,
|
||
scratch=None, balance: bool = True):
|
||
# dec_chunk_lat > 0 captures the decode as a sequence of windows instead of one call,
|
||
# which is what lets the decoder engine be built with a low profile ceiling: scratch is
|
||
# linear in the ceiling and committed at create_execution_context(), so a 4096-ceiling
|
||
# engine reserves ~5.5 GB whether it decodes 32 latents or 4096. Chunking inside the graph
|
||
# keeps every enqueue within `dec_chunk_lat` while the render length stays free.
|
||
# Addresses are re-pointed per window between enqueues; set_tensor_address is host-side, so
|
||
# each recorded enqueue bakes the offsets that were current when it was captured.
|
||
self.dec_chunk_lat = int(dec_chunk_lat or 0)
|
||
self.dec_balance = bool(balance)
|
||
self.dec_win_lat = self.dec_chunk_lat # actual window width (<= chunk once balanced)
|
||
self.t5_runner = t5_runner
|
||
self.dit = dit
|
||
self.dec_runner = dec_runner
|
||
self.L = L
|
||
self.steps = steps
|
||
self.requested_samples = requested_samples
|
||
# Will be allocated in build().
|
||
self.input_ids_buf = None # (1, 256) int64, device
|
||
self.attn_mask_buf = None # (1, 256) int64, device
|
||
self.t5_hidden_internal = None # (1, 256, 768) fp32, device — T5's output
|
||
self.initial_noise_buf = None # (1, 256, L) fp32, device
|
||
self._noise_bufs = None # list of (1, 256, L) fp32, device
|
||
self._sigma_curr_bufs = None
|
||
self._sigma_next_bufs = None
|
||
self.latents_out_buf = None # (1, 256, L) fp32, device — final latent
|
||
self.decoder_in_buf = None # (1, 256, L) fp32, device — what dec reads
|
||
self.pcm_int32_buf = None # (1, T_lat*4096, 2) int32, device
|
||
self.pcm_int16_buf = None # (T_lat*4096, 2) int16, device
|
||
self.pinned_host_pcm = None # (T_lat*4096, 2) int16, pinned host
|
||
self.local_add_cond_buf = None # (1, 257, L) fp32, device (kept zero)
|
||
self._dec_plan = None # [(start, size)] windows, when chunking
|
||
self._shared_scratch = scratch # SA3Inference._scratch, shared by every engine+graph
|
||
self._dec_ctx_main = None # THIS graph's decoder context (never the runner's -- see
|
||
# build(); destroying a context a graph was captured from
|
||
# segfaults that graph's replay)
|
||
self._dec_ctx_tail = None # 2nd context for a short tail window (odd L)
|
||
self._dec_win_i32 = None # one window's int32 PCM, reused by every window
|
||
self._dec_win_lat = None # one window's latents, CONTIGUOUS (see _decode_windows)
|
||
self._ceiling_buf = None # persistent scalar for a runtime limiter ceiling
|
||
self._dec_scratch = None # shared scratch buffer for the decoder contexts
|
||
self._graph = None
|
||
self._built = False
|
||
|
||
def set_limiter_ceiling(self, value: float):
|
||
"""Change the limiter ceiling between graph replays -- a 4-byte write, no recapture.
|
||
|
||
canon.LIMITER_BYPASS disables limiting; the graph's +-1.0 clamp still bounds the output.
|
||
"""
|
||
if self._ceiling_buf is None:
|
||
raise ValueError("this decoder engine bakes the ceiling in; rebuild with "
|
||
"build_ship_capped.py --ceiling-input")
|
||
self._ceiling_buf.fill_(float(value))
|
||
|
||
@staticmethod
|
||
def _plan_windows(N: int, chunk: int, trim: int = 16, balance: bool = True):
|
||
"""Even-aligned windows covering N. `balance` shrinks them to the smallest uniform size
|
||
that still covers N, instead of always using the full `chunk`.
|
||
|
||
Why balance matters: the stitch requires consecutive starts to be exactly `w - trim`
|
||
apart, so a fixed-width plan covers only `chunk - trim` = 240 new latents per window. Just
|
||
above the chunk size that is brutally wasteful -- L=257 pads to 258 and needs TWO full
|
||
256-windows, decoding 512 latents to produce 258. Measured 25.9 ms against 12.3 ms at
|
||
L=256: one extra latent for 2.1x the time. The waste decays as 256*n/N (1.99x at 257,
|
||
1.50x at 512, 1.25x at 1024, 1.06x at 4096), which is exactly the shape of the measured
|
||
latency curve -- predicted-from-waste tracks measured within 3-7% at every length.
|
||
|
||
Balancing picks the smallest n with an even w <= chunk satisfying n*w - (n-1)*trim >= N,
|
||
which pins the waste at a flat ~1.07x everywhere. Measured recovery, eager, vs the
|
||
fixed-width plan: L=257 1.68x, 512 1.30x, 1024 1.14x, 1291 1.07x, and nothing at 2048+
|
||
where the fixed plan was already optimal. Accuracy is UNCHANGED (+-0.1 dB vs the shipped
|
||
full-length decode at every length) -- the extra SWA edges cost nothing measurable.
|
||
|
||
⚠ `balance` is a property of the CALLER, not of the decoder. A balanced window size varies
|
||
with N, and on the eager path a varying decoder shape costs ~34-40 ms per switch (which
|
||
made balanced decoding 22% WORSE on a mixed-length stream -- hence balance=False in
|
||
canon.decode_chunked). Inside a captured graph the shape is resolved once at capture and
|
||
replay never rebinds, so the switch is free and only the saved work remains.
|
||
"""
|
||
w = chunk
|
||
if balance:
|
||
for n in range(1, N // 2 + 2):
|
||
cand = -(-(N + (n - 1) * trim) // n)
|
||
cand += cand & 1 # keep every start even (SAME-S parity)
|
||
if cand <= chunk:
|
||
w = cand
|
||
break
|
||
step = (w - trim) // 2 * 2
|
||
starts = list(range(0, N - w, step))
|
||
last = -(-(N - w) // 2) * 2
|
||
if not starts or starts[-1] != last:
|
||
starts.append(last)
|
||
return [(st, (N - st) if i == len(starts) - 1 else w)
|
||
for i, st in enumerate(starts)]
|
||
|
||
def _decode_windows(self, capture_stream):
|
||
"""Enqueue the windowed decode, narrowing each window into pcm_int16_buf.
|
||
|
||
Interiors are trimmed by half the overlap on each neighbour-facing edge, matching
|
||
canon.decode_chunked and the repo's encode_chunked. Addresses are re-pointed per window;
|
||
each enqueue records the offsets current at capture time.
|
||
"""
|
||
SPL = SAMPLES_PER_LATENT
|
||
trim = 16
|
||
half = trim // 2
|
||
n = len(self._dec_plan)
|
||
for i, (st, sz) in enumerate(self._dec_plan):
|
||
ctx = (self._dec_ctx_tail if (self._dec_ctx_tail is not None
|
||
and sz != self.dec_win_lat)
|
||
else (self._dec_ctx_main or self.dec_runner.context))
|
||
# Gather the window into a CONTIGUOUS buffer. Offsetting the latent pointer instead
|
||
# does not work: decoder_in_buf is (1, 256, L) with the 256 channels strided by L, so
|
||
# a pointer bumped by `st` latents only shifts within the first channel's row and the
|
||
# engine reads 256*chunk contiguous elements of the wrong data (measured: cos 0.0027
|
||
# against the unwindowed decode, 80% of samples wrong). The copy is a kernel, so it
|
||
# captures fine, and the engine's bound address stays fixed across all windows.
|
||
self._dec_win_lat[..., :sz].copy_(self.decoder_in_buf[..., st:st + sz])
|
||
ctx.set_tensor_address("latent", self._dec_win_lat.data_ptr())
|
||
ctx.set_tensor_address("pcm", self._dec_win_i32.data_ptr())
|
||
if self._ceiling_buf is not None:
|
||
# Engines with a runtime limiter ceiling: bind it here, because this loop enqueues
|
||
# the context directly and never goes through decoder_decode (where the default
|
||
# lives). Unbound, TRT refuses the enqueue -- and under capture that would drop the
|
||
# decode stage silently.
|
||
ctx.set_input_shape("limiter_ceiling", ())
|
||
ctx.set_tensor_address("limiter_ceiling", self._ceiling_buf.data_ptr())
|
||
_enqueue(ctx, capture_stream, f"decoder window {i} (mega-graph)")
|
||
kl = 0 if i == 0 else half
|
||
kr = sz if i == n - 1 else sz - (trim - half)
|
||
self.pcm_int16_buf[(st + kl) * SPL:(st + kr) * SPL] = \
|
||
self._dec_win_i32[0, kl * SPL:kr * SPL, :].clamp(-32767, 32767).to(
|
||
canon.torch.int16)
|
||
|
||
def build(self, sigmas, seconds: float, sigma_max: float):
|
||
"""Allocate persistent buffers, prime + warm all engines, capture the graph.
|
||
|
||
sigmas: 1-D tensor of length steps+1 (pre-shifted schedule).
|
||
seconds: float — duration condition for DiT.
|
||
sigma_max: float — used to initialize seeding semantics (informational).
|
||
"""
|
||
torch = canon.torch
|
||
L = self.L
|
||
ns = self.steps
|
||
T_full = L * SAMPLES_PER_LATENT
|
||
# ── Padded DECODE length ──
|
||
# An odd L cannot be tiled by even-start windows of `dec_chunk_lat` (every start must be
|
||
# even or that window decodes at cos 0.978), so the tail comes out one latent short and
|
||
# needs a SECOND bound shape -- a second context, and an in-loop shape transition that
|
||
# costs ~54 ms on SAME-L. Replicating the final latent makes the length even, at which
|
||
# point every window is exactly dec_chunk_lat and the tail disappears: measured over
|
||
# L=257..4096, even lengths need no shape but 256, odd lengths need 255 as well.
|
||
# The extra latent's audio is decoded and then never returned (run() slices to
|
||
# requested_samples), exactly as canon.decode_chunked already does on the eager path.
|
||
self._L_dec = L + 1 if (self.dec_chunk_lat and L > self.dec_chunk_lat and L % 2) else L
|
||
T_dec = self._L_dec * SAMPLES_PER_LATENT
|
||
|
||
# We force all three engines to execute on the same stream and capture
|
||
# the whole pipeline on that one stream. Easiest: take the DiT runner's
|
||
# stream as the canonical "capture stream" and just hand its
|
||
# cuda_stream to every execute_async_v3 call.
|
||
capture_stream = self.dit.runner.stream
|
||
|
||
# ── T5 engine: set static shapes + bind addresses ──
|
||
t5_ctx = self.t5_runner.context
|
||
t5_ctx.set_input_shape("input_ids", (1, T5_MAX_LEN))
|
||
t5_ctx.set_input_shape("attention_mask", (1, T5_MAX_LEN))
|
||
# Allocate I/O buffers.
|
||
ids_dt = self.t5_runner.in_dtype["input_ids"] # int64
|
||
mask_dt = self.t5_runner.in_dtype["attention_mask"] # int64
|
||
self.input_ids_buf = canon.torch.zeros(1, T5_MAX_LEN, dtype=ids_dt, device="cuda")
|
||
self.attn_mask_buf = canon.torch.zeros(1, T5_MAX_LEN, dtype=mask_dt, device="cuda")
|
||
hid_dt = self.t5_runner.out_dtype["hidden_states"] # fp32
|
||
self.t5_hidden_internal = canon.torch.zeros(1, T5_MAX_LEN, COND_DIM, dtype=hid_dt, device="cuda")
|
||
t5_ctx.set_tensor_address("input_ids", self.input_ids_buf.data_ptr())
|
||
t5_ctx.set_tensor_address("attention_mask", self.attn_mask_buf.data_ptr())
|
||
t5_ctx.set_tensor_address("hidden_states", self.t5_hidden_internal.data_ptr())
|
||
|
||
# ── DiT: bind persistent buffers (this also configures shapes/output) ──
|
||
self.dit.bind_persistent(L)
|
||
# The DiT consumes (1, 257, L) local_add_cond. For the text-to-audio
|
||
# fast path it's all zeros — pre-fill once and never touch again.
|
||
self.local_add_cond_buf = self.dit._local_add_cond_buf
|
||
self.local_add_cond_buf.zero_()
|
||
# Seconds is per-inference (small) but constant during a graph replay.
|
||
# We write to dit._sec_buf BEFORE replay (outside the graph) so the
|
||
# captured kernels read the new value.
|
||
self.dit._sec_buf[0] = float(seconds)
|
||
|
||
# ── Initial noise buffer (the graph copies it into dit._x_buf) ──
|
||
self.initial_noise_buf = canon.torch.empty(1, IO_CHANNELS, L,
|
||
dtype=canon.torch.float32, device="cuda")
|
||
|
||
# ── Per-step sigma buffers (graph reads from these by pointer) ──
|
||
self._sigma_curr_bufs = [
|
||
canon.torch.tensor(float(sigmas[i]), dtype=canon.torch.float32, device="cuda")
|
||
for i in range(ns)
|
||
]
|
||
self._sigma_next_bufs = [
|
||
canon.torch.tensor(float(sigmas[i + 1]), dtype=canon.torch.float32, device="cuda")
|
||
for i in range(ns)
|
||
]
|
||
self._noise_bufs = [
|
||
canon.torch.empty(1, IO_CHANNELS, L, dtype=canon.torch.float32, device="cuda")
|
||
for _ in range(ns - 1)
|
||
]
|
||
|
||
# ── Final-latent buffer (graph writes here after last step) ──
|
||
self.latents_out_buf = canon.torch.empty(1, IO_CHANNELS, L,
|
||
dtype=canon.torch.float32, device="cuda")
|
||
|
||
# ── Decoder: bind addresses ──
|
||
dec_ctx = self.dec_runner.context
|
||
# A runtime limiter ceiling must be bound BEFORE capture: the window loop enqueues the
|
||
# context directly rather than through decoder_decode (where the default lives), and an
|
||
# unbound input makes TRT refuse the enqueue -- which under capture would drop the decode
|
||
# stage silently. Persistent buffer => set_limiter_ceiling() works between replays with no
|
||
# recapture, the same way seconds_total does.
|
||
if "limiter_ceiling" in self.dec_runner.in_dtype:
|
||
self._ceiling_buf = canon.torch.tensor(
|
||
canon.DEFAULT_LIMITER_CEILING, dtype=canon.torch.float32, device="cuda")
|
||
if self.dec_chunk_lat and L > self.dec_chunk_lat:
|
||
self._dec_plan = self._plan_windows(self._L_dec, self.dec_chunk_lat,
|
||
balance=self.dec_balance)
|
||
# Every window in a balanced plan is the same width, but that width is <= chunk, so
|
||
# the buffers and bound shape follow the PLAN rather than the engine ceiling.
|
||
self.dec_win_lat = max(sz for _, sz in self._dec_plan)
|
||
sizes = sorted({sz for _, sz in self._dec_plan})
|
||
# SAME-S tiles latents in PAIRS (17 tokens/latent, 34-token chunks, midpoint-shifted by
|
||
# one latent), so every window START must be even or that window decodes at cos 0.978.
|
||
assert all(st % 2 == 0 for st, _ in self._dec_plan), "odd window start"
|
||
# An odd L cannot be covered by uniform even-start windows, so the tail is shorter and
|
||
# needs its own shape. A second context is cheap because both share one scratch buffer
|
||
# (validated bit-identical, and it survives graph capture).
|
||
assert len(sizes) <= 2, f"expected at most 2 window shapes, got {sizes}"
|
||
dec_ctx.set_input_shape("latent", (1, IO_CHANNELS, self.dec_win_lat))
|
||
else:
|
||
self._dec_plan = None
|
||
dec_ctx.set_input_shape("latent", (1, IO_CHANNELS, L))
|
||
dec_in_dt = self.dec_runner.in_dtype["latent"]
|
||
self.decoder_in_buf = canon.torch.empty(1, IO_CHANNELS, self._L_dec,
|
||
dtype=dec_in_dt, device="cuda")
|
||
# Auto-detect output flavor like decoder_decode does.
|
||
if "pcm" in self.dec_runner.out_dtype:
|
||
self._dec_out_name = "pcm"
|
||
pcm_dt = self.dec_runner.out_dtype["pcm"]
|
||
if self._dec_plan is not None:
|
||
# One window-sized int32 buffer, reused by every window, instead of a
|
||
# full-length one: at L=16384 that is 8 MB rather than 537 MB. The int16
|
||
# destination already exists (pcm_int16_buf), and the graph clamps to
|
||
# +-32767 before its int32 cast, so narrowing per window is bit-exact.
|
||
self._dec_win_i32 = canon.torch.empty(
|
||
1, self.dec_win_lat * SAMPLES_PER_LATENT, 2, dtype=pcm_dt, device="cuda")
|
||
self._dec_win_lat = canon.torch.zeros(
|
||
1, IO_CHANNELS, self.dec_win_lat,
|
||
dtype=self.dec_runner.in_dtype["latent"], device="cuda")
|
||
self.pcm_int32_buf = self._dec_win_i32 # stage-5 code reads this name
|
||
sizes = sorted({sz for _, sz in self._dec_plan})
|
||
tail = [sz for sz in sizes if sz != self.dec_win_lat]
|
||
if tail:
|
||
# An odd L cannot be tiled by uniform even-start windows, so the tail window is
|
||
# shorter and needs its own bound shape -- i.e. a second context.
|
||
#
|
||
# Both decoder contexts must then SHARE one scratch buffer, and the main
|
||
# context has to be REPLACED rather than merely re-pointed: TRTRunner created
|
||
# it with the default (STATIC) strategy, so it already owns a full private
|
||
# scratch allocation. Adding a shared buffer beside it pays 2x instead of 1x,
|
||
# which showed up as odd lengths costing ~350 MiB more than their even
|
||
# neighbours (L=323 -> 3004 MiB vs L=256 -> 2650). Building both contexts
|
||
# USER_MANAGED against one buffer and dropping the STATIC one leaves a single
|
||
# allocation. Sharing is safe because the two run sequentially on one stream --
|
||
# verified bit-identical and stable under graph capture.
|
||
# get_device_memory_size_for_profile_v2, NOT device_memory_size: on a
|
||
# multi-profile engine the latter is the MAX ACROSS PROFILES, so a low-band
|
||
# deployment would allocate the wide band's 5778 MB here. That showed up as
|
||
# odd lengths peaking at 7854 MiB against 2342 for even ones -- the tail
|
||
# context is only created when L is odd.
|
||
#
|
||
# Both contexts belong to THIS GRAPH and are built here rather than reusing
|
||
# (and formerly REPLACING) self.dec_runner.context. A captured CUDA graph's
|
||
# recorded launches reference the context they were enqueued from, so handing
|
||
# the runner a new context destroys the old one underneath every graph already
|
||
# in the LRU cache: the next replay of an earlier graph segfaults. That is not
|
||
# hypothetical -- two cached graphs whose lengths both need a tail (e.g. 645
|
||
# then 1291) crashed on the first replay after the second capture. Per-graph
|
||
# contexts are cheap (a few hundred KB of host state each) because they carry
|
||
# no scratch of their own.
|
||
eng = self.dec_runner.engine
|
||
if self._shared_scratch is not None:
|
||
# Preferred: the one buffer _share_scratch() already sized for every
|
||
# engine. T5 -> DiT -> decoder run sequentially on one stream, so the
|
||
# decoder's windows reuse it with no extra allocation at all.
|
||
sptr, sbytes = (self._shared_scratch.data_ptr(),
|
||
self._shared_scratch.numel())
|
||
else:
|
||
# Standalone use (no shared buffer): one per-graph allocation, sized from
|
||
# get_device_memory_size_for_profile_v2, NOT device_memory_size -- on a
|
||
# multi-profile engine the latter is the MAX ACROSS PROFILES, so a
|
||
# low-band deployment would allocate the wide band's 5778 MB here. That
|
||
# showed up as odd lengths peaking at 7854 MiB against 2342 for even ones.
|
||
_pi = getattr(self.dec_runner, "profile", 0) or 0
|
||
sbytes = eng.get_device_memory_size_for_profile_v2(_pi)
|
||
self._dec_scratch = canon.torch.empty(sbytes, dtype=canon.torch.uint8,
|
||
device="cuda")
|
||
sptr = self._dec_scratch.data_ptr()
|
||
self._dec_ctx_main = eng.create_execution_context(
|
||
trt.ExecutionContextAllocationStrategy.USER_MANAGED)
|
||
self._dec_ctx_main.set_device_memory(sptr, sbytes)
|
||
self._dec_ctx_main.set_input_shape("latent",
|
||
(1, IO_CHANNELS, self.dec_win_lat))
|
||
dec_ctx = self._dec_ctx_main
|
||
self._dec_ctx_tail = eng.create_execution_context(
|
||
trt.ExecutionContextAllocationStrategy.USER_MANAGED)
|
||
self._dec_ctx_tail.set_device_memory(sptr, sbytes)
|
||
self._dec_ctx_tail.set_input_shape("latent", (1, IO_CHANNELS, tail[0]))
|
||
self._dec_ctx_tail.set_tensor_address("latent", self._dec_win_lat.data_ptr())
|
||
self._dec_ctx_tail.set_tensor_address("pcm", self._dec_win_i32.data_ptr())
|
||
dec_ctx.set_tensor_address("latent", self.decoder_in_buf.data_ptr())
|
||
dec_ctx.set_tensor_address("pcm", self._dec_win_i32.data_ptr())
|
||
else:
|
||
self.pcm_int32_buf = canon.torch.empty(1, T_full, 2, dtype=pcm_dt, device="cuda")
|
||
dec_ctx.set_tensor_address("latent", self.decoder_in_buf.data_ptr())
|
||
dec_ctx.set_tensor_address("pcm", self.pcm_int32_buf.data_ptr())
|
||
if self._ceiling_buf is not None:
|
||
dec_ctx.set_input_shape("limiter_ceiling", ())
|
||
dec_ctx.set_tensor_address("limiter_ceiling", self._ceiling_buf.data_ptr())
|
||
else:
|
||
# Legacy fp output — graph-capture supported but the postprocess
|
||
# path is different. Not the production target.
|
||
self._dec_out_name = "audio"
|
||
au_dt = self.dec_runner.out_dtype["audio"]
|
||
# (1, 2, T_full) audio in [-1, 1]
|
||
self._audio_legacy_buf = canon.torch.empty(1, 2, T_full, dtype=au_dt, device="cuda")
|
||
dec_ctx.set_tensor_address("latent", self.decoder_in_buf.data_ptr())
|
||
dec_ctx.set_tensor_address("audio", self._audio_legacy_buf.data_ptr())
|
||
|
||
# ── int16 narrow + pinned-host output buffers ──
|
||
# Full (T_full, 2) int16 GPU buffer — kept resident; we narrow with a
|
||
# tensor view at replay time. The DtoH copy goes only to the requested
|
||
# window (first requested_samples).
|
||
self.pcm_int16_buf = canon.torch.empty(T_dec, 2, dtype=canon.torch.int16, device="cuda")
|
||
# Pinned host — pre-pinned, persistent. Exactly the right size for
|
||
# this T_lat. The captured DMA copies from device to here.
|
||
self.pinned_host_pcm = canon.torch.empty(T_full, 2, dtype=canon.torch.int16,
|
||
pin_memory=True)
|
||
|
||
# ── Warmup: replicate canonical's warmup sequence exactly so that the
|
||
# decoder context's internal state at the time of the captured call
|
||
# matches canonical's state at its "real" call.
|
||
#
|
||
# The decoder engine carries internal state across calls (TRT 10
|
||
# workspace + perhaps cached cuBLAS handles). The state evolves
|
||
# with each call and depends on the inputs. To get bit-exact
|
||
# parity with canonical we MUST follow the same call sequence:
|
||
# 3 warmups with zero T5/DiT/decoder inputs,
|
||
# then the real captured call.
|
||
with canon.torch.cuda.stream(capture_stream):
|
||
# Pre-fill the input buffers with zeros to mimic canonical warmup.
|
||
self.input_ids_buf.zero_()
|
||
self.attn_mask_buf.zero_()
|
||
# decoder_in_buf will be loaded with zero latents below.
|
||
zero_latents = canon.torch.zeros_like(self.decoder_in_buf)
|
||
zero_t = canon.torch.tensor([0.5], device="cuda")
|
||
zero_h = canon.torch.zeros_like(self.dit._t5_hidden_buf)
|
||
zero_m = canon.torch.zeros_like(self.dit._t5_mask_buf)
|
||
zero_l = canon.torch.zeros_like(self.dit._local_add_cond_buf)
|
||
zero_x = canon.torch.zeros_like(self.dit._x_buf)
|
||
for _w in range(3):
|
||
# T5 with zero ids (matches t5gemma_encode(" ") closely enough
|
||
# that the engine state evolution is the same — what matters
|
||
# for downstream determinism is that EACH engine's call count
|
||
# and rough input pattern matches canonical).
|
||
_enqueue(self.t5_runner.context, capture_stream, "t5 (mega-graph)")
|
||
# DiT single step with zero inputs (matches canon's warmup
|
||
# `dit.step(_w_x, _w_t, _w_h, _w_m, 30.0, _w_l)`).
|
||
self.dit._t_buf[0] = 0.5
|
||
self.dit._x_buf.copy_(zero_x)
|
||
self.dit._t5_hidden_buf.copy_(zero_h)
|
||
self.dit._t5_mask_buf.copy_(zero_m)
|
||
self.dit._local_add_cond_buf.copy_(zero_l)
|
||
self.dit._sec_buf[0] = 30.0
|
||
_enqueue(self.dit.runner.context, capture_stream, "dit (mega-graph)")
|
||
# Decoder with zero latents (matches canon's _w_lat = zeros).
|
||
self.decoder_in_buf.copy_(zero_latents)
|
||
if self._dec_plan is not None:
|
||
# Same window sequence the capture will record -- this warmup exists so the
|
||
# decoder context's state at capture matches what it saw while warming.
|
||
self._decode_windows(capture_stream)
|
||
else:
|
||
_enqueue(self.dec_runner.context, capture_stream, "decoder (mega-graph)")
|
||
# Then canon runs ONE full sampling pass (8 DiT calls, no decoder).
|
||
# Mirror that here so the DiT context's state also matches.
|
||
self.dit._x_buf.copy_(zero_x)
|
||
for i in range(ns):
|
||
self.dit._t_buf[0] = float(sigmas[i])
|
||
_enqueue(self.dit.runner.context, capture_stream, "dit (mega-graph)")
|
||
# Then canon's GraphPingpongSampler.build() does 2 more 8-step
|
||
# passes of the DiT (still no decoder). Mirror that too.
|
||
for _w in range(2):
|
||
self.dit._x_buf.copy_(zero_x)
|
||
for i in range(ns):
|
||
self.dit._t_buf[0] = float(sigmas[i])
|
||
_enqueue(self.dit.runner.context, capture_stream, "dit (mega-graph)")
|
||
v = self.dit._vel_buf.float()
|
||
denoised = self.dit._x_buf - self._sigma_curr_bufs[i] * v
|
||
if i < ns - 1:
|
||
nb = self._noise_bufs[i]
|
||
nb.normal_()
|
||
new_x = (1.0 - self._sigma_next_bufs[i]) * denoised + self._sigma_next_bufs[i] * nb
|
||
else:
|
||
new_x = denoised
|
||
if i < ns - 1:
|
||
self.dit._x_buf.copy_(new_x)
|
||
else:
|
||
self.latents_out_buf.copy_(new_x)
|
||
# Now restore the DiT's persistent input buffers to zeros (they
|
||
# got mutated above) — the captured graph will refill them anyway,
|
||
# but be clean.
|
||
self.dit._t5_hidden_buf.zero_()
|
||
self.dit._t5_mask_buf.zero_()
|
||
self.dit._local_add_cond_buf.zero_()
|
||
self.dit._x_buf.zero_()
|
||
capture_stream.synchronize()
|
||
|
||
# ── Now capture ──
|
||
self._graph = canon.torch.cuda.CUDAGraph()
|
||
with canon.torch.cuda.stream(capture_stream):
|
||
with canon.torch.cuda.graph(self._graph, stream=capture_stream):
|
||
# Stage 1: T5
|
||
_enqueue(self.t5_runner.context, capture_stream, "t5 (mega-graph)")
|
||
# Cast attn_mask → fp32 → dit's t5_mask_buf
|
||
self.dit._t5_mask_buf.copy_(self.attn_mask_buf.float())
|
||
# T5 hidden (fp32) → dit's t5_hidden_buf
|
||
self.dit._t5_hidden_buf.copy_(self.t5_hidden_internal)
|
||
# Initial noise → dit's x_buf
|
||
self.dit._x_buf.copy_(self.initial_noise_buf)
|
||
# Stage 3: DiT loop
|
||
for i in range(ns):
|
||
self.dit._t_buf[0] = self._sigma_curr_bufs[i]
|
||
_enqueue(self.dit.runner.context, capture_stream, "dit (mega-graph)")
|
||
v = self.dit._vel_buf.float()
|
||
denoised = self.dit._x_buf - self._sigma_curr_bufs[i] * v
|
||
if i < ns - 1:
|
||
nb = self._noise_bufs[i]
|
||
new_x = (1.0 - self._sigma_next_bufs[i]) * denoised + self._sigma_next_bufs[i] * nb
|
||
else:
|
||
new_x = denoised
|
||
if i < ns - 1:
|
||
self.dit._x_buf.copy_(new_x)
|
||
else:
|
||
self.latents_out_buf.copy_(new_x)
|
||
# Stage 4: Decoder
|
||
self.decoder_in_buf[..., :L].copy_(self.latents_out_buf)
|
||
if self._L_dec > L:
|
||
# Edge-replicate into the pad slot. A device copy, so it captures like any
|
||
# other stage and costs one extra latent of decode.
|
||
self.decoder_in_buf[..., L:].copy_(self.latents_out_buf[..., -1:])
|
||
if self._dec_plan is not None:
|
||
self._decode_windows(capture_stream)
|
||
else:
|
||
_enqueue(self.dec_runner.context, capture_stream, "decoder (mega-graph)")
|
||
# Stage 5a: narrow + cast int32 → int16 (or legacy fp32 → int16)
|
||
if self._dec_plan is not None:
|
||
# Already done per window: _decode_windows narrows each window straight into
|
||
# pcm_int16_buf, so there is no full-length int32 buffer to convert here
|
||
# (that buffer is window-sized in this path -- 8 MB instead of 537 MB at
|
||
# L=16384). Running the copy below would try to write a window over the
|
||
# whole output.
|
||
pass
|
||
elif self._dec_out_name == "pcm":
|
||
# Belt-and-suspenders int16 clamp. New engines (from the
|
||
# FP32 clip+scale fix in the ONNX producer) already bound
|
||
# the int32 output to ±32767, so this is a no-op for them.
|
||
# Kept for backwards-compat with any older engine still in
|
||
# use, which has BF16 trunk rounding 32767 → 32768 and
|
||
# wrapping on int16 downcast (audible clicks).
|
||
self.pcm_int16_buf.copy_(self.pcm_int32_buf[0].clamp(-32767, 32767))
|
||
else:
|
||
a = self._audio_legacy_buf[0].clamp(-1.0, 1.0) * 32767.0
|
||
self.pcm_int16_buf.copy_(a.to(canon.torch.int16).T)
|
||
# Stage 5b: DtoH (captured non_blocking copy into pinned host)
|
||
self.pinned_host_pcm[:self.requested_samples].copy_(
|
||
self.pcm_int16_buf[:self.requested_samples], non_blocking=True)
|
||
|
||
self._built = True
|
||
|
||
def run(self, input_ids_cpu, attn_mask_cpu, seed: int, seconds: float):
|
||
"""Execute the captured pipeline. Returns a numpy view of the pinned PCM.
|
||
|
||
input_ids_cpu / attn_mask_cpu: (1, 256) int64 CPU tensors (tokenizer output).
|
||
seed: int — used to seed the in-loop noise + the initial latent randn.
|
||
seconds: float — duration condition. Written into dit._sec_buf before replay.
|
||
|
||
Returns: numpy.ndarray (requested_samples, 2) int16 — a view into the
|
||
pinned host buffer (zero-copy; valid until next run() overwrites it).
|
||
"""
|
||
assert self._built, "call build() first"
|
||
torch = canon.torch
|
||
stream = self.dit.runner.stream
|
||
|
||
# All host writes must happen on the SAME stream the graph replays on,
|
||
# otherwise the replay races the input copies.
|
||
with torch.cuda.stream(stream):
|
||
# Update T5 input buffers via HtoD copy (captured by the graph
|
||
# as raw device buffers — but the HtoD copy itself runs here,
|
||
# OUTSIDE the graph, before replay).
|
||
self.input_ids_buf.copy_(input_ids_cpu.to(self.input_ids_buf.dtype),
|
||
non_blocking=True)
|
||
self.attn_mask_buf.copy_(attn_mask_cpu.to(self.attn_mask_buf.dtype),
|
||
non_blocking=True)
|
||
# Update scalar inputs.
|
||
self.dit._sec_buf[0] = float(seconds)
|
||
# Refresh noise + initial latent (matches sa3_trt.py:
|
||
# noise = randn(seed=args.seed); replay-seed = args.seed + 1)
|
||
g_init = torch.Generator(device="cuda")
|
||
g_init.manual_seed(int(seed))
|
||
self.initial_noise_buf.normal_(generator=g_init)
|
||
g_loop = torch.Generator(device="cuda")
|
||
g_loop.manual_seed(int(seed) + 1)
|
||
for nb in self._noise_bufs:
|
||
nb.normal_(generator=g_loop)
|
||
# Replay the full pipeline.
|
||
self._graph.replay()
|
||
stream.synchronize()
|
||
# pinned_host_pcm has been written by the DtoH; return a view.
|
||
return self.pinned_host_pcm[:self.requested_samples].numpy()
|
||
|
||
|
||
# ─── Reusable inference class (CLI + gradio share this) ─────────────────
|
||
class SA3Inference:
|
||
"""SA3 TRT inference, set up once and reused.
|
||
|
||
Loads engines + warms up + builds a CUDA graph at __init__. Subsequent
|
||
generate() calls are graph replays (~30 ms on H100 at L=324). Different
|
||
(T_lat, steps) combos build their own graphs on demand and are cached
|
||
(LRU evict at MAX_GRAPHS).
|
||
|
||
MVP scope: cfg=1.0, no init_audio, no inpaint, sigma_max=1.0 — i.e. the
|
||
mega-graph fast path. Other configs raise NotImplementedError; the
|
||
eager-fallback wiring is the next step.
|
||
|
||
Thread-safe: serialize_generate() acquires an internal Lock so two
|
||
callers can't trample the shared CUDA stream / persistent buffers.
|
||
"""
|
||
|
||
MAX_GRAPHS = 4
|
||
DEFAULT_SIGMA_MAX = 1.0 # mega-graph path requires this
|
||
# Hard latent-length bounds. Every DiT engine's TRT profile is L=1..4096, the
|
||
# SAME-L decoder caps at 4096, and for --precision fp8 the baked RoPE cos/sin
|
||
# tables are sized to exactly 4096 (a longer render needs a re-bake). Enforced
|
||
# in get_graph() so every library caller (gradio / API / generate) is covered,
|
||
# not just the CLI.
|
||
DIT_MIN_L = 1
|
||
DIT_MAX_L = 4096
|
||
# The decoders' own floor is DECODER_MIN_L (32) from sa3_trt_core — the DiT's
|
||
# profile starts at 1, so nothing used to enforce it and a short render passed
|
||
# validation, ran the DiT happily, then died or fell silent inside the decoder
|
||
# instead of being rejected up front.
|
||
|
||
def __init__(self, dit: str, decoder: str, *,
|
||
precision: str | None = None,
|
||
dec_precision: str = "canonical",
|
||
default_T_lat: int = 324, default_steps: int = 8,
|
||
default_seconds: float = 30.0,
|
||
models_dir: Path | None = None,
|
||
with_encoder: bool = False,
|
||
chunking: bool = True,
|
||
dit_engine: "Path | str | None" = None,
|
||
dec_engine: "Path | str | None" = None,
|
||
quiet: bool = False):
|
||
"""Load engines + build a warmup graph.
|
||
|
||
Args:
|
||
dit: one of DIT_CHOICES — "sm-music" / "sm-sfx" / "medium"
|
||
decoder: one of DECODER_PATHS — "same-s" / "same-l"
|
||
precision: None (default → "fp16" for every model), or an
|
||
explicit "fp16" (canonical: FP16 trunk, FP32
|
||
RMSNorm/RoPE islands, FMHA-fused FP16 attention core),
|
||
"fp8" (medium only; max-speed clean tier — fp8 linears
|
||
+ bf16 fused FMHA + baked fp32 RoPE, ~1.3× faster than
|
||
fp16 and clean at long sequence, capped L≤4096),
|
||
or "fp32" (bit-equiv
|
||
PyTorch eager, ~2× slower). ("fp16mixed" and the
|
||
retired medium-only "bf16" alias to "fp16".)
|
||
Engines auto-download
|
||
from HF if the requested file is missing.
|
||
default_T_lat: latent length to build the initial graph at
|
||
default_steps: pingpong steps for the initial graph
|
||
default_seconds: duration condition for the initial graph (used for
|
||
the warmup pass; per-call generate() overrides)
|
||
models_dir: override the canonical models/ root (optional)
|
||
with_encoder: also load the audio encoder TRT engine (needed for
|
||
future audio-to-audio / inpaint modes)
|
||
dit_engine: explicit DiT engine path, overriding the precision-derived default.
|
||
⚠ Needed by any caller that picks an engine FILE rather than a
|
||
precision: get_dit_engine_path() rebuilds the path from
|
||
<arch>/<subdir>/dit_<precision>.trt and ignores DIT_CHOICES[...]
|
||
["engine"], so mutating that global silently had no effect -- the
|
||
gradio variant dropdown was a no-op and every render used the fp16
|
||
engine whatever the user chose.
|
||
dec_engine: same, for the decoder.
|
||
chunking: True (default) selects the engines' LOW band -- decode in
|
||
256-latent windows, encode in 64 -- which is the only thing that
|
||
reduces the scratch reservation, because TRT commits scratch from
|
||
the profile ceiling before any shape is bound. 509 MB instead of
|
||
8143 for the SAME-L decoder. False selects the WIDE band: one
|
||
single-shot call, ~10-20% faster above L=256 and ~5.4 GB more
|
||
resident. Ignored by single-profile engines.
|
||
quiet: suppress per-stage print() output from canon helpers
|
||
"""
|
||
if dit not in DIT_CHOICES:
|
||
raise ValueError(f"unknown dit={dit!r}; valid: {list(DIT_CHOICES)}")
|
||
if decoder not in DECODER_PATHS:
|
||
raise ValueError(f"unknown decoder={decoder!r}; valid: {list(DECODER_PATHS)}")
|
||
# Resolve precision default per model: fp16 everywhere. Retired tokens
|
||
# (fp16mixed, bf16) alias to fp16.
|
||
if precision is None:
|
||
precision = canon.default_precision(dit)
|
||
precision = canon.normalize_precision(precision)
|
||
if precision not in canon.PRECISIONS:
|
||
raise ValueError(f"unknown precision={precision!r}; valid: {canon.PRECISIONS}")
|
||
if precision == "fp8" and dit != "medium":
|
||
raise ValueError("precision='fp8' is only available for dit='medium' "
|
||
"(fp8 linears + bf16 FMHA + baked fp32 RoPE; sm-music/sm-sfx ship fp16 only)")
|
||
|
||
# Quiet: patch canon's stage/sub/_stage_vram to no-ops so loading
|
||
# doesn't spam stdout (gradio in particular wants a clean log).
|
||
if quiet:
|
||
canon._USE_COLOR = False
|
||
canon.stage = lambda *a, **kw: None
|
||
canon.sub = lambda *a, **kw: None
|
||
canon._stage_vram = lambda *a, **kw: 0
|
||
|
||
# models-dir override (mirrors main()'s logic).
|
||
if models_dir is not None and str(models_dir) != str(canon.MODELS_DIR):
|
||
new_root = Path(models_dir).resolve()
|
||
new_arch_dir = new_root / ARCH
|
||
canon.T5GEMMA_PATH = new_arch_dir / "t5gemma" / "t5gemma_fp16.trt"
|
||
for kk in DIT_CHOICES:
|
||
DIT_CHOICES[kk]["engine"] = new_arch_dir / DIT_CHOICES[kk]["engine"].relative_to(canon.ARCH_DIR)
|
||
for kk in DECODER_PATHS:
|
||
DECODER_PATHS[kk] = new_arch_dir / DECODER_PATHS[kk].relative_to(canon.ARCH_DIR)
|
||
for kk in ENCODER_PATHS:
|
||
ENCODER_PATHS[kk] = new_arch_dir / ENCODER_PATHS[kk].relative_to(canon.ARCH_DIR)
|
||
canon.MODELS_DIR = new_root
|
||
canon.ARCH_DIR = new_arch_dir
|
||
|
||
# chunking picks the optimization profile (band), not a runtime code path: profile 0 is
|
||
# the low-ceiling band, 1 the wide one. A single-profile engine only has 0.
|
||
self.chunking = bool(chunking)
|
||
self.dec_profile = 0 if chunking else 1
|
||
self._scratch = None
|
||
self._scratch_bytes = 0
|
||
self._eager = None
|
||
self._eager_dit = None
|
||
self._cfg_batched = None
|
||
self.dit_name = dit
|
||
self.decoder_name = decoder
|
||
self.precision = precision
|
||
# _encoder_runner() resolves the encoder by this tier; without it audio-to-audio
|
||
# and inpainting died on the first attribute read.
|
||
self.dec_precision = dec_precision
|
||
self.with_encoder = with_encoder
|
||
self.quiet = quiet
|
||
|
||
# 1. Lazy-download any missing engines (precision-aware).
|
||
needed = canon.get_engine_files(dit, decoder, precision, with_encoder=with_encoder, dec_tier=dec_precision)
|
||
_ensure_files(needed)
|
||
|
||
# 2. Heavy imports (torch + tensorrt + plugin).
|
||
sub(dim("Loading..."))
|
||
t0 = time.time()
|
||
_import_heavy()
|
||
sub(f"{dim('heavy imports')} {(time.time()-t0)*1000:.0f} ms")
|
||
_init_nvml()
|
||
torch = canon.torch
|
||
|
||
# 3. Tokenizer + dist_shift (from bundled runtime.py).
|
||
t0 = time.time()
|
||
import runtime as rt
|
||
rt.MODELS_DIR = str(canon.MODELS_DIR)
|
||
rt.ARCH_DIR = str(canon.ARCH_DIR)
|
||
state = rt.load()
|
||
self.tokenizer = state["tokenizer"]
|
||
self.dist_shift = state["dist_shift"]
|
||
sub(f"{dim('tokenizer + dist-shift')} {(time.time()-t0)*1000:.0f} ms")
|
||
|
||
# 4. Engine load (parallel).
|
||
import concurrent.futures
|
||
engine_specs = {
|
||
"t5": canon.T5GEMMA_PATH,
|
||
"dit": Path(dit_engine) if dit_engine else canon.get_dit_engine_path(dit, precision),
|
||
"dec": (Path(dec_engine) if dec_engine else
|
||
canon.resolve_decoder_engine(decoder, precision, dec_precision)),
|
||
}
|
||
if with_encoder:
|
||
engine_specs["enc"] = canon.resolve_encoder_engine(decoder, dec_precision)
|
||
t0 = time.time()
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=len(engine_specs)) as ex:
|
||
# user_managed=True -- _share_scratch() immediately below owns the allocation.
|
||
futs = {name: ex.submit(TRTRunner, path, None, True,
|
||
self.dec_profile if name in ("dec", "enc") else 0)
|
||
for name, path in engine_specs.items()}
|
||
self.runners = {name: fut.result() for name, fut in futs.items()}
|
||
sub(f"{dim(f'engine load (parallel, {len(self.runners)} engines)')} {(time.time()-t0)*1000:.0f} ms")
|
||
|
||
self._share_scratch()
|
||
self.dit = DiTRunner(self.runners["dit"])
|
||
|
||
# 5. Graph cache + lock.
|
||
self._graphs: dict[tuple[int, int], FullPipelineGraph] = {}
|
||
self._graph_lru: list[tuple[int, int]] = []
|
||
self._lock = threading.Lock()
|
||
|
||
# 6. Build the initial graph at the default config.
|
||
t0 = time.time()
|
||
self.get_graph(default_T_lat, default_steps, default_seconds)
|
||
torch.cuda.synchronize()
|
||
sub(f"{dim(f'warmup + capture (T_lat={default_T_lat}, steps={default_steps})')} "
|
||
f"{(time.time()-t0)*1000:.0f} ms")
|
||
|
||
@staticmethod
|
||
def resolve_T_lat(seconds: float, decoder: str) -> int:
|
||
"""seconds → T_lat (natural ceil; no even-bump).
|
||
|
||
Earlier code bumped T_lat to even for same-s under the assumption that
|
||
the BF16 trunk's internal chunking required it. Empirically, on real
|
||
DiT-output latents the engine produces correct audio (cos ≥ 0.99 vs
|
||
PT eager) at any L in [32, 4096], odd or even. Removing the bump keeps
|
||
the DiT sequence length at exactly what the user requested.
|
||
"""
|
||
return max(1, math.ceil(seconds * SAMPLE_RATE / SAMPLES_PER_LATENT))
|
||
|
||
def _share_scratch(self):
|
||
"""One scratch buffer for every engine loaded at construction (see canon.share_scratch).
|
||
|
||
⚠ Call this ONCE, from __init__, before any execution context or CUDA graph exists.
|
||
canon.share_scratch() REPLACES each runner's execution context and frees the previous
|
||
scratch allocation, so calling it again after a render would leave the captured
|
||
mega-graph and the eager path's contexts pointing at destroyed objects -- a segfault
|
||
with no Python traceback. That is why the lazily-loaded encoder gets its OWN scratch
|
||
(see _encoder_runner) instead of being folded in here.
|
||
"""
|
||
buf = canon.share_scratch(self.runners)
|
||
if buf is not None:
|
||
self._scratch, self._scratch_bytes = buf, buf.numel()
|
||
|
||
def close(self):
|
||
"""Tear the instance down in the ONE order that is safe on the GPU.
|
||
|
||
A captured mega-graph holds raw pointers into the execution contexts it was
|
||
captured against and into the shared scratch buffer; so do the eager path's
|
||
_SecondContext objects. Free the runners first and destroying those graphs
|
||
afterwards touches memory that is already gone -- which surfaces as
|
||
`CUDA error: an illegal memory access was encountered` on the NEXT render, in
|
||
a different instance, because a poisoned CUDA context stays poisoned. (That is
|
||
what the gradio LRU eviction used to do, so switching models twice bricked the
|
||
server.) Python's own GC order over __dict__ is arbitrary, so relying on
|
||
refcounting to get this right is not an option: it has to be explicit.
|
||
|
||
Idempotent — safe to call twice, and safe on a half-constructed instance.
|
||
"""
|
||
t = canon.torch
|
||
try:
|
||
t.cuda.synchronize() # nothing may still be in flight
|
||
except Exception:
|
||
pass
|
||
# 1. graphs, 2. eager contexts, 3. runner contexts + engines, 4. the buffer they shared
|
||
for attr in ("_graphs", "_graph_lru"):
|
||
try:
|
||
getattr(self, attr).clear()
|
||
except Exception:
|
||
pass
|
||
self._eager = self._eager_dit = None
|
||
for _, runner in list(getattr(self, "runners", {}).items()):
|
||
try:
|
||
runner.free()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
self.runners.clear()
|
||
except Exception:
|
||
pass
|
||
self._scratch, self._scratch_bytes = None, 0
|
||
try:
|
||
t.cuda.synchronize()
|
||
t.cuda.empty_cache()
|
||
except Exception:
|
||
pass
|
||
|
||
def engine_footprint(self) -> dict:
|
||
"""Per-engine resident VRAM: weights (the deserialised engine) + scratch.
|
||
|
||
Scratch is read per PROFILE, never device_memory_size_v2 -- on the two-profile
|
||
autoencoders the latter reports the MAX ACROSS profiles, so a chunked deployment
|
||
would be blamed for the wide band's 7.8 GB it never reserves.
|
||
|
||
`shared` marks the engines folded into one buffer by canon.share_scratch: their
|
||
scratch is the SAME allocation, so summing the column double-counts.
|
||
"""
|
||
out = {}
|
||
for name, r in self.runners.items():
|
||
try:
|
||
prof = getattr(r, "profile", 0) or 0
|
||
scratch = r.engine.get_device_memory_size_for_profile_v2(prof)
|
||
except Exception:
|
||
scratch = 0
|
||
# TRT exposes no API for an engine's weight allocation; the serialised file
|
||
# size is the faithful proxy, since weights dominate it.
|
||
try:
|
||
weights = Path(r.engine_path).stat().st_size
|
||
except Exception:
|
||
weights = 0
|
||
out[name] = {"weights": weights, "scratch": scratch,
|
||
"profile": getattr(r, "profile", 0) or 0,
|
||
"shared": r._scratch is None}
|
||
out["_shared_scratch"] = (self._scratch.numel() if self._scratch is not None else 0)
|
||
return out
|
||
|
||
def _dec_min_L(self) -> int:
|
||
"""The loaded decoder engine's real profile floor.
|
||
|
||
DECODER_MIN_L (32) is the SHIPPED engines' floor and was a build-time convention, not a
|
||
model limit -- the DiT's floor is 1, and an engine rebuilt with `--min-l 1` decodes L=1
|
||
natively at cos 0.9993 (SAME-L) against eager, where padding up to 32 is out-of-
|
||
distribution and measures 0.82-0.92. So the guard has to ask the engine, not a constant.
|
||
"""
|
||
try:
|
||
# ⚠ profile_bounds returns (min, max) as plain ints, not shape tuples. This read
|
||
# `lo[-1]`, which raises TypeError on an int, so EVERY call fell into the except
|
||
# and returned the stale DECODER_MIN_L=32 — rejecting every render under ~2.97 s
|
||
# even though the shipped chunkable engines all report a floor of 1.
|
||
lo, _ = canon.profile_bounds(self.runners["dec"], "latent")
|
||
return int(lo)
|
||
except Exception:
|
||
return DECODER_MIN_L
|
||
|
||
def get_graph(self, T_lat: int, steps: int, seconds: float) -> FullPipelineGraph:
|
||
"""Return the cached graph for (T_lat, steps), building on miss.
|
||
|
||
LRU-evicts when the cache exceeds MAX_GRAPHS to keep VRAM bounded.
|
||
Called under the inference lock when reached from generate().
|
||
"""
|
||
# Hard length cap (see DIT_MAX_L). Reject cleanly here so the library path
|
||
# fails with a clear message rather than a cryptic TRT profile error deep
|
||
# in graph build. For fp8 this is a real correctness limit (baked RoPE table).
|
||
dec_min = self._dec_min_L()
|
||
if T_lat < dec_min:
|
||
min_s = dec_min * SAMPLES_PER_LATENT / SAMPLE_RATE
|
||
raise ValueError(
|
||
f"T_lat={T_lat} is below this decoder engine's profile minimum of "
|
||
f"{dec_min} (~{min_s:.2f}s). The DiT accepts it but the decoder refuses the "
|
||
f"shape, so it is rejected here instead of failing obscurely further in. "
|
||
f"Rebuild the decoder with a lower floor (build_capped.py --min-l 1) to go below.")
|
||
if not (self.DIT_MIN_L <= T_lat <= self.DIT_MAX_L):
|
||
max_s = self.DIT_MAX_L * SAMPLES_PER_LATENT / SAMPLE_RATE
|
||
reason = ("the fp8 baked RoPE table (sized to L=4096; a longer render needs a "
|
||
"re-bake — build/make_rope_baked_onnx.py --max-t)"
|
||
if self.precision == "fp8" else
|
||
"the DiT engine's optimization profile (built for L=1..4096)")
|
||
raise ValueError(
|
||
f"T_lat={T_lat} is out of range [{self.DIT_MIN_L}, {self.DIT_MAX_L}] "
|
||
f"(~{max_s:.0f}s max), bounded by {reason} and the SAME-L decoder's 4096 cap.")
|
||
key = (T_lat, steps)
|
||
if key in self._graphs:
|
||
self._graph_lru.remove(key)
|
||
self._graph_lru.append(key)
|
||
return self._graphs[key]
|
||
|
||
sigma_max = self.DEFAULT_SIGMA_MAX
|
||
sigmas = build_pingpong_schedule(steps, sigma_max=sigma_max,
|
||
dist_shift=self.dist_shift, latent_len=T_lat)
|
||
# Size output buffer to MAX possible for this T_lat. generate() slices
|
||
# to the actual requested samples, so one graph per (T_lat, steps)
|
||
# serves any seconds within that T_lat's range.
|
||
max_samples = T_lat * SAMPLES_PER_LATENT
|
||
_, dec_hi = canon.profile_bounds(self.runners["dec"], "latent")
|
||
chunk = dec_hi if T_lat > dec_hi else 0
|
||
graph = FullPipelineGraph(self.runners["t5"], self.dit, self.runners["dec"],
|
||
T_lat, steps, max_samples, dec_chunk_lat=chunk,
|
||
scratch=self._scratch)
|
||
graph.build(sigmas, seconds, sigma_max)
|
||
canon.torch.cuda.synchronize()
|
||
|
||
self._graphs[key] = graph
|
||
self._graph_lru.append(key)
|
||
# LRU eviction. Python GC + torch's caching allocator reclaim GPU
|
||
# memory once the FullPipelineGraph reference is dropped.
|
||
while len(self._graphs) > self.MAX_GRAPHS:
|
||
evict = self._graph_lru.pop(0)
|
||
del self._graphs[evict]
|
||
return graph
|
||
|
||
# ── Eager path: CFG / negative prompt / audio-to-audio / inpaint ──────
|
||
#
|
||
# Same math as sa3_trt_core.main()'s eager branch (sequential dual-pass CFG with
|
||
# APG, linear init-latent mix for audio-to-audio, keep-mask local_add_cond +
|
||
# paste-back for inpaint), but driven from the already-loaded engines instead of
|
||
# a second process, and on its own execution contexts so the captured mega-graph
|
||
# stays valid (see _SecondContext).
|
||
def _eager_runners(self):
|
||
if self._eager is None:
|
||
self._eager = {n: _SecondContext(self.runners[n], self._scratch)
|
||
for n in ("t5", "dit", "dec")}
|
||
self._eager_dit = DiTRunner(self._eager["dit"])
|
||
# Var-batch engine? Then CFG's cond+uncond pair goes through as one batch-2 call.
|
||
self._cfg_batched = BatchedDiT.attach(self.runners["dit"], batch=2)
|
||
if not self.quiet:
|
||
sub(dim("CFG dual-pass: " + ("batched (var-batch engine, profile "
|
||
f"{self._cfg_batched.profile_index})"
|
||
if self._cfg_batched else
|
||
"sequential (engine is batch-1 only)")))
|
||
return self._eager
|
||
|
||
def _encoder_runner(self):
|
||
"""The SAME-S/L encoder engine — only needed for audio-to-audio / inpaint."""
|
||
if "enc" not in self.runners:
|
||
_ensure_files(ENCODER_FILES[self.decoder_name])
|
||
# The encoder arrives long after construction -- the first audio-to-audio or
|
||
# inpaint of a session -- by which point the mega-graph is captured and the
|
||
# eager contexts exist. Re-running canon.share_scratch() to fold it in would
|
||
# swap those contexts out and free the buffer underneath them: a segfault with
|
||
# no traceback, reachable as "generate, audio-to-audio, generate".
|
||
#
|
||
# So bind the encoder ALONE. If the buffer the other engines already share is
|
||
# big enough it costs nothing extra -- the encoder runs before them on the same
|
||
# stream under the same lock, exactly the sequential-reuse assumption
|
||
# canon.share_scratch() is built on -- and no other context is touched. Only
|
||
# when the encoder needs more does it fall back to its own allocation, which is
|
||
# still strictly additive rather than disruptive.
|
||
enc = TRTRunner(
|
||
canon.resolve_encoder_engine(self.decoder_name, self.dec_precision),
|
||
None, True)
|
||
need = enc.engine.get_device_memory_size_for_profile_v2(enc.profile)
|
||
if self._scratch is not None and self._scratch.numel() >= need:
|
||
enc.context.set_device_memory(self._scratch.data_ptr(), self._scratch.numel())
|
||
else:
|
||
enc._scratch = canon.torch.empty(need, dtype=canon.torch.uint8, device="cuda")
|
||
enc.context.set_device_memory(enc._scratch.data_ptr(), need)
|
||
if not self.quiet:
|
||
sub(dim(f"encoder scratch {need / 2**20:.0f} MiB (own buffer — larger than "
|
||
f"the {self._scratch.numel() / 2**20:.0f} MiB shared one)"
|
||
if self._scratch is not None else
|
||
f"encoder scratch {need / 2**20:.0f} MiB"))
|
||
self.runners["enc"] = enc
|
||
return self.runners["enc"]
|
||
|
||
@staticmethod
|
||
def _read_audio_any(path: str) -> np.ndarray:
|
||
"""Any soundfile-readable audio → (2, T) float32 @ 44.1 kHz, in [-1, 1].
|
||
|
||
canon.read_wav() only accepts 16-bit PCM at 44.1 kHz; the gradio file picker
|
||
will happily hand over 24-bit, 48 kHz, mono, or mp3.
|
||
"""
|
||
import soundfile as sf
|
||
data, sr = sf.read(path, dtype="float32", always_2d=True) # (T, C)
|
||
a = data.T # (C, T)
|
||
if a.shape[0] == 1:
|
||
a = np.repeat(a, 2, axis=0)
|
||
elif a.shape[0] > 2:
|
||
a = a[:2]
|
||
if sr != SAMPLE_RATE:
|
||
import torch as _t
|
||
import torchaudio
|
||
a = torchaudio.functional.resample(_t.from_numpy(np.ascontiguousarray(a)),
|
||
sr, SAMPLE_RATE).numpy()
|
||
return np.ascontiguousarray(a, dtype=np.float32)
|
||
|
||
def generate_eager(self, prompt: str, *,
|
||
seconds: float = 30.0, steps: int = 8,
|
||
seed: int | None = None,
|
||
cfg: float = 1.0, apg: float = 1.0,
|
||
negative_prompt: str | None = None,
|
||
init_audio_path: str | None = None,
|
||
init_noise_level: float = 1.0,
|
||
inpaint_range: tuple[float, float] | None = None,
|
||
) -> tuple[np.ndarray, dict]:
|
||
"""Advanced-mode render (no mega-graph). Returns (pcm_int16 (T,2), timing)."""
|
||
torch = canon.torch
|
||
T_lat = self.resolve_T_lat(seconds, self.decoder_name)
|
||
dec_min = self._dec_min_L()
|
||
if T_lat < dec_min:
|
||
raise ValueError(f"T_lat={T_lat} is below the {self.decoder_name} decoder engine's "
|
||
f"profile minimum of {dec_min} "
|
||
f"(~{dec_min * SAMPLES_PER_LATENT / SAMPLE_RATE:.2f}s)")
|
||
if not (self.DIT_MIN_L <= T_lat <= self.DIT_MAX_L):
|
||
raise ValueError(f"T_lat={T_lat} out of range [{self.DIT_MIN_L}, {self.DIT_MAX_L}]")
|
||
sigma_max = float(init_noise_level)
|
||
if not (0.0 < sigma_max <= 1.0):
|
||
raise ValueError(f"init_noise_level must be in (0, 1], got {sigma_max}")
|
||
if inpaint_range is not None and init_audio_path is None:
|
||
raise ValueError("inpaint needs init audio to keep the un-inpainted part of")
|
||
if seed is None:
|
||
seed = random.randint(0, 2**31 - 1)
|
||
|
||
with self._lock:
|
||
t_start = time.time()
|
||
# Per-stage peaks for the debug view. reset_peak_memory_stats is free when
|
||
# nobody reads it, and the numbers are torch-side allocations only: TRT weights
|
||
# and scratch are reported separately by engine_footprint(), because they are
|
||
# resident for the whole process rather than attributable to one stage.
|
||
_vram = {}
|
||
def _mark(k, base):
|
||
# transient working set for this stage: everything above the resting
|
||
# allocation. Without subtracting the base, every stage just reports
|
||
# the persistent shared scratch and they all look identical.
|
||
_vram[k] = max(0, torch.cuda.max_memory_allocated() - base)
|
||
def _base():
|
||
torch.cuda.reset_peak_memory_stats()
|
||
return torch.cuda.memory_allocated()
|
||
R = self._eager_runners()
|
||
dit = self._eager_dit
|
||
|
||
# 1. T5 (+ the uncond branch when CFG is on)
|
||
t0 = time.time(); _b = _base()
|
||
embeds, mask = t5gemma_encode(R["t5"], self.tokenizer, prompt)
|
||
null_embeds = null_mask = None
|
||
if cfg != 1.0:
|
||
null_embeds, null_mask = t5gemma_encode(
|
||
R["t5"], self.tokenizer, negative_prompt or "")
|
||
if not (negative_prompt or "").strip():
|
||
# No negative prompt → all-zero hidden + all-zero mask, which the
|
||
# engine's bundled conditioner turns into the learned padding
|
||
# embedding at every position (canon's "zero cross-attn" uncond).
|
||
null_embeds = torch.zeros_like(embeds)
|
||
null_mask = torch.zeros_like(mask)
|
||
t5_ms = (time.time() - t0) * 1000
|
||
_mark('t5', _b)
|
||
|
||
# 2. init audio → latents (audio-to-audio / inpaint)
|
||
init_latents = None
|
||
encode_ms = 0.0
|
||
if init_audio_path:
|
||
t0 = time.time(); _b = _base()
|
||
a = self._read_audio_any(init_audio_path)
|
||
target = T_lat * SAMPLES_PER_LATENT
|
||
a = (a[:, :target] if a.shape[-1] >= target
|
||
else np.pad(a, ((0, 0), (0, target - a.shape[-1]))))
|
||
init_latents = canon.encode_chunked(
|
||
self._encoder_runner(),
|
||
torch.from_numpy(np.ascontiguousarray(a)).unsqueeze(0).cuda())
|
||
encode_ms = (time.time() - t0) * 1000
|
||
_mark('encode', _b)
|
||
|
||
# 3. schedule + initial latent
|
||
sigmas = build_pingpong_schedule(steps, sigma_max=sigma_max,
|
||
dist_shift=self.dist_shift, latent_len=T_lat)
|
||
g = torch.Generator(device="cuda")
|
||
g.manual_seed(int(seed))
|
||
pure_noise = torch.randn(1, IO_CHANNELS, T_lat, device="cuda",
|
||
dtype=torch.float32, generator=g)
|
||
lat_range = None
|
||
if inpaint_range is not None:
|
||
s_sec, e_sec = inpaint_range
|
||
if not (0 <= s_sec < e_sec <= seconds):
|
||
raise ValueError(f"inpaint range {s_sec},{e_sec} must satisfy "
|
||
f"0 <= start < end <= {seconds}")
|
||
lat_range = (max(0, int(round(s_sec * SAMPLE_RATE / SAMPLES_PER_LATENT))),
|
||
min(T_lat, int(round(e_sec * SAMPLE_RATE / SAMPLES_PER_LATENT))))
|
||
if init_latents is not None and lat_range is None:
|
||
noise = init_latents * (1.0 - sigma_max) + pure_noise * sigma_max
|
||
else:
|
||
noise = pure_noise
|
||
|
||
# 4. local_add_cond: zeros, or the inpaint keep-mask + masked input
|
||
if lat_range is not None:
|
||
s0, s1 = lat_range
|
||
keep = torch.ones((1, 1, T_lat), device="cuda", dtype=torch.float32)
|
||
keep[:, :, s0:s1] = 0.0
|
||
local_add_cond = torch.cat([keep, init_latents * keep], dim=1).contiguous()
|
||
paste_back = (init_latents, keep)
|
||
else:
|
||
local_add_cond = torch.zeros((1, 257, T_lat), device="cuda", dtype=torch.float32)
|
||
paste_back = None
|
||
|
||
batched = self._cfg_batched if cfg != 1.0 else None
|
||
if batched is not None:
|
||
cat = torch.cat
|
||
h2 = cat([embeds.float(), null_embeds.float()], 0)
|
||
m2 = cat([mask.float(), null_mask.float()], 0)
|
||
lac2 = cat([local_add_cond, local_add_cond], 0)
|
||
sec2 = torch.full((2,), float(seconds), device="cuda")
|
||
|
||
def model_fn(x, t):
|
||
if cfg == 1.0:
|
||
return dit.step(x, t, embeds, mask, seconds, local_add_cond)
|
||
if batched is not None:
|
||
tv = float(t.item() if torch.is_tensor(t) else t)
|
||
v2 = batched.step(cat([x, x], 0), torch.full((2,), tv, device="cuda"),
|
||
h2, m2, sec2, lac2)
|
||
v_cond, v_uncond = v2[0:1], v2[1:2]
|
||
else:
|
||
v_cond = dit.step(x, t, embeds, mask, seconds, local_add_cond)
|
||
v_uncond = dit.step(x, t, null_embeds, null_mask, seconds, local_add_cond)
|
||
sigma = t.reshape(-1, 1, 1).float()
|
||
cond_d = x.float() - v_cond.float() * sigma
|
||
uncond_d = x.float() - v_uncond.float() * sigma
|
||
diff = cond_d - uncond_d
|
||
if apg <= 0.0:
|
||
cfg_diff = diff
|
||
else:
|
||
norm = torch.sqrt((cond_d * cond_d).sum(dim=(-2, -1), keepdim=True))
|
||
unit = cond_d / torch.clamp(norm, min=1e-8)
|
||
parallel = (diff * unit).sum(dim=(-2, -1), keepdim=True) * unit
|
||
orth = diff - parallel
|
||
cfg_diff = orth if apg >= 1.0 else (apg * orth + (1.0 - apg) * diff)
|
||
cfg_d = cond_d + (cfg - 1.0) * cfg_diff
|
||
return ((x.float() - cfg_d) / sigma).to(x.dtype)
|
||
|
||
t0 = time.time(); _b = _base()
|
||
latents = sample_flow_pingpong(model_fn, noise, sigmas, seed=int(seed) + 1,
|
||
paste_back=paste_back)
|
||
torch.cuda.synchronize()
|
||
sampling_ms = (time.time() - t0) * 1000
|
||
_mark('dit', _b)
|
||
|
||
# 5. decode
|
||
t0 = time.time(); _b = _base()
|
||
# Eager path serves CFG / APG / a2a / inpaint, which the mega-graph does not cover.
|
||
# It must window the decode for the same reason the graph does: a capped decoder engine
|
||
# refuses any length above its profile ceiling, and that refusal used to surface as the
|
||
# previous request's buffer rather than an error.
|
||
_, _dec_hi = canon.profile_bounds(R["dec"], "latent")
|
||
out = (canon.decode_chunked(R["dec"], latents, chunk_lat=_dec_hi)
|
||
if latents.shape[-1] > _dec_hi else decoder_decode(R["dec"], latents))
|
||
# ⚠ int16 too: decode_chunked accumulates in int16 (see its docstring). Testing only
|
||
# for int32 sends chunked output down the float branch, where it is clamped to ±1 and
|
||
# multiplied by 32767 a second time -- saturating every sample AND transposing to
|
||
# (2, T). Any render past the decoder's chunk ceiling hit this.
|
||
if out.dtype in (torch.int32, torch.int16):
|
||
pcm_full = out[0].to(torch.int16).cpu().numpy()
|
||
else:
|
||
pcm_full = (out.clamp(-1, 1) * 32767.0).to(torch.int16
|
||
).squeeze(0).T.contiguous().cpu().numpy()
|
||
decode_ms = (time.time() - t0) * 1000
|
||
_mark('decode', _b)
|
||
|
||
actual_samples = int(round(seconds * SAMPLE_RATE))
|
||
pcm = pcm_full[:actual_samples].copy()
|
||
inference_ms = (time.time() - t_start) * 1000
|
||
mode = ("inpaint" if lat_range is not None else
|
||
"audio-to-audio" if init_latents is not None else "text-to-audio")
|
||
return pcm, {
|
||
"inference_ms": inference_ms,
|
||
"graph_build_ms": 0.0,
|
||
"realtime": (seconds * 1000.0) / inference_ms if inference_ms > 0 else 0.0,
|
||
"seed": int(seed), "T_lat": T_lat, "samples": len(pcm),
|
||
"t5_ms": t5_ms, "encode_ms": encode_ms,
|
||
"sampling_ms": sampling_ms, "decode_ms": decode_ms,
|
||
"mode": mode, "cfg": cfg, "sigma_max": sigma_max, "eager": True,
|
||
"stage_vram": _vram, "steps": int(steps),
|
||
}
|
||
|
||
def generate(self, prompt: str, *,
|
||
seconds: float = 30.0, steps: int = 8,
|
||
seed: int | None = None,
|
||
# Below: MVP raises NotImplementedError; wiring planned.
|
||
init_noise_level: float = 1.0,
|
||
negative_prompt: str | None = None,
|
||
cfg: float = 1.0,
|
||
apg: float = 1.0,
|
||
init_audio_path: str | None = None,
|
||
inpaint_range: tuple[float, float] | None = None,
|
||
) -> tuple[np.ndarray, dict]:
|
||
"""Generate one audio clip. Returns (pcm_int16, timing_dict).
|
||
|
||
Returns:
|
||
pcm: (T_samples, 2) int16 numpy array, T_samples = round(seconds*44100)
|
||
timing: dict with 'inference_ms', 'graph_build_ms' (0 if cache hit),
|
||
'realtime', 'seed', 'T_lat', 'samples'
|
||
"""
|
||
# Anything the mega-graph cannot serve -- CFG's dual pass, an init latent to mix or paste
|
||
# back, sigma_max != 1 -- routes to the eager sampler. It runs on its OWN execution
|
||
# contexts (see _SecondContext) so an already-captured graph stays valid.
|
||
if (cfg != 1.0 or init_audio_path is not None or inpaint_range is not None
|
||
or init_noise_level != 1.0):
|
||
return self.generate_eager(
|
||
prompt, seconds=seconds, steps=steps, seed=seed,
|
||
cfg=cfg, apg=apg, negative_prompt=negative_prompt,
|
||
init_audio_path=init_audio_path, init_noise_level=init_noise_level,
|
||
inpaint_range=inpaint_range)
|
||
|
||
T_lat = self.resolve_T_lat(seconds, self.decoder_name)
|
||
if not (1 <= T_lat <= 4096):
|
||
raise ValueError(f"T_lat={T_lat} out of engine range [1, 4096]")
|
||
if seed is None:
|
||
seed = random.randint(0, 2**31 - 1)
|
||
|
||
with self._lock:
|
||
t0 = time.time()
|
||
graph = self.get_graph(T_lat, steps, seconds)
|
||
graph_build_ms = (time.time() - t0) * 1000
|
||
|
||
# Tokenize on CPU (outside graph).
|
||
tok = self.tokenizer(prompt, return_tensors="pt", max_length=T5_MAX_LEN,
|
||
padding="max_length", truncation=True)
|
||
ids_cpu = tok["input_ids"]
|
||
mask_cpu = tok["attention_mask"]
|
||
|
||
t0 = time.time()
|
||
pcm_full = graph.run(ids_cpu, mask_cpu, seed=int(seed), seconds=seconds)
|
||
inference_ms = (time.time() - t0) * 1000
|
||
|
||
# Slice and detach from the pinned-host buffer (the next run() would
|
||
# overwrite it — we want a stable copy to hand to the caller).
|
||
actual_samples = int(round(seconds * SAMPLE_RATE))
|
||
pcm = pcm_full[:actual_samples].copy()
|
||
|
||
return pcm, {
|
||
"inference_ms": inference_ms,
|
||
"graph_build_ms": graph_build_ms,
|
||
"realtime": (seconds * 1000.0) / inference_ms if inference_ms > 0 else 0.0,
|
||
"seed": int(seed),
|
||
"T_lat": T_lat,
|
||
"samples": actual_samples,
|
||
}
|
||
|
||
|
||
# ─── Main (mirrors sa3_trt.main() but routes through FullPipelineGraph) ─
|
||
def main():
|
||
# Update module-level engine paths if --models-dir is overridden. We keep
|
||
# canon's globals authoritative so we go through its path-fixup logic.
|
||
ap = _HelpfulParser(
|
||
description="SA3 text-to-audio via TensorRT — full-pipeline CUDA-graph variant",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog=(
|
||
"modes\n"
|
||
" text-to-audio --prompt P (mega-graph fast path)\n"
|
||
" audio-to-audio --prompt P --init-audio IN.wav (falls back to eager)\n"
|
||
" inpainting --prompt P --init-audio IN.wav --inpaint-range A,B (eager)\n"
|
||
" negative CFG --prompt P --cfg N --negative-prompt P_NEG (eager)\n"
|
||
),
|
||
)
|
||
ap.add_argument("--prompt", default=None)
|
||
ap.add_argument("--negative-prompt", default=None)
|
||
ap.add_argument("--init-audio", default=None)
|
||
ap.add_argument("--inpaint-range", default=None)
|
||
ap.add_argument("--dit", choices=list(DIT_CHOICES.keys()), default=None)
|
||
ap.add_argument("--decoder", choices=list(DECODER_PATHS.keys()), default=None)
|
||
ap.add_argument("--dec-precision", choices=list(canon.DECODER_TIERS), default="canonical",
|
||
help="Decoder/encoder quantization tier (orthogonal to --precision, which is the "
|
||
"DiT). 'canonical' (bf16, default); 'fp8' (~1.14x, near-transparent); "
|
||
"'fp8_fast' (SAME-S only; fp8 attention projections too, ~1.22x, lossier). "
|
||
"Engine auto-downloads from HF. (For max-fidelity decode use --precision fp32.)")
|
||
ap.add_argument("--precision", choices=list(canon.PRECISIONS) + list(canon._PRECISION_ALIAS), default=None,
|
||
help="DiT engine precision. Default is 'fp16' for every model: "
|
||
"canonical FP16 trunk + FP32 RMSNorm/RoPE islands + FMHA-fused FP16 "
|
||
"attention core, fp32-accurate at every length. 'fp8' (medium only) "
|
||
"is the max-speed clean tier — fp8 linears + bf16 fused FMHA + baked "
|
||
"fp32 RoPE, ~1.3x faster than fp16 at every length and clean at "
|
||
"long sequence (L<=4096). 'fp32' = bit-equiv PyTorch eager, slower. "
|
||
"('fp16mixed' and the retired medium-only 'bf16' are accepted as "
|
||
"deprecated aliases for 'fp16'.) Auto-downloads from HF.")
|
||
ap.add_argument("--models-dir", default=str(canon.MODELS_DIR))
|
||
ap.add_argument("--seconds", type=float, default=30.0)
|
||
ap.add_argument("--steps", type=int, default=8)
|
||
ap.add_argument("--seed", type=int, default=None)
|
||
ap.add_argument("--init-noise-level", type=float, default=1.0)
|
||
ap.add_argument("--cfg", type=float, default=1.0)
|
||
ap.add_argument("--apg", type=float, default=1.0)
|
||
ap.add_argument("--free-models", action=argparse.BooleanOptionalAction, default=False)
|
||
ap.add_argument("--pinned-copy", action=argparse.BooleanOptionalAction, default=True)
|
||
ap.add_argument("--quiet", action="store_true")
|
||
ap.add_argument("--out", default="out.wav")
|
||
ap.add_argument("--chunking", action=argparse.BooleanOptionalAction, default=True,
|
||
help="SAME-L only. On (default): decode in 256-latent windows and encode in "
|
||
"64, using the engines' LOW optimization profile -- 509 MB of decoder "
|
||
"scratch instead of 8143, because TRT commits scratch from the profile "
|
||
"ceiling before any shape is bound. Off: one single-shot call on the WIDE "
|
||
"profile, ~10-20%% faster above L=256 and ~5.4 GB more resident. Chunking "
|
||
"is a MEMORY mechanism, not a speed one: it only wins on time at exactly "
|
||
"L=256, where the whole request is one window.")
|
||
ap.add_argument("--limiter-ceiling", type=float, default=None,
|
||
help="Peak ceiling the decoder's limiter holds, linear (0.977 = -0.2021 dBFS "
|
||
"everywhere by default). ⚠ The SHIPPED engines BAKE this so their IO "
|
||
"stays ('latent','pcm') and they drop in for the pre-limiter engines; "
|
||
"passing this flag against them is an error. Rebuild with "
|
||
"build_samel_chunkable.py --ceiling-input for a settable ceiling, or use "
|
||
"--dec-precision legacy for the pre-limiter hard-clip behaviour.")
|
||
ap.add_argument("--mega-graph", action=argparse.BooleanOptionalAction, default=True,
|
||
help="Capture the entire pipeline in one CUDA graph (T5+DiT+decoder+narrow+DtoH). "
|
||
"On by default. Falls back to eager path for cfg≠1.0, inpaint, or audio-to-audio.")
|
||
args = ap.parse_args()
|
||
if args.steps < 1:
|
||
ap.error(f"--steps must be ≥ 1 (got {args.steps})")
|
||
if args.precision in canon._PRECISION_ALIAS: # retired token → canonical, with a note
|
||
_canon_p = canon._PRECISION_ALIAS[args.precision]
|
||
print(f" note: --precision {args.precision} is retired → using {_canon_p}", file=sys.stderr)
|
||
args.precision = _canon_p
|
||
|
||
# Mute display in quiet mode — match sa3_trt's behavior.
|
||
if args.quiet:
|
||
canon._USE_COLOR = False
|
||
# The imports above are by name; we have to re-bind the names that
|
||
# other code calls.
|
||
def _noop_stage(idx_total, label, ms=None): pass
|
||
def _noop_sub(text): pass
|
||
def _noop_vram(label): return 0
|
||
canon.stage = _noop_stage
|
||
canon.sub = _noop_sub
|
||
canon._stage_vram = _noop_vram
|
||
|
||
# Engine dir override — fix up canon's globals if user passed --models-dir.
|
||
if args.models_dir != str(canon.MODELS_DIR):
|
||
new_root = Path(args.models_dir).resolve()
|
||
new_arch_dir = new_root / ARCH
|
||
canon.T5GEMMA_PATH = new_arch_dir / "t5gemma" / "t5gemma_fp16.trt"
|
||
for kk in DIT_CHOICES:
|
||
DIT_CHOICES[kk]["engine"] = new_arch_dir / DIT_CHOICES[kk]["engine"].relative_to(canon.ARCH_DIR)
|
||
for kk in DECODER_PATHS:
|
||
DECODER_PATHS[kk] = new_arch_dir / DECODER_PATHS[kk].relative_to(canon.ARCH_DIR)
|
||
for kk in ENCODER_PATHS:
|
||
ENCODER_PATHS[kk] = new_arch_dir / ENCODER_PATHS[kk].relative_to(canon.ARCH_DIR)
|
||
canon.MODELS_DIR = new_root
|
||
canon.ARCH_DIR = new_arch_dir
|
||
|
||
args = prompt_user_if_missing(args)
|
||
if args.prompt is None:
|
||
args.prompt = input("Prompt: ").strip()
|
||
|
||
# T_lat
|
||
T_lat = max(1, math.ceil(args.seconds * SAMPLE_RATE / SAMPLES_PER_LATENT))
|
||
target_dur = T_lat * SAMPLES_PER_LATENT / SAMPLE_RATE
|
||
|
||
DIT_MIN_L, DIT_MAX_L = 1, 4096
|
||
if T_lat < DECODER_MIN_L:
|
||
min_s = DECODER_MIN_L * SAMPLES_PER_LATENT / SAMPLE_RATE
|
||
sys.exit(f"error: --seconds {args.seconds} gives T_lat={T_lat}, below the "
|
||
f"{args.decoder} decoder's minimum of {DECODER_MIN_L} latents "
|
||
f"(~{min_s:.2f}s). The DiT would accept it, but the decoder's "
|
||
f"optimization profile starts at {DECODER_MIN_L}; it would refuse the "
|
||
f"shape and the render would come back empty rather than erroring. "
|
||
f"Use --seconds {min_s:.2f} or more.")
|
||
if T_lat < DIT_MIN_L or T_lat > DIT_MAX_L:
|
||
max_s = DIT_MAX_L * SAMPLES_PER_LATENT / SAMPLE_RATE
|
||
why = ("the fp8 baked RoPE table (sized to L=4096; a longer render needs a re-bake — "
|
||
"build/make_rope_baked_onnx.py --max-t)"
|
||
if args.precision == "fp8" else
|
||
"the DiT engine's L=1..4096 profile") + " and the SAME-L decoder's 4096 cap"
|
||
sys.exit(f"error: T_lat={T_lat} (= {args.seconds}s) out of [{DIT_MIN_L}, {DIT_MAX_L}] "
|
||
f"(~{max_s:.0f}s max), bounded by {why}.")
|
||
|
||
# Inpaint validation (unchanged from sa3_trt)
|
||
inpaint_range = None
|
||
inp_start_sec = inp_end_sec = None
|
||
if args.inpaint_range is not None:
|
||
if args.init_audio is None:
|
||
sys.exit("error: --inpaint-range requires --init-audio")
|
||
s_str, e_str = args.inpaint_range.split(",")
|
||
inp_start_sec = float(s_str.strip()); inp_end_sec = float(e_str.strip())
|
||
if not (0 <= inp_start_sec < inp_end_sec <= args.seconds):
|
||
sys.exit(f"error: invalid inpaint range")
|
||
inp_start_lat = max(0, int(round(inp_start_sec * SAMPLE_RATE / SAMPLES_PER_LATENT)))
|
||
inp_end_lat = min(T_lat, int(round(inp_end_sec * SAMPLE_RATE / SAMPLES_PER_LATENT)))
|
||
inpaint_range = (inp_start_lat, inp_end_lat)
|
||
|
||
sigma_max = float(args.init_noise_level)
|
||
mode = ("inpaint" if inpaint_range else
|
||
"audio-to-audio" if args.init_audio else "text-to-audio")
|
||
|
||
# Can we use the mega-graph fast path?
|
||
use_mega = (args.mega_graph and args.cfg == 1.0 and
|
||
inpaint_range is None and args.init_audio is None and
|
||
sigma_max == 1.0)
|
||
# sigma_max==1 → noise = pure_noise (no init_latents mixing); plus the
|
||
# negative cases above. Anything else → eager.
|
||
|
||
# Banner
|
||
_STAGE_VRAMS.clear()
|
||
print()
|
||
banner(f"SA3 → TRT {mode} {'[mega-graph]' if use_mega else '[eager]'}")
|
||
k = lambda s: dim(f"{s:>10}")
|
||
v = lambda s, w=10: f"{s:<{w}}"
|
||
print(f" {k('prompt')} {bold(repr(args.prompt))}")
|
||
if args.negative_prompt:
|
||
suffix = "" if args.cfg != 1.0 else dim(" (ignored: --cfg=1.0)")
|
||
print(f" {k('neg prompt')} {bold(repr(args.negative_prompt))}{suffix}")
|
||
_dtier = f" {k('dec-tier')} {magenta(v(args.dec_precision))}" if args.dec_precision != "canonical" else ""
|
||
print(f" {k('dit')} {magenta(v(args.dit))} {k('decoder')} {magenta(v(args.decoder))}{_dtier} {k('precision')} {v(args.precision)}")
|
||
print(f" {k('σmax')} {bold(f'{sigma_max:.2f}')}")
|
||
print(f" {k('seconds')} {v(f'{args.seconds}s')} {k('steps')} {v(args.steps)} {k('seed')} {args.seed}")
|
||
print(f" {k('cfg')} {v(args.cfg)} {k('mega-graph')} {v('on' if use_mega else 'off (fallback)')}")
|
||
print(f" {k('T_lat')} {T_lat} {dim(f'({target_dur:.2f}s → trimmed to {args.seconds}s)')}")
|
||
print()
|
||
|
||
# Everything below this point allocates. If the mega-graph cannot serve this
|
||
# configuration, hand off NOW — before the lazy-download, the heavy imports and
|
||
# above all the engine load — so only one set of engines is ever resident.
|
||
if not use_mega:
|
||
_delegate_to_eager(args)
|
||
return
|
||
|
||
# Lazy-download
|
||
needed = canon.get_engine_files(args.dit, args.decoder, args.precision,
|
||
with_encoder=bool(args.init_audio), dec_tier=args.dec_precision)
|
||
_ensure_files(needed)
|
||
|
||
# Heavy imports
|
||
sub(dim("Loading..."))
|
||
t0 = time.time()
|
||
_import_heavy()
|
||
sub(f"{dim('heavy imports')} {(time.time()-t0)*1000:.0f} ms")
|
||
|
||
_init_nvml()
|
||
|
||
# Tokenizer + dist_shift
|
||
t0 = time.time()
|
||
import runtime as rt
|
||
rt.MODELS_DIR = str(canon.MODELS_DIR)
|
||
rt.ARCH_DIR = str(canon.ARCH_DIR)
|
||
state = rt.load()
|
||
tokenizer = state["tokenizer"]
|
||
dist_shift = state["dist_shift"]
|
||
sub(f"{dim('tokenizer + dist-shift')} {(time.time()-t0)*1000:.0f} ms")
|
||
|
||
# Engine load
|
||
import concurrent.futures
|
||
engine_specs = {
|
||
"t5": canon.T5GEMMA_PATH,
|
||
"dit": canon.get_dit_engine_path(args.dit, args.precision),
|
||
"dec": canon.resolve_decoder_engine(args.decoder, args.precision, args.dec_precision),
|
||
}
|
||
if args.init_audio:
|
||
engine_specs["enc"] = canon.resolve_encoder_engine(args.decoder, args.dec_precision)
|
||
t0 = time.time()
|
||
# The decoder/encoder carry two profiles: 0 = low ceiling (chunked, minimum scratch),
|
||
# 1 = wide (single-shot). Selecting it at context creation is what makes the saving real --
|
||
# a context created on the default profile reserves the MAX across profiles.
|
||
_band = 0 if args.chunking else 1
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=len(engine_specs)) as ex:
|
||
# user_managed=True: allocate NO private scratch. share_scratch() below hands every
|
||
# engine one buffer. Letting each runner commit its own first and freeing it afterwards
|
||
# leaves both resident at peak -- measured 11924 MiB where this path takes 8946.
|
||
futs = {name: ex.submit(TRTRunner, path, None, True,
|
||
_band if name in ("dec", "enc") else 0)
|
||
for name, path in engine_specs.items()}
|
||
runners = {name: fut.result() for name, fut in futs.items()}
|
||
# One scratch buffer for every engine: they run sequentially on a single stream, so holding
|
||
# a private allocation each costs several GB for nothing.
|
||
_shared_scratch = canon.share_scratch(runners) # keep alive for the process lifetime
|
||
sub(f"{dim(f'engine load (parallel, {len(runners)} engines)')} {(time.time()-t0)*1000:.0f} ms")
|
||
|
||
# ── Warmup + graph build ──
|
||
torch = canon.torch
|
||
WARMUP_PASSES = 3
|
||
t0 = time.time()
|
||
dit = DiTRunner(runners["dit"])
|
||
|
||
if use_mega:
|
||
# ── Mega-graph fast path ──
|
||
sigmas = build_pingpong_schedule(args.steps, sigma_max=sigma_max,
|
||
dist_shift=dist_shift, latent_len=T_lat)
|
||
requested_samples = int(round(args.seconds * SAMPLE_RATE))
|
||
# The graph copies the full pinned-host buffer up to requested_samples;
|
||
# the WAV save reads the same buffer.
|
||
# Windowed decode when the request exceeds the decoder's bound ceiling. profile_bounds
|
||
# reports the ACTIVE profile, so --no-chunking (wide band) reports 4096 and takes the
|
||
# single-shot path while the default (low band) reports 256 and windows.
|
||
_, _dec_hi = canon.profile_bounds(runners["dec"], "latent")
|
||
mega = FullPipelineGraph(runners["t5"], dit, runners["dec"],
|
||
T_lat, args.steps, requested_samples,
|
||
dec_chunk_lat=(_dec_hi if T_lat > _dec_hi else 0))
|
||
mega.build(sigmas, args.seconds, sigma_max)
|
||
if args.limiter_ceiling is not None:
|
||
mega.set_limiter_ceiling(args.limiter_ceiling)
|
||
torch.cuda.synchronize()
|
||
sub(f"{dim(f'warmup + capture')} {(time.time()-t0)*1000:.0f} ms")
|
||
print()
|
||
|
||
# ── Inference ──
|
||
t_wall_start = time.time()
|
||
# Tokenize (CPU only — outside the graph).
|
||
tok = tokenizer(args.prompt, return_tensors="pt", max_length=T5_MAX_LEN,
|
||
padding="max_length", truncation=True)
|
||
ids_cpu = tok["input_ids"] # (1, 256) int64
|
||
mask_cpu = tok["attention_mask"]
|
||
# One replay does everything (T5 + DiT loop + dec + narrow + DtoH).
|
||
pcm = mega.run(ids_cpu, mask_cpu, seed=int(args.seed), seconds=args.seconds)
|
||
t_inference = time.time() - t_wall_start
|
||
stage("[1/1]", f"Full pipeline (mega-graph: T5+DiT+dec+DtoH)", t_inference * 1000)
|
||
_stage_vram("Mega-graph inference")
|
||
|
||
# WAV save (excluded from inference timing — disk I/O).
|
||
out_path = Path(args.out)
|
||
if not out_path.is_absolute():
|
||
out_path = canon.OUTPUT_DIR / out_path
|
||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||
t0w = time.time()
|
||
with wave.open(str(out_path), "wb") as w:
|
||
w.setnchannels(2); w.setsampwidth(2); w.setframerate(SAMPLE_RATE)
|
||
w.writeframes(pcm.tobytes())
|
||
t_disk = (time.time() - t0w) * 1000
|
||
try:
|
||
out_display = out_path.relative_to(canon.REPO).as_posix()
|
||
except (ValueError, AttributeError):
|
||
out_display = str(out_path)
|
||
stage("[2/2]", f"WAV → {out_display}", t_disk)
|
||
|
||
t_full = time.time() - t_wall_start
|
||
audio_dur = pcm.shape[0] / SAMPLE_RATE
|
||
print()
|
||
rule()
|
||
def _fmt_t(s):
|
||
return f"{s*1000:.0f} ms" if s < 1 else f"{s:.2f} s"
|
||
col1 = f"Inference {bold(_fmt_t(t_inference))}"
|
||
col2 = f"Inference + Saving the WAV {bold(_fmt_t(t_full))}"
|
||
col1_rt = dim(f"{audio_dur/t_inference:.0f}× realtime")
|
||
col2_rt = dim(f"{audio_dur/t_full:.0f}× realtime")
|
||
import re
|
||
_ANSI = re.compile(r'\x1b\[[0-9;]*m')
|
||
def _vlen(s): return len(_ANSI.sub('', s))
|
||
gap = max(_vlen(col1), _vlen(col1_rt)) + 4
|
||
indent = " " + " " * len("done ")
|
||
print(f" {bold(green('done'))} {col1}{' ' * (gap - _vlen(col1))}{col2}")
|
||
print(f"{indent}{col1_rt}{' ' * (gap - _vlen(col1_rt))}{col2_rt}")
|
||
rule()
|
||
print(f" {bold(green('▸ saved'))} {bold(cyan(out_display))} {dim(f'({out_path.resolve()})')}")
|
||
return
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|