"""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 //dit_.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()