huchukato
9e6d4c195d
cuda_graph_instance is a cudaGraphExec_t from cudaGraphInstantiate() and must be released with cudaGraphExecDestroy(), not cudaGraphDestroy(). The handle-type mismatch corrupted the CUDA context on CUDA 13, causing a silent process death after the node finished (keep_model_loaded=False) or cudaErrorInvalidValue on the next run (keep_model_loaded=True).
528 lines
20 KiB
Python
528 lines
20 KiB
Python
import torch
|
|
from torch.cuda import nvtx
|
|
from collections import OrderedDict
|
|
import numpy as np
|
|
from polygraphy.backend.common import bytes_from_path
|
|
from polygraphy import util
|
|
from polygraphy.backend.trt import ModifyNetworkOutputs, Profile
|
|
from polygraphy.backend.trt import (
|
|
engine_from_bytes,
|
|
engine_from_network,
|
|
network_from_onnx_path,
|
|
save_engine,
|
|
)
|
|
import time
|
|
import threading
|
|
from tqdm import tqdm
|
|
import copy
|
|
from polygraphy.logger import G_LOGGER
|
|
from logging import error, warning
|
|
import cuda.bindings.runtime as cudart
|
|
import subprocess
|
|
import sys
|
|
|
|
# Lazy import tensorrt to avoid import conflicts
|
|
_trt = None
|
|
_trt_available = False
|
|
|
|
def get_trt():
|
|
global _trt, _trt_available
|
|
if _trt is None:
|
|
try:
|
|
import tensorrt as trt
|
|
_trt = trt
|
|
_trt_available = True
|
|
except ImportError:
|
|
print("[ComfyUI-RIFE-TensorRT] Warning: TensorRT not available")
|
|
_trt = None
|
|
_trt_available = False
|
|
return _trt
|
|
|
|
def is_trt_available():
|
|
global _trt_available
|
|
return _trt_available
|
|
|
|
def get_trt_logger():
|
|
trt = get_trt()
|
|
if trt:
|
|
return trt.Logger(trt.Logger.ERROR)
|
|
else:
|
|
# Create a simple fallback logger
|
|
class SimpleLogger:
|
|
def __init__(self, level):
|
|
self.level = level
|
|
ERROR = 0
|
|
WARNING = 1
|
|
return SimpleLogger(SimpleLogger.ERROR)
|
|
|
|
TRT_LOGGER = get_trt_logger()
|
|
G_LOGGER.module_severity = G_LOGGER.ERROR
|
|
|
|
# Map of numpy dtype -> torch dtype
|
|
numpy_to_torch_dtype_dict = {
|
|
np.uint8: torch.uint8,
|
|
np.int8: torch.int8,
|
|
np.int16: torch.int16,
|
|
np.int32: torch.int32,
|
|
np.int64: torch.int64,
|
|
np.float16: torch.float16,
|
|
np.float32: torch.float32,
|
|
np.float64: torch.float64,
|
|
np.complex64: torch.complex64,
|
|
np.complex128: torch.complex128,
|
|
}
|
|
if np.version.full_version >= "1.24.0":
|
|
numpy_to_torch_dtype_dict[np.bool_] = torch.bool
|
|
else:
|
|
numpy_to_torch_dtype_dict[np.bool] = torch.bool
|
|
|
|
# Map of torch dtype -> numpy dtype
|
|
torch_to_numpy_dtype_dict = {
|
|
value: key for (key, value) in numpy_to_torch_dtype_dict.items()
|
|
}
|
|
|
|
# https://github.com/Jeff-LiangF/streamv2v/blob/18c1a3bd56ff348d54a3300605936980bb13b03c/src/streamv2v/acceleration/tensorrt/utilities.py
|
|
def CUASSERT(cuda_ret):
|
|
err = cuda_ret[0]
|
|
if err != cudart.cudaError_t.cudaSuccess:
|
|
# Special handling for CUDA ERROR 35 (CUDA_ERROR_NO_DEVICE)
|
|
if err == 35:
|
|
raise RuntimeError(
|
|
f"CUDA ERROR: {err} (Device not available or CUDA state corrupted). "
|
|
f"Try: 1) Restart ComfyUI, 2) Check GPU availability with nvidia-smi, "
|
|
f"3) Ensure no other processes are using the GPU. "
|
|
f"Error reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t"
|
|
)
|
|
else:
|
|
raise RuntimeError(
|
|
f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t"
|
|
)
|
|
if len(cuda_ret) > 1:
|
|
return cuda_ret[1]
|
|
return None
|
|
|
|
def safe_cuda_call(func, *args, max_retries=3, **kwargs):
|
|
"""Wrapper for CUDA calls with retry mechanism for ERROR 35"""
|
|
for attempt in range(max_retries):
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except RuntimeError as e:
|
|
if "CUDA ERROR: 35" in str(e) and attempt < max_retries - 1:
|
|
print(f"⚠️ CUDA ERROR 35 detected, attempt {attempt + 1}/{max_retries}")
|
|
print("Trying to reset CUDA context...")
|
|
|
|
# Try to reset CUDA context
|
|
try:
|
|
import torch
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.synchronize()
|
|
print("✅ CUDA cache cleared and synchronized")
|
|
except Exception as reset_error:
|
|
print(f"⚠️ Could not reset CUDA: {reset_error}")
|
|
|
|
# Small delay before retry
|
|
import time
|
|
time.sleep(2) # Increased delay for cloud environments
|
|
|
|
if attempt == max_retries - 2: # Last attempt
|
|
print("🔄 Final retry attempt...")
|
|
else:
|
|
# Re-raise the error if it's not ERROR 35 or we've exhausted retries
|
|
raise e
|
|
|
|
def safe_cuda_call_with_graph_fallback(func, *args, **kwargs):
|
|
"""Wrapper with CUDA graph fallback for persistent ERROR 35"""
|
|
try:
|
|
# First try with CUDA graph enabled
|
|
return safe_cuda_call(func, *args, **kwargs)
|
|
except RuntimeError as e:
|
|
if "CUDA ERROR: 35" in str(e):
|
|
print("⚠️ CUDA ERROR 35 persists, trying without CUDA graph...")
|
|
# Force disable CUDA graph and retry
|
|
if 'use_cuda_graph' in kwargs:
|
|
kwargs['use_cuda_graph'] = False
|
|
print("🔄 Retrying inference without CUDA graph optimization...")
|
|
return safe_cuda_call(func, *args, max_retries=2, **kwargs)
|
|
raise e
|
|
|
|
def build_progress_feedback(stop_event):
|
|
"""Simple progress feedback during TensorRT engine build"""
|
|
phases = [
|
|
"🔧 Analyzing ONNX model...",
|
|
"⚡ Optimizing layers...",
|
|
"🏗️ Building TensorRT engine...",
|
|
"🔧 Tuning performance...",
|
|
"✅ Finalizing engine..."
|
|
]
|
|
|
|
start_time = time.time()
|
|
|
|
for i, phase in enumerate(phases):
|
|
if stop_event.is_set():
|
|
break
|
|
|
|
elapsed = int(time.time() - start_time)
|
|
print(f"{phase} ({elapsed}s elapsed)")
|
|
|
|
# Wait between phases (simulate progress)
|
|
for j in range(10): # Check every 0.5 seconds for 5 seconds per phase
|
|
if stop_event.is_set():
|
|
break
|
|
time.sleep(0.5)
|
|
|
|
if not stop_event.is_set():
|
|
elapsed = int(time.time() - start_time)
|
|
print(f"🎯 Build completed in {elapsed}s!")
|
|
|
|
class TQDMProgressMonitor:
|
|
def __init__(self):
|
|
trt = get_trt()
|
|
# Initialize attributes regardless of TensorRT availability
|
|
self._active_phases = {}
|
|
self._step_result = True
|
|
self.max_indent = 5
|
|
|
|
# Only inherit from real TensorRT IProgressMonitor if available
|
|
if is_trt_available():
|
|
try:
|
|
trt.IProgressMonitor.__init__(self)
|
|
except Exception as e:
|
|
print(f"Warning: Could not initialize IProgressMonitor: {e}")
|
|
# If dummy, don't try to inherit - just work with our methods
|
|
|
|
def phase_start(self, phase_name, parent_phase, num_steps):
|
|
leave = False
|
|
try:
|
|
if parent_phase is not None:
|
|
nbIndents = (
|
|
self._active_phases.get(parent_phase, {}).get(
|
|
"nbIndents", self.max_indent
|
|
)
|
|
+ 1
|
|
)
|
|
if nbIndents >= self.max_indent:
|
|
return
|
|
else:
|
|
nbIndents = 0
|
|
leave = True
|
|
self._active_phases[phase_name] = {
|
|
"tq": tqdm(
|
|
total=num_steps, desc=phase_name, leave=leave, position=nbIndents
|
|
),
|
|
"nbIndents": nbIndents,
|
|
"parent_phase": parent_phase,
|
|
}
|
|
except KeyboardInterrupt:
|
|
# The phase_start callback cannot directly cancel the build, so request the cancellation from within step_complete.
|
|
_step_result = False
|
|
|
|
def phase_finish(self, phase_name):
|
|
try:
|
|
if phase_name in self._active_phases.keys():
|
|
self._active_phases[phase_name]["tq"].update(
|
|
self._active_phases[phase_name]["tq"].total
|
|
- self._active_phases[phase_name]["tq"].n
|
|
)
|
|
|
|
parent_phase = self._active_phases[phase_name].get("parent_phase", None)
|
|
while parent_phase is not None:
|
|
self._active_phases[parent_phase]["tq"].refresh()
|
|
parent_phase = self._active_phases[parent_phase].get(
|
|
"parent_phase", None
|
|
)
|
|
if (
|
|
self._active_phases[phase_name]["parent_phase"]
|
|
in self._active_phases.keys()
|
|
):
|
|
self._active_phases[
|
|
self._active_phases[phase_name]["parent_phase"]
|
|
]["tq"].refresh()
|
|
del self._active_phases[phase_name]
|
|
pass
|
|
except KeyboardInterrupt:
|
|
_step_result = False
|
|
|
|
def step_complete(self, phase_name, step):
|
|
try:
|
|
if phase_name in self._active_phases.keys():
|
|
self._active_phases[phase_name]["tq"].update(
|
|
step - self._active_phases[phase_name]["tq"].n
|
|
)
|
|
return self._step_result
|
|
except KeyboardInterrupt:
|
|
# There is no need to propagate this exception to TensorRT. We can simply cancel the build.
|
|
return False
|
|
|
|
class Engine:
|
|
def __init__(
|
|
self,
|
|
engine_path,
|
|
):
|
|
self.engine_path = engine_path
|
|
self.engine = None
|
|
self.context = None
|
|
self.buffers = OrderedDict()
|
|
self.tensors = OrderedDict()
|
|
self.inputs = {} # Initialize here
|
|
self.outputs = {} # Initialize here
|
|
self.cuda_graph_instance = None # cuda graph
|
|
self.graph = None
|
|
|
|
def __del__(self):
|
|
# Clean up CUDA graph resources
|
|
if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
|
|
try:
|
|
cudart.cudaGraphExecDestroy(self.cuda_graph_instance)
|
|
except:
|
|
pass
|
|
if hasattr(self, 'graph') and self.graph is not None:
|
|
try:
|
|
cudart.cudaGraphDestroy(self.graph)
|
|
except:
|
|
pass
|
|
|
|
del self.engine
|
|
del self.context
|
|
del self.buffers
|
|
del self.tensors
|
|
|
|
def reset(self, engine_path=None):
|
|
# Clean up CUDA graph resources first
|
|
if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
|
|
try:
|
|
cudart.cudaGraphExecDestroy(self.cuda_graph_instance)
|
|
except:
|
|
pass
|
|
self.cuda_graph_instance = None
|
|
if hasattr(self, 'graph') and self.graph is not None:
|
|
try:
|
|
cudart.cudaGraphDestroy(self.graph)
|
|
except:
|
|
pass
|
|
self.graph = None
|
|
|
|
if hasattr(self, 'engine') and self.engine is not None:
|
|
del self.engine
|
|
if hasattr(self, 'context') and self.context is not None:
|
|
del self.context
|
|
if hasattr(self, 'buffers'):
|
|
del self.buffers
|
|
if hasattr(self, 'tensors'):
|
|
del self.tensors
|
|
|
|
self.engine = None
|
|
self.context = None
|
|
self.engine_path = engine_path if engine_path else self.engine_path
|
|
|
|
self.buffers = OrderedDict()
|
|
self.tensors = OrderedDict()
|
|
self.inputs = {}
|
|
self.outputs = {}
|
|
|
|
def build(
|
|
self,
|
|
onnx_path,
|
|
fp16,
|
|
input_profile=None,
|
|
enable_refit=False,
|
|
enable_preview=False,
|
|
enable_all_tactics=False,
|
|
timing_cache=None,
|
|
update_output_names=None,
|
|
):
|
|
print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")
|
|
|
|
# Simple check - if TensorRT is not available, we can't build
|
|
if not is_trt_available():
|
|
print("❌ TensorRT is not available. Cannot build engine.")
|
|
raise RuntimeError("TensorRT not available - please install TensorRT first")
|
|
|
|
print("✅ TensorRT is available, proceeding with build...")
|
|
|
|
# Start progress feedback thread
|
|
stop_event = threading.Event()
|
|
progress_thread = threading.Thread(target=build_progress_feedback, args=(stop_event,))
|
|
progress_thread.daemon = True
|
|
progress_thread.start()
|
|
|
|
try:
|
|
p = [Profile()]
|
|
if input_profile:
|
|
p = [Profile() for i in range(len(input_profile))]
|
|
for _p, i_profile in zip(p, input_profile):
|
|
for name, dims in i_profile.items():
|
|
assert len(dims) == 3
|
|
_p.add(name, min=dims[0], opt=dims[1], max=dims[2])
|
|
|
|
config_kwargs = {}
|
|
if not enable_all_tactics:
|
|
config_kwargs["tactic_sources"] = []
|
|
|
|
try:
|
|
trt = get_trt()
|
|
network = network_from_onnx_path(
|
|
onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]
|
|
)
|
|
except Exception as e:
|
|
print(f"❌ Failed to load ONNX network: {e}")
|
|
raise RuntimeError(f"ONNX loading failed: {e}")
|
|
|
|
if update_output_names:
|
|
print(f"Updating network outputs to {update_output_names}")
|
|
network = ModifyNetworkOutputs(network, update_output_names)
|
|
|
|
builder = network[0]
|
|
config = builder.create_builder_config()
|
|
|
|
# Skip progress monitor for simplicity - avoid interface issues
|
|
# config.progress_monitor = TQDMProgressMonitor()
|
|
|
|
trt = get_trt()
|
|
config.set_flag(trt.BuilderFlag.FP16) if fp16 else None
|
|
config.set_flag(trt.BuilderFlag.REFIT) if enable_refit else None
|
|
|
|
profiles = copy.deepcopy(p)
|
|
for profile in profiles:
|
|
# Last profile is used for set_calibration_profile.
|
|
calib_profile = profile.fill_defaults(network[1]).to_trt(
|
|
builder, network[1]
|
|
)
|
|
config.add_optimization_profile(calib_profile)
|
|
|
|
try:
|
|
engine = engine_from_network(
|
|
network,
|
|
config,
|
|
)
|
|
except Exception as e:
|
|
error(f"Failed to build engine: {e}")
|
|
return 1
|
|
try:
|
|
save_engine(engine, path=self.engine_path)
|
|
print(f"✅ Engine saved successfully to: {self.engine_path}")
|
|
except Exception as e:
|
|
error(f"Failed to save engine: {e}")
|
|
return 1
|
|
return 0
|
|
finally:
|
|
# Stop progress feedback thread
|
|
stop_event.set()
|
|
progress_thread.join(timeout=1)
|
|
|
|
def load(self):
|
|
self.engine = engine_from_bytes(bytes_from_path(self.engine_path))
|
|
|
|
def activate(self, reuse_device_memory=None):
|
|
# If engine was reset, reload it
|
|
if self.engine is None:
|
|
self.load()
|
|
|
|
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 get_input_profile_bounds(self):
|
|
"""Return min/opt/max shapes for the first optimization profile.
|
|
|
|
Returns a dict: {tensor_name: (min_shape, opt_shape, max_shape)}.
|
|
Uses engine.get_tensor_profile_shape (TensorRT 10+). Returns empty
|
|
dict if the API is unavailable (older TensorRT builds).
|
|
"""
|
|
bounds = {}
|
|
try:
|
|
trt_instance = get_trt()
|
|
for idx in range(self.engine.num_io_tensors):
|
|
name = self.engine.get_tensor_name(idx)
|
|
if self.engine.get_tensor_mode(name) == trt_instance.TensorIOMode.INPUT:
|
|
# get_tensor_profile_shape returns [min, opt, max] as trt.Dims
|
|
profile_shapes = self.engine.get_tensor_profile_shape(name, 0)
|
|
min_shape = tuple(profile_shapes[0])
|
|
opt_shape = tuple(profile_shapes[1])
|
|
max_shape = tuple(profile_shapes[2])
|
|
bounds[name] = (min_shape, opt_shape, max_shape)
|
|
except Exception as e:
|
|
print(f"⚠️ Could not retrieve profile bounds: {e}")
|
|
return bounds
|
|
|
|
def allocate_buffers(self, shape_dict=None, device="cuda"):
|
|
# Clean up CUDA graph resources since tensors will be recreated
|
|
if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
|
|
try:
|
|
cudart.cudaGraphExecDestroy(self.cuda_graph_instance)
|
|
except:
|
|
pass
|
|
self.cuda_graph_instance = None
|
|
if hasattr(self, 'graph') and self.graph is not None:
|
|
try:
|
|
cudart.cudaGraphDestroy(self.graph)
|
|
except:
|
|
pass
|
|
self.graph = None
|
|
|
|
nvtx.range_push("allocate_buffers")
|
|
for idx in range(self.engine.num_io_tensors):
|
|
name = self.engine.get_tensor_name(idx)
|
|
binding = self.engine[idx]
|
|
if shape_dict and binding in shape_dict:
|
|
shape = shape_dict[binding]["shape"]
|
|
else:
|
|
shape = self.context.get_tensor_shape(name)
|
|
|
|
trt_instance = get_trt()
|
|
dtype = trt_instance.nptype(self.engine.get_tensor_dtype(name))
|
|
if self.engine.get_tensor_mode(name) == trt_instance.TensorIOMode.INPUT:
|
|
self.context.set_input_shape(name, shape)
|
|
tensor = torch.empty(
|
|
tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]
|
|
).to(device)
|
|
|
|
self.buffers[name] = tensor
|
|
self.tensors[name] = tensor
|
|
|
|
if self.engine.get_tensor_mode(name) == trt_instance.TensorIOMode.INPUT:
|
|
self.inputs[name] = tensor
|
|
else:
|
|
self.outputs[name] = tensor
|
|
nvtx.range_pop()
|
|
|
|
def infer(self, feed_dict, stream, use_cuda_graph=False):
|
|
"""Inference with CUDA error recovery and graph fallback"""
|
|
def _do_inference():
|
|
for name, buf in feed_dict.items():
|
|
self.tensors[name].copy_(buf)
|
|
|
|
nvtx.range_push("TensorRT Inference")
|
|
|
|
# Set tensor addresses explicitly for all bindings
|
|
for name, tensor in self.tensors.items():
|
|
self.context.set_tensor_address(name, tensor.data_ptr())
|
|
|
|
if use_cuda_graph:
|
|
if self.cuda_graph_instance is not None:
|
|
CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream.ptr))
|
|
CUASSERT(cudart.cudaStreamSynchronize(stream.ptr))
|
|
else:
|
|
# do inference before CUDA graph capture
|
|
noerror = self.context.execute_async_v3(stream.ptr)
|
|
if not noerror:
|
|
raise ValueError("ERROR: inference failed.")
|
|
# capture cuda graph
|
|
CUASSERT(
|
|
cudart.cudaStreamBeginCapture(stream.ptr, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal)
|
|
)
|
|
self.context.execute_async_v3(stream.ptr)
|
|
self.graph = CUASSERT(cudart.cudaStreamEndCapture(stream.ptr))
|
|
self.cuda_graph_instance = CUASSERT(cudart.cudaGraphInstantiate(self.graph, 0))
|
|
else:
|
|
noerror = self.context.execute_async_v3(stream.ptr)
|
|
if not noerror:
|
|
raise ValueError("ERROR: inference failed.")
|
|
|
|
nvtx.range_pop()
|
|
return self.tensors
|
|
|
|
# Use safe wrapper with CUDA graph fallback for cloud environments
|
|
return safe_cuda_call_with_graph_fallback(_do_inference)
|