Upload tensorrt_convert.py

tensorrt_convert.py

@@ -29,7 +29,7 @@ def trtlog(msg: str):
 # Opset handling:
 #   - If COMFY_TRT_ONNX_OPSET is set, use that integer.
 #   - Otherwise, leave opset_version=None so torch.onnx uses the
-#     recommended opset for this PyTorch version (e.g. 20 on 2.9).
+#     recommended opset for this PyTorch version (e.g. 20 in 2.9+).
 DEFAULT_ONNX_OPSET = None
 _env_opset = os.getenv("COMFY_TRT_ONNX_OPSET")
 if _env_opset is not None:
@@ -403,28 +403,78 @@ class TRT_MODEL_CONVERSION_BASE:
 
         # -----------------------------------------------------------------
         # Build dynamic_shapes spec for torch.export / dynamo=True
+        #   - STATIC node: no dynamic_shapes at all (fully static export)
+        #   - DYNAMIC node: only create Dim if max > min
         # -----------------------------------------------------------------
-        B = Dim("batch", min=batch_size_min, max=batch_size_max)
-        H = Dim("height", min=height_min // 8, max=height_max // 8)
-        W = Dim("width", min=width_min // 8, max=width_max // 8)
-        T = Dim(
-            "tokens",
-            min=context_len_min * context_min,
-            max=context_len * context_max,
-        )
-
-        dynamic_shapes = {
-            "x": {0: B, 2: H, 3: W},
-            "timesteps": {0: B},
-            "context": {0: B, 1: T},
-        }
+        dynamic_shapes = None
 
-        if "y" in input_names:
-            dynamic_shapes["y"] = {0: B}
-        if "guidance" in input_names:
-            dynamic_shapes["guidance"] = {0: B}
-
-        trtlog(f"dynamic_shapes spec: {dynamic_shapes}")
+        def _maybe_dim(name: str, min_v: int, max_v: int):
+            """Create Dim only if there is real dynamism (max > min)."""
+            if max_v < min_v:
+                trtlog(
+                    f"WARNING: Dim {name} has min>{max_v}>{min_v}, swapping to fix."
+                )
+                min_v, max_v = max_v, min_v
+            if max_v > min_v:
+                trtlog(f"Dim {name}: dynamic range [{min_v}, {max_v}]")
+                return Dim(name, min=min_v, max=max_v)
+            else:
+                trtlog(f"Dim {name}: static value {min_v}, not using Dim.")
+                return None
+
+        if not is_static:
+            # Only build dynamic_shapes for the DYNAMIC node
+            B = _maybe_dim("batch", batch_size_min, batch_size_max)
+            H = _maybe_dim("height", height_min // 8, height_max // 8)
+            W = _maybe_dim("width", width_min // 8, width_max // 8)
+            tokens_min = context_len_min * context_min
+            tokens_max = context_len * context_max
+            T = _maybe_dim("tokens", tokens_min, tokens_max)
+
+            dynamic_shapes = {}
+
+            # x: [B, C, H, W]
+            x_dyn = {}
+            if B is not None:
+                x_dyn[0] = B
+            if H is not None:
+                x_dyn[2] = H
+            if W is not None:
+                x_dyn[3] = W
+            if x_dyn:
+                dynamic_shapes["x"] = x_dyn
+
+            # timesteps: [B]
+            if B is not None:
+                dynamic_shapes["timesteps"] = {0: B}
+
+            # context: [B, T, context_dim]
+            ctx_dyn = {}
+            if B is not None:
+                ctx_dyn[0] = B
+            if T is not None:
+                ctx_dyn[1] = T
+            if ctx_dyn:
+                dynamic_shapes["context"] = ctx_dyn
+
+            # y: [B, y_dim]
+            if "y" in input_names and B is not None:
+                dynamic_shapes["y"] = {0: B}
+
+            # guidance: [B, ...]
+            if "guidance" in input_names and B is not None:
+                dynamic_shapes["guidance"] = {0: B}
+
+            if not dynamic_shapes:
+                trtlog(
+                    "No dimensions are actually dynamic for DYNAMIC node. "
+                    "Export will effectively be static."
+                )
+                dynamic_shapes = None
+            else:
+                trtlog(f"dynamic_shapes spec: {dynamic_shapes}")
+        else:
+            trtlog("STATIC node: skipping torch.export.Dim and dynamic_shapes entirely.")
 
         # -----------------------------------------------------------------
         # Build example inputs (using OPT shapes)
@@ -440,7 +490,9 @@ class TRT_MODEL_CONVERSION_BASE:
             )
 
         # -----------------------------------------------------------------
-        # ONNX export with Dynamo + dynamic_shapes
+        # ONNX export with Dynamo (dynamo=True)
+        #   - For static: dynamic_shapes=None, so shapes are fully specialized.
+        #   - For dynamic: dynamic_shapes guides symbolic shapes.
         # -----------------------------------------------------------------
         os.makedirs(os.path.dirname(output_onnx), exist_ok=True)
 
@@ -449,6 +501,10 @@ class TRT_MODEL_CONVERSION_BASE:
             f"opset_version={DEFAULT_ONNX_OPSET}, dtype={dtype}, "
             f"output={output_onnx}"
         )
+        if dynamic_shapes is None:
+            trtlog("ONNX export will be STATIC (no dynamic_shapes).")
+        else:
+            trtlog("ONNX export will use dynamic_shapes (see spec above).")
 
         try:
             torch.onnx.export(
@@ -463,8 +519,7 @@ class TRT_MODEL_CONVERSION_BASE:
                 dynamic_shapes=dynamic_shapes,
                 # NOTE:
                 #   - We intentionally do NOT pass dynamic_axes here.
-                #     dynamic_axes is for the legacy TorchScript exporter,
-                #     dynamic_shapes + dynamo=True is the modern path.
+                #     dynamic_axes is for the legacy TorchScript exporter.
             )
             trtlog("torch.onnx.export completed successfully.")
         except Exception as e:
@@ -828,6 +883,7 @@ class STATIC_TRT_MODEL_CONVERSION(TRT_MODEL_CONVERSION_BASE):
         context_opt,
         num_video_frames,
     ):
+        # STATIC: all min/opt/max are identical
         return super()._convert(
             model,
             filename_prefix,