From 611fe410d07481b09b8cfc28802b772974f1771e Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 19 Jul 2026 13:59:31 -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: guard Optimizer.fold_constants() against a 0-node graph. This is the single call site shared by both BaseModel.optimize() and CLIP's optimize() override (which calls fold_constants() directly, bypassing BaseModel.optimize()), so placing the guard here -- rather than only in optimize() -- covers every model type (CLIP/UNet/VAE encoder/decoder). Co-Authored-By: Claude Sonnet 5 --- .../acceleration/tensorrt/builder.py | 32 ++++++++++++++++++- .../acceleration/tensorrt/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 69de1bcdb..3f2094377 100644 --- a/src/streamdiffusion/acceleration/tensorrt/builder.py +++ b/src/streamdiffusion/acceleration/tensorrt/builder.py @@ -2,6 +2,7 @@ import os from typing import * +import onnx import torch from .models import BaseModel @@ -16,6 +17,29 @@ def create_onnx_path(name, onnx_dir, opt=True): return os.path.join(onnx_dir, name + (".opt" if opt else "") + ".onnx") +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 + + class EngineBuilder: def __init__( self, @@ -47,9 +71,15 @@ def build( force_onnx_export: bool = False, force_onnx_optimize: bool = False, ): - 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}") else: + if not force_onnx_export and os.path.exists(onnx_path): + print( + f"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}") export_onnx( self.network, diff --git a/src/streamdiffusion/acceleration/tensorrt/models.py b/src/streamdiffusion/acceleration/tensorrt/models.py index 3c631aec0..dbbe6626d 100644 --- a/src/streamdiffusion/acceleration/tensorrt/models.py +++ b/src/streamdiffusion/acceleration/tensorrt/models.py @@ -46,6 +46,15 @@ def select_outputs(self, keep, names=None): self.graph.outputs[i].name = name def fold_constants(self, return_onnx=False): + if len(self.graph.nodes) == 0: + # Catches a corrupt/truncated source ONNX (e.g. left behind by an interrupted + # export). CLIP's optimize() override calls fold_constants() directly without + # going through BaseModel.optimize(), so the guard has to live here -- the one + # call site both paths share -- rather than in optimize() itself. + raise RuntimeError( + "Optimizer.fold_constants: input ONNX graph has 0 nodes -- the source ONNX is " + "empty or corrupt. Delete the cached .onnx file for this model and rebuild." + ) onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True) self.graph = gs.import_onnx(onnx_graph) if return_onnx: