Merge pull request #1 from ComfyAssets/fix/pr14-review-fixes

vitosans · web-flow · commit e7c216d172e3 · 2026-03-12T16:37:48.000-07:00
Fix error handling and CUDA cleanup from PR yuvraj108c#14
diff --git a/__init__.py b/__init__.py
@@ -3,23 +3,153 @@
 from comfy.model_management import get_torch_device
 from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife, logger
 from .trt_utilities import Engine
+from .utilities import download_file, ColoredLogger
 import folder_paths
 import time
 from polygraphy import cuda
+import comfy.model_management as mm
+import tensorrt
+import json
 
 ENGINE_DIR = os.path.join(folder_paths.models_dir, "tensorrt", "rife")
 
+# Image dimensions for TensorRT engine building
+IMAGE_DIM_MIN = 256
+IMAGE_DIM_OPT = 512
+IMAGE_DIM_MAX = 3840
+
+# Logger for this module
+rife_logger = ColoredLogger("ComfyUI-Rife-Tensorrt")
+
+# Function to load configuration
+def load_node_config(config_filename="load_rife_config.json"):
+    """Loads node configuration from a JSON file."""
+    current_dir = os.path.dirname(__file__)
+    config_path = os.path.join(current_dir, config_filename)
+
+    default_config = {
+        "model": {
+            "options": ["rife49_ensemble_True_scale_1_sim"],
+            "default": "rife49_ensemble_True_scale_1_sim",
+            "tooltip": "Default model (fallback from code)"
+        },
+        "precision": {
+            "options": ["fp16", "fp32"],
+            "default": "fp16",
+            "tooltip": "Default precision (fallback from code)"
+        }
+    }
+
+    try:
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+        rife_logger.info(f"Successfully loaded configuration from {config_filename}")
+        return config
+    except FileNotFoundError:
+        rife_logger.warning(f"Configuration file '{config_path}' not found. Using default fallback configuration.")
+        return default_config
+    except json.JSONDecodeError:
+        rife_logger.error(f"Error decoding JSON from '{config_path}'. Using default fallback configuration.")
+        return default_config
+    except Exception as e:
+        rife_logger.error(f"An unexpected error occurred while loading '{config_path}': {e}. Using default fallback.")
+        return default_config
+
+# Load the configuration once when the module is imported
+LOAD_RIFE_NODE_CONFIG = load_node_config()
+
+class LoadRifeTensorrtModel:
+    @classmethod
+    def INPUT_TYPES(cls):
+        # Use the pre-loaded configuration
+        model_config = LOAD_RIFE_NODE_CONFIG.get("model", {})
+        precision_config = LOAD_RIFE_NODE_CONFIG.get("precision", {})
+
+        # Provide sensible defaults if keys are missing in the config
+        model_options = model_config.get("options", ["rife49_ensemble_True_scale_1_sim"])
+        model_default = model_config.get("default", "rife49_ensemble_True_scale_1_sim")
+        model_tooltip = model_config.get("tooltip", "Select a RIFE model.")
+
+        precision_options = precision_config.get("options", ["fp16", "fp32"])
+        precision_default = precision_config.get("default", "fp16")
+        precision_tooltip = precision_config.get("tooltip", "Select precision.")
+
+        return {
+            "required": {
+                "model": (model_options, {"default": model_default, "tooltip": model_tooltip}),
+                "precision": (precision_options, {"default": precision_default, "tooltip": precision_tooltip}),
+            }
+        }
+
+    RETURN_NAMES = ("rife_trt_model",)
+    RETURN_TYPES = ("RIFE_TRT_MODEL",)
+    CATEGORY = "tensorrt"
+    DESCRIPTION = "Load RIFE tensorrt models, they will be built automatically if not found."
+    FUNCTION = "load_rife_tensorrt_model"
+
+    def load_rife_tensorrt_model(self, model, precision):
+        tensorrt_models_dir = os.path.join(folder_paths.models_dir, "tensorrt", "rife")
+        onnx_models_dir = os.path.join(folder_paths.models_dir, "onnx")
+
+        os.makedirs(tensorrt_models_dir, exist_ok=True)
+        os.makedirs(onnx_models_dir, exist_ok=True)
+
+        onnx_model_path = os.path.join(onnx_models_dir, f"{model}.onnx")
+
+        # Build tensorrt model path with detailed naming
+        engine_channel = 3
+        engine_min_batch, engine_opt_batch, engine_max_batch = 1, 1, 1
+        engine_min_h, engine_opt_h, engine_max_h = IMAGE_DIM_MIN, IMAGE_DIM_OPT, IMAGE_DIM_MAX
+        engine_min_w, engine_opt_w, engine_max_w = IMAGE_DIM_MIN, IMAGE_DIM_OPT, IMAGE_DIM_MAX
+        tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{model}_{precision}_{engine_min_batch}x{engine_channel}x{engine_min_h}x{engine_min_w}_{engine_opt_batch}x{engine_channel}x{engine_opt_h}x{engine_opt_w}_{engine_max_batch}x{engine_channel}x{engine_max_h}x{engine_max_w}_{tensorrt.__version__}.trt")
+
+        if not os.path.exists(tensorrt_model_path):
+            if not os.path.exists(onnx_model_path):
+                onnx_model_download_url = f"https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/{model}.onnx"
+                rife_logger.info(f"Downloading {onnx_model_download_url}")
+                download_file(url=onnx_model_download_url, save_path=onnx_model_path)
+            else:
+                rife_logger.info(f"ONNX model found at: {onnx_model_path}")
+
+            rife_logger.info(f"Building TensorRT engine for {onnx_model_path}: {tensorrt_model_path}")
+            mm.soft_empty_cache()
+            s = time.time()
+            engine = Engine(tensorrt_model_path)
+            ret = engine.build(
+                onnx_path=onnx_model_path,
+                fp16=True if precision == "fp16" else False,
+                input_profile=[
+                    {
+                        "img0": [(engine_min_batch, engine_channel, engine_min_h, engine_min_w), (engine_opt_batch, engine_channel, engine_opt_h, engine_opt_w), (engine_max_batch, engine_channel, engine_max_h, engine_max_w)],
+                        "img1": [(engine_min_batch, engine_channel, engine_min_h, engine_min_w), (engine_opt_batch, engine_channel, engine_opt_h, engine_opt_w), (engine_max_batch, engine_channel, engine_max_h, engine_max_w)],
+                    }
+                ],
+            )
+            if ret != 0:
+                if os.path.exists(tensorrt_model_path):
+                    os.remove(tensorrt_model_path)
+                raise RuntimeError(f"Failed to build TensorRT engine from {onnx_model_path}")
+            e = time.time()
+            rife_logger.info(f"Time taken to build: {(e-s)} seconds")
+
+        rife_logger.info(f"Loading TensorRT engine: {tensorrt_model_path}")
+        mm.soft_empty_cache()
+        engine = Engine(tensorrt_model_path)
+        engine.load()
+
+        return (engine,)
+
 class RifeTensorrt:
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
-                "frames": ("IMAGE", ),
-                "engine": (os.listdir(ENGINE_DIR),),
-                "clear_cache_after_n_frames": ("INT", {"default": 100, "min": 1, "max": 1000}),
-                "multiplier": ("INT", {"default": 2, "min": 1}),
-                "use_cuda_graph": ("BOOLEAN", {"default": True}),
-                "keep_model_loaded": ("BOOLEAN", {"default": False}),
+                "frames": ("IMAGE", {"tooltip": "Input frames for video frame interpolation"}),
+                "rife_trt_model": ("RIFE_TRT_MODEL", {"tooltip": "Tensorrt model built and loaded"}),
+                "clear_cache_after_n_frames": ("INT", {"default": 100, "min": 1, "max": 1000, "tooltip": "Clear CUDA cache after processing this many frames"}),
+                "multiplier": ("INT", {"default": 2, "min": 1, "tooltip": "Frame interpolation multiplier"}),
+                "use_cuda_graph": ("BOOLEAN", {"default": True, "tooltip": "Use CUDA graph for better performance"}),
+                "keep_model_loaded": ("BOOLEAN", {"default": False, "tooltip": "Keep model loaded in memory after processing"}),
             },
         }
 
@@ -31,7 +161,7 @@ def INPUT_TYPES(s):
     def vfi(
         self,
         frames,
-        engine,
+        rife_trt_model,
         clear_cache_after_n_frames=100,
         multiplier=2,
         use_cuda_graph=True,
@@ -45,24 +175,21 @@ def vfi(
         }
 
         cudaStream = cuda.Stream()
-        engine_path = os.path.join(ENGINE_DIR, engine)
-        if (not hasattr(self, 'engine') or self.engine_label != engine):
-            self.engine = Engine(engine_path)
-            logger(f"Loading TensorRT engine: {engine_path}")
-            self.engine.load()
-            self.engine.activate()
-            self.engine_label = engine
-        else:
-            logger(f"Using cached TensorRT engine: {engine_path}")
-
-        self.engine.allocate_buffers(shape_dict=shape_dict)
+
+        # Use the provided model directly
+        engine = rife_trt_model
+        logger(f"Using loaded TensorRT engine")
+
+        # Activate and allocate buffers for the engine
+        engine.activate()
+        engine.allocate_buffers(shape_dict=shape_dict)
 
         frames = preprocess_frames(frames)
 
         def return_middle_frame(frame_0, frame_1, timestep):
             timestep_t = torch.tensor([timestep], dtype=torch.float32).to(get_torch_device())
             # s = time.time()
-            output = self.engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream, use_cuda_graph)
+            output = engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream, use_cuda_graph)
             # e = time.time()
             # print(f"Time taken to infer: {(e-s)*1000} ms")
 
@@ -71,19 +198,21 @@ def return_middle_frame(frame_0, frame_1, timestep):
 
         result = generate_frames_rife(frames, clear_cache_after_n_frames, multiplier, return_middle_frame)
         out = postprocess_frames(result)
-        
+
         if not keep_model_loaded:
-            del self.engine, self.engine_label
+            engine.reset()
 
         return (out,)
 
 
 NODE_CLASS_MAPPINGS = {
     "RifeTensorrt": RifeTensorrt,
+    "LoadRifeTensorrtModel": LoadRifeTensorrtModel,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "RifeTensorrt": "⚡ Rife Tensorrt",
+    "LoadRifeTensorrtModel": "Load Rife Tensorrt Model",
 }
 
 __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
diff --git a/load_rife_config.json b/load_rife_config.json
@@ -0,0 +1,16 @@
+{
+  "model": {
+    "options": [
+      "rife47_ensemble_True_scale_1_sim",
+      "rife48_ensemble_True_scale_1_sim",
+      "rife49_ensemble_True_scale_1_sim"
+    ],
+    "default": "rife49_ensemble_True_scale_1_sim",
+    "tooltip": "RIFE models for video frame interpolation. These models have been tested with tensorrt. Loaded from config."
+  },
+  "precision": {
+    "options": ["fp16", "fp32"],
+    "default": "fp16",
+    "tooltip": "Precision to build the tensorrt engines. Loaded from config."
+  }
+}
diff --git a/readme.md b/readme.md
@@ -2,9 +2,9 @@
 
 # ComfyUI Rife TensorRT ⚡
 
-[![python](https://img.shields.io/badge/python-3.10.12-green)](https://www.python.org/downloads/release/python-31012/)
-[![cuda](https://img.shields.io/badge/cuda-12.4-green)](https://developer.nvidia.com/cuda-downloads)
-[![trt](https://img.shields.io/badge/TRT-10.4.0-green)](https://developer.nvidia.com/tensorrt)
+[![python](https://img.shields.io/badge/python-3.12.11-green)](https://www.python.org/downloads/release/python-31211/)
+[![cuda](https://img.shields.io/badge/cuda-12.9-green)](https://developer.nvidia.com/cuda-downloads)
+[![trt](https://img.shields.io/badge/TRT-10.13.3.9-green)](https://developer.nvidia.com/tensorrt)
 [![by-nc-sa/4.0](https://img.shields.io/badge/license-CC--BY--NC--SA--4.0-lightgrey)](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.en)
 
 ![node](https://github.com/user-attachments/assets/5fd6d529-300c-42a5-b9cf-46e031f0bcb5)
@@ -40,27 +40,40 @@ cd ./ComfyUI-Rife-Tensorrt
 pip install -r requirements.txt
 ```
 
-## 🛠️ Building Tensorrt Engine
+## 🛠️ Supported Models
 
-1. Download one of the following onnx models:
-   - [rife49_ensemble_True_scale_1_sim.onnx](https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/rife49_ensemble_True_scale_1_sim.onnx)
-   - [rife48_ensemble_True_scale_1_sim.onnx](https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/rife48_ensemble_True_scale_1_sim.onnx)
-   - [rife47_ensemble_True_scale_1_sim.onnx](https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/rife47_ensemble_True_scale_1_sim.onnx)
-2. Edit onnx/trt paths inside [export_trt.py](./export_trt.py) and build tensorrt engine by running:
-   - `python export_trt.py`
+The following RIFE models are supported and will be automatically downloaded and built:
+   - **rife49_ensemble_True_scale_1_sim** (default) - Latest and most accurate
+   - **rife48_ensemble_True_scale_1_sim** - Good balance of speed and quality
+   - **rife47_ensemble_True_scale_1_sim** - Fastest option
 
-3. Place the exported engine inside ComfyUI `/models/tensorrt/rife` directory
+Models are automatically downloaded from [HuggingFace](https://huggingface.co/yuvraj108c/rife-onnx) and TensorRT engines are built on first use.
 
 ## ☀️ Usage
 
-- Insert node by `Right Click -> tensorrt -> Rife Tensorrt`
-- Image resolutions between `256x256` and `3840x3840` will work with the tensorrt engines 
+1. **Load Model**: Insert `Right Click -> Add Node -> tensorrt -> Load Rife Tensorrt Model`
+   - Choose your preferred RIFE model (rife47, rife48, or rife49)
+   - Select precision (fp16 recommended for speed, fp32 for maximum accuracy)
+   - The model will be automatically downloaded and TensorRT engine built on first use
+
+2. **Process Frames**: Insert `Right Click -> Add Node -> tensorrt -> Rife Tensorrt`
+   - Connect the loaded model from step 1
+   - Input your video frames
+   - Configure interpolation settings (multiplier, CUDA graph, etc.)
+   - Image resolutions between `256x256` and `3840x3840` are supported 
 
 ## 🤖 Environment tested
 
-- Ubuntu 22.04 LTS, Cuda 12.4, Tensorrt 10.4.0, Python 3.10, RTX 3070 GPU
+- WSL Ubuntu 24.04.03 LTS, Cuda 12.9, Tensorrt 10.13.3.9, Python 3.12.11, RTX 5080 GPU
 - Windows (Not tested, but should work)
 
+## 🚨 Updates
+
+### December 2025
+- **Automatic Model Management**: No more manual downloads! Models are automatically downloaded from HuggingFace and TensorRT engines are built on demand
+- **Improved Workflow**: New two-node system with `Load Rife Tensorrt Model` + `Rife Tensorrt` for better organization
+- **Updated Dependencies**: TensorRT updated to 10.13.3.9 for better performance and compatibility
+
 ## 👏 Credits
 
 - https://github.com/styler00dollar/VSGAN-tensorrt-docker
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,9 @@
 einops
 colored
 polygraphy
-tensorrt==10.4.0
-cuda-python
+#tensorrt==10.13.3.9
+tensorrt
+cuda-python
+requests
+
+tqdm
diff --git a/trt_utilities.py b/trt_utilities.py
diff --git a/utilities.py b/utilities.py
