diff --git a/.gitignore b/.gitignore index 236f609..9881d64 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ models __pycache__ -.vscode \ No newline at end of file +.vscode +CLAUDE.md diff --git a/__init__.py b/__init__.py index c13155b..0ead840 100644 --- a/__init__.py +++ b/__init__.py @@ -1,89 +1,14 @@ -import torch -import os -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 -import folder_paths -import time -from polygraphy import cuda - -ENGINE_DIR = os.path.join(folder_paths.models_dir, "tensorrt", "rife") - -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}), - }, - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "tensorrt" - OUTPUT_NODE=True - - def vfi( - self, - frames, - engine, - clear_cache_after_n_frames=100, - multiplier=2, - use_cuda_graph=True, - keep_model_loaded=False, - ): - B, H, W, C = frames.shape - shape_dict = { - "img0": {"shape": (1, 3, H, W)}, - "img1": {"shape": (1, 3, H, W)}, - "output": {"shape": (1, 3, H, W)}, - } - - 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) - - 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) - # e = time.time() - # print(f"Time taken to infer: {(e-s)*1000} ms") - - result = output['output'] - return result - - 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 - - return (out,) - +from .nodes.load_rife_tensorrt import LoadRifeTensorrtModel +from .nodes.rife_tensorrt import RifeTensorrt 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 new file mode 100644 index 0000000..3c01bd0 --- /dev/null +++ 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." + } +} \ No newline at end of file diff --git a/nodes/load_rife_tensorrt.py b/nodes/load_rife_tensorrt.py new file mode 100644 index 0000000..3be2d8b --- /dev/null +++ b/nodes/load_rife_tensorrt.py @@ -0,0 +1,91 @@ +from ..trt_utilities import Engine +from ..utilities import download_file, load_node_config, rife_logger +import folder_paths +import time +import comfy.model_management as mm +import tensorrt +import os + +# Image dimensions for TensorRT engine building +IMAGE_DIM_MIN = 256 +IMAGE_DIM_OPT = 512 +IMAGE_DIM_MAX = 3840 + +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) + 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)], + } + ], + ) + 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,) diff --git a/nodes/rife_tensorrt.py b/nodes/rife_tensorrt.py new file mode 100644 index 0000000..f85d529 --- /dev/null +++ b/nodes/rife_tensorrt.py @@ -0,0 +1,56 @@ +import torch +import os +from comfy.model_management import get_torch_device +from ..vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife +from ..trt_utilities import Engine +import folder_paths +import time +import comfy.model_management as mm + +class RifeTensorrt: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "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"}), + }, + } + + RETURN_TYPES = ("IMAGE", ) + FUNCTION = "vfi" + CATEGORY = "tensorrt" + + def vfi( + self, + frames, + rife_trt_model, + clear_cache_after_n_frames=100, + multiplier=2, + ): + B, H, W, C = frames.shape + shape_dict = { + "img0": {"shape": (1, 3, H, W)}, + "img1": {"shape": (1, 3, H, W)}, + "output": {"shape": (1, 3, H, W)}, + } + + cudaStream = torch.cuda.current_stream().cuda_stream + engine = rife_trt_model + 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()) + output = engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream) + result = output['output'] + return result + + result = generate_frames_rife(frames, clear_cache_after_n_frames, multiplier, return_middle_frame) + out = postprocess_frames(result) + + return (out,) \ No newline at end of file diff --git a/readme.md b/readme.md index 6374357..48323a0 100644 --- a/readme.md +++ b/readme.md @@ -2,21 +2,36 @@ # 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.3-green)](https://www.python.org/downloads/release/python-3123//) +[![cuda](https://img.shields.io/badge/cuda-13.0-green)](https://developer.nvidia.com/cuda-downloads) +[![trt](https://img.shields.io/badge/TRT-10.14.1.48-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) - - - This project provides a [TensorRT](https://github.com/NVIDIA/TensorRT) implementation of [RIFE](https://github.com/hzwer/ECCV2022-RIFE) for ultra fast frame interpolation inside ComfyUI -This project is licensed under [CC BY-NC-SA](https://creativecommons.org/licenses/by-nc-sa/4.0/), everyone is FREE to access, use, modify and redistribute with the same license. +**Last tested**: 08 June 2026 (ComfyUI v0.23.0 | Torch 2.12.0 | Python 3.12.3 | L40S | CUDA 13.0 | Ubuntu 24.04) -If you like the project, please give me a star! ⭐ +Screenshot 2026-06-08 at 09 21 51 + + + +## ⭐ Support +If you like my projects and wish to see updates and new features, please consider supporting me. It helps a lot! + +[![ComfyUI-Depth-Anything-Tensorrt](https://img.shields.io/badge/ComfyUI--Depth--Anything--Tensorrt-blue?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Depth-Anything-Tensorrt) +[![ComfyUI-Upscaler-Tensorrt](https://img.shields.io/badge/ComfyUI--Upscaler--Tensorrt-blue?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Upscaler-Tensorrt) +[![ComfyUI-Dwpose-Tensorrt](https://img.shields.io/badge/ComfyUI--Dwpose--Tensorrt-blue?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Dwpose-Tensorrt) +[![ComfyUI-Rife-Tensorrt](https://img.shields.io/badge/ComfyUI--Rife--Tensorrt-blue?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Rife-Tensorrt) + +[![ComfyUI-Whisper](https://img.shields.io/badge/ComfyUI--Whisper-gray?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Whisper) +[![ComfyUI_InvSR](https://img.shields.io/badge/ComfyUI__InvSR-gray?style=flat-square)](https://github.com/yuvraj108c/ComfyUI_InvSR) +[![ComfyUI-Thera](https://img.shields.io/badge/ComfyUI--Thera-gray?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Thera) +[![ComfyUI-Video-Depth-Anything](https://img.shields.io/badge/ComfyUI--Video--Depth--Anything-gray?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Video-Depth-Anything) +[![ComfyUI-PiperTTS](https://img.shields.io/badge/ComfyUI--PiperTTS-gray?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-PiperTTS) + +[![buy-me-coffees](https://i.imgur.com/3MDbAtw.png)](https://www.buymeacoffee.com/yuvraj108cZ) +[![paypal-donation](https://i.imgur.com/w5jjubk.png)](https://paypal.me/yuvraj108c) --- @@ -40,26 +55,35 @@ 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 -## 🤖 Environment tested +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, etc.) + - Image resolutions between `256x256` and `3840x3840` are supported -- Ubuntu 22.04 LTS, Cuda 12.4, Tensorrt 10.4.0, Python 3.10, RTX 3070 GPU -- Windows (Not tested, but should work) + +## 🚨 Updates + +### 08 June 2026 +- **Automatic Model Management**: No more manual downloads! Models are automatically downloaded from HuggingFace and TensorRT engines are built on demand. [PR#14](https://github.com/yuvraj108c/ComfyUI-Rife-Tensorrt/pull/14) by [@reaperhammer](https://github.com/reaperhammer) +- **Improved Workflow + Codebase**: New two-node system with `Load Rife Tensorrt Model` + `Rife Tensorrt` for better organization +- **Remove cuda-python**: No more cuda installation issues on windows ## 👏 Credits diff --git a/requirements.txt b/requirements.txt index b479d99..c222890 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,4 @@ -einops -colored -polygraphy -tensorrt==10.4.0 -cuda-python \ No newline at end of file +einops +colored +polygraphy +tensorrt \ No newline at end of file diff --git a/export_onnx.py b/scripts/export_onnx.py similarity index 98% rename from export_onnx.py rename to scripts/export_onnx.py index 97f2948..62b7d54 100644 --- a/export_onnx.py +++ b/scripts/export_onnx.py @@ -87,7 +87,7 @@ def export_onnx(ckpt_name, ensemble, scale_factor): model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) arch_ver = CKPT_NAME_VER_DICT[ckpt_name] interpolation_model = IFNet(arch_ver=arch_ver) - interpolation_model.load_state_dict(torch.load(model_path)) + interpolation_model.load_state_dict(torch.load(model_path, map_location=TORCH_DEVICE)) interpolation_model.eval().to(TORCH_DEVICE) # # dummy data diff --git a/export_trt.py b/scripts/export_trt.py similarity index 100% rename from export_trt.py rename to scripts/export_trt.py diff --git a/trt_utilities.py b/trt_utilities.py index eaa680b..3ea5567 100644 --- a/trt_utilities.py +++ b/trt_utilities.py @@ -1,3 +1,20 @@ +# +# Copyright 2022 The HuggingFace Inc. team. +# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# import torch from torch.cuda import nvtx from collections import OrderedDict @@ -16,7 +33,6 @@ import tensorrt as trt from logging import error, warning from tqdm import tqdm import copy -import cuda.bindings.runtime as cudart TRT_LOGGER = trt.Logger(trt.Logger.ERROR) G_LOGGER.module_severity = G_LOGGER.ERROR @@ -44,17 +60,6 @@ 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: - 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 - class TQDMProgressMonitor(trt.IProgressMonitor): def __init__(self): trt.IProgressMonitor.__init__(self) @@ -125,6 +130,7 @@ class TQDMProgressMonitor(trt.IProgressMonitor): # There is no need to propagate this exception to TensorRT. We can simply cancel the build. return False + class Engine: def __init__( self, @@ -166,7 +172,7 @@ class Engine: timing_cache=None, update_output_names=None, ): - print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}") + # print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}") p = [Profile()] if input_profile: p = [Profile() for i in range(len(input_profile))] @@ -217,6 +223,7 @@ class Engine: return 0 def load(self): + # print(f"Loading TensorRT engine: {self.engine_path}") self.engine = engine_from_bytes(bytes_from_path(self.engine_path)) def activate(self, reuse_device_memory=None): @@ -246,32 +253,25 @@ class Engine: nvtx.range_pop() def infer(self, feed_dict, stream, use_cuda_graph=False): + nvtx.range_push("set_tensors") for name, buf in feed_dict.items(): self.tensors[name].copy_(buf) 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() + nvtx.range_push("execute") + noerror = self.context.execute_async_v3(stream) + if not noerror: + raise ValueError("ERROR: inference failed.") + nvtx.range_pop() return self.tensors + def __str__(self): + out = "" + for opt_profile in range(self.engine.num_optimization_profiles): + for binding_idx in range(self.engine.num_bindings): + name = self.engine.get_binding_name(binding_idx) + shape = self.engine.get_profile_shape(opt_profile, name) + out += f"\t{name} = {shape}\n" + return out \ No newline at end of file diff --git a/utilities.py b/utilities.py new file mode 100644 index 0000000..4d805d0 --- /dev/null +++ b/utilities.py @@ -0,0 +1,141 @@ +import requests +from tqdm import tqdm +import logging +import sys +import json +import os + +class ColoredLogger: + COLORS = { + 'RED': '\033[91m', + 'GREEN': '\033[92m', + 'YELLOW': '\033[93m', + 'BLUE': '\033[94m', + 'MAGENTA': '\033[95m', + 'RESET': '\033[0m' + } + + LEVEL_COLORS = { + 'DEBUG': COLORS['BLUE'], + 'INFO': COLORS['GREEN'], + 'WARNING': COLORS['YELLOW'], + 'ERROR': COLORS['RED'], + 'CRITICAL': COLORS['MAGENTA'] + } + + def __init__(self, name="MY-APP"): + self.logger = logging.getLogger(name) + self.logger.setLevel(logging.DEBUG) + self.app_name = name + + # Prevent message propagation to parent loggers + self.logger.propagate = False + + # Clear existing handlers + self.logger.handlers = [] + + # Create console handler + handler = logging.StreamHandler(sys.stdout) + handler.setLevel(logging.DEBUG) + + # Custom formatter class to handle colored components + class ColoredFormatter(logging.Formatter): + def format(self, record): + # Color the level name according to severity + level_color = ColoredLogger.LEVEL_COLORS.get(record.levelname, '') + colored_levelname = f"{level_color}{record.levelname}{ColoredLogger.COLORS['RESET']}" + + # Color the logger name in blue + colored_name = f"{ColoredLogger.COLORS['BLUE']}{record.name}{ColoredLogger.COLORS['RESET']}" + + # Set the colored components + record.levelname = colored_levelname + record.name = colored_name + + return super().format(record) + + # Create formatter with the new format + formatter = ColoredFormatter('[%(name)s|%(levelname)s] - %(message)s') + handler.setFormatter(formatter) + + self.logger.addHandler(handler) + + + def debug(self, message): + self.logger.debug(f"{self.COLORS['BLUE']}{message}{self.COLORS['RESET']}") + + def info(self, message): + self.logger.info(f"{self.COLORS['GREEN']}{message}{self.COLORS['RESET']}") + + def warning(self, message): + self.logger.warning(f"{self.COLORS['YELLOW']}{message}{self.COLORS['RESET']}") + + def error(self, message): + self.logger.error(f"{self.COLORS['RED']}{message}{self.COLORS['RESET']}") + + def critical(self, message): + self.logger.critical(f"{self.COLORS['MAGENTA']}{message}{self.COLORS['RESET']}") + +rife_logger = ColoredLogger("ComfyUI-Rife-Tensorrt") + +def download_file(url, save_path): + """ + Download a file from URL with progress bar + + Args: + url (str): URL of the file to download + save_path (str): Path to save the file as + """ + GREEN = '\033[92m' + RESET = '\033[0m' + response = requests.get(url, stream=True) + response.raise_for_status() + total_size = int(response.headers.get('content-length', 0)) + + with open(save_path, 'wb') as file, tqdm( + desc=save_path, + total=total_size, + unit='iB', + unit_scale=True, + unit_divisor=1024, + colour='green', + bar_format=f'{GREEN}{{l_bar}}{{bar}}{RESET}{GREEN}{{r_bar}}{RESET}' + ) as progress_bar: + for data in response.iter_content(chunk_size=1024): + size = file.write(data) + progress_bar.update(size) + + +# 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 \ No newline at end of file diff --git a/vfi_utilities.py b/vfi_utilities.py index f9ec616..88d7b16 100644 --- a/vfi_utilities.py +++ b/vfi_utilities.py @@ -7,7 +7,9 @@ import einops from comfy.model_management import soft_empty_cache, get_torch_device import numpy as np from comfy.utils import ProgressBar -from colored import Fore, Back, Style +from colored import Fore, Back, Style +from .utilities import rife_logger +from tqdm import tqdm DEVICE = get_torch_device() @@ -25,9 +27,6 @@ def load_file_from_github_release(model_type, ckpt_name): error_str = '\n\n'.join(error_strs) raise Exception(f"Tried all GitHub base urls to download {ckpt_name} but no suceess. Below is the error log:\n\n{error_str}") -def logger(msg): - print(f'{Style.reset}{Fore.cyan}⚡ [Rife Tensorrt] - {msg}{Style.reset}') - def preprocess_frames(frames): return einops.rearrange(frames[..., :3], "n h w c -> n c h w") @@ -45,7 +44,15 @@ def generate_frames_rife( out_len = 0 number_of_frames_processed_since_last_cleared_cuda_cache = 0 - pbar = ProgressBar(len(frames)) + pbar = ProgressBar(len(frames)-1) + + bar_format = "[\033[94mComfyUI-Rife-Tensorrt\033[0m|\033[92mINFO\033[0m] - \033[92m{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]" + progress_bar = tqdm( + total=len(frames)-1, + desc="Interpolating", + bar_format=bar_format, + disable=((len(frames)-1) == 1) + ) for frame_itr in range(len(frames) - 1): # Skip the final frame since there are no frames after it @@ -67,19 +74,25 @@ def generate_frames_rife( if number_of_frames_processed_since_last_cleared_cuda_cache >= clear_cache_after_n_frames: soft_empty_cache() number_of_frames_processed_since_last_cleared_cuda_cache = 0 - logger("Clearing cache...") + rife_logger.info("Clearing cache...") pbar.update(1) + progress_bar.update(1) - + progress_bar.refresh() + progress_bar.close() + # Append final frame output_frames[out_len] = frames[-1:] - logger(f"done! - {(len(frames) -1) * (multiplier-1)} new frames generated at resolution: {output_frames[0].shape}") + # Get actual frame shape from first interpolated frame (CHW format) + actual_frame = output_frames[0] + h, w = actual_frame.shape[1], actual_frame.shape[2] + rife_logger.info(f"done! - {out_len} total frames output at resolution: {h}x{w}") out_len += 1 # clear cache for courtesy soft_empty_cache() - logger("Final clearing cache done ...") + rife_logger.info("Final clearing cache done ...") # res = output_frames[:out_len] return res \ No newline at end of file