From 323ebdcc492d8a599e56c636d5648378e62253db Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 19 Jul 2026 13:46:25 -0400 Subject: [PATCH] fix: validate cached ONNX before reuse to prevent empty-graph TensorRT build crash An interrupted/failed ONNX export leaves a syntactically-valid-but-empty protobuf at onnx_path. The export cache check was existence-only, so the empty file was reused, feeding a 0-node/opset-less graph into polygraphy fold_constants -> ORT shape inference bails on opset<7 -> returns None -> get_num_nodes(None) crashes as "'NoneType' object has no attribute 'graph'". builder.py: add _onnx_cache_valid() (size>0, >=1 node, opset>=7, loaded structure-only) and gate the export cache-hit on it; delete + re-export a bad cache entry. models.py: fail fast in BaseModel.optimize() on a 0-node graph. Applies to every engine type (UNet/VAE enc+dec/ControlNet/safety). Co-Authored-By: Claude Sonnet 5 --- .../acceleration/tensorrt/builder.py | 32 ++++++++++++++++++- .../acceleration/tensorrt/models/models.py | 9 ++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/src/streamdiffusion/acceleration/tensorrt/builder.py b/src/streamdiffusion/acceleration/tensorrt/builder.py index 0db23310..5d5fb4a8 100644 --- a/src/streamdiffusion/acceleration/tensorrt/builder.py +++ b/src/streamdiffusion/acceleration/tensorrt/builder.py @@ -8,6 +8,7 @@ from pathlib import Path from typing import * +import onnx import torch from .models.models import BaseModel @@ -27,6 +28,29 @@ class StageStatus: FAILED = "failed" +def _onnx_cache_valid(path: str) -> bool: + """Check a cached source ONNX is usable before reusing it, skipping export. + + A prior run killed mid-export (OOM, crash, manual kill) can leave a + syntactically-valid-but-empty protobuf at onnx_path. os.path.exists() alone + can't tell that apart from a real export, and a 0-node / opset-less graph + makes polygraphy's constant folding blow up downstream with an opaque + 'NoneType' object has no attribute 'graph' (ORT symbolic shape inference + refuses opset < 7 and returns None). Load structure-only (no external + weight data) and reject anything that looks truncated. + """ + try: + if os.path.getsize(path) == 0: + return False + model = onnx.load(path, load_external_data=False) + if len(model.graph.node) == 0: + return False + opset = max((o.version for o in model.opset_import), default=0) + return opset >= 7 + except Exception: + return False + + def _write_build_stats(engine_path: str, stats: dict): """Append build stats to a JSON-lines file next to the engine directory.""" try: @@ -235,10 +259,16 @@ def build( _fp8_onnx_path = os.path.join(engine_dir_early, f"{artifact_prefix}.fp8.onnx") # --- ONNX Export --- - if not force_onnx_export and os.path.exists(onnx_path): + if not force_onnx_export and os.path.exists(onnx_path) and _onnx_cache_valid(onnx_path): print(f"Found cached model: {onnx_path}") stats["stages"]["onnx_export"] = {"status": "cached"} else: + if not force_onnx_export and os.path.exists(onnx_path): + _build_logger.warning( + f"[BUILD] Cached ONNX at {onnx_path} is empty/corrupt (likely from an " + f"interrupted prior export) -- discarding and re-exporting." + ) + os.remove(onnx_path) print(f"Exporting model: {onnx_path}") t0 = time.perf_counter() _export_kwargs = { diff --git a/src/streamdiffusion/acceleration/tensorrt/models/models.py b/src/streamdiffusion/acceleration/tensorrt/models/models.py index 9035998c..4d3a13f7 100644 --- a/src/streamdiffusion/acceleration/tensorrt/models/models.py +++ b/src/streamdiffusion/acceleration/tensorrt/models/models.py @@ -203,6 +203,15 @@ def get_shape_dict(self, batch_size, image_height, image_width): def optimize(self, onnx_graph): opt = Optimizer(onnx_graph, verbose=self.verbose) opt.info(self.name + ": original") + if len(opt.graph.nodes) == 0: + # Guards against a corrupt/truncated source ONNX (e.g. left behind by an + # interrupted export) slipping into fold_constants, where it would otherwise + # surface as an opaque polygraphy "'NoneType' object has no attribute 'graph'" + # once ORT shape inference bails on the missing/invalid opset. + raise RuntimeError( + f"{self.name}: input ONNX graph has 0 nodes -- the source ONNX is empty or " + f"corrupt. Delete the cached .onnx file for this engine and rebuild." + ) opt.cleanup() opt.info(self.name + ": cleanup") opt.fold_constants()