diff --git a/docs/api_conformance_audit_2026-07-11.md b/docs/api_conformance_audit_2026-07-11.md new file mode 100644 index 000000000..ba00eba9b --- /dev/null +++ b/docs/api_conformance_audit_2026-07-11.md @@ -0,0 +1,415 @@ +# API Conformance Audit — CUDA Python / TensorRT / ONNX (2026-07-11) + +> **Untracked working document — never stage or commit** (same convention as +> `docs/perf_bestpractices_audit_2026-07-10.md`). +> +> Audit of the codebase's usage of the CUDA Python, TensorRT, and ONNX-family APIs +> against official documentation and the *installed* package versions. Report only — +> no code was changed this phase. Branch `refactor/py-code-review-remediation`, +> HEAD `d5efd4d`. + +--- + +## 1. Method and ground truth + +Two evidence sources, used together: + +1. **Online docs (browser, pulled this session).** NVIDIA's "latest" TensorRT docs + are version **11.1.0** — one major version ahead of the installed 10.16. They are + used as the authoritative statement of semantic contracts (which are stable across + 10.x→11.x) and as the *direction of travel* for deprecations/removals. +2. **Installed-package introspection (venv, authoritative for pinned versions).** + Enum membership, method presence, signatures, and docstrings were introspected from + the actual installed wheels (`venv/Scripts/python.exe`). Where 11.1 docs and 10.16 + bindings disagree, the introspection result is treated as ground truth for today's + behavior and the 11.1 docs as ground truth for what an upgrade will break. + +### Pinned versions → doc sources + +| Package | Installed | Doc source used | +|---|---|---| +| tensorrt_cu12 | 10.16.1.11 | docs.nvidia.com TensorRT 11.1.0 (`/latest/_static/python-api/...`, `/latest/_static/c-api/...`) + venv introspection | +| cuda-python | 12.9.7 (dist 12.9.0) | venv introspection (docstrings/warnings); nvidia.github.io/cuda-python | +| onnx | 1.19.1 | venv introspection (`onnx.save_model` signature) | +| torch | 2.8.0+cu128 | venv introspection (`torch.onnx.export` signature + docstring) | +| polygraphy | 0.49.26 | venv introspection (`engine_from_bytes` signature) | +| nvidia-modelopt | 0.43.0 | venv introspection (`modelopt.onnx.quantization.quantize` signature) | +| onnxruntime-gpu | 1.24.4 | (used only indirectly by modelopt calibration — no direct call sites) | + +### Key doc pages cited below + +- **[TRT-CTX-PY]** `https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/ExecutionContext.html` +- **[TRT-CTX-CPP]** `https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/c-api/classnvinfer1_1_1_i_execution_context.html` +- **[TRT-BCFG]** `https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/BuilderConfig.html` +- **[TRT-ONNX]** `https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/parsers/Onnx/pyOnnx.html` +- **[TRT-REFIT]** `https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/Refitter.html` +- **[INTROSPECT]** venv introspection of installed packages, this session. + +### Installed TRT 10.16.1.11 enum/API ground truth [INTROSPECT] + +- `trt.TacticSource` members: `{CUBLAS, CUBLAS_LT, CUDNN, EDGE_MASK_CONVOLUTIONS, JIT_CONVOLUTIONS}` + (11.1 docs list **only** `EDGE_MASK_CONVOLUTIONS` and `JIT_CONVOLUTIONS`, both "Enabled by default" [TRT-BCFG]). +- `trt.BuilderFlag` includes `FP16, FP8, TF32, REFIT, SPARSE_WEIGHTS`; **no `STRONGLY_TYPED`** (removed in 10.12). +- `trt.PreviewFeature` members: `{PROFILE_SHARING_0806, ALIASED_PLUGIN_IO_10_03, RUNTIME_ACTIVATION_RESIZE_10_10, MULTIDEVICE_RUNTIME_10_16}`. +- `trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH` still present (deprecated). +- `trt.OnnxParserFlag.NATIVE_INSTANCENORM` present. +- `IExecutionContext.set_device_memory(memory, size)` **exists** (two-arg, V2 semantics: + docstring requires "256-byte aligned device memory" and size ≥ `get_device_memory_size_v2`); + the legacy one-arg `device_memory` property setter also still exists. There is **no** + Python attribute named `set_device_memory_v2` — the V2 API surfaces under the plain name. +- `ICudaEngine.device_memory_size_v2` exists. +- `Refitter` has `get_all`, `set_weights`, `set_named_weights`, `refit_cuda_engine`, + `refit_cuda_engine_async`; `trt.WeightsRole` has 6 members. +- `trt.nptype(trt.DataType.FP8)` raises `TypeError` ("Could not resolve TensorRT datatype + to an equivalent numpy datatype"). +- Accessing deprecated enum members emits **no** runtime `DeprecationWarning` — deprecation + is silent at runtime; only the docs/headers say so. + +--- + +## 2. Verdict summary + +**No CRITICAL findings.** Every documented contract with a runtime consequence +(tensor-address alignment, input-lifetime around `execute_async_v3`, timing-cache +lifetime, graph-capture rules, refit-vs-enqueued-work) was checked and the code +conforms — in several cases with explicit, comment-documented mitigations. The +findings below are deprecation debt (breaks on a future TRT 11 upgrade), missing +defensive checks, and minor idiom/perf notes. + +| ID | Severity | File:line | Summary | +|---|---|---|---| +| H1 | HIGH | `src/streamdiffusion/acceleration/tensorrt/utilities.py:388-401` | `TacticSource.CUBLAS/CUBLAS_LT` removed in TRT 11 → whole tactic-scoping block silently no-ops on upgrade | +| H2 | HIGH | `utilities.py:1002` | Deprecated `context.device_memory` property setter (superseded since TRT 10.1); path is currently dead code | +| M1 | MEDIUM | `preprocessing/processors/trt_base.py:135,144,195,200,235` | `set_input_shape` / `set_tensor_address` bool returns unchecked → silent failure risk | +| M2 | MEDIUM | `utilities.py:544-570` | Zero-copy bind path has no 256-byte alignment guard (contract-safe under all default configs — proven in §5 — but unguarded against config drift) | +| M3 | MEDIUM | `src/streamdiffusion/tools/compile_depth_anything_tensorrt.py:129` | Deprecated `NetworkDefinitionCreationFlag.EXPLICIT_BATCH` (sibling RAFT tool already fixed) | +| L1 | LOW | `utilities.py:1525` | `do_constant_folding=True` — "Deprecated option" in torch 2.8 | +| L2 | LOW | `utilities.py:1519-1530` | Legacy TorchScript ONNX exporter (`dynamo=False`); torch 2.8 recommends `dynamo=True` but keeps `False` as default | +| L3 | LOW | `utilities.py:1579` | External-data detection by `.pb` suffix scan is heuristic (safe in current flow, brittle to reordering) | +| I1–I5 | INFO | various | Redundant flag, CPU-alloc idiom, refit sync-scope note, legacy-import fallback, ctypes cudart | + +Conformance confirmations (checked, no finding) are listed in §4. + +--- + +## 3. Findings + +### H1 — `TacticSource.CUBLAS/CUBLAS_LT` are gone in TRT 11; the SM_120 tactic-scoping block will silently vanish on upgrade + +**File:** [utilities.py:388-401](../src/streamdiffusion/acceleration/tensorrt/utilities.py) + +```python +if gpu_profile.compute_capability >= (12, 0): + try: + sources = ( + (1 << int(trt.TacticSource.CUBLAS)) + | (1 << int(trt.TacticSource.CUBLAS_LT)) + | (1 << int(trt.TacticSource.JIT_CONVOLUTIONS)) + | (1 << int(trt.TacticSource.EDGE_MASK_CONVOLUTIONS)) + ) + config.set_tactic_sources(sources) + ... + except (AttributeError, TypeError) as e: + logger.debug(f"[TRT Config] set_tactic_sources not available: {e}") +``` + +**Docs:** [TRT-BCFG] documents `tensorrt.TacticSource` with exactly two members — +`EDGE_MASK_CONVOLUTIONS` and `JIT_CONVOLUTIONS`, each annotated "Enabled by default". +`CUBLAS`, `CUBLAS_LT`, and `CUDNN` no longer exist in 11.1. On installed 10.16 all five +members still exist [INTROSPECT], so the block works today. + +**Consequence on TRT 11 upgrade:** the first expression, `trt.TacticSource.CUBLAS`, +raises `AttributeError`; the `except` swallows it at DEBUG level. Net effect: no call to +`set_tactic_sources` at all. Because the code's *primary intent* on SM_120 is to **exclude +CUDNN** (per the block comment) and CUDNN tactics no longer exist in TRT 11, the silent +no-op is behaviorally harmless there — but the comment's rationale becomes stale, the +INFO log ("CUDNN excluded for SM_120+") never fires, and cuBLAS/cuBLAS_LT scoping intent +is silently dropped rather than consciously retired. + +**Remediation:** gate on `hasattr(trt.TacticSource, "CUBLAS")` instead of try/except so +the TRT 11 path is an explicit, logged branch (e.g. "TRT ≥11: default tactic sources +already exclude cuDNN/cuBLAS — nothing to scope"), and update the comment. No behavior +change on 10.16. + +--- + +### H2 — deprecated `context.device_memory` property setter (dead path today) + +**File:** [utilities.py:999-1002](../src/streamdiffusion/acceleration/tensorrt/utilities.py) + +```python +def activate(self, reuse_device_memory=None): + if reuse_device_memory: + ... + self.context.device_memory = reuse_device_memory +``` + +**Docs:** [TRT-CTX-CPP] `setDeviceMemory`: *"Deprecated in TensorRT 10.1. Superseded by +setDeviceMemoryV2()."* — with the additional caveat *"Weight streaming related scratch +memory will be allocated by TensorRT if the memory is set by this API"* (i.e. the legacy +setter's size assumption excludes weight-streaming scratch). The installed 10.16 Python +binding already exposes the successor as two-arg `set_device_memory(memory, size)` whose +docstring requires *"256-byte aligned device memory"* and *"size ... at least as large as +CudaEngine.get_device_memory_size_v2"* [INTROSPECT]. + +**Liveness:** grep of all 15 `activate()` call sites (runtime_engines, preprocessing +processors) shows **none pass `reuse_device_memory`** — the deprecated line is +unreachable in practice. `activate()` therefore always takes the +`create_execution_context()` branch. + +**Remediation (pick one):** +1. Delete the `reuse_device_memory` parameter and branch (dead code); or +2. Migrate to `self.context.set_device_memory(ptr, self.engine.device_memory_size_v2)` + if the memory-reuse feature is intended to come back. + +Either way the deprecated one-arg setter should not survive to a TRT 11 upgrade. + +--- + +### M1 — `trt_base.TensorRTEngine` ignores `set_input_shape` / `set_tensor_address` return values + +**File:** [trt_base.py:135](../src/streamdiffusion/preprocessing/processors/trt_base.py) +(also :144, :195, :200 for `set_input_shape`; :235 for `set_tensor_address`) + +```python +self.context.set_input_shape(name, input_shape) # :135, :144, :195, :200 +... +self.context.set_tensor_address(name, tensor.data_ptr()) # :235 +``` + +**Docs:** [TRT-CTX-PY] `set_input_shape(self, name, shape) -> bool` — returns whether the +shape was set successfully (False for out-of-profile shapes, wrong rank, unknown name); +`set_tensor_address` likewise returns `bool`. A `False` return is TRT's *only* signal — +no exception is raised. + +**Consequence:** an out-of-profile input shape (e.g. a preprocessing resolution outside +the engine's optimization profile) is silently accepted by Python; the failure then +surfaces downstream as an `execute_async_v3` failure or — worse — garbage output with +stale shapes. The core `Engine` class in utilities.py already does this correctly +(bool check at :1055; raises on `set_tensor_address` failure at :1209-1210); the second, +independent `TensorRTEngine` in trt_base.py (used by depth/pose/realesrgan/temporal-net +processors, i.e. live per-frame paths) does not. + +**Remediation:** mirror the utilities.py pattern — raise `RuntimeError` with tensor name +and shape when either call returns `False`. + +--- + +### M2 — zero-copy bind path (`_staging_action`) has no 256-byte alignment guard + +**File:** [utilities.py:544-570](../src/streamdiffusion/acceleration/tensorrt/utilities.py) +(decision helper), bind at :1194/:1208. + +```python +if name not in zero_copy_names or not is_contiguous or not dtype_match: + return "copy" +``` + +**Docs:** [TRT-CTX-CPP] `setTensorAddress`: *"The pointer must have at least 256-byte +alignment."* The Sub-phase 5.6 zero-copy path binds caller-owned kvo/fio per-layer +*views* (bucket base + slot offset) directly to TRT, so view offsets must preserve +256-byte alignment. + +**Current status: conformant.** §5 proves arithmetically that every per-layer view +offset is a multiple of 256 bytes under all default configurations (kvo unconditionally; +fio for every FI-eligible hidden dim under `max_fi_up_blocks ≤ 3`). No violation exists +today. + +**Why still a finding:** the guarantee is *implicit* — it rests on (a) fp16 dtype, +(b) hidden dims being multiples of 64, (c) the default FI eligibility mask, and +(d) PyTorch's allocator alignment. None of these is asserted anywhere. A config drift +(`max_fi_up_blocks=4` exposing H=320 layers, combined with an odd +`maxframes·batch·seq` product — see §5) or an atypical model would silently violate the +documented contract, and TRT does not reliably diagnose misaligned addresses (undefined +behavior territory). + +**Remediation:** add `cur_ptr % 256 == 0` to `_staging_action`'s eligibility test +(misaligned → fall back to `"copy"`, which is always safe). One integer modulo per input +per frame on the CPU; keeps the pure-function unit-testability. Optionally log once at +WARNING when the fallback triggers. + +--- + +### M3 — deprecated `EXPLICIT_BATCH` flag in the Depth-Anything compile tool + +**File:** [compile_depth_anything_tensorrt.py:129](../src/streamdiffusion/tools/compile_depth_anything_tensorrt.py) + +```python +network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) +``` + +**Docs:** `EXPLICIT_BATCH` is deprecated since TRT 10.0 (networks are always +explicit-batch); the flag still exists in installed 10.16 [INTROSPECT] and is ignored. +The sibling tool already carries the fix and the rationale: +[compile_raft_tensorrt.py:152](../src/streamdiffusion/tools/compile_raft_tensorrt.py) — +`network = builder.create_network() # EXPLICIT_BATCH deprecated/ignored in TRT 10.x`. + +**Remediation:** same one-line change as the RAFT tool. Low blast radius (offline tool), +but it breaks outright when TRT removes the enum member. + +--- + +### L1 — `do_constant_folding=True` is a deprecated `torch.onnx.export` option + +**File:** [utilities.py:1525](../src/streamdiffusion/acceleration/tensorrt/utilities.py) + +**Docs [INTROSPECT, torch 2.8.0 docstring]:** *"do_constant_folding: Deprecated +option."* (alongside `training`, `operator_export_type`, `custom_opsets`). Harmless +today; drop it whenever the export call is next touched. + +### L2 — legacy TorchScript exporter (`dynamo=False`) + +**File:** [utilities.py:1519-1530](../src/streamdiffusion/acceleration/tensorrt/utilities.py) + +**Docs [INTROSPECT, torch 2.8.0]:** `dynamo=False` **is still the default** in 2.8 and is +not deprecated; the docstring says `dynamo=True` "is the recommended way to export models +to ONNX." No action required now — but the custom attention processors / export wrappers +were built against TS-exporter tracing semantics, so treat any future move to +`dynamo=True` as its own validation project, not a flag flip. Being explicit +(`dynamo=False`) is *better* than relying on the default, since the default is the likely +thing to change in a future torch. + +### L3 — external-data detection by `.pb` suffix scan + +**File:** [utilities.py:1579](../src/streamdiffusion/acceleration/tensorrt/utilities.py) + +```python +external_data_files = [f for f in os.listdir(onnx_dir) if f.endswith(".pb")] +``` + +Safe in the current flow **only because** `export_onnx` consolidates torch's external +tensor files (named `onnx__*`, no `.pb` suffix) into a single `weights.pb` and deletes +the originals (:1545-1562) before `optimize_onnx` runs. If a >2GB model ever reached +`optimize_onnx` straight from `torch.onnx.export` (which writes extension-less external +files), the scan would report "no external data" and `onnx.load(onnx_path)` at :1611 +would still succeed (external data loads by default) but the subsequent plain +`onnx.save` at :1614 would fail on the 2GB protobuf limit. Prefer +`onnx.external_data_helper`-based detection (check `TensorProto.data_location == +EXTERNAL` on the loaded model) or document the ordering invariant at the scan site. + +### INFO notes + +- **I1 — redundant parser flag.** `parser.set_flag(trt.OnnxParserFlag.NATIVE_INSTANCENORM)` + at utilities.py:769/:893 — [TRT-ONNX]: *"This flag is ON by default"* (and required for + version-/hardware-compatible engines). Not deprecated; setting it is a no-op. Keep or + drop, cosmetic either way. +- **I2 — CPU-alloc-then-copy.** utilities.py:1066 + `torch.empty(tuple(shape), dtype=torch_dtype).to(device=device)` allocates on CPU then + copies; trt_base.py:165 already uses the direct `device=` form (with a comment saying + why). One-time cost at buffer (re)allocation only — perf-cosmetic. +- **I3 — refit sync scope.** [TRT-REFIT]: *"The behavior is undefined if the engine has + pending enqueued work."* The refit path synchronizes `torch.cuda.current_stream()` at + utilities.py:723 before `refit_cuda_engine()` — correct for the load-time-only refit + flow (comment there documents this as defensive). Note it syncs torch's current stream, + not the engine's polygraphy stream; if refit were ever made reachable mid-stream, sync + the engine stream too. +- **I4 — legacy cuda-python import fallback.** utilities.py:37-41 prefers + `from cuda.bindings import runtime as cudart` and falls back to `from cuda import cudart`. + Installed 12.9.7 emits on the legacy import: *"The cuda.cudart module is deprecated and + will be removed in a future release, please switch to use the cuda.bindings.runtime + module instead."* [INTROSPECT]. The preference order is already correct; the fallback + only fires on cuda-python < 12.x installs. Fine as is. +- **I5 — ctypes cudart in `tools/cuda_l2_cache.py`** (`cudaDeviceSetLimit`, + `cudaStreamSetAttribute` via ctypes): bypasses cuda-python entirely; works but carries + its own struct-layout risk. Out of scope for remediation; noted for completeness. + +--- + +## 4. Conformance confirmations (checked, no finding) + +| Contract | Doc | Code | Verdict | +|---|---|---|---| +| `execute_async_v3` input lifetime (inputs must not be modified/freed before stream sync) | [TRT-CTX-PY] | Inputs staged via `copy_()` into persistent engine buffers *on the engine stream* (utilities.py:1158-1196); zero-copy names are persistent, address-stable caches by construction | ✔ | +| Default-stream warning (`enqueueV3` on default stream ⇒ extra device sync) | [TRT-CTX-CPP] | Core engine runs on polygraphy-owned non-default stream; trt_base engines create a dedicated `torch.cuda.Stream()` with pre/post event barriers + `record_stream()` | ✔ | +| `setTensorAddress` 256-byte alignment | [TRT-CTX-CPP] | Proven for all bound tensors: full buffers are torch base allocations (512-B aligned); kvo/fio view offsets ≡ 0 mod 256 under defaults (§5). Guard recommended (M2) | ✔ | +| CUDA graph capture: quiesce streams, ThreadLocal mode, instantiate signature | cuda-python 12.9 [INTROSPECT] | 3 warmup passes, engine-stream + legacy-stream drain before `cudaStreamBeginCapture(..., ThreadLocal)`; `cudaGraphInstantiate(graph, 0)` matches documented `(graph, unsigned long long flags)`; `CUASSERT` correctly unpacks `(err, result)` tuples; graph-launch failure falls back to `execute_async_v3` and re-captures | ✔ | +| Rebind-after-capture rule (addresses baked into graph) | [TRT-CTX-CPP] | `_staging_action` returns `bind_and_reset` on pointer change with live graph → `reset_cuda_graph()`; addresses only rebound when no graph instance exists (utilities.py:1198-1212) | ✔ | +| Timing cache lifetime ("must not be destroyed until after the engine is built") | [TRT-BCFG] | Local `trt_cache` created/loaded at utilities.py:807-817 stays in scope through `build_serialized_network` | ✔ | +| `builder_optimization_level` valid range 0–5 (default 3) | [TRT-BCFG] | Profile uses 3 or 4 | ✔ | +| `avg_timing_iterations` (default 1), `max_num_tactics` (default −1) | [TRT-BCFG] | 8/4 iterations, cap 64 — valid values, guarded by try/except for older TRT | ✔ | +| `PreviewFeature.RUNTIME_ACTIVATION_RESIZE_10_10` | [TRT-BCFG] + [INTROSPECT] | Present in installed 10.16 *and* still documented in 11.1 | ✔ | +| `BuilderFlag.STRONGLY_TYPED` removed in 10.12 | [INTROSPECT] | `hasattr(trt.BuilderFlag, "STRONGLY_TYPED")` guard at utilities.py:913-919 correctly skips on 10.16; FP8 path uses the network-creation flag instead | ✔ | +| `trt.nptype` has no FP8 mapping | [INTROSPECT] | `TypeError` → FP8 fallback at utilities.py:1043-1051 | ✔ | +| Refitter API (`get_all`, `set_weights`, `WeightsRole`, `refit_cuda_engine`) | [TRT-REFIT] | All still current in 11.1 docs, zero deprecation notes on the page; weight lifetime rule ("can be unset and released ... after refit_cuda_engine returns") satisfied — numpy arrays outlive the call | ✔ | +| `set_input_shape`/`set_tensor_address` return checks (core engine) | [TRT-CTX-PY] | utilities.py:1055 (bool check), :1209-1210 (raise) | ✔ (trt_base gap = M1) | +| `onnx.save_model` external-data kwargs | [INTROSPECT] onnx 1.19.1 | `save_as_external_data=True, all_tensors_to_one_file=True, location="weights.pb", convert_attribute=False` all match the installed signature | ✔ | +| `modelopt.onnx.quantization.quantize` kwargs | [INTROSPECT] 0.43.0 | `quantize_mode`, `use_external_data_format`, `high_precision_dtype` all valid; fp8_quantize.py additionally guards every kwarg via `inspect.signature` | ✔ | +| `polygraphy engine_from_bytes(serialized_engine, runtime=None)` | [INTROSPECT] 0.49.26 | Called with a single positional arg | ✔ | +| cuda-python module layout (`cuda.bindings.runtime` preferred) | [INTROSPECT] 12.9.7 | Preferred import first, legacy fallback second (utilities.py:37-41) | ✔ | + +--- + +## 5. Appendix — kvo/fio zero-copy alignment arithmetic (checkpoint #1) + +**Contract:** [TRT-CTX-CPP] `setTensorAddress`: *"The pointer must have at least 256-byte +alignment."* + +**What gets bound:** `Engine.infer` binds `buf.data_ptr()` for every feed-dict entry in +`zero_copy_names` (utilities.py:1180-1194). `zero_copy_names` is exactly the kvo + fio +input names (unet_engine.py:100/:108). Those tensors are the per-layer views created in +[models/utils.py](../src/streamdiffusion/acceleration/tensorrt/models/utils.py): + +- kvo (`create_kvo_cache`, :118-122): bucket `torch.zeros(L, 2, mf, B, S, H, dtype=fp16)`; + per-layer view = `bucket[slot]`. +- fio (`create_fi_cache`, :239-251): bucket `torch.zeros(L, mf, B, S, H, dtype=fp16)`; + per-layer view = `bucket[slot]`. + +(`L` = layers in bucket, `mf` = cache_maxframes, `B` = batch, `S` = sequence length, +`H` = attention hidden dim; fp16 ⇒ 2 bytes/element.) + +**Base pointers:** each bucket is a fresh `torch.zeros` CUDA allocation. PyTorch's CUDA +caching allocator returns block-granular base pointers (512-byte rounding), so +`bucket.data_ptr() % 256 == 0`. (This is allocator behavior, not a documented API +contract — one more reason for the M2 guard, which makes the whole argument +assumption-free at runtime.) + +**View offsets:** + +- kvo: `offset(slot) = slot · (2·mf·B·S·H elements) · 2 B = 4·slot·mf·B·S·H` bytes. + Every attention hidden dim in the supported UNets (SD1.5/SD-Turbo/SDXL: 320, 640, 1280) + is a multiple of 64, so `4·H ≡ 0 (mod 256)` ⇒ **every kvo view offset is a multiple of + 256 unconditionally** (independent of `mf`, `B`, `S`, `slot`). +- fio: `offset(slot) = 2·slot·mf·B·S·H` bytes ⇒ multiple of 256 iff `slot·mf·B·S·(H/64)` + is even. + - Default FI eligibility (`max_fi_up_blocks=2`, models/utils.py:127-164) admits only + mid-block + first two up-blocks ⇒ `H ∈ {1280}` (SD1.5: up_blocks[0] has no + attentions, up_blocks[1] is H=1280) or `H ∈ {1280, 640}` (SDXL). `H/64 ∈ {20, 10}` + is even ⇒ **aligned unconditionally**. + - `max_fi_up_blocks=3` adds H=640 layers (SD1.5) — still even ⇒ aligned. + - Only a non-default `max_fi_up_blocks=4` (which also contradicts the function's + documented "final up-block always excluded" rule) would admit H=320 (`H/64 = 5`, + odd). Then alignment requires `slot·mf·B·S` even. At the H=320 level + `S = (h/8)·(w/8)`, odd only when *both* `h/8` and `w/8` are odd (resolutions + ≡ 8 mod 16, e.g. 968×968) — so a violation additionally needs odd `mf`, odd `B`, + and an odd slot. Narrow, but reachable ⇒ finding M2. + +**Rebound caches:** `stream_parameter_updater.py:1089` replaces `kvo_cache[i]` with a +standalone tensor — a fresh base allocation (≥256-aligned) — and the pointer change is +detected by `_staging_action` → `bind_and_reset` → graph reset. Conformant. + +**Verdict: no alignment violation exists in any default configuration.** The M2 guard +converts this from an arithmetic argument into a runtime invariant. + +--- + +## 6. Ranked fix shortlist (candidate follow-up gated commits) + +1. **M1 — check `set_input_shape`/`set_tensor_address` returns in trt_base.py** + (raise like utilities.py does). Cheapest, closes a real silent-failure hole on live + per-frame paths. +2. **M2 — add `cur_ptr % 256 == 0` to `_staging_action` eligibility** (+ unit test rows + in the existing pure-function test). Makes the documented alignment contract a runtime + invariant instead of an arithmetic proof. +3. **H1 — replace try/except with `hasattr(trt.TacticSource, "CUBLAS")` version gate** + and refresh the comment/log so the TRT 11 behavior is explicit, not accidental. +4. **H2 — delete the dead `reuse_device_memory` branch** (or migrate to + `set_device_memory(ptr, engine.device_memory_size_v2)` if reuse is planned). +5. **M3 — `create_network()` in compile_depth_anything_tensorrt.py:129**, same one-liner + as the RAFT tool. +6. **L1 — drop `do_constant_folding=True`** next time the export call is touched. +7. **L3 — harden external-data detection** in `optimize_onnx` (data_location check or a + comment documenting the `export_onnx`-consolidates-first invariant). + +Items 1–5 are behavior-preserving on the installed stack (TRT 10.16 / cuda-python 12.9 / +torch 2.8) and would each pass the standard gate (pytest 117+/0, ruff, pyrefly ≤72 +baseline) independently. diff --git a/docs/perf_bestpractices_audit_2026-07-10.md b/docs/perf_bestpractices_audit_2026-07-10.md new file mode 100644 index 000000000..a330adedc --- /dev/null +++ b/docs/perf_bestpractices_audit_2026-07-10.md @@ -0,0 +1,358 @@ +# StreamDiffusion Performance & Best-Practices Audit — 2026-07-10 + +Read-only audit against CUDA / PyTorch / TensorRT best practices, combining static code review +with measured profiling of the real-time denoising loop. No source files were changed — this +document is the only artifact produced. + +**Scope:** `src/streamdiffusion/{pipeline.py,wrapper.py}`, `acceleration/tensorrt/*`, `modules/*`, +`preprocessing/*`. **Config profiled:** StreamDiffusionTD "Quality/FP16" preset — SDXL-turbo, +512×512, img2img, FP8 TensorRT engines, CUDA graphs enabled. + +**Repo state at audit time:** branch `refactor/py-code-review-remediation` (2 commits ahead of +`origin/SDTD_032_dev`), HEAD `6718511`. The immediately preceding commit (`5a44f34`, same day) +already remediated exception-hygiene issues in `pipeline.py`, `ipadapter_module.py`, +`controlnet_module.py`, `image_processing_module.py`, `latent_processing_module.py`, +`acceleration/tensorrt/__init__.py`, `engine_manager.py`, and the TRT export wrappers — but **not** +`wrapper.py`, which was untouched. Every finding below was checked against this HEAD; where a +finding sits in a file/region that commit touched, its exact lines were diffed to confirm the +issue is distinct from what was already fixed (see "General code quality" §, findings #4–#10). + +--- + +## Prioritized quick wins + +Ranked by (impact × how directly it affects the measured per-frame budget) ÷ effort. "Measured" +column links to the Phase 2 data below; "Static" findings are code-review-only. + +| # | Finding | Impact | Effort | Evidence | +|---|---|---|---|---| +| 1 | **~13.6ms/frame (66% of the entire 20.8ms budget) inside `unet_step` is unattributed by any profiler region** — only ~4.6ms of its 18.2ms median is explained by measured children (`trt_infer`×3 + staging + scheduler). The unmeasured remainder is CFG-buffer prep, hook dispatch, and conditioning-lookup Python/CUDA logic that today has zero instrumentation. | **High** — largest single lever on real-time throughput | Small (instrument first, then optimize what's found) | `profiler_logs/*_stats.json`; `pipeline.py:820-1041` (`unet_step`) | +| 2 | `set_tensor_address` is reissued for every I/O tensor on **every frame**, even on the already-captured CUDA-graph replay fast path, defeating part of the point of graph capture (eliminate repeated host-side API calls) | High — recurring per-frame host overhead × every TRT engine in the pipeline | Small — guard with `if not (use_cuda_graph and self.cuda_graph_instance is not None):` | `acceleration/tensorrt/utilities.py:1118-1125` | +| 3 | Engine / timing-cache files are written non-atomically (`open(...,"wb").write()` straight to the final path); an interrupted multi-minute TRT build leaves a truncated file that the next run's `os.path.exists()` check silently treats as a valid cache hit | High — silent corruption / wasted rebuild risk | Small-medium — write to `.tmp` + `os.replace()` | `acceleration/tensorrt/utilities.py:779-789,907`; `builder.py:284-285` | +| 4 | RNG/seed silently resets to the hardcoded default (`seed=2`) on **every runtime prompt change** — `StreamDiffusion.prepare()`'s `generator` parameter defaults to a module-load-time-shared `torch.Generator()` mutable default, and neither `wrapper.py` `prepare()` branch (single-prompt or blending) forwards `self.generator`/`self.current_seed` when calling `stream.prepare()` | High — visible, reproducible loss of user-configured determinism in production | Small — thread `self.generator`/seed through both wrapper `prepare()` call sites; fix the mutable default | `pipeline.py:402-403`; `wrapper.py:493-498,519-524` | +| 5 | `__call__` / `postprocess_image("latent")` return aliased, persistent internal buffers (`_image_decode_buf`, `_prev_image_buf`) that get `.copy_()`'d in place next frame — any caller that retains the reference (queues it, hands to async consumer) sees silent mutation | High — silent data corruption for any downstream buffering consumer | Medium — `.clone()` on public returns of these buffers, or document the aliasing contract loudly | `pipeline.py:1272-1285`; `wrapper.py:1013-1014` | +| 6 | ControlNet's dynamic-shape profile floor (384px) is narrower than UNet's (256px) — a resolution in [256,384) that's valid for UNet will hard-fail specifically on the ControlNet engine when dynamic shapes + ControlNet are combined | Medium-high — real runtime failure risk in a reachable config | Small — lower `min_ctrl_h`/`min_ctrl_w` to 256, or share one constant | `acceleration/tensorrt/models/controlnet_models.py:68-73` vs `models/models.py:160-161` | +| 7 | NSFW safety-checker `.item()` blocks the frame on a synchronous GPU→host readback every frame it's enabled — no async pattern, despite the codebase already having the correct template one file away | High when the (opt-in, default-off) feature is enabled | Small-medium — reuse `image_filter.py`'s 1-frame-delay pinned-buffer pattern | `acceleration/tensorrt/runtime_engines/unet_engine.py:515`; `wrapper.py:1308` | +| 8 | `output_type="pil"` path does a blocking, unpinned `.cpu()` into freshly-allocated pageable memory every frame — the `"np"` path 30 lines above already does this correctly (pinned buffer + `non_blocking=True` + single `Event.record()/synchronize()`) | Medium-high for `"pil"` consumers | Small — route `"pil"` through the same `_output_pin_buf`/`Event` machinery | `wrapper.py:1264` vs `wrapper.py:1019-1037` | +| 9 | Stringly-typed OOM detection (`"out of memory" in str(e).lower()`) duplicated in two engine exception handlers, instead of `isinstance(e, torch.cuda.OutOfMemoryError)` (already used correctly elsewhere in the same file) | High — a PyTorch/TensorRT message-wording change silently breaks the OOM→CPU-fallback path | Small-medium — extract one `_is_oom_error(e)` helper, `isinstance` check first | `wrapper.py:2342-2349,2400-2407` | +| 10 | SDXL pipeline-type-mismatch retry failure is caught and only logged as a warning — execution continues with the wrong (non-SDXL) pipeline object rather than failing fast | High — garbled output for SDXL models on this edge case, with no error surfaced | Small — raise, or fail fast at engine-build time | `wrapper.py:1504-1514` | +| 11 | No recovery path if `cudaGraphLaunch` fails on the fast-replay path — an uncaught `CUASSERT` takes down the entire real-time loop instead of falling back to `execute_async_v3` for one frame | Medium — rare trigger, severe blast radius | Medium — try/except → `reset_cuda_graph()` → fall back, re-capture next call | `acceleration/tensorrt/utilities.py:1125` | +| 12 | `_load_model` is a ~1390-line function combining pipeline load, LoRA fusion, VAE setup, and TRT/xformers/sfast acceleration dispatch for UNet+VAE+ControlNet+IPAdapter+safety-checker — the exact logic this audit was asked to verify is concentrated in one untestable function | High — maintainability/testability of the acceleration-mode selection logic itself | Large — split into `_load_pipe()`, `_setup_lora()`, `_setup_vae()`, `_build_tensorrt_engines()`, `_install_hook_modules()` | `wrapper.py:1321-2712` | + +--- + +## Phase 2 — Measured per-frame budget + +**Method:** `scripts/profiling/profile_nsys.py --target benchmark` against the cached SDXL-turbo +FP8 TRT engine, 3 extra warmup + 11 `torch.profiler`-captured + 10 `nsys`-gated frames (24 total +measured `predict_x0_batch` calls). Steady-state wall-clock per frame from the nsys `frame_nsys*` +NVTX spans: **~20.5–21.5ms (~47–49 FPS)**. + +**Authoritative timing source.** The repo's own CUDA-event-based profiler +(`src/streamdiffusion/tools/gpu_profiler.py`, exported to `profiler_logs/*_stats.json`) is used +for every number below, **not** the raw `nsys stats --report nvtx_pushpop_sum` output. +Reason: NVTX push/pop ranges nested *inside* a captured CUDA graph fire only at capture time (the +first ~3 warmup passes), not on each graph replay — so their statistics wildly understate +steady-state cost for graph-replayed regions (e.g. `nvtx_pushpop_sum` shows `predict_x0_batch` +med=1.845ms, vs. the CUDA-event stats' correct **18.254ms**). Anyone profiling this pipeline again +should use the `profiler_logs/*_stats.json` export, or wrap regions of interest so they sit outside +the captured graph, not the raw NVTX summary. + +### Per-region breakdown (p50, `profiler_logs/sdtd_benchmark_20260710_092608_stats.json`, n=24 unless noted) + +| Region | p50 (ms) | p95 (ms) | Total (ms) | % of frame | +|---|---|---|---|---| +| `predict_x0_batch` (= `unet_step`, 1 call/frame) | 18.254 | 31.25 | 633.8 | ~88% | +| ┗ `trt_infer` (n=72, 3/frame) | 1.42 | 30.34 | 688.5 | (nested) | +| ┗ `trt.input_staging` (n=72, 3/frame) | 0.075 | 0.169 | 5.2 | (nested) | +| ┗ `sched.step_batch` (n=48, 2/frame) | 0.06 | 0.097 | 3.4 | (nested) | +| ┗ `sched.rebuild` (n=24) | 0.031 | 0.063 | 0.9 | (nested) | +| **Unattributed remainder inside `unet_step`** | **~13.6** | — | — | **~66%** | +| `encode_image` | 1.222 | 2.093 | 48.9 | ~6% | +| `decode_image` | 1.134 | 1.749 | 34.9 | ~5% | +| `glue.ipc_pack_rgba` | 0.086 | 0.118 | 22.0 | <1% | + +`trt_infer`'s three calls (3×1.42ms) plus staging/scheduler sum to **~4.6ms**, against an +`unet_step` median of **18.2ms** — a **~13.6ms/frame gap with no nested `profiler.region()` +call inside it today**. This is the single largest lever in the pipeline: it's 66% of the entire +frame budget, and is exactly the span covered by "General code quality" finding #12 below +(`unet_step`'s CFG-buffer-prep, hook-dispatch loop, and SDXL-conditioning-cache lookup) — +recommend adding `profiler.region()` spans around those three sub-phases before attempting any +optimization, since right now it's not possible to tell whether the cost is Python dispatch +overhead, un-graphed CUDA kernels (RCFG blend, buffer copies), or something else. + +### Kernel summary (`nsys stats --report cuda_gpu_kern_sum`, top contributors) + +FP8 (`e4m3`) TensorRT GEMM/conv kernels dominate as expected for an FP8-quantized UNet: +`sm89_xmma_gemm_e4m3...` entries account for ~17.6% + 13.8% + 3.7% + 3.3% + ... (roughly 45%+ of +total kernel time across the top 20 rows) — confirms FP8 quantization is effectively applied to +the primary UNet path, no action needed there. + +One notable outlier: **`cutlass_80_wmma_tensorop_f16_s161616gemm_f16...` is the single +second-largest kernel at 16.5% (1260 instances)** — this is an **FP16** tensor-core GEMM, not FP8. +Combined with several other `sm80`/`sm75` `f16f16` conv kernels and `fmha_cutlassF_f16` (flash +attention, 2.1%), roughly a quarter of total kernel time runs in FP16 rather than FP8. This audit +could not attribute these kernels to a specific engine (ControlNet vs. VAE vs. a text encoder) — +`ncu`'s source-correlation was blocked (see Limitations) — but given `fp8_quantize.py` exists in +this repo, it's worth checking whether ControlNet and/or the VAE are intentionally left at FP16 or +are simply not yet covered by the FP8 quantization pass. **Follow-up:** run `ncu` with +`--import-source yes` once GPU performance-counter access is available, to attribute this 16.5% +to a specific engine and confirm whether it's a legitimate FP8 candidate. + +### Memory-copy summary (`nsys stats --report cuda_gpu_mem_time_sum`) + +| Operation | Time (%) | Total (ms) | Count | Max single (ms) | +|---|---|---|---|---| +| H2D memcpy | 97.2 | 2167.3 | 1742 | 338.7 | +| memset | 2.6 | 58.6 | 166 | 51.6 | +| D2D memcpy | 0.1 | 2.4 | 502 | 0.03 | +| D2H memcpy | 0.1 | 1.9 | 1865 | 0.01 | + +H2D dominates nsys's GPU-op-time view, but the total (2.17s) **exceeds** the entire measured +24-frame NVTX capture window (~1.7s) — this is almost certainly TRT engine deserialization / +model-weight upload during process startup (`prepare()`), captured because `nsys profile` wraps +the whole Python process, not a per-frame steady-state cost. **Not corroborated against +timestamps in this pass** — treat as a cold-start/engine-load latency note (relevant to +TouchDesigner scene-load / engine-switch UX), not a per-frame throughput finding. D2H per-frame +traffic (1865 copies, avg <1μs) is negligible and consistent with the pinned-buffer output path. + +The first frame (`frame_warmup0`) costs ~1.16s — expected CUDA-graph-capture + cuBLAS/cuDNN +heuristic-search tax, paid once at startup per the `WARMUP_RUNS` mechanism. Not a runtime finding; +worth confirming this doesn't re-trigger on prompt/LoRA/engine hot-swap in the TD integration, as +that would surface as a UX stutter rather than a steady-state FPS regression. + +### Limitations of this profiling pass + +- **`ncu` (Nsight Compute) did not produce results.** First attempt hit `ERR_NVGPUCTRPERM` (GPU + performance-counter access disabled); after enabling "Allow access to the GPU performance + counters to all users" in NVIDIA Control Panel → Developer Settings, a retry (`--set basic`, + `--launch-count 30`) was attempted but **did not bound itself** — it ran for ~105 CPU-minutes + with the GPU pegged at full utilization and no completion in sight, and was terminated. Root + cause: `ncu`'s counter collection replays each kernel serially per metric pass, and this + benchmark loop launches its UNet/ControlNet/VAE work via **captured CUDA graphs** + (`acceleration/tensorrt/utilities.py:1123-1125`) — profiling graph-launched kernels with `ncu` + is known to multiply replay cost far beyond a normal (non-graphed) kernel count, so a naive + `--set basic` pass over the whole benchmark loop doesn't terminate in practical time. Per-kernel + SOL%, roofline, and occupancy for the FP8 GEMM kernels above are **not measured** — only the + nsys kernel-time ranking (by wall-clock share) is available. +- `nsys stats --report gpu_gaps` doesn't exist in the installed Nsight Systems 2026.3.1 (the + profiling README's report name is stale/version-mismatched); no direct GPU-idle-gap report was + substituted. +- `--cn-scale` defaulted to `0.0` in this capture — ControlNet TRT engine structure is visible in + the kernel/layer traces, but active per-frame ControlNet compute contribution was **not + confirmed** in this run. A follow-up capture with `--cn-scale > 0` would be needed to isolate its + cost. + +--- + +## TensorRT acceleration stack (static review) + +Overall maturity: this is a mature, incident-hardened hand-built TensorRT stack (not +`torch_tensorrt`) — GPU-tiered builder config, version/compute-capability/config-aware engine +cache-key naming, a documented (not accidental) single-CUDA-stream correctness fix, FP8 +finite-scale validation, and graph-reset-on-buffer-reallocation logic all reflect real +incident-driven engineering. Remaining gaps are concentrated and mostly small/isolated: + +1. **Non-atomic engine/timing-cache writes** — see quick win #3. +2. **`set_tensor_address` reissued every frame on the graph-replay fast path** — see quick win #2. +3. **ControlNet/UNet dynamic-shape profile-floor mismatch (384px vs 256px)** — see quick win #6. +4. **No recovery path on `cudaGraphLaunch` failure** — see quick win #11. +5. **`Engine.refit()` has no stream/context synchronization guard** (`utilities.py:565-665`, + `refitter.refit_cuda_engine()` at `:663`) — violates TensorRT's REFIT requirement that no + `IExecutionContext` be actively executing during a refit. Confirmed currently unreachable + (`enable_refit`/`build_enable_refit` default `False`, zero callers of `.refit(` repo-wide) — + a landmine only if REFIT-based weight hot-swap (e.g. live LoRA switching) is wired up later + without adding the guard. Impact: low today / high if activated unguarded. Effort: small. +6. **Single shared CUDA stream across UNet/ControlNet/VAE/aux engines** (`wrapper.py:2263`, + explicitly documented as a deliberate fix for a prior race in `controlnet_engine.py:155-157`). + Standard multi-stream overlap (e.g. decode frame N-1 while denoising frame N) is traded away + for correctness simplicity — do not revert casually. Impact: low (correctness fine, pure + missed-optimization). Effort: large (would need careful multi-stream + `cudaEvent` + cross-stream sync redesign to reclaim overlap without reintroducing the race). +7. **`opt_batch_size` defaults to `1`** in the public builder entry points (`builder.py:96` and 5 + call sites in `acceleration/tensorrt/__init__.py`), which would mistune the optimization + profile if ever called without an explicit value. Verified **not currently triggered** — the + real call path always threads `stream.trt_unet_batch_size` through correctly + (`wrapper.py:2315` → `engine_manager.py:224-233` → `EngineBuilder.build()`). Dormant footgun + for a future direct/test-harness caller. Impact: low (dormant). Effort: small — drop the + default, make it a required kwarg. +8. **FP8 sanity-gate comment mismatches the actual threshold** (`builder.py:307` comment says + "< 100", code checks `< 500` at `:323`) — cosmetic, update the comment. + +--- + +## PyTorch hot-path / sync-free review (static review) + +The per-frame diffusion core is in genuinely good shape: buffers are pre-allocated and reused +(`x_t_latent_buffer`, `_combined_latent_buf`, `_cfg_latent_buf`/`_cfg_t_buf`, +`_stock_noise_bufs` ping-pong, `_image_decode_buf`, `_prev_image_buf`), `pin_memory`/ +`non_blocking=True` are applied correctly at every H2D/D2H boundary that matters (input staging, +similarity-filter skip-probability readback, the `"np"` output path), `torch.inference_mode()` +wraps every hot entry point, and the frame-skip similarity filter uses a genuinely sync-free +1-frame-delay pinned-buffer pattern (`image_filter.py`) that should be the template for the two +findings below, not the exception. + +- **NSFW safety-checker blocking `.item()` every frame** — see quick win #7. +- **`"pil"` output path's blocking, unpinned `.cpu()`** — see quick win #8. +- **`_tensor_to_pil_safe` does two GPU-tensor truthiness syncs before its eventual `.cpu()`** + (`preprocessing/preprocessing_orchestrator.py:601,606,610`) — three host syncs where the + codebase's own `processors/base.py:tensor_to_pil()` (a near-duplicate helper) correctly does + the `.cpu()` transfer first and compares ranges after, at line `:175-208`. This is a regression + of a fix that already exists elsewhere in the same codebase. Impact: medium (gated to + ControlNet/image-hook configs using non-tensor-native processors). Effort: small — reorder to + match `base.py`'s pattern, or dedupe the two implementations. +- **Systemic, already self-tracked: ~20 of 28 preprocessor classes fall back to a per-frame + PIL/CPU round-trip** (`preprocessing/processors/base.py:87-109`) — only 8 processors set + `gpu_native = True`. The code already emits a one-time-per-class warning pointing at this exact + gap ("[GPU-residency]... Set gpu_native=True..."), so this is not a new discovery, but it is the + single largest concentration of avoidable per-frame syncs in the codebase, gated behind which + preprocessor a user selects rather than always active. Worth surfacing to whoever owns that + backlog item. +- **`_ipc_pack_rgba`/`_ipc_pack_unit_rgba` allocate fresh tensors every frame** (`wrapper.py:1052- + 1061,1092-1099`, `torch.full_like`/`torch.cat`) rather than writing into a persistent buffer like + every other output path in the same file. Impact: low (small tensors, cheap kernels, but + avoidable). Effort: small — hoist a persistent BGRA buffer, write in-place. +- **`channels_last` is absent repo-wide, and correctly so for the primary path** — the supported + real-time deployment runs UNet and VAE through TensorRT, so there are no PyTorch convs in the + hot path for `channels_last` to attach to. Only relevant if `acceleration="none"`/`"xformers"` + (PyTorch VAE fallback) or one of the ~20 non-TRT preprocessors above is selected — not a gap in + the intended architecture. + +--- + +## General code quality & best practices (static review) + +Two systemic patterns stand out across `pipeline.py`/`wrapper.py`/`modules/*`: (1) broad, +repeated tolerance for silently-swallowed exceptions that will make production failures invisible +until they surface as a worse downstream symptom, and (2) a handful of very large, +multi-responsibility functions that concentrate most of the acceleration-mode-selection logic this +audit was asked to review, making it hard to verify or safely extend. Findings already covered +above as quick wins are cross-referenced rather than repeated. + +- **RNG/seed mutable-default reset** — quick win #4. +- **Aliased buffer returns (`"latent"` output, `prev_image_result`)** — quick win #5. +- **Stringly-typed OOM detection duplicated** — quick win #9. +- **SDXL pipeline-mismatch retry failure silently swallowed** — quick win #10. +- **`_load_model` god function (~1390 lines, 30+ params)** — quick win #12. +- **`cleanup_engines_and_rebuild` targets a hardcoded `"engines"` path**, ignoring the configured + `engine_dir` (`wrapper.py:3062-3069`) — for any non-default `engine_dir` this OOM-recovery + method either silently no-ops or, worse, could delete an unrelated `engines/` folder in the cwd. + Impact: medium (breaks the one method whose entire purpose is OOM recovery, for any custom + `engine_dir`). Effort: small — `engines_dir = str(getattr(self, "_engine_dir", "engines"))`. +- **Manual `__del__()` invocation in `cleanup_gpu_memory()`** (`wrapper.py:2907-2912,2939-2943`) + — calling a dunder destructor explicitly does not stop the GC from invoking it again once the + object is actually collected; native TRT/CUDA cleanup logic inside `__del__` can run twice, + risking a double-free or a silently swallowed second exception. Impact: medium-high (double + cleanup of native handles is a classic source of intermittent, hard-to-reproduce crashes). + Effort: small — call a proper `close()`/`release()` if the wrapper exposes one, or just `del` + the reference. +- **`_load_lora_with_offline_fallback` swallows every per-candidate exception with zero logging** + (`pipeline.py:349-371`) — if candidate #2 of 5 fails for an unrelated reason (permissions, + corruption) while the rest fail with "not found", the caller only ever sees the final "not + found" — actively misleading during debugging. Impact: medium. Effort: small — add + `logger.debug` inside the loop. **Not touched by the 2026-07-09 remediation commit.** +- **Bare `except Exception: pass` with zero logging in "Advanced model detection failed" fallback** + (`wrapper.py:1855-1856`), inconsistent with every sibling handler within 50 lines that does log. + Impact: medium (both detection paths failing leaves stale `model_type`/`is_sdxl` with no + diagnostic trail). Effort: small. **In `wrapper.py`, which the 2026-07-09 remediation commit did + not touch at all.** +- **Per-frame dynamic import + broad try/except-pass chains in the UNet hot-path hook** + (`modules/ipadapter_module.py:434-521`, `_unet_hook`, invoked once per denoising step) — a + `from diffusers_ipadapter...` import executes on every call instead of being hoisted to module + scope, plus five separate `try/except Exception: pass` fallbacks, including one explicitly + commented "Do not add fallback mechanisms" that nonetheless silently swallows and continues. + Note: the 2026-07-09 remediation commit touched this same file but at different lines + (`:360-372`, `set_scale`/attribute-attach try/excepts) — this hot-path hook block was not part + of that pass. Impact: medium (per-step import overhead when IPAdapter is enabled; transient + failures invisible in production). Effort: small — hoist the import, add debug logging. +- **Pointless `try/except Exception: pass` wrapping a single `logger.info()` call** + (`wrapper.py:2103-2106`) — dead defensive code. Impact: low. Effort: small — delete it. +- **Duplicated latent-cache/decode/output-buffer logic between `__call__` and `txt2img`(_sd_turbo)** + (`pipeline.py:1272-1285` vs `1367-1394`) — the `_image_decode_buf` lazy-allocate-then-`.copy_()` + pattern is copy-pasted near-verbatim across 2-3 entry points; any fix to it (e.g. quick win #5) + must be applied in every copy or silently regresses one call path. Impact: medium. Effort: + small-medium — extract a shared `_decode_and_cache()` helper. +- **`unet_step` mixes five responsibilities in ~220 lines**: CFG buffer prep, SDXL conditioning + cache, generic hook dispatch, ControlNet/IPAdapter kwarg extraction, TRT-vs-PyTorch calling + convention (`pipeline.py:820-1041`). This is the exact span identified as the ~13.6ms/frame + profiling blind spot (quick win #1) — splitting it into `_prepare_cfg_batch()`, + `_get_sdxl_conditioning()`, `_dispatch_unet_hooks()`, `_call_unet()` would both fix the + maintainability issue and naturally create the instrumentation boundaries needed to localize + that cost. Impact: medium (maintainability) / high (unlocks quick win #1). Effort: large. +- **Parameter sprawl**: `_load_model` (30 params), `__init__` (50+ params) — plain + positional/keyword lists rather than a config object; makes the god-function split above harder + until addressed. Impact: medium. Effort: large — introduce a `StreamDiffusionConfig` dataclass, + can be done incrementally alongside the `_load_model` split. +- **Minor/cosmetic**: `__call__`'s `x` parameter is typed without `Optional` despite a `None` + default (`pipeline.py:1191`); `_get_scheduler_scalings` has no type hints (`pipeline.py:736`). + Impact: low. Effort: small. + +--- + +## Follow-ups not completed by this audit + +1. GPU performance-counter access is now enabled, but a naive `ncu --set basic` pass over the full + benchmark loop does not terminate in practical time because of CUDA-graph replay cost (see + Limitations above). To get SOL%/occupancy for the FP8 GEMM kernels and attribute the 16.5% + FP16 cutlass kernel to a specific engine, retry `ncu` with the workload scoped down — e.g. + disable CUDA-graph capture for the profiled process (`use_cuda_graph=False` on the engine + wrappers, if exposed as a config knob) so `ncu` sees ungraphed kernel launches, and/or bound + the pass tightly with `--launch-skip N --launch-count <10-20>` targeting only the steady-state + frames, and/or `--kernel-name regex:sm89_xmma_gemm|cutlass` to profile just the GEMM kernels of + interest instead of the entire per-frame kernel set. +2. Add `profiler.region()` spans inside `unet_step` around CFG-buffer prep, hook dispatch, and + conditioning lookup to localize the ~13.6ms/frame currently unattributed (quick win #1) — + this is the highest-value next step and doesn't require any of the above. +3. Correlate H2D memcpy timestamps against `prepare()` vs. per-frame windows to confirm the 2.17s + H2D total is a one-time load cost, not a steady-state one. +4. Re-run `profile_nsys.py` with `--cn-scale > 0` to isolate ControlNet's actual per-frame compute + contribution (this run had `--cn-scale 0.0`, the script's default). + +--- + +## Verification + +- Repo `git status` shows only this new file as untracked; no source under `src/streamdiffusion/` + was modified. +- Profiling artifacts referenced above (`profiles/*.txt`, `profiles/*.nsys-rep`, + `profiler_logs/*_stats.json`) were generated by the repo's own + `scripts/profiling/profile_nsys.py`/`nsys stats`, not hand-authored. +- Four citations (`utilities.py:1118-1125`, `pipeline.py:402-403`, + `controlnet_models.py:68-73`, `wrapper.py:1264`) were re-read directly against source and match + the findings exactly. + +### MCP cross-check (2026-07-10, post-publication) + +A second, independent verification pass used the repo's semantic code-search index +(`D:\dev\SDTD_032_dev`, 4707 chunks, `last_indexed 2026-07-10T08:16`Z — after the +`5a44f34`/`6718511` commits above, so the index reflects the exact HEAD this report was written +against). Method: `search_code` (k=7) to locate each claim's region, then exact `Grep`/`Read` to +pin the specific cited line(s). ~20 of the ~30 path:line citations in this report were checked, +covering all 12 quick wins plus a sample of the general-findings section. + +**Result: no hallucinated citations.** Every checked claim resolves to real code at (or within a +line or two of) its cited location, including the highest-leverage ones: the `unet_step` god-method +(`pipeline.py:826-1041`, cyclomatic complexity 52 — corroborates quick win #1/#12's "unattributed +~13.6ms" and "220-line, five-responsibility function" characterizations independently of the +profiler data); `Engine.refit()` (`utilities.py:565-665`); the graph-replay `set_tensor_address` +loop (`utilities.py:1118-1125`); the single shared `cuda.Stream()` (`wrapper.py:2263`); and the +UNet-vs-ControlNet shape-floor mismatch (`models.py:160` = 256 vs. ControlNet's 384). + +Two findings check out as understated, not overstated: +- **Quick win #9 (stringly-typed OOM):** `wrapper.py:2345` and `:2403` do match `"out of memory" + in error_msg` as described — but the same file already catches the typed + `torch.cuda.OutOfMemoryError` correctly at `wrapper.py:2029`. The typed exception is known to the + codebase, making the two string-matched sites more clearly a regression/inconsistency than an + unavoidable gap. +- **TensorRT finding #8 (FP8 sanity-gate comment mismatch):** confirmed exactly — `builder.py:307` + comment reads "< 100 means quantization is inactive" while `builder.py:323` gates on `_qdq < 500`. + +One characterization was not independently confirmed at the same precision as the rest: the +per-frame *invocation frequency* claimed for the IPAdapter hot-path hook +(`modules/ipadapter_module.py:434-521`) — the function-local imports at `:462`/`:479` are confirmed +present, but this pass did not separately verify the hook fires on every denoising step versus, +e.g., once per `generate()` call. + +Also independently re-confirmed via direct grep of `src/streamdiffusion/`: zero occurrences of +`channels_last`/`memory_format` repo-wide, matching the "absent repo-wide" claim in the PyTorch +hot-path section exactly. diff --git a/src/streamdiffusion/acceleration/tensorrt/builder.py b/src/streamdiffusion/acceleration/tensorrt/builder.py index 045dd7390..1e767aba4 100644 --- a/src/streamdiffusion/acceleration/tensorrt/builder.py +++ b/src/streamdiffusion/acceleration/tensorrt/builder.py @@ -2,6 +2,7 @@ import json import logging import os +import shutil import time from datetime import datetime, timezone from pathlib import Path @@ -70,6 +71,101 @@ def _run_fp8_stage(name: str, fn, stats: dict, allow_fallback: bool, engine_file ) from err +def _check_fp8_disk_space(onnx_opt_path: str, allow_fallback: bool) -> bool: + """Preflight disk-space check before FP8 ONNX quantization. + + ModelOpt keeps several full-size copies of the external-data ONNX weights on disk + at once (opt, opt_named, opt_named_extended, fp8 output) plus calibration tensors. + Require free space >= 5x the source ONNX's on-disk footprint (external data + included), with a ~28 GB floor for SDXL-scale models, so we fail fast before the + ~40 min export/optimize/calibrate pipeline instead of after it. + """ + onnx_dir = os.path.dirname(onnx_opt_path) + source_size = os.path.getsize(onnx_opt_path) + weights_pb = os.path.join(onnx_dir, "weights.pb") + if os.path.exists(weights_pb): + source_size += os.path.getsize(weights_pb) + + required = max(5 * source_size, 28 * 1024**3) + free = shutil.disk_usage(onnx_dir).free + if free >= required: + return True + + _build_logger.warning( + f"[BUILD] Low disk space for FP8 quantization on {onnx_dir}: " + f"{free / 1024**3:.1f} GB free, need ~{required / 1024**3:.1f} GB." + ) + if allow_fallback: + _build_logger.warning("[BUILD] Falling back to FP16 (fp8_allow_fp16_fallback=True).") + return False + raise RuntimeError( + f"Insufficient disk space for FP8 quantization on {onnx_dir}: " + f"{free / 1024**3:.1f} GB free, need ~{required / 1024**3:.1f} GB. " + "Free up space, or set fp8_allow_fp16_fallback=True in TRT_PROFILES to build FP16 instead." + ) + + +def _cleanup_intermediates(engine_dir: str, fp8_ok: bool): + """Delete intermediate ONNX/build artifacts, preserving .engine, .cache, calib_data.npz, + build_stats.json, and (only when fp8_ok) the cached unet.fp8.onnx* artifact. + + Two-pass deletion handles Windows file locks (gc.collect releases Python handles). + Runs from a `finally` block so it also fires when a build stage raises, instead of + orphaning tens of GB of external-data ONNX copies on failure. + """ + _keep_suffixes = (".engine", ".cache") + _keep_exact = {"build_stats.json", "timing.cache", "calib_data.npz"} + _to_delete = [] + for file in os.listdir(engine_dir): + # Keep the FP8 quantized ONNX artifact only if quantization actually succeeded + # (marked by the ".ok" sentinel) -- a partial file from a failed run must be swept. + if fp8_ok and "fp8.onnx" in file: + continue + if file in _keep_exact or any(file.endswith(s) for s in _keep_suffixes): + continue + _to_delete.append(os.path.join(engine_dir, file)) + + if not _to_delete: + return + + _failed = [] + for fpath in _to_delete: + try: + os.remove(fpath) + except OSError: + _failed.append(fpath) + + # Release Python-held file handles (ONNX model refs), retry locked files. + # Per-file poll with 50ms backoff instead of a single global sleep -- most + # handles release within 1-2 retries on Windows; worst case ~0.5s same as before. + if _failed: + gc.collect() + torch.cuda.empty_cache() + _still_failed = [] + for fpath in _failed: + _last_err = None + for _attempt in range(10): + try: + os.remove(fpath) + _last_err = None + break + except OSError as _e: + _last_err = _e + time.sleep(0.05) + if _last_err is not None: + _still_failed.append(os.path.basename(fpath)) + _build_logger.warning(f"[BUILD] Could not delete temp file {os.path.basename(fpath)}: {_last_err}") + if _still_failed: + _build_logger.warning( + f"[BUILD] {len(_still_failed)} intermediate files could not be cleaned. " + f"Manual cleanup: delete all files except *.engine, calib_data.npz, unet.fp8.onnx from {engine_dir}" + ) + cleaned = len(_to_delete) - len(_still_failed) + else: + cleaned = len(_to_delete) + _build_logger.info(f"[BUILD] Cleaned {cleaned}/{len(_to_delete)} intermediate files") + + def create_onnx_path(name, onnx_dir, opt=True): return os.path.join(onnx_dir, name + (".opt" if opt else "") + ".onnx") @@ -201,157 +297,168 @@ def build( ) _build_logger.info(f"Verified ONNX opt file: {onnx_opt_path} ({opt_file_size / (1024**2):.1f} MB)") - # --- FP8: Capture calibration tensors (once, cached in calib_data.npz) --- - if fp8 and pipe_ref is not None: - if os.path.exists(_calib_data_path): - _build_logger.info(f"[BUILD] FP8 calibration data cached: {_calib_data_path}") - stats["stages"]["fp8_calib_capture"] = {"status": StageStatus.CACHED} - else: + try: + # --- FP8: Capture calibration tensors (once, cached in calib_data.npz) --- + if fp8 and pipe_ref is not None: + if os.path.exists(_calib_data_path): + _build_logger.info(f"[BUILD] FP8 calibration data cached: {_calib_data_path}") + stats["stages"]["fp8_calib_capture"] = {"status": StageStatus.CACHED} + else: - def _calib_fn(): - if is_controlnet: - from .fp8_quantize import capture_calibration_data_controlnet + def _calib_fn(): + if is_controlnet: + from .fp8_quantize import capture_calibration_data_controlnet + + _build_logger.info( + f"[BUILD] FP8 CN calibration: {calibration_steps} synthetic passes, " + f"res={opt_image_width}x{opt_image_height}" + ) + capture_calibration_data_controlnet( + cn_model=pipe_ref, + n_calibration_steps=calibration_steps, + image_height=opt_image_height, + image_width=opt_image_width, + batch_size=opt_batch_size, + save_path=_calib_data_path, + ) + else: + from .fp8_quantize import _load_calibration_prompts, capture_calibration_data + + prompts = calibration_prompts or _load_calibration_prompts() + _build_logger.info( + f"[BUILD] FP8 activation capture: {len(prompts)} prompts × " + f"{calibration_steps} steps, guidance_scale={fp8_guidance_scale}" + ) + capture_calibration_data( + pipe_ref, + prompts, + num_inference_steps=calibration_steps, + save_path=_calib_data_path, + guidance_scale=fp8_guidance_scale, + onnx_path=onnx_opt_path, + use_cached_attn=fp8_use_cached_attn, + use_controlnet=fp8_use_controlnet, + num_ip_layers=fp8_num_ip_layers, + ) + + if not _run_fp8_stage( + "fp8_calib_capture", _calib_fn, stats, fp8_allow_fp16_fallback, engine_filename + ): + fp8 = False + elif fp8 and pipe_ref is None: + _build_logger.warning( + "[BUILD] fp8=True but pipe_ref not provided — FP8 calibration skipped. " + "Pass pipe_ref in engine_build_options for proper activation capture." + ) + fp8 = False + + # --- FP8: Inject native FLOAT8E4M3FN Q/DQ into the ONNX (cached in unet.fp8.onnx) --- + if fp8: + if os.path.exists(_fp8_onnx_path + ".ok"): + _build_logger.info(f"[BUILD] FP8 ONNX cached: {_fp8_onnx_path}") + stats["stages"]["fp8_onnx_quantize"] = {"status": StageStatus.CACHED} + elif not _check_fp8_disk_space(onnx_opt_path, fp8_allow_fp16_fallback): + fp8 = False + else: - _build_logger.info( - f"[BUILD] FP8 CN calibration: {calibration_steps} synthetic passes, " - f"res={opt_image_width}x{opt_image_height}" - ) - capture_calibration_data_controlnet( - cn_model=pipe_ref, - n_calibration_steps=calibration_steps, - image_height=opt_image_height, - image_width=opt_image_width, - batch_size=opt_batch_size, - save_path=_calib_data_path, - ) - else: - from .fp8_quantize import _load_calibration_prompts, capture_calibration_data + def _quant_fn(): + from .fp8_quantize import load_calibration_data, quantize_onnx_fp8 - prompts = calibration_prompts or _load_calibration_prompts() - _build_logger.info( - f"[BUILD] FP8 activation capture: {len(prompts)} prompts × " - f"{calibration_steps} steps, guidance_scale={fp8_guidance_scale}" - ) - capture_calibration_data( - pipe_ref, - prompts, - num_inference_steps=calibration_steps, - save_path=_calib_data_path, - guidance_scale=fp8_guidance_scale, + calib_data = load_calibration_data(_calib_data_path) + if calib_data is None: + raise RuntimeError(f"Calibration data missing after capture step: {_calib_data_path}") + quantize_onnx_fp8( onnx_path=onnx_opt_path, + output_path=_fp8_onnx_path, + calibration_data=calib_data, use_cached_attn=fp8_use_cached_attn, + use_feature_injection=fp8_use_feature_injection, use_controlnet=fp8_use_controlnet, num_ip_layers=fp8_num_ip_layers, ) - if not _run_fp8_stage("fp8_calib_capture", _calib_fn, stats, fp8_allow_fp16_fallback, engine_filename): - fp8 = False - elif fp8 and pipe_ref is None: - _build_logger.warning( - "[BUILD] fp8=True but pipe_ref not provided — FP8 calibration skipped. " - "Pass pipe_ref in engine_build_options for proper activation capture." - ) - fp8 = False + if not _run_fp8_stage( + "fp8_onnx_quantize", _quant_fn, stats, fp8_allow_fp16_fallback, engine_filename + ): + fp8 = False - # --- FP8: Inject native FLOAT8E4M3FN Q/DQ into the ONNX (cached in unet.fp8.onnx) --- - if fp8: - if os.path.exists(_fp8_onnx_path + ".ok"): - _build_logger.info(f"[BUILD] FP8 ONNX cached: {_fp8_onnx_path}") - stats["stages"]["fp8_onnx_quantize"] = {"status": StageStatus.CACHED} - else: + # Select the ONNX to feed into TRT: FP8-quantized when available, else plain opt. + _trt_onnx_path = _fp8_onnx_path if (fp8 and os.path.exists(_fp8_onnx_path + ".ok")) else onnx_opt_path - def _quant_fn(): - from .fp8_quantize import load_calibration_data, quantize_onnx_fp8 - - calib_data = load_calibration_data(_calib_data_path) - if calib_data is None: - raise RuntimeError(f"Calibration data missing after capture step: {_calib_data_path}") - quantize_onnx_fp8( - onnx_path=onnx_opt_path, - output_path=_fp8_onnx_path, - calibration_data=calib_data, - use_cached_attn=fp8_use_cached_attn, - use_feature_injection=fp8_use_feature_injection, - use_controlnet=fp8_use_controlnet, - num_ip_layers=fp8_num_ip_layers, - ) - - if not _run_fp8_stage("fp8_onnx_quantize", _quant_fn, stats, fp8_allow_fp16_fallback, engine_filename): - fp8 = False - - # Select the ONNX to feed into TRT: FP8-quantized when available, else plain opt. - _trt_onnx_path = _fp8_onnx_path if (fp8 and os.path.exists(_fp8_onnx_path + ".ok")) else onnx_opt_path - - # --- TRT Engine Build --- - if not force_engine_build and os.path.exists(engine_path): - print(f"Found cached engine: {engine_path}") - stats["stages"]["trt_build"] = {"status": "cached"} - else: - t0 = time.perf_counter() - build_engine( - engine_path=engine_path, - onnx_opt_path=_trt_onnx_path, - model_data=self.model, - opt_image_height=opt_image_height, - opt_image_width=opt_image_width, - opt_batch_size=opt_batch_size, - build_static_batch=build_static_batch, - build_dynamic_shape=build_dynamic_shape, - build_all_tactics=build_all_tactics, - build_enable_refit=build_enable_refit, - fp8=fp8, - builder_optimization_level=builder_optimization_level, - ) - elapsed = time.perf_counter() - t0 - stats["stages"]["trt_build"] = {"status": "built", "elapsed_s": round(elapsed, 2)} - _build_logger.info(f"[BUILD] TRT engine build ({engine_filename}): {elapsed:.1f}s") - - # --- FP8 Q/DQ layer count (sanity gate: < 500 means quantization is inactive) --- - if fp8 and os.path.exists(engine_path): - try: - import json as _json - import re as _re - - import tensorrt as trt - - _rt = trt.Runtime(BUILD_TRT_LOGGER) - with open(engine_path, "rb") as _f: - _eng = _rt.deserialize_cuda_engine(_f.read()) - _insp = _eng.create_engine_inspector() - _info = _insp.get_engine_information(trt.LayerInformationFormat.JSON) - _qdq = _info.count("QuantizeLinear") + _info.count("DequantizeLinear") - stats["fp8_qdq_layers"] = _qdq - _build_logger.info(f"[BUILD] FP8 engine Q/DQ layer count: {_qdq}") - if _qdq < 500: - _build_logger.warning( - f"[BUILD] Low Q/DQ count ({_qdq} < 500) — FP8 quantization likely inactive or incomplete" - ) - - # Fused-MHA check: count attention layers TRT fused into a single kernel. - # Pattern is empirical — FLUX uses "_gemm_mha_v2"; SDXL on Ada may differ. - # First build logs sample names so the regex can be confirmed or tightened. - _MHA_RE = _re.compile(r"mha|fmha|MultiHead|FlashAttn", _re.IGNORECASE) + # --- TRT Engine Build --- + if not force_engine_build and os.path.exists(engine_path): + print(f"Found cached engine: {engine_path}") + stats["stages"]["trt_build"] = {"status": "cached"} + else: + t0 = time.perf_counter() + build_engine( + engine_path=engine_path, + onnx_opt_path=_trt_onnx_path, + model_data=self.model, + opt_image_height=opt_image_height, + opt_image_width=opt_image_width, + opt_batch_size=opt_batch_size, + build_static_batch=build_static_batch, + build_dynamic_shape=build_dynamic_shape, + build_all_tactics=build_all_tactics, + build_enable_refit=build_enable_refit, + fp8=fp8, + builder_optimization_level=builder_optimization_level, + ) + elapsed = time.perf_counter() - t0 + stats["stages"]["trt_build"] = {"status": "built", "elapsed_s": round(elapsed, 2)} + _build_logger.info(f"[BUILD] TRT engine build ({engine_filename}): {elapsed:.1f}s") + + # --- FP8 Q/DQ layer count (sanity gate: < 500 means quantization is inactive) --- + if fp8 and os.path.exists(engine_path): try: - _layers = _json.loads(_info).get("Layers", []) - except Exception: - _layers = [] - _total = len(_layers) - _mha_names = [_l.get("Name", "") for _l in _layers if _MHA_RE.search(_l.get("Name", ""))] - _mha_count = len(_mha_names) - stats["mha_fused_kernels"] = _mha_count - stats["total_engine_layers"] = _total - _build_logger.info(f"[BUILD] FP8 engine fused MHA layers: {_mha_count} / {_total} total") - if _mha_count == 0 and _total > 0: - _build_logger.warning( - "[BUILD] No fused MHA layers detected — attention may be running decomposed " - "(slower). Sample layer names (first 5): " + str([_l.get("Name", "") for _l in _layers[:5]]) - ) - else: - _build_logger.info(f"[BUILD] Sample fused-MHA layer names: {_mha_names[:3]}") - except Exception as _e: - _build_logger.warning(f"[BUILD] FP8 inspector check skipped: {_e}") + import json as _json + import re as _re + + import tensorrt as trt + + _rt = trt.Runtime(BUILD_TRT_LOGGER) + with open(engine_path, "rb") as _f: + _eng = _rt.deserialize_cuda_engine(_f.read()) + _insp = _eng.create_engine_inspector() + _info = _insp.get_engine_information(trt.LayerInformationFormat.JSON) + _qdq = _info.count("QuantizeLinear") + _info.count("DequantizeLinear") + stats["fp8_qdq_layers"] = _qdq + _build_logger.info(f"[BUILD] FP8 engine Q/DQ layer count: {_qdq}") + if _qdq < 500: + _build_logger.warning( + f"[BUILD] Low Q/DQ count ({_qdq} < 500) — FP8 quantization likely inactive or incomplete" + ) - # Record totals (before cleanup so build_stats.json is preserved) + # Fused-MHA check: count attention layers TRT fused into a single kernel. + # Pattern is empirical — FLUX uses "_gemm_mha_v2"; SDXL on Ada may differ. + # First build logs sample names so the regex can be confirmed or tightened. + _MHA_RE = _re.compile(r"mha|fmha|MultiHead|FlashAttn", _re.IGNORECASE) + try: + _layers = _json.loads(_info).get("Layers", []) + except Exception: + _layers = [] + _total = len(_layers) + _mha_names = [_l.get("Name", "") for _l in _layers if _MHA_RE.search(_l.get("Name", ""))] + _mha_count = len(_mha_names) + stats["mha_fused_kernels"] = _mha_count + stats["total_engine_layers"] = _total + _build_logger.info(f"[BUILD] FP8 engine fused MHA layers: {_mha_count} / {_total} total") + if _mha_count == 0 and _total > 0: + _build_logger.warning( + "[BUILD] No fused MHA layers detected — attention may be running decomposed " + "(slower). Sample layer names (first 5): " + + str([_l.get("Name", "") for _l in _layers[:5]]) + ) + else: + _build_logger.info(f"[BUILD] Sample fused-MHA layer names: {_mha_names[:3]}") + except Exception as _e: + _build_logger.warning(f"[BUILD] FP8 inspector check skipped: {_e}") + finally: + _fp8_ok = fp8 and os.path.exists(_fp8_onnx_path + ".ok") + _cleanup_intermediates(engine_dir_early, _fp8_ok) + + # Cleanup already ran in the `finally` above (also covers the failure path). total_elapsed = time.perf_counter() - build_total_start stats["total_elapsed_s"] = round(total_elapsed, 2) stats["build_end"] = datetime.now(timezone.utc).isoformat() @@ -362,58 +469,3 @@ def _quant_fn(): _build_logger.info(f"[BUILD] {engine_filename} complete: {total_elapsed:.1f}s total") _write_build_stats(engine_path, stats) - - # Cleanup ONNX artifacts — preserve .engine, calib_data.npz, unet.fp8.onnx* files, timing.cache, build_stats.json - # Two-pass deletion to handle Windows file locks (gc.collect releases Python handles) - _keep_suffixes = (".engine", ".cache") - _keep_exact = {"build_stats.json", "timing.cache", "calib_data.npz"} - engine_dir = os.path.dirname(engine_path) - _to_delete = [] - for file in os.listdir(engine_dir): - # Keep all files that are part of the FP8 quantized ONNX artifact - # (unet.fp8.onnx and any external data companion like unet.fp8.onnx.data) - if "fp8.onnx" in file: - continue - if file in _keep_exact or any(file.endswith(s) for s in _keep_suffixes): - continue - _to_delete.append(os.path.join(engine_dir, file)) - - if _to_delete: - _failed = [] - for fpath in _to_delete: - try: - os.remove(fpath) - except OSError: - _failed.append(fpath) - - # Release Python-held file handles (ONNX model refs), retry locked files. - # Per-file poll with 50ms backoff instead of a single global sleep — most - # handles release within 1-2 retries on Windows; worst case ~0.5s same as before. - if _failed: - gc.collect() - torch.cuda.empty_cache() - _still_failed = [] - for fpath in _failed: - _last_err = None - for _attempt in range(10): - try: - os.remove(fpath) - _last_err = None - break - except OSError as _e: - _last_err = _e - time.sleep(0.05) - if _last_err is not None: - _still_failed.append(os.path.basename(fpath)) - _build_logger.warning( - f"[BUILD] Could not delete temp file {os.path.basename(fpath)}: {_last_err}" - ) - if _still_failed: - _build_logger.warning( - f"[BUILD] {len(_still_failed)} intermediate files could not be cleaned. " - f"Manual cleanup: delete all files except *.engine, calib_data.npz, unet.fp8.onnx from {engine_dir}" - ) - cleaned = len(_to_delete) - len(_still_failed) - else: - cleaned = len(_to_delete) - _build_logger.info(f"[BUILD] Cleaned {cleaned}/{len(_to_delete)} intermediate files") diff --git a/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py b/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py index de317065d..85d081c7a 100644 --- a/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py +++ b/src/streamdiffusion/acceleration/tensorrt/runtime_engines/unet_engine.py @@ -55,10 +55,39 @@ def __init__(self, filepath: str, stream: "cuda.Stream", use_cuda_graph: bool = self._fio_cache_len: int = -1 # -1 = not yet initialized # Sub-phase 5.6: kvo/fio cache inputs are persistent, address-stable, - # TRT-contiguous tensors (see models/utils.py create_kvo_cache/create_fi_cache), - # so Engine.infer() can bind TensorRT directly to them instead of copying into - # its own staging buffer every frame. Rebuilt only when the kvo/fio lazy-init - # branches below fire (cache length change), never per-frame. + # TRT-contiguous tensors in steady state (models/utils.py + # create_kvo_cache/create_fi_cache), so Engine.infer() can bind + # TensorRT directly to them instead of copying into its own staging + # buffer every frame. `_staging_action` falls back to a copy on + # non-contiguous/dtype-mismatch/mis-aligned inputs, and to + # bind_and_reset on a pointer change — see utilities.py. Rebuilt only + # when the kvo/fio lazy-init branches below fire (cache length + # change), never per-frame. + # + # ControlNet input_control_* residuals are deliberately, PERMANENTLY + # EXCLUDED here (reverted from Phase-2 D2 / f8ff50f) and stay on the + # DtoD-copy staging path. Two reasons, confirmed against code, not + # just bisection: + # 1. The residual source toggles between the ControlNet engine's + # live output and idle cached-dummy-zero tensors + # (_add_controlnet_residuals vs _add_cached_dummy_inputs below), + # and a missed bind_and_reset on that toggle silently pins the + # UNet's CUDA graph to whichever buffer was bound at capture + # time. + # 2. copy_() runs on the engine stream and orders the ControlNet + # engine's residual write before the UNet graph's read of it; a + # zero-copy bind drops that ordering guarantee with no + # replacement synchronization (see the "copy_() executes on + # engine stream to guarantee ordering" log line in utilities.py). + # Re-enabling zero-copy for these names reproduced the "ControlNet + # produces no visual change" regression on the rig for both a + # single-scale (Canny) and ramped-scale (scribble) config, even + # though single-ControlNet residuals are themselves contiguous, + # 256-aligned, persistent buffers (controlnet_module.py + # build_unet_hook single-CN path) — i.e. the failure is NOT explained + # by the copy-fallback path ever firing; it is the ordering guarantee + # above. Do not re-add these names to _zero_copy_names without a + # rig-verified CUDA-event ordering fix. self._zero_copy_names: FrozenSet[str] = frozenset() self._shape_dict: Dict[str, Any] = {} diff --git a/src/streamdiffusion/acceleration/tensorrt/utilities.py b/src/streamdiffusion/acceleration/tensorrt/utilities.py index 7cea0b456..caea2daa2 100644 --- a/src/streamdiffusion/acceleration/tensorrt/utilities.py +++ b/src/streamdiffusion/acceleration/tensorrt/utilities.py @@ -385,8 +385,11 @@ def _apply_gpu_profile_to_config( # + EDGE_MASK_CONVOLUTIONS — the sources that produce valid SM_120 kernels. # TRT 10.16 exposes TacticSource as an int enum (not IntFlag), so the bitmask # is built via (1 << int(source)). No-op on Ada/Ampere. + # CUBLAS/CUBLAS_LT are removed from TacticSource in TRT 11 (cuBLAS tactics + # dropped entirely) — gate on hasattr so that's a deliberate, logged branch + # rather than a silently-swallowed AttributeError. if gpu_profile.compute_capability >= (12, 0): - try: + if hasattr(trt.TacticSource, "CUBLAS"): sources = ( (1 << int(trt.TacticSource.CUBLAS)) | (1 << int(trt.TacticSource.CUBLAS_LT)) @@ -397,8 +400,11 @@ def _apply_gpu_profile_to_config( logger.info( "[TRT Config] tactic sources = CUBLAS|CUBLAS_LT|JIT_CONV|EDGE_MASK (CUDNN excluded for SM_120+)" ) - except (AttributeError, TypeError) as e: - logger.debug(f"[TRT Config] set_tactic_sources not available: {e}") + else: + logger.info( + "[TRT Config] TRT >=11: CUBLAS/CUBLAS_LT removed from TacticSource — " + "default tactic sources already exclude cuDNN/cuBLAS, nothing to scope" + ) # max_num_tactics: cap profiling candidates per layer to reduce build time. # Available since TRT 10.x; -1 (default) lets TRT decide heuristically. 64 is a @@ -558,12 +564,28 @@ def _staging_action( Returns one of: "copy" - copy into self.tensors[name] (today's behavior; always safe) + "copy_and_reset" - copy into self.tensors[name], but reset the CUDA graph first + because this name was previously bound zero-copy into a live + graph, which still reads the old bound address rather than + self.tensors[name] "bind" - skip the copy; bind TensorRT directly to the caller's tensor "bind_and_reset" - skip the copy and bind directly, but reset the CUDA graph first because the caller's address changed while a graph built from the old address was still live + + A pointer that is not 256-byte aligned falls back to "copy": TensorRT's + setTensorAddress contract requires ≥256-byte alignment, and real torch CUDA + allocations are always ≥512-byte aligned, so this guard is a defensive + invariant against config drift rather than a path exercised in production. """ - if name not in zero_copy_names or not is_contiguous or not dtype_match: + if name not in zero_copy_names or not is_contiguous or not dtype_match or cur_ptr % 256 != 0: + # Falling back to copy. If this name was bound zero-copy into a LIVE graph, + # that graph still reads the stale bound address, not self.tensors[name] — + # force one reset so it re-captures reading the staging buffer. prev_ptr is + # only non-None for names actually bound before, so plain copy-path inputs + # (never in zero_copy_names) never trip this. + if graph_exists and prev_ptr is not None: + return "copy_and_reset" return "copy" if graph_exists and cur_ptr != prev_ptr: return "bind_and_reset" @@ -999,12 +1021,8 @@ def get_input_profile_bounds(self, name: str): except Exception: return None - def activate(self, reuse_device_memory=None): - if reuse_device_memory: - self.context = self.engine.create_execution_context_without_device_memory() - self.context.device_memory = reuse_device_memory - else: - self.context = self.engine.create_execution_context() + def activate(self): + self.context = self.engine.create_execution_context() # Attach per-layer profiler when STREAMDIFFUSION_PROFILE_TRT is set. # Requires engines built with profiling_verbosity=DETAILED for meaningful names. @@ -1191,8 +1209,14 @@ def infer(self, feed_dict, stream, use_cuda_graph=False, zero_copy_names: frozen cur_ptr, graph_exists, ) - if action == "copy": + if action in ("copy", "copy_and_reset"): self.tensors[name].copy_(buf) + if action == "copy_and_reset": + needs_reset = True + # Next frame this name has no prior bind, so + # _staging_action sees prev_ptr=None and returns + # plain "copy" — converges after one reset. + self._bound_ptrs.pop(name, None) else: bind_ptrs[name] = cur_ptr if action == "bind_and_reset": @@ -1525,7 +1549,6 @@ def export_onnx( onnx_path, export_params=True, opset_version=onnx_opset, - do_constant_folding=True, input_names=model_data.get_input_names(), output_names=model_data.get_output_names(), dynamic_axes=model_data.get_dynamic_axes(), @@ -1577,15 +1600,17 @@ def optimize_onnx( ): import os - # Check if external data files exist (indicating external data format was used) onnx_dir = os.path.dirname(onnx_path) - external_data_files = [f for f in os.listdir(onnx_dir) if f.endswith(".pb")] - uses_external_data = len(external_data_files) > 0 + # Inspect TensorProto.data_location on the raw (unloaded) model rather than + # scanning the directory for ".pb" filenames — load_external_data_for_model + # resets data_location back to DEFAULT once external data is loaded, so this + # check must happen before loading. + onnx_model = onnx.load(onnx_path, load_external_data=False) + uses_external_data = any(onnx.external_data_helper.uses_external_data(t) for t in onnx_model.graph.initializer) if uses_external_data: - logger.info(f"Optimizing ONNX with external data (found: {external_data_files})") - # Load model with external data - onnx_model = onnx.load(onnx_path, load_external_data=True) + logger.info("Optimizing ONNX with external data") + onnx.external_data_helper.load_external_data_for_model(onnx_model, onnx_dir) onnx_opt_graph = model_data.optimize(onnx_model) # Create output directory @@ -1610,8 +1635,7 @@ def optimize_onnx( logger.info("ONNX optimization complete with external data") else: - # Standard optimization for smaller models - onnx_model = onnx.load(onnx_path) + # No external data to load — the model loaded above is already complete. onnx_opt_graph = model_data.optimize(onnx_model) onnx.save(onnx_opt_graph, onnx_opt_path) diff --git a/src/streamdiffusion/config.py b/src/streamdiffusion/config.py index fc1d5756c..bd97afcf6 100644 --- a/src/streamdiffusion/config.py +++ b/src/streamdiffusion/config.py @@ -1,16 +1,16 @@ -import logging -import os -import sys -import yaml import json +import logging from pathlib import Path from typing import Any, Dict, List, Tuple, Union import yaml +from .param_schema import DEFAULTS + logger = logging.getLogger(__name__) + def load_config(config_path: Union[str, Path]) -> Dict[str, Any]: """Load StreamDiffusion configuration from YAML or JSON file""" config_path = Path(config_path) @@ -96,7 +96,7 @@ def create_wrapper_from_config(config: Dict[str, Any], **overrides) -> Any: seed_blend_config = final_config["seed_blending"] wrapper.update_stream_params( seed_list=seed_blend_config.get("seed_list", []), - interpolation_method=seed_blend_config.get("interpolation_method", "linear"), + seed_interpolation_method=seed_blend_config.get("interpolation_method", "linear"), ) return wrapper @@ -106,7 +106,7 @@ def _extract_wrapper_params(config: Dict[str, Any]) -> Dict[str, Any]: """Extract parameters for StreamDiffusionWrapper.__init__() from config""" param_map = { "model_id_or_path": config.get("model_id", "stabilityai/sd-turbo"), - "t_index_list": config.get("t_index_list", [0, 16, 32, 45]), + "t_index_list": config.get("t_index_list", list(DEFAULTS["t_index_list"])), "lora_dict": config.get("lora_dict"), "mode": config.get("mode", "img2img"), "output_type": config.get("output_type", "pil"), @@ -128,12 +128,12 @@ def _extract_wrapper_params(config: Dict[str, Any]) -> Dict[str, Any]: "similar_filter_sleep_fraction": config.get("similar_filter_sleep_fraction", 0.025), "use_denoising_batch": config.get("use_denoising_batch", True), "cfg_type": config.get("cfg_type", "self"), - "seed": config.get("seed", 2), - "use_safety_checker": config.get("use_safety_checker", False), + "seed": config.get("seed", DEFAULTS["seed"]), + "use_safety_checker": config.get("use_safety_checker", DEFAULTS["use_safety_checker"]), "skip_diffusion": config.get("skip_diffusion", False), "engine_dir": config.get("engine_dir", "engines"), - "normalize_prompt_weights": config.get("normalize_prompt_weights", True), - "normalize_seed_weights": config.get("normalize_seed_weights", True), + "normalize_prompt_weights": config.get("normalize_prompt_weights", DEFAULTS["normalize_prompt_weights"]), + "normalize_seed_weights": config.get("normalize_seed_weights", DEFAULTS["normalize_seed_weights"]), "scheduler": config.get("scheduler", "lcm"), "sampler": config.get("sampler", "normal"), "compile_engines_only": config.get("compile_engines_only", False), @@ -161,12 +161,12 @@ def _extract_wrapper_params(config: Dict[str, Any]) -> Dict[str, Any]: param_map["use_cached_attn"] = config.get("use_cached_attn", False) - param_map["cache_maxframes"] = config.get("cache_maxframes", 1) - param_map["cache_interval"] = config.get("cache_interval", 1) + param_map["cache_maxframes"] = config.get("cache_maxframes", DEFAULTS["cache_maxframes"]) + param_map["cache_interval"] = config.get("cache_interval", DEFAULTS["cache_interval"]) # cn_cache_interval: ControlNet residual reuse interval. # 1 (default) = disabled, run CN every frame. # N > 1 = run CN once every N frames; reuse residuals between (control latency = N-1 frames). - param_map["cn_cache_interval"] = config.get("cn_cache_interval", 1) + param_map["cn_cache_interval"] = config.get("cn_cache_interval", DEFAULTS["cn_cache_interval"]) # Feature Injection (StreamV2V §3.4.2) — requires use_cached_attn=True if "use_feature_injection" not in config: @@ -178,8 +178,8 @@ def _extract_wrapper_params(config: Dict[str, Any]) -> Dict[str, Any]: param_map["use_feature_injection"] = config.get("use_feature_injection", True) # fi_strength: blend weight α (thesis §3.4.2 Eq 3.2 α=0.75; default matches thesis). # fi_threshold: cosine-similarity gate below which injection is skipped (default 0.98). - param_map["fi_strength"] = config.get("fi_strength", 0.75) - param_map["fi_threshold"] = config.get("fi_threshold", 0.98) + param_map["fi_strength"] = config.get("fi_strength", DEFAULTS["fi_strength"]) + param_map["fi_threshold"] = config.get("fi_threshold", DEFAULTS["fi_threshold"]) # max_cache_maxframes: allocation cap for the KVO/FI cache ring buffers (VRAM). # cache_maxframes is the live logical write window; this is the hard upper bound. param_map["max_cache_maxframes"] = config.get("max_cache_maxframes", 4) @@ -206,10 +206,10 @@ def _extract_prepare_params(config: Dict[str, Any]) -> Dict[str, Any]: """Extract parameters for wrapper.prepare() from config""" prepare_params = { "prompt": config.get("prompt", ""), - "negative_prompt": config.get("negative_prompt", ""), - "num_inference_steps": config.get("num_inference_steps", 50), - "guidance_scale": config.get("guidance_scale", 1.2), - "delta": config.get("delta", 1.0), + "negative_prompt": config.get("negative_prompt", DEFAULTS["negative_prompt"]), + "num_inference_steps": config.get("num_inference_steps", DEFAULTS["num_inference_steps"]), + "guidance_scale": config.get("guidance_scale", DEFAULTS["guidance_scale"]), + "delta": config.get("delta", DEFAULTS["delta"]), } # Handle prompt blending configuration @@ -335,67 +335,6 @@ def _prepare_single_hook_config(hook_config: Dict[str, Any], hook_type: str) -> } -def _validate_pipeline_hook_configs(config: Dict[str, Any]) -> None: - """Validate pipeline hook configurations following ControlNet/IPAdapter validation pattern""" - hook_types = ["image_preprocessing", "image_postprocessing", "latent_preprocessing", "latent_postprocessing"] - - for hook_type in hook_types: - if hook_type in config: - hook_config = config[hook_type] - if not isinstance(hook_config, dict): - raise ValueError(f"_validate_config: '{hook_type}' must be a dictionary") - - # Validate enabled field - if "enabled" in hook_config: - enabled = hook_config["enabled"] - if not isinstance(enabled, bool): - raise ValueError(f"_validate_config: '{hook_type}.enabled' must be a boolean") - - # Validate processors field - if "processors" in hook_config: - processors = hook_config["processors"] - if not isinstance(processors, list): - raise ValueError(f"_validate_config: '{hook_type}.processors' must be a list") - - for i, processor in enumerate(processors): - if not isinstance(processor, dict): - raise ValueError(f"_validate_config: '{hook_type}.processors[{i}]' must be a dictionary") - - # Validate processor type (required) - if "type" not in processor: - raise ValueError( - f"_validate_config: '{hook_type}.processors[{i}]' missing required 'type' field" - ) - - if not isinstance(processor["type"], str): - raise ValueError(f"_validate_config: '{hook_type}.processors[{i}].type' must be a string") - - # Validate enabled field (optional, defaults to True) - if "enabled" in processor: - enabled = processor["enabled"] - if not isinstance(enabled, bool): - raise ValueError( - f"_validate_config: '{hook_type}.processors[{i}].enabled' must be a boolean" - ) - - # Validate order field (optional) - if "order" in processor: - order = processor["order"] - if not isinstance(order, int): - raise ValueError( - f"_validate_config: '{hook_type}.processors[{i}].order' must be an integer" - ) - - # Validate params field (optional, coerce None to empty dict) - if "params" in processor: - if processor["params"] is None: - processor["params"] = {} - elif not isinstance(processor["params"], dict): - raise ValueError( - f"_validate_config: '{hook_type}.processors[{i}].params' must be a dictionary" - ) - - def create_prompt_blending_config( base_config: Dict[str, Any], prompt_list: List[Tuple[str, float]], @@ -548,4 +487,4 @@ def _validate_config(config: Dict[str, Any]) -> None: raise ValueError(f"_validate_config: Seed value {i} must be a non-negative integer") if not isinstance(weight, (int, float)) or weight < 0: - raise ValueError(f"_validate_config: Seed weight {i} must be a non-negative number") \ No newline at end of file + raise ValueError(f"_validate_config: Seed weight {i} must be a non-negative number") diff --git a/src/streamdiffusion/modules/controlnet_module.py b/src/streamdiffusion/modules/controlnet_module.py index 2fc6ef520..03c57345d 100644 --- a/src/streamdiffusion/modules/controlnet_module.py +++ b/src/streamdiffusion/modules/controlnet_module.py @@ -13,7 +13,6 @@ from streamdiffusion.preprocessing.preprocessing_orchestrator import ( PreprocessingOrchestrator, ) -from streamdiffusion.preprocessing.orchestrator_user import OrchestratorUser from streamdiffusion.tools.gpu_profiler import profiler @@ -89,6 +88,15 @@ def __init__(self, device: str = "cuda", dtype: torch.dtype = torch.float16) -> self._cn_cache_images_version: int = -1 self._cn_cache_scale_hash: Optional[tuple] = None + # Persistent multi-ControlNet residual merge buffers (Phase-2 prep). The naive + # `merged_down[j] = merged_down[j] + ds[j]` allocates a fresh tensor every frame, + # which would thrash a CUDA graph if these residuals are later bound zero-copy into + # the UNet's input_control_* inputs. Lazily (re)allocated only when residual + # shape/dtype/device changes, so pointers stay stable across frames. + self._cn_merged_down: Optional[List[torch.Tensor]] = None + self._cn_merged_mid: Optional[torch.Tensor] = None + self._cn_merged_shape_key: Optional[tuple] = None + # ---------- Public API (used by wrapper in a later step) ---------- def install(self, stream) -> None: self._stream = stream @@ -116,6 +124,9 @@ def install(self, stream) -> None: self._cn_cached_residuals = None self._cn_cache_images_version = -1 self._cn_cache_scale_hash = None + self._cn_merged_down = None + self._cn_merged_mid = None + self._cn_merged_shape_key = None def add_controlnet( self, cfg: ControlNetConfig, control_image: Optional[Union[str, Any, torch.Tensor]] = None @@ -501,13 +512,6 @@ def _unet_hook(ctx: StepCtx) -> UnetKwargsDelta: encoder_hidden_states = self._stream.prompt_embeds[:, : self._expected_text_len, :] - base_kwargs: Dict[str, Any] = { - "sample": x_t, - "timestep": t_list, - "encoder_hidden_states": encoder_hidden_states, - "return_dict": False, - } - down_samples_list: List[List[torch.Tensor]] = [] mid_samples_list: List[torch.Tensor] = [] @@ -623,16 +627,34 @@ def _unet_hook(ctx: StepCtx) -> UnetKwargsDelta: mid_block_additional_residual=mid_samples_list[0], ) else: - # Merge multiple ControlNet residuals - merged_down = down_samples_list[0] - merged_mid = mid_samples_list[0] + # Merge multiple ControlNet residuals into persistent buffers, in place. + # Do NOT alias down_samples_list[0]/mid_samples_list[0] — those are + # engine-A's own persistent output buffers. Reallocate only when the + # residual shape/dtype/device changes (resolution/batch/model swap), + # so the merged tensors stay pointer-stable across frames. + shape_key = ( + tuple(t.shape for t in down_samples_list[0]), + mid_samples_list[0].shape, + down_samples_list[0][0].dtype, + down_samples_list[0][0].device, + ) + if self._cn_merged_down is None or self._cn_merged_shape_key != shape_key: + self._cn_merged_down = [torch.empty_like(t) for t in down_samples_list[0]] + self._cn_merged_mid = torch.empty_like(mid_samples_list[0]) + self._cn_merged_shape_key = shape_key + + for j in range(len(self._cn_merged_down)): + self._cn_merged_down[j].copy_(down_samples_list[0][j]) + self._cn_merged_mid.copy_(mid_samples_list[0]) + for ds, ms in zip(down_samples_list[1:], mid_samples_list[1:]): - for j in range(len(merged_down)): - merged_down[j] = merged_down[j] + ds[j] - merged_mid = merged_mid + ms + for j in range(len(self._cn_merged_down)): + self._cn_merged_down[j].add_(ds[j]) + self._cn_merged_mid.add_(ms) + _result = UnetKwargsDelta( - down_block_additional_residuals=merged_down, - mid_block_additional_residual=merged_mid, + down_block_additional_residuals=self._cn_merged_down, + mid_block_additional_residual=self._cn_merged_mid, ) # Residual cache write: store result for reuse on upcoming intermediate frames. diff --git a/src/streamdiffusion/param_schema.py b/src/streamdiffusion/param_schema.py new file mode 100644 index 000000000..6ecaa222a --- /dev/null +++ b/src/streamdiffusion/param_schema.py @@ -0,0 +1,136 @@ +"""Single source of truth for the runtime-tunable StreamDiffusion parameter set. + +This module owns parameter *identity* — the 25 names accepted by +:meth:`StreamDiffusionWrapper.update_stream_params`, which of those 23 are +forwarded to :meth:`StreamParameterUpdater.update_stream_params`, and each +param's *construction-time* default (the value a fresh wrapper gets when the +key is absent from config — see ``config._extract_wrapper_params`` and +``config._extract_prepare_params``). + +It deliberately does **not** own order-dependent *apply* logic — the real +``update_stream_params`` implementations stay hand-written in +``wrapper.py`` / ``stream_parameter_updater.py``. A signature-parity test +(``tests/unit/test_param_schema.py``) asserts those signatures stay in sync +with ``PARAM_NAMES`` / ``UPDATER_PARAM_NAMES``, so drift is *caught*, not +*prevented*. + +Note on ``default`` vs. the runtime signatures: every parameter in +``update_stream_params`` defaults to ``None`` at the call-site (meaning +"leave the current value unchanged"), except the two interpolation-method +Literals. ``ParamSpec.default`` here is a *different* concept — the +concrete construction-time value — and intentionally does not mirror the +``None`` sentinels. + +This module has no torch import and no dependency on the rest of the +``streamdiffusion`` package, so it stays cheap to import in isolation +(e.g. from a lightweight test or tool). In practice ``streamdiffusion``'s +own ``__init__.py`` eagerly imports ``.pipeline``/``.wrapper`` (torch-heavy) +before this module would ever be reached via ``from streamdiffusion...``, +so the "cheap import" property is good hygiene rather than a load-time win +for current consumers. +""" + +from dataclasses import dataclass +from typing import Any, Dict, List, Literal, Tuple + + +# Interpolation-method aliases shared by the wrapper and updater signatures. +PromptInterpolationMethod = Literal["linear", "slerp", "cosine_weighted"] +SeedInterpolationMethod = Literal["linear", "slerp"] + + +@dataclass(frozen=True) +class ParamSpec: + """Identity of one runtime-tunable parameter. + + Attributes + ---------- + name: + Keyword name, shared verbatim by the wrapper and (when ``updater`` + is True) the updater signatures. + default: + Construction-time default (see module docstring) — NOT the + ``update_stream_params`` runtime default, which is ``None`` for + all but the two interpolation-method params. + updater: + Whether this param is forwarded to + ``StreamParameterUpdater.update_stream_params``. False only for + ``use_safety_checker`` / ``safety_checker_threshold``, which the + wrapper handles itself (wrapper.py update_stream_params tail). + """ + + name: str + default: Any + updater: bool = True + + +# Order matches StreamDiffusionWrapper.update_stream_params exactly +# (wrapper.py:681-712). +PARAMS: Tuple[ParamSpec, ...] = ( + ParamSpec("num_inference_steps", 50), + ParamSpec("guidance_scale", 1.2), + ParamSpec("delta", 1.0), + # Stored as a tuple so DEFAULTS never hands out a mutable list that a + # caller could alias-mutate; consumers that need a list should do + # list(DEFAULTS["t_index_list"]). + ParamSpec("t_index_list", (0, 16, 32, 45)), + ParamSpec("seed", 2), + ParamSpec("prompt_list", None), + ParamSpec("negative_prompt", ""), + ParamSpec("prompt_interpolation_method", "slerp"), + ParamSpec("normalize_prompt_weights", True), + ParamSpec("seed_list", None), + ParamSpec("seed_interpolation_method", "linear"), + ParamSpec("normalize_seed_weights", True), + ParamSpec("controlnet_config", None), + ParamSpec("ipadapter_config", None), + ParamSpec("image_preprocessing_config", None), + ParamSpec("image_postprocessing_config", None), + ParamSpec("latent_preprocessing_config", None), + ParamSpec("latent_postprocessing_config", None), + ParamSpec("use_safety_checker", False, updater=False), + ParamSpec("safety_checker_threshold", 0.5, updater=False), + ParamSpec("cache_maxframes", 1), + ParamSpec("cache_interval", 1), + ParamSpec("cn_cache_interval", 1), + ParamSpec("fi_strength", 0.75), + ParamSpec("fi_threshold", 0.98), +) + +# All 25 params the wrapper accepts, in wrapper signature order. +PARAM_NAMES: Tuple[str, ...] = tuple(p.name for p in PARAMS) + +# The 23 params forwarded to the updater, in updater signature order +# (a contiguous subsequence of PARAM_NAMES once use_safety_checker / +# safety_checker_threshold are removed). +UPDATER_PARAM_NAMES: Tuple[str, ...] = tuple(p.name for p in PARAMS if p.updater) + +# name -> construction-time default. +DEFAULTS: Dict[str, Any] = {p.name: p.default for p in PARAMS} + + +def floor_num_inference_steps(num_inference_steps: int, max_t_index: int) -> int: + """Raise ``num_inference_steps`` to ``max_t_index + 1`` if it's too small to + hold the largest t_index value. Never lowers it. + + Extracted 1:1 from stream_parameter_updater.py's two + ``if num_inference_steps <= max_t_index`` branches (~:328-348); callers + keep their own branch-specific warning text — this helper only owns the + arithmetic. + """ + return max(num_inference_steps, max_t_index + 1) + + +def rescale_t_index_list(old_t_list: List[int], old_num_steps: int, new_num_steps: int) -> List[int]: + """Proportionally rescale t_index values from an old step-count space to a + new one, clamped to the new space's valid range. + + Extracted 1:1 from stream_parameter_updater.py:358-359. + + Example: rescale_t_index_list([0, 16, 32, 45], 50, 9) == [0, 3, 5, 7]. + (Not [0, 3, 6, 8] — that is the result for new_num_steps=10. The source + comment this was extracted from had the same off-by-one; see + tests/unit/test_param_schema.py for the corrected golden.) + """ + scale_factor = (new_num_steps - 1) / (old_num_steps - 1) if old_num_steps > 1 else 1.0 + return [min(round(t * scale_factor), new_num_steps - 1) for t in old_t_list] diff --git a/src/streamdiffusion/pipeline.py b/src/streamdiffusion/pipeline.py index cdbd86570..7ebe828dc 100644 --- a/src/streamdiffusion/pipeline.py +++ b/src/streamdiffusion/pipeline.py @@ -153,7 +153,11 @@ def __init__( self.add_time_ids = None # Initialize parameter updater - self._param_updater = StreamParameterUpdater(self, normalize_prompt_weights, normalize_seed_weights) + self._param_updater = StreamParameterUpdater( + self, + normalize_prompt_weights=normalize_prompt_weights, + normalize_seed_weights=normalize_seed_weights, + ) # Hook containers (step 1: introduced but initially no-op) self.embedding_hooks: List[EmbeddingHook] = [] diff --git a/src/streamdiffusion/preprocessing/processors/trt_base.py b/src/streamdiffusion/preprocessing/processors/trt_base.py index 56b8852a6..b97778b21 100644 --- a/src/streamdiffusion/preprocessing/processors/trt_base.py +++ b/src/streamdiffusion/preprocessing/processors/trt_base.py @@ -132,7 +132,13 @@ def allocate_buffers(self, device: str = "cuda", input_shape: tuple = None): if self.engine.get_tensor_mode(name) != trt.TensorIOMode.INPUT: continue if input_shape is not None: - self.context.set_input_shape(name, input_shape) + if not self.context.set_input_shape(name, input_shape): + raise RuntimeError( + f"TensorRTEngine.allocate_buffers: set_input_shape failed for " + f"'{name}' with shape {input_shape}. The engine was built for a " + f"fixed shape range — revert the parameter change or rebuild the " + f"engine for the new shape." + ) else: static_shape = tuple(self.context.get_tensor_shape(name)) if any(d < 0 for d in static_shape): @@ -141,7 +147,11 @@ def allocate_buffers(self, device: str = "cuda", input_shape: tuple = None): f"shape {static_shape} but no input_shape was provided. " "Pass input_shape=(N, C, H, W) when using a dynamic engine." ) - self.context.set_input_shape(name, static_shape) + if not self.context.set_input_shape(name, static_shape): + raise RuntimeError( + f"TensorRTEngine.allocate_buffers: set_input_shape failed for " + f"'{name}' with shape {static_shape}." + ) # Pass 2: allocate buffers for ALL tensors (output shapes resolved by TRT now). for idx in range(self.engine.num_io_tensors): @@ -192,12 +202,18 @@ def infer(self, feed_dict: dict, stream=None) -> OrderedDict: self.tensors[name] = cached[name] # Re-apply input shapes to TRT context (context state is NOT cached). for name, shape in new_input_shapes.items(): - self.context.set_input_shape(name, shape) + if not self.context.set_input_shape(name, shape): + raise RuntimeError( + f"TensorRTEngine.infer: set_input_shape failed for '{name}' with shape {shape}." + ) else: # LRU miss: reallocate changed inputs, re-derive output shapes. for name, fed_shape in new_input_shapes.items(): if fed_shape != tuple(self.tensors[name].shape): - self.context.set_input_shape(name, fed_shape) + if not self.context.set_input_shape(name, fed_shape): + raise RuntimeError( + f"TensorRTEngine.infer: set_input_shape failed for '{name}' with shape {fed_shape}." + ) self.tensors[name] = torch.empty( fed_shape, dtype=self.tensors[name].dtype, @@ -232,7 +248,8 @@ def infer(self, feed_dict: dict, stream=None) -> OrderedDict: self.tensors[name].copy_(buf) for name, tensor in self.tensors.items(): - self.context.set_tensor_address(name, tensor.data_ptr()) + if not self.context.set_tensor_address(name, tensor.data_ptr()): + raise RuntimeError(f"TensorRTEngine.infer: set_tensor_address failed for '{name}'") # --- Cross-stream synchronization --- # The input copy_() calls above ran on the CURRENT (default) stream. diff --git a/src/streamdiffusion/stream_parameter_updater.py b/src/streamdiffusion/stream_parameter_updater.py index 13ca8cd1a..42b9c20da 100644 --- a/src/streamdiffusion/stream_parameter_updater.py +++ b/src/streamdiffusion/stream_parameter_updater.py @@ -5,6 +5,12 @@ import torch import torch.nn.functional as F +from .param_schema import ( + PromptInterpolationMethod, + SeedInterpolationMethod, + floor_num_inference_steps, + rescale_t_index_list, +) from .preprocessing.orchestrator_user import OrchestratorUser @@ -29,6 +35,7 @@ class StreamParameterUpdater(OrchestratorUser): def __init__( self, stream_diffusion, + *, wrapper=None, normalize_prompt_weights: bool = True, normalize_seed_weights: bool = True, @@ -281,10 +288,10 @@ def update_stream_params( seed: Optional[int] = None, prompt_list: Optional[List[Tuple[str, float]]] = None, negative_prompt: Optional[str] = None, - prompt_interpolation_method: Literal["linear", "slerp", "cosine_weighted"] = "slerp", + prompt_interpolation_method: PromptInterpolationMethod = "slerp", normalize_prompt_weights: Optional[bool] = None, seed_list: Optional[List[Tuple[int, float]]] = None, - seed_interpolation_method: Literal["linear", "slerp"] = "linear", + seed_interpolation_method: SeedInterpolationMethod = "linear", normalize_seed_weights: Optional[bool] = None, controlnet_config: Optional[List[Dict[str, Any]]] = None, ipadapter_config: Optional[Dict[str, Any]] = None, @@ -330,21 +337,23 @@ def update_stream_params( if t_index_list is None: # Check against current t_list max_t_index = max(self.stream.t_list) if self.stream.t_list else 0 - if num_inference_steps <= max_t_index: + floored = floor_num_inference_steps(num_inference_steps, max_t_index) + if floored != num_inference_steps: logger.warning( f"update_stream_params: num_inference_steps ({num_inference_steps}) is too small for " - f"current t_list (max index: {max_t_index}). Adjusting to {max_t_index + 1}." + f"current t_list (max index: {max_t_index}). Adjusting to {floored}." ) - num_inference_steps = max_t_index + 1 + num_inference_steps = floored else: # Check against provided t_index_list max_t_index = max(t_index_list) if t_index_list else 0 - if num_inference_steps <= max_t_index: + floored = floor_num_inference_steps(num_inference_steps, max_t_index) + if floored != num_inference_steps: logger.warning( f"update_stream_params: num_inference_steps ({num_inference_steps}) is too small for " - f"provided t_index_list (max index: {max_t_index}). Adjusting to {max_t_index + 1}." + f"provided t_index_list (max index: {max_t_index}). Adjusting to {floored}." ) - num_inference_steps = max_t_index + 1 + num_inference_steps = floored old_num_steps = len(self.stream.timesteps) self.stream.scheduler.set_timesteps(num_inference_steps, self.stream.device) @@ -353,9 +362,8 @@ def update_stream_params( # If t_index_list wasn't explicitly provided, rescale existing t_list proportionally if t_index_list is None and old_num_steps > 0: # Rescale each index proportionally to the new number of steps - # e.g., if t_list = [0, 16, 32, 45] with 50 steps -> [0, 3, 6, 8] with 9 steps - scale_factor = (num_inference_steps - 1) / (old_num_steps - 1) if old_num_steps > 1 else 1.0 - t_index_list = [min(round(t * scale_factor), num_inference_steps - 1) for t in self.stream.t_list] + # e.g., if t_list = [0, 16, 32, 45] with 50 steps -> [0, 3, 5, 7] with 9 steps + t_index_list = rescale_t_index_list(self.stream.t_list, old_num_steps, num_inference_steps) # Now update timestep-dependent parameters with the correct t_index_list if t_index_list is not None: diff --git a/src/streamdiffusion/tools/compile_depth_anything_tensorrt.py b/src/streamdiffusion/tools/compile_depth_anything_tensorrt.py index ac0ed7465..b842ec101 100644 --- a/src/streamdiffusion/tools/compile_depth_anything_tensorrt.py +++ b/src/streamdiffusion/tools/compile_depth_anything_tensorrt.py @@ -9,14 +9,15 @@ python -m streamdiffusion.tools.compile_depth_anything_tensorrt --model_size small --resolution 518 """ -import torch +import importlib.util import logging -import os from pathlib import Path -from typing import Optional + import fire +import torch -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') + +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) try: @@ -24,18 +25,13 @@ from streamdiffusion.acceleration.tensorrt.utilities import BUILD_TRT_LOGGER - TENSORRT_AVAILABLE = True except ImportError: TENSORRT_AVAILABLE = False BUILD_TRT_LOGGER = None logger.warning("TensorRT not available. Please install it first.") -try: - import onnx - ONNX_AVAILABLE = True -except ImportError: - ONNX_AVAILABLE = False +ONNX_AVAILABLE = importlib.util.find_spec("onnx") is not None # Depth Anything model configs DEPTH_ANYTHING_MODELS = { @@ -55,10 +51,7 @@ def export_depth_anything_to_onnx( - onnx_path: Path, - model_size: str = "small", - resolution: int = 518, - device: str = "cuda" + onnx_path: Path, model_size: str = "small", resolution: int = 518, device: str = "cuda" ) -> bool: """Export Depth Anything model to ONNX format""" try: @@ -107,6 +100,7 @@ def export_depth_anything_to_onnx( except Exception as e: logger.error(f"Failed to export ONNX: {e}") import traceback + traceback.print_exc() return False @@ -126,11 +120,11 @@ def build_tensorrt_engine( try: builder = trt.Builder(BUILD_TRT_LOGGER) - network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + network = builder.create_network() # EXPLICIT_BATCH deprecated/ignored in TRT 10.x parser = trt.OnnxParser(network, BUILD_TRT_LOGGER) # Parse ONNX - with open(onnx_path, 'rb') as f: + with open(onnx_path, "rb") as f: if not parser.parse(f.read()): for i in range(parser.num_errors): logger.error(f"ONNX parse error: {parser.get_error(i)}") @@ -146,10 +140,12 @@ def build_tensorrt_engine( # Set optimization profile for fixed resolution profile = builder.create_optimization_profile() - profile.set_shape("input", - (1, 3, resolution, resolution), # min - (1, 3, resolution, resolution), # opt - (1, 3, resolution, resolution)) # max + profile.set_shape( + "input", + (1, 3, resolution, resolution), # min + (1, 3, resolution, resolution), # opt + (1, 3, resolution, resolution), + ) # max config.add_optimization_profile(profile) # Build engine @@ -162,7 +158,7 @@ def build_tensorrt_engine( # Save engine engine_path.parent.mkdir(parents=True, exist_ok=True) - with open(engine_path, 'wb') as f: + with open(engine_path, "wb") as f: f.write(serialized_engine) logger.info(f"TensorRT engine saved: {engine_path}") @@ -171,6 +167,7 @@ def build_tensorrt_engine( except Exception as e: logger.error(f"Failed to build TensorRT engine: {e}") import traceback + traceback.print_exc() return False @@ -210,7 +207,7 @@ def compile_depth_anything( # Check if engine already exists if engine_path.exists(): logger.info(f"Engine already exists: {engine_path}") - overwrite = input("Overwrite? (y/N): ").lower().strip() == 'y' + overwrite = input("Overwrite? (y/N): ").lower().strip() == "y" if not overwrite: return @@ -232,9 +229,9 @@ def compile_depth_anything( logger.info("Removed intermediate ONNX file") logger.info(f"\nSuccess! Engine saved to: {engine_path}") - logger.info(f"\nTo use in config:") - logger.info(f' preprocessor: "depth_tensorrt"') - logger.info(f' preprocessor_params:') + logger.info("\nTo use in config:") + logger.info(' preprocessor: "depth_tensorrt"') + logger.info(" preprocessor_params:") logger.info(f' engine_path: "{engine_path}"') diff --git a/src/streamdiffusion/wrapper.py b/src/streamdiffusion/wrapper.py index c75149704..e8b00ed58 100644 --- a/src/streamdiffusion/wrapper.py +++ b/src/streamdiffusion/wrapper.py @@ -10,6 +10,7 @@ from .image_utils import postprocess_image from .model_detection import detect_model +from .param_schema import PromptInterpolationMethod, SeedInterpolationMethod from .pipeline import StreamDiffusion from .tools.gpu_profiler import configure as _configure_profiler from .tools.gpu_profiler import profiler @@ -690,11 +691,11 @@ def update_stream_params( # Prompt blending parameters prompt_list: Optional[List[Tuple[str, float]]] = None, negative_prompt: Optional[str] = None, - prompt_interpolation_method: Literal["linear", "slerp", "cosine_weighted"] = "slerp", + prompt_interpolation_method: PromptInterpolationMethod = "slerp", normalize_prompt_weights: Optional[bool] = None, # Seed blending parameters seed_list: Optional[List[Tuple[int, float]]] = None, - seed_interpolation_method: Literal["linear", "slerp"] = "linear", + seed_interpolation_method: SeedInterpolationMethod = "linear", normalize_seed_weights: Optional[bool] = None, # ControlNet configuration controlnet_config: Optional[List[Dict[str, Any]]] = None, diff --git a/tests/unit/test_config_extraction_golden.py b/tests/unit/test_config_extraction_golden.py new file mode 100644 index 000000000..b4b8f749a --- /dev/null +++ b/tests/unit/test_config_extraction_golden.py @@ -0,0 +1,102 @@ +"""Golden-snapshot test for config.py's construction-time extraction functions. + +Captures the exact output of _extract_wrapper_params / _extract_prepare_params +for a minimal config, BEFORE Stage 2 Increment 3 delegates their literal +defaults to param_schema.DEFAULTS. If Inc 3 changes any default value (instead +of just its source), this test catches the divergence — the output must stay +byte-identical. + +t_index_list is asserted to remain a `list` (not param_schema.DEFAULTS' +internal tuple) since StreamDiffusionWrapper.__init__ types it List[int] and +downstream consumers (e.g. save_config/JSON) expect a list. +""" + +import torch + +from streamdiffusion.config import _extract_prepare_params, _extract_wrapper_params + + +MINIMAL_CONFIG = {"model_id": "stabilityai/sd-turbo"} + +EXPECTED_WRAPPER_PARAMS = { + "model_id_or_path": "stabilityai/sd-turbo", + "t_index_list": [0, 16, 32, 45], + "mode": "img2img", + "output_type": "pil", + "device": "cuda", + "dtype": torch.float16, + "frame_buffer_size": 1, + "width": 512, + "height": 512, + "warmup": 10, + "acceleration": "tensorrt", + "do_add_noise": True, + "use_tiny_vae": True, + "enable_similar_image_filter": False, + "similar_image_filter_threshold": 0.98, + "similar_image_filter_max_skip_frame": 10, + "similar_filter_sleep_fraction": 0.025, + "use_denoising_batch": True, + "cfg_type": "self", + "seed": 2, + "use_safety_checker": False, + "skip_diffusion": False, + "engine_dir": "engines", + "normalize_prompt_weights": True, + "normalize_seed_weights": True, + "scheduler": "lcm", + "sampler": "normal", + "compile_engines_only": False, + "static_shapes": False, + "fp8": False, + "vae_builder_optimization_level": 3, + "build_engines_if_missing": True, + "fp8_allow_fp16_fallback": False, + "use_controlnet": False, + "use_ipadapter": False, + "use_cached_attn": False, + "cache_maxframes": 1, + "cache_interval": 1, + "cn_cache_interval": 1, + "use_feature_injection": True, + "fi_strength": 0.75, + "fi_threshold": 0.98, + "max_cache_maxframes": 4, + "use_cuda_ipc_output": False, + "cuda_ipc_num_slots": 2, + "controlnet_preview_passthrough": False, + "debug_mode": False, +} + +EXPECTED_PREPARE_PARAMS = { + "prompt": "", + "negative_prompt": "", + "num_inference_steps": 50, + "guidance_scale": 1.2, + "delta": 1.0, +} + + +def test_extract_wrapper_params_minimal_config_byte_identical(): + result = _extract_wrapper_params(MINIMAL_CONFIG) + assert result == EXPECTED_WRAPPER_PARAMS + + +def test_extract_wrapper_params_t_index_list_is_a_list(): + """Must stay a `list`, not param_schema.DEFAULTS' internal immutable tuple.""" + result = _extract_wrapper_params(MINIMAL_CONFIG) + assert isinstance(result["t_index_list"], list) + + +def test_extract_wrapper_params_t_index_list_not_aliased_to_schema_default(): + """Mutating the returned list must not corrupt param_schema.DEFAULTS or a + second call's output (guards against a shared-mutable-default regression).""" + result = _extract_wrapper_params(MINIMAL_CONFIG) + result["t_index_list"].append(999) + second_result = _extract_wrapper_params(MINIMAL_CONFIG) + assert second_result["t_index_list"] == [0, 16, 32, 45] + + +def test_extract_prepare_params_minimal_config_byte_identical(): + result = _extract_prepare_params(MINIMAL_CONFIG) + assert result == EXPECTED_PREPARE_PARAMS diff --git a/tests/unit/test_controlnet_residual_merge.py b/tests/unit/test_controlnet_residual_merge.py new file mode 100644 index 000000000..d9b1224e1 --- /dev/null +++ b/tests/unit/test_controlnet_residual_merge.py @@ -0,0 +1,186 @@ +""" +Regression tests for ControlNetModule's multi-ControlNet residual merge (Phase-2 prep, D1). + +CPU-only and model-free: fake ControlNet callables stand in for real engines so the +merge logic in build_unet_hook()'s closure can be exercised without CUDA/TRT. + +Root cause being guarded: the old merge (`merged_down[j] = merged_down[j] + ds[j]`) +allocated a fresh tensor every frame, so the UNet's input_control_* residuals would +never be pointer-stable across frames — a prerequisite for zero-copy binding those +inputs (Phase-2 D2). This also verifies the merge never aliases engine A's own +persistent output buffer. +""" + +import torch + +from streamdiffusion.hooks import StepCtx +from streamdiffusion.modules.controlnet_module import ControlNetModule + + +class _FakeCN: + """Stands in for a ControlNet engine: returns fixed-value residuals, records calls.""" + + def __init__(self, down_value: float, mid_value: float, down_shapes, mid_shape): + self.calls = 0 + self._down_value = down_value + self._mid_value = mid_value + self._down_shapes = down_shapes + self._mid_shape = mid_shape + + def __call__( + self, sample, timestep, encoder_hidden_states, controlnet_cond, conditioning_scale, return_dict=False + ): + self.calls += 1 + down = [torch.full(shape, self._down_value, dtype=torch.float32) for shape in self._down_shapes] + mid = torch.full(self._mid_shape, self._mid_value, dtype=torch.float32) + return down, mid + + +class _FakeStream: + def __init__(self, text_len: int = 77, batch: int = 1): + self.prompt_embeds = torch.randn(batch, text_len, 8) + + +def _make_module_with_two_controlnets(down_shapes, mid_shape) -> tuple: + module = ControlNetModule(device="cpu", dtype=torch.float32) + module._stream = _FakeStream() + + cn_a = _FakeCN(down_value=1.0, mid_value=2.0, down_shapes=down_shapes, mid_shape=mid_shape) + cn_b = _FakeCN(down_value=3.0, mid_value=4.0, down_shapes=down_shapes, mid_shape=mid_shape) + + module.controlnets = [cn_a, cn_b] + module.controlnet_images = [torch.randn(1, 3, 8, 8), torch.randn(1, 3, 8, 8)] + module.controlnet_scales = [1.0, 1.0] + module.enabled_list = [True, True] + + return module, cn_a, cn_b + + +def _make_ctx(batch: int = 1) -> StepCtx: + return StepCtx( + x_t_latent=torch.randn(batch, 4, 8, 8), + t_list=torch.tensor([0]), + step_index=0, + guidance_mode="none", + sdxl_cond=None, + ) + + +class TestControlNetResidualMerge: + def test_merge_is_numerically_correct(self): + down_shapes = [(1, 4, 8, 8), (1, 4, 4, 4)] + mid_shape = (1, 4, 2, 2) + module, cn_a, cn_b = _make_module_with_two_controlnets(down_shapes, mid_shape) + hook = module.build_unet_hook() + + result = hook(_make_ctx()) + + assert cn_a.calls == 1 and cn_b.calls == 1 + for merged, shape in zip(result.down_block_additional_residuals, down_shapes): + assert torch.allclose(merged, torch.full(shape, 4.0)), "1.0 + 3.0 == 4.0 per down block" + assert torch.allclose(result.mid_block_additional_residual, torch.full(mid_shape, 6.0)), "2.0 + 4.0 == 6.0" + + def test_merge_buffers_are_pointer_stable_across_frames(self): + """Same shape on consecutive frames must reuse the same tensor objects + (required before the UNet's input_control_* inputs can be zero-copy bound).""" + down_shapes = [(1, 4, 8, 8)] + mid_shape = (1, 4, 2, 2) + module, _cn_a, _cn_b = _make_module_with_two_controlnets(down_shapes, mid_shape) + hook = module.build_unet_hook() + + result1 = hook(_make_ctx()) + down_ptr1 = [t.data_ptr() for t in result1.down_block_additional_residuals] + mid_ptr1 = result1.mid_block_additional_residual.data_ptr() + + result2 = hook(_make_ctx()) + down_ptr2 = [t.data_ptr() for t in result2.down_block_additional_residuals] + mid_ptr2 = result2.mid_block_additional_residual.data_ptr() + + assert down_ptr1 == down_ptr2, "merged down-block buffers must be reused, not reallocated" + assert mid_ptr1 == mid_ptr2, "merged mid-block buffer must be reused, not reallocated" + # Values still correct on frame 2 (buffer was overwritten, not just left stale) + assert torch.allclose(result2.down_block_additional_residuals[0], torch.full(down_shapes[0], 4.0)) + + def test_merge_reallocates_on_shape_change(self): + """A resolution/batch change must produce a new buffer, not corrupt-reuse the old one.""" + down_shapes = [(1, 4, 8, 8)] + mid_shape = (1, 4, 2, 2) + module, cn_a, cn_b = _make_module_with_two_controlnets(down_shapes, mid_shape) + hook = module.build_unet_hook() + + result1 = hook(_make_ctx()) + down_ptr1 = result1.down_block_additional_residuals[0].data_ptr() + + # Simulate a resolution change: new shapes for both fake engines. + new_down_shapes = [(1, 4, 16, 16)] + new_mid_shape = (1, 4, 4, 4) + cn_a._down_shapes = new_down_shapes + cn_a._mid_shape = new_mid_shape + cn_b._down_shapes = new_down_shapes + cn_b._mid_shape = new_mid_shape + + result2 = hook(_make_ctx()) + assert result2.down_block_additional_residuals[0].shape == new_down_shapes[0] + down_ptr2 = result2.down_block_additional_residuals[0].data_ptr() + assert down_ptr2 != down_ptr1, "shape change must trigger reallocation" + + def test_merge_does_not_alias_engine_output_buffer(self): + """The merge must not mutate/alias down_samples_list[0] — that's engine A's + own persistent output buffer, reused by the engine on the next frame.""" + down_shapes = [(1, 4, 8, 8)] + mid_shape = (1, 4, 2, 2) + module, cn_a, _cn_b = _make_module_with_two_controlnets(down_shapes, mid_shape) + hook = module.build_unet_hook() + + result = hook(_make_ctx()) + + # Re-invoke cn_a directly (as the engine would on its own next call) and confirm + # its freshly-returned buffer is untouched by the merge (still all 1.0, not 4.0). + fresh_down, fresh_mid = cn_a( + sample=None, timestep=None, encoder_hidden_states=None, controlnet_cond=None, conditioning_scale=1.0 + ) + assert torch.allclose(fresh_down[0], torch.full(down_shapes[0], 1.0)) + assert torch.allclose(fresh_mid, torch.full(mid_shape, 2.0)) + # Sanity: the merge result itself is still the summed value. + assert torch.allclose(result.down_block_additional_residuals[0], torch.full(down_shapes[0], 4.0)) + + def test_single_controlnet_bypasses_merge_buffers(self): + """With only one active ControlNet, the engine's own output is returned + directly — no merge buffer should be allocated.""" + module = ControlNetModule(device="cpu", dtype=torch.float32) + module._stream = _FakeStream() + down_shapes = [(1, 4, 8, 8)] + mid_shape = (1, 4, 2, 2) + cn_a = _FakeCN(down_value=5.0, mid_value=6.0, down_shapes=down_shapes, mid_shape=mid_shape) + module.controlnets = [cn_a] + module.controlnet_images = [torch.randn(1, 3, 8, 8)] + module.controlnet_scales = [1.0] + module.enabled_list = [True] + + hook = module.build_unet_hook() + hook(_make_ctx()) + + assert module._cn_merged_down is None + assert module._cn_merged_mid is None + + def test_install_resets_merge_buffers(self): + down_shapes = [(1, 4, 8, 8)] + mid_shape = (1, 4, 2, 2) + module, _cn_a, _cn_b = _make_module_with_two_controlnets(down_shapes, mid_shape) + module.build_unet_hook()(_make_ctx()) + assert module._cn_merged_down is not None + + class _MinimalStream: + unet_hooks = [] + controlnets = None + controlnet_scales = None + preprocessors = None + + # attach_orchestrator requires a preprocessing orchestrator; reuse the existing one + # (install() only touches it when _preprocessing_orchestrator is None). + module._preprocessing_orchestrator = object() + module.install(_MinimalStream()) + + assert module._cn_merged_down is None + assert module._cn_merged_mid is None + assert module._cn_merged_shape_key is None diff --git a/tests/unit/test_param_schema.py b/tests/unit/test_param_schema.py new file mode 100644 index 000000000..cb5271044 --- /dev/null +++ b/tests/unit/test_param_schema.py @@ -0,0 +1,151 @@ +"""Unit tests for src/streamdiffusion/param_schema.py. + +Covers: + - PARAM_NAMES / UPDATER_PARAM_NAMES counts and ordering + - DEFAULTS golden values (construction-time defaults, not the + update_stream_params None-sentinel defaults) + - floor_num_inference_steps / rescale_t_index_list vs hand-computed values + - signature parity: PARAM_NAMES / UPDATER_PARAM_NAMES must match the real + StreamDiffusionWrapper.update_stream_params / StreamParameterUpdater. + update_stream_params signatures — this is the regression lock that + catches drift the moment either signature changes. + +CPU-only, no CUDA required (importing streamdiffusion pulls in torch, but +no tensor is ever created). +""" + +import inspect + +from streamdiffusion.param_schema import ( + DEFAULTS, + PARAM_NAMES, + UPDATER_PARAM_NAMES, + floor_num_inference_steps, + rescale_t_index_list, +) +from streamdiffusion.stream_parameter_updater import StreamParameterUpdater +from streamdiffusion.wrapper import StreamDiffusionWrapper + + +WRAPPER_ONLY_PARAMS = {"use_safety_checker", "safety_checker_threshold"} + + +class TestParamNames: + def test_param_names_count(self): + assert len(PARAM_NAMES) == 25 + + def test_updater_param_names_count(self): + assert len(UPDATER_PARAM_NAMES) == 23 + + def test_updater_param_names_is_ordered_subsequence_of_param_names(self): + """Dropping the two wrapper-only names from PARAM_NAMES, in place, + must yield exactly UPDATER_PARAM_NAMES (order preserved).""" + filtered = tuple(n for n in PARAM_NAMES if n not in WRAPPER_ONLY_PARAMS) + assert filtered == UPDATER_PARAM_NAMES + + def test_wrapper_only_params_excluded_from_updater(self): + assert not (WRAPPER_ONLY_PARAMS & set(UPDATER_PARAM_NAMES)) + assert WRAPPER_ONLY_PARAMS <= set(PARAM_NAMES) + + def test_no_duplicate_names(self): + assert len(PARAM_NAMES) == len(set(PARAM_NAMES)) + + +class TestSignatureParity: + """Regression lock: PARAM_NAMES / UPDATER_PARAM_NAMES must track the real + signatures. If either signature changes without updating param_schema.py, + these fail immediately.""" + + def test_wrapper_signature_matches_param_names(self): + sig = inspect.signature(StreamDiffusionWrapper.update_stream_params) + params = [name for name in sig.parameters if name != "self"] + assert tuple(params) == PARAM_NAMES + + def test_updater_signature_matches_updater_param_names(self): + sig = inspect.signature(StreamParameterUpdater.update_stream_params) + params = [name for name in sig.parameters if name != "self"] + assert tuple(params) == UPDATER_PARAM_NAMES + + +class TestDefaultsGolden: + """Spot-check construction-time defaults against the literals confirmed + identical in both config.py (_extract_wrapper_params / + _extract_prepare_params) and StreamDiffusionWrapper.__init__.""" + + def test_prepare_time_defaults(self): + assert DEFAULTS["num_inference_steps"] == 50 + assert DEFAULTS["guidance_scale"] == 1.2 + assert DEFAULTS["delta"] == 1.0 + + def test_t_index_list_default_and_immutability(self): + assert list(DEFAULTS["t_index_list"]) == [0, 16, 32, 45] + # Must not be a mutable list a caller could alias-mutate. + assert not isinstance(DEFAULTS["t_index_list"], list) + + def test_scalar_defaults(self): + assert DEFAULTS["seed"] == 2 + assert DEFAULTS["negative_prompt"] == "" + assert DEFAULTS["use_safety_checker"] is False + assert DEFAULTS["safety_checker_threshold"] == 0.5 + assert DEFAULTS["normalize_prompt_weights"] is True + assert DEFAULTS["normalize_seed_weights"] is True + assert DEFAULTS["cache_maxframes"] == 1 + assert DEFAULTS["cache_interval"] == 1 + assert DEFAULTS["cn_cache_interval"] == 1 + assert DEFAULTS["fi_strength"] == 0.75 + assert DEFAULTS["fi_threshold"] == 0.98 + + def test_interpolation_method_defaults(self): + assert DEFAULTS["prompt_interpolation_method"] == "slerp" + assert DEFAULTS["seed_interpolation_method"] == "linear" + + def test_config_only_params_default_none(self): + for name in ( + "prompt_list", + "seed_list", + "controlnet_config", + "ipadapter_config", + "image_preprocessing_config", + "image_postprocessing_config", + "latent_preprocessing_config", + "latent_postprocessing_config", + ): + assert DEFAULTS[name] is None + + +class TestFloorNumInferenceSteps: + def test_no_change_when_already_large_enough(self): + assert floor_num_inference_steps(50, 45) == 50 + + def test_raises_when_too_small(self): + assert floor_num_inference_steps(9, 45) == 46 + + def test_boundary_equal_to_max_t_index_is_too_small(self): + # Original code: `if num_inference_steps <= max_t_index: ... = max_t_index + 1` + assert floor_num_inference_steps(45, 45) == 46 + + def test_boundary_one_above_max_t_index_is_fine(self): + assert floor_num_inference_steps(46, 45) == 46 + + +class TestRescaleTIndexList: + def test_golden_50_to_9(self): + """Corrected golden — the code's scale_factor=(new-1)/(old-1) gives + [0,3,5,7] for 50->9, NOT [0,3,6,8] (that off-by-one was in the + original source comment at stream_parameter_updater.py:357 and was + copied into an earlier draft of this extraction).""" + assert rescale_t_index_list([0, 16, 32, 45], 50, 9) == [0, 3, 5, 7] + + def test_golden_50_to_10(self): + """[0,3,6,8] is the correct result for new_num_steps=10, not 9.""" + assert rescale_t_index_list([0, 16, 32, 45], 50, 10) == [0, 3, 6, 8] + + def test_single_old_step_no_division_by_zero(self): + assert rescale_t_index_list([0], 1, 9) == [0] + + def test_same_step_count_is_identity(self): + assert rescale_t_index_list([0, 16, 32, 45], 50, 50) == [0, 16, 32, 45] + + def test_result_clamped_to_new_range(self): + result = rescale_t_index_list([0, 49], 50, 5) + assert max(result) <= 4 diff --git a/tests/unit/test_param_updater_binding.py b/tests/unit/test_param_updater_binding.py new file mode 100644 index 000000000..a87630a9c --- /dev/null +++ b/tests/unit/test_param_updater_binding.py @@ -0,0 +1,134 @@ +"""Regression tests for two parameter-binding bugs fixed in Stage 1: + + 1. ``StreamDiffusion.__init__`` constructed ``StreamParameterUpdater`` with three + *positional* args (``pipeline.py:156``), so the normalize flags landed on the + wrong fields (``wrapper`` <- prompt flag, ``normalize_prompt_weights`` <- seed + flag, ``normalize_seed_weights`` stuck at its default ``True``). The updater + ``__init__`` is now keyword-only past ``stream_diffusion`` so this mis-binding + is impossible to reintroduce silently. + + 2. ``create_wrapper_from_config`` passed ``interpolation_method=`` to + ``wrapper.update_stream_params`` on the seed-only-blending path + (``config.py:99``); that kwarg does not exist (the real name is + ``seed_interpolation_method``), so any config with ``seed_blending`` but no + ``prompt_blending`` raised ``TypeError`` at construction time. + +All tests run on CPU and construct no real pipeline. +""" + +import types +from unittest.mock import patch + +import pytest +import torch + +from streamdiffusion.stream_parameter_updater import StreamParameterUpdater + + +# --------------------------------------------------------------------------- +# Fixtures — mirror the minimal-fake-stream pattern from +# tests/unit/test_prompt_interpolation.py so __init__ runs without a real pipeline. +# --------------------------------------------------------------------------- + + +def _fake_stream(): + """Return a minimal namespace that looks like a StreamDiffusion instance.""" + stream = types.SimpleNamespace() + stream.device = torch.device("cpu") + stream.dtype = torch.float32 + stream.batch_size = 1 + stream.cfg_type = "none" + stream.guidance_scale = 1.0 + stream.prompt_embeds = None + stream.negative_prompt_embeds = None + stream._preprocessing_orchestrator = None + stream.embedding_hooks = [] + return stream + + +def _make_updater(**kwargs) -> StreamParameterUpdater: + """Construct an updater with a fake stream, no-op'ing attach_orchestrator.""" + stream = _fake_stream() + orig_attach = StreamParameterUpdater.attach_orchestrator + + def _noop_attach(self, s): # noqa: ANN001 + self._preprocessing_orchestrator = None + + StreamParameterUpdater.attach_orchestrator = _noop_attach + try: + updater = StreamParameterUpdater(stream, **kwargs) + finally: + StreamParameterUpdater.attach_orchestrator = orig_attach + + updater._embedding_orchestrator = None + return updater + + +# --------------------------------------------------------------------------- +# Bug 1 — normalize flags must bind to the correct fields. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "prompt_flag, seed_flag", + [(True, False), (False, True), (False, False), (True, True)], +) +def test_normalize_flags_bind_to_correct_fields(prompt_flag, seed_flag): + updater = _make_updater( + normalize_prompt_weights=prompt_flag, + normalize_seed_weights=seed_flag, + ) + assert updater.normalize_prompt_weights is prompt_flag + assert updater.normalize_seed_weights is seed_flag + # wrapper stays None at construction — the real reference is attached later + # by StreamDiffusionWrapper (wrapper.py:454), never by a mis-bound bool. + assert updater.wrapper is None + + +def test_getters_echo_constructed_flags(): + updater = _make_updater(normalize_prompt_weights=False, normalize_seed_weights=True) + assert updater.get_normalize_prompt_weights() is False + assert updater.get_normalize_seed_weights() is True + + +def test_positional_flag_binding_is_rejected(): + """The keyword-only barrier makes the original bug a hard TypeError.""" + stream = _fake_stream() + with pytest.raises(TypeError): + StreamParameterUpdater(stream, False, False) # noqa: F841 + + +# --------------------------------------------------------------------------- +# Bug 2 — seed-only blending config must reach update_stream_params with the +# correct kwarg name (seed_interpolation_method), not interpolation_method. +# --------------------------------------------------------------------------- + + +class _FakeWrapper: + """Stub with the *real* update_stream_params kwarg name so the wrong name + (interpolation_method) would raise TypeError, faithfully reproducing the bug.""" + + def __init__(self, **kwargs): + self.update_calls = [] + + def prepare(self, **kwargs): # not exercised on the seed-only path, present for safety + pass + + def update_stream_params(self, *, seed_list=None, seed_interpolation_method="linear"): + self.update_calls.append({"seed_list": seed_list, "seed_interpolation_method": seed_interpolation_method}) + + +def test_seed_only_blending_uses_seed_interpolation_method(): + from streamdiffusion import config as config_mod + + seed_list = [(1, 0.5), (2, 0.5)] + cfg = {"seed_blending": {"seed_list": seed_list, "interpolation_method": "linear"}} + + with patch("streamdiffusion.StreamDiffusionWrapper", _FakeWrapper): + wrapper = config_mod.create_wrapper_from_config(cfg) + + assert isinstance(wrapper, _FakeWrapper) + assert len(wrapper.update_calls) == 1 + call = wrapper.update_calls[0] + assert call["seed_list"] == seed_list + assert call["seed_interpolation_method"] == "linear" diff --git a/tests/unit/test_td_pending_params.py b/tests/unit/test_td_pending_params.py index 4ea0e43b5..23bf4c6ae 100644 --- a/tests/unit/test_td_pending_params.py +++ b/tests/unit/test_td_pending_params.py @@ -14,10 +14,11 @@ """ import threading -import types import unittest from typing import Any, Dict, Optional +from streamdiffusion.param_schema import PARAM_NAMES + # --------------------------------------------------------------------------- # Minimal faithful replica of the three methods under test. @@ -25,12 +26,14 @@ # td_manager.py will break these tests and alert the developer. # --------------------------------------------------------------------------- + class _FakeStream: cfg_type = "none" class _FakeWrapper: """Records calls made to update_stream_params.""" + def __init__(self): self.stream = _FakeStream() self.calls: list = [] @@ -50,15 +53,31 @@ class _Manager: """ VALID_PARAMS = { - 'num_inference_steps', 'guidance_scale', 'delta', 't_index_list', 'seed', - 'prompt_list', 'negative_prompt', 'prompt_interpolation_method', - 'normalize_prompt_weights', 'seed_list', 'seed_interpolation_method', - 'normalize_seed_weights', 'controlnet_config', 'ipadapter_config', - 'image_preprocessing_config', 'image_postprocessing_config', - 'latent_preprocessing_config', 'latent_postprocessing_config', - 'use_safety_checker', 'safety_checker_threshold', - 'cache_maxframes', 'cache_interval', 'fi_strength', 'fi_threshold', - 'cn_cache_interval', + "num_inference_steps", + "guidance_scale", + "delta", + "t_index_list", + "seed", + "prompt_list", + "negative_prompt", + "prompt_interpolation_method", + "normalize_prompt_weights", + "seed_list", + "seed_interpolation_method", + "normalize_seed_weights", + "controlnet_config", + "ipadapter_config", + "image_preprocessing_config", + "image_postprocessing_config", + "latent_preprocessing_config", + "latent_postprocessing_config", + "use_safety_checker", + "safety_checker_threshold", + "cache_maxframes", + "cache_interval", + "fi_strength", + "fi_threshold", + "cn_cache_interval", } def __init__(self): @@ -71,33 +90,30 @@ def __init__(self): # --- Replica of td_manager.py _apply_parameters --- def _apply_parameters(self, params: Dict[str, Any]) -> None: - import logging, random + import random + filtered_params = {k: v for k, v in params.items() if k in self.VALID_PARAMS} - if 'guidance_scale' in filtered_params: - cfg_type = getattr(self.wrapper.stream, 'cfg_type', None) - if cfg_type in ("full", "initialize") and filtered_params['guidance_scale'] <= 1.0: - filtered_params['guidance_scale'] = 1.2 + if "guidance_scale" in filtered_params: + cfg_type = getattr(self.wrapper.stream, "cfg_type", None) + if cfg_type in ("full", "initialize") and filtered_params["guidance_scale"] <= 1.0: + filtered_params["guidance_scale"] = 1.2 - if 'seed_list' in filtered_params: + if "seed_list" in filtered_params: self._randomize_seed_indices = [] new_seed_list = [] - for idx, (seed, weight) in enumerate(filtered_params['seed_list']): + for idx, (seed, weight) in enumerate(filtered_params["seed_list"]): if seed == -1: self._randomize_seed_indices.append(idx) seed = random.randint(0, 2**32 - 1) new_seed_list.append((seed, weight)) - filtered_params['seed_list'] = new_seed_list + filtered_params["seed_list"] = new_seed_list self.wrapper.update_stream_params(**filtered_params) # --- Replica of td_manager.py update_parameters --- def update_parameters(self, params: Dict[str, Any]) -> None: - render_alive = ( - self.streaming - and self.stream_thread is not None - and self.stream_thread.is_alive() - ) + render_alive = self.streaming and self.stream_thread is not None and self.stream_thread.is_alive() if render_alive: with self._pending_params_lock: self._pending_params.update(params) @@ -134,6 +150,7 @@ def _stop_thread(t: threading.Thread): # Test cases # --------------------------------------------------------------------------- + class TestUpdateParametersDefer(unittest.TestCase): """ (a) update_parameters defers when streaming, applies directly when not. @@ -283,5 +300,21 @@ def test_invalid_keys_filtered_out(self): self.assertAlmostEqual(mgr.wrapper.calls[0]["guidance_scale"], 1.5) +class TestValidParamsMatchesSchema(unittest.TestCase): + """ + Drift lock (Stage 2 Increment 4): td_manager.py's runtime whitelist (now + `set(param_schema.PARAM_NAMES)` in the real module -- see + StreamDiffusionTD/td_manager.py::_apply_parameters) must stay exactly + the 25-name set param_schema.py owns. _Manager.VALID_PARAMS above is a + frozen replica of the *pre-refactor* literal list, kept here so this + file never needs to import the real td_manager.py (CUDA/TD deps -- see + module docstring). If param_schema.PARAM_NAMES ever adds/removes/renames + a param, this test fails and both td_manager.py mirrors need updating. + """ + + def test_schema_param_names_matches_frozen_replica(self): + self.assertEqual(set(PARAM_NAMES), _Manager.VALID_PARAMS) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_zero_copy_staging_5_6.py b/tests/unit/test_zero_copy_staging_5_6.py index 56bbe341f..dbba00bbf 100644 --- a/tests/unit/test_zero_copy_staging_5_6.py +++ b/tests/unit/test_zero_copy_staging_5_6.py @@ -34,8 +34,9 @@ ) -PTR_A = 0xDEAD_BEEF -PTR_B = 0xFEED_FACE +PTR_A = 0x1000_0000 # 256-byte aligned (audit M2 guard) +PTR_B = 0x2000_0000 # 256-byte aligned (audit M2 guard) +PTR_MISALIGNED = PTR_A + 1 # not 256-byte aligned class TestStagingActionNotZeroCopy: @@ -54,13 +55,18 @@ def test_name_not_in_zero_copy_names_copies(self): assert result == "copy" def test_empty_zero_copy_names_always_copies(self): - """Default frozenset() → today's behavior exactly, for every name.""" + """Default frozenset() → today's behavior exactly, for every name. + + prev_ptr=None is the realistic value here: a name that is never in + zero_copy_names is never bound, so Engine._bound_ptrs never has an + entry for it and _bound_ptrs.get(name) is always None. + """ result = _staging_action( "kvo_cache_in_0", frozenset(), True, True, - PTR_A, + None, PTR_A, True, ) @@ -179,3 +185,149 @@ def test_first_bind_with_no_prior_ptr_recorded(self): False, ) assert result == "bind" + + +class TestStagingActionAlignmentGuard: + """Audit M2: setTensorAddress requires >=256-byte alignment. A mis-aligned + cur_ptr must fall back to "copy" even when every other eligibility check + would otherwise allow a bind — never silently bind an unaligned address.""" + + def test_misaligned_pointer_falls_back_to_copy_no_graph(self): + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + True, + True, + None, + PTR_MISALIGNED, + False, + ) + assert result == "copy" + + def test_misaligned_pointer_falls_back_to_copy_with_live_graph(self): + """A name previously bound zero-copy (prev_ptr set from an earlier, + valid aligned bind) whose current address is invalid for + setTensorAddress must copy — but a live graph is still replaying the + stale bound address, so this must ALSO force a reset. + + This is the regression lock: f8ff50f's bind->copy fallback left the + live UNet CUDA graph replaying a frozen buffer forever once a + zero-copy input fell back to copy mid-stream (see unet_engine.py's + header comment for the full incident writeup). Before the + copy_and_reset outcome existed, this exact transition asserted plain + "copy" here — i.e. the test documented the bug as correct behavior. + """ + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + True, + True, + PTR_A, # previously bound at a valid, aligned address + PTR_MISALIGNED, # now invalid -> must fall back to copy + True, + ) + assert result == "copy_and_reset" + + +class TestStagingActionCopyAndReset: + """copy_and_reset is the general-purpose safety net for ANY zero-copy name + (currently kvo/fio only — see unet_engine.py) that falls back to copy while + a graph built from its previously-bound address is still live. Exercises + the non-contiguous fallback path specifically (the misaligned-pointer path + is covered by TestStagingActionAlignmentGuard above).""" + + def test_non_contiguous_previously_bound_with_live_graph_copies_and_resets(self): + """A previously-bound name (e.g. a cache tensor that got sliced/viewed + this frame) that is no longer contiguous, with a live graph, must copy + AND force a reset — the graph is still reading the old bound address.""" + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + False, # is_contiguous + True, + PTR_A, # previously bound + PTR_A, + True, + ) + assert result == "copy_and_reset" + + def test_non_contiguous_previously_bound_no_live_graph_just_copies(self): + """Same fallback, but no graph exists yet — nothing to invalidate, so a + plain copy is correct; forcing a reset here would be a needless no-op + reset on every non-graphed frame.""" + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + False, + True, + PTR_A, + PTR_A, + False, + ) + assert result == "copy" + + def test_never_bound_name_falls_back_with_live_graph_just_copies(self): + """prev_ptr=None means this name has never been bound zero-copy, so no + graph could possibly be reading a stale address for it — plain copy, + no spurious reset.""" + result = _staging_action( + "kvo_cache_in_0", + frozenset({"kvo_cache_in_0"}), + False, + True, + None, + PTR_A, + True, + ) + assert result == "copy" + + +class TestStagingActionControlNetResiduals: + """Phase-2 D2 attempted to opt input_control_* UNet inputs into this same + zero-copy path; that was reverted (see unet_engine.py header comment) after + it reproduced a "ControlNet produces no visual change" regression on the + rig, driven by a lost engine-stream copy ordering guarantee — not by the + bind->copy fallback covered above. input_control_* names are PERMANENTLY + excluded from the real UNet engine's _zero_copy_names. The tests below + exercise _staging_action as a pure function using ControlNet-shaped names + only as convenient example zero-copy names; they do not describe current + runtime behavior for actual ControlNet residuals.""" + + def test_steady_state_control_residual_binds_without_reset(self): + """Same persistent merge/output buffer across frames — plain bind.""" + result = _staging_action( + "input_control_00", + frozenset({"input_control_00"}), + True, + True, + PTR_A, + PTR_A, + True, + ) + assert result == "bind" + + def test_idle_active_toggle_binds_and_resets(self): + """Idle dummy-zero buffer -> active ControlNet residual buffer address + flip while a graph is live must force a re-capture.""" + result = _staging_action( + "input_control_00", + frozenset({"input_control_00"}), + True, + True, + PTR_A, + PTR_B, + True, + ) + assert result == "bind_and_reset" + + def test_middle_control_residual_binds_without_reset(self): + result = _staging_action( + "input_control_middle", + frozenset({"input_control_middle"}), + True, + True, + PTR_A, + PTR_A, + True, + ) + assert result == "bind"