746 lines
34 KiB
Python
746 lines
34 KiB
Python
import torch
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
import re
|
|
from pathlib import Path
|
|
from comfy.model_management import get_torch_device
|
|
from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife, logger, InterpolationStateList, infer_tiled
|
|
from .trt_utilities import Engine
|
|
from .utilities import download_file, ColoredLogger
|
|
import folder_paths
|
|
import time
|
|
|
|
def detect_cloud_environment():
|
|
"""Detect if running in cloud environment with CUDA 12 where CUDA graph may have issues"""
|
|
# Check for common cloud environment indicators
|
|
cloud_indicators = [
|
|
'RUNPOD', 'COLAB', 'KAGGLE', 'PAPERSPACE', 'GOOGLE_COLAB',
|
|
'JUPYTER_RUNTIME', 'DOCKER', 'CONTAINER'
|
|
]
|
|
|
|
# Check environment variables
|
|
for indicator in cloud_indicators:
|
|
if indicator in os.environ:
|
|
print(f"🌩️ Detected cloud environment: {indicator}")
|
|
|
|
# Check CUDA version - only disable CUDA graph for CUDA 12
|
|
try:
|
|
import torch
|
|
cuda_version = torch.version.cuda
|
|
if cuda_version and cuda_version.startswith('12'):
|
|
print(f"🌩️ CUDA {cuda_version} detected in cloud environment - disabling CUDA graph for stability")
|
|
return True
|
|
else:
|
|
print(f"✅ CUDA {cuda_version} detected - keeping CUDA graph enabled")
|
|
return False
|
|
except:
|
|
print("⚠️ Could not detect CUDA version, disabling CUDA graph for safety")
|
|
return True
|
|
|
|
# Check if running in Docker/container
|
|
if os.path.exists('/.dockerenv'):
|
|
print("🌩️ Detected Docker environment")
|
|
try:
|
|
import torch
|
|
cuda_version = torch.version.cuda
|
|
if cuda_version and cuda_version.startswith('12'):
|
|
print(f"🌩️ CUDA {cuda_version} in Docker - disabling CUDA graph")
|
|
return True
|
|
else:
|
|
print(f"✅ CUDA {cuda_version} in Docker - keeping CUDA graph enabled")
|
|
return False
|
|
except:
|
|
print("⚠️ Could not detect CUDA version in Docker, disabling CUDA graph")
|
|
return True
|
|
|
|
# Check for common cloud provider files
|
|
cloud_files = [
|
|
'/run/cloud-init/result.json',
|
|
'/var/lib/cloud/data/result.json'
|
|
]
|
|
|
|
for file_path in cloud_files:
|
|
if os.path.exists(file_path):
|
|
print("🌩️ Detected cloud environment via cloud-init")
|
|
try:
|
|
import torch
|
|
cuda_version = torch.version.cuda
|
|
if cuda_version and cuda_version.startswith('12'):
|
|
print(f"🌩️ CUDA {cuda_version} in cloud - disabling CUDA graph")
|
|
return True
|
|
else:
|
|
print(f"✅ CUDA {cuda_version} in cloud - keeping CUDA graph enabled")
|
|
return False
|
|
except:
|
|
print("⚠️ Could not detect CUDA version, disabling CUDA graph for safety")
|
|
return True
|
|
|
|
return False
|
|
|
|
# Auto-detect CUDA and install appropriate TensorRT packages
|
|
def _auto_install_tensorrt():
|
|
"""Auto-detect CUDA version and install the matching TensorRT wheels.
|
|
|
|
The NVIDIA CUDA Toolkit must already be installed on the system.
|
|
This function installs only the TensorRT packages via pip.
|
|
A marker file prevents repeated install attempts on every ComfyUI startup.
|
|
"""
|
|
disable_auto_install = os.environ.get("DISABLE_TENSORRT_AUTO_INSTALL", "false").lower() == "true"
|
|
if disable_auto_install:
|
|
print("[ComfyUI-RIFE-TensorRT] Auto-installation disabled via DISABLE_TENSORRT_AUTO_INSTALL")
|
|
return True
|
|
|
|
node_dir = Path(__file__).resolve().parent
|
|
installed_marker = node_dir / ".tensorrt_auto_installed"
|
|
failed_marker = node_dir / ".tensorrt_auto_install_failed"
|
|
|
|
# Skip if we already installed successfully in a previous run AND tensorrt is importable.
|
|
if installed_marker.exists():
|
|
try:
|
|
import tensorrt
|
|
return True
|
|
except ImportError:
|
|
print("[ComfyUI-RIFE-TensorRT] Install marker exists but tensorrt is not importable; retrying.")
|
|
try:
|
|
installed_marker.unlink()
|
|
except Exception:
|
|
pass
|
|
|
|
# Avoid retrying too often after a failure (do not block every startup).
|
|
if failed_marker.exists():
|
|
try:
|
|
last_fail = failed_marker.stat().st_mtime
|
|
if time.time() - last_fail < 3600:
|
|
print("[ComfyUI-RIFE-TensorRT] Recent failed install attempt; skipping auto-install.")
|
|
print("To retry now, delete the .tensorrt_auto_install_failed marker file or wait 1 hour.")
|
|
return False
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
# Check if TensorRT is already installed.
|
|
try:
|
|
import tensorrt
|
|
print(f"[ComfyUI-RIFE-TensorRT] TensorRT already installed (version: {tensorrt.__version__})")
|
|
installed_marker.touch()
|
|
return True
|
|
except ImportError:
|
|
print("[ComfyUI-RIFE-TensorRT] TensorRT not found, detecting CUDA version...")
|
|
|
|
# Detect CUDA version
|
|
cuda_version = None
|
|
|
|
# Try PyTorch CUDA version first (most reliable on cloud platforms)
|
|
try:
|
|
if hasattr(torch, 'version') and hasattr(torch.version, 'cuda'):
|
|
pytorch_cuda = torch.version.cuda
|
|
if pytorch_cuda:
|
|
cuda_version = pytorch_cuda
|
|
print(f"[ComfyUI-RIFE-TensorRT] Detected CUDA version from PyTorch: {cuda_version}")
|
|
except Exception as e:
|
|
print(f"[ComfyUI-RIFE-TensorRT] Could not detect CUDA from PyTorch: {e}")
|
|
|
|
# Try nvcc command as fallback
|
|
if not cuda_version:
|
|
try:
|
|
result = subprocess.run(["nvcc", "--version"], capture_output=True, text=True)
|
|
if result.returncode == 0:
|
|
match = re.search(r"release (\d+\.\d+)", result.stdout)
|
|
if match:
|
|
cuda_version = match.group(1)
|
|
print(f"[ComfyUI-RIFE-TensorRT] Detected CUDA version from nvcc: {cuda_version}")
|
|
except Exception:
|
|
pass
|
|
|
|
# Try CUDA_PATH
|
|
if not cuda_version and os.environ.get("CUDA_PATH"):
|
|
nvcc_path = os.path.join(os.environ["CUDA_PATH"], "bin", "nvcc")
|
|
if os.path.exists(nvcc_path):
|
|
try:
|
|
result = subprocess.run([nvcc_path, "--version"], capture_output=True, text=True)
|
|
if result.returncode == 0:
|
|
match = re.search(r"release (\d+\.\d+)", result.stdout)
|
|
if match:
|
|
cuda_version = match.group(1)
|
|
print(f"[ComfyUI-RIFE-TensorRT] Detected CUDA via CUDA_PATH: {cuda_version}")
|
|
except Exception:
|
|
pass
|
|
|
|
# Try CUDA_HOME
|
|
if not cuda_version and os.environ.get("CUDA_HOME"):
|
|
nvcc_path = os.path.join(os.environ["CUDA_HOME"], "bin", "nvcc")
|
|
if os.path.exists(nvcc_path):
|
|
try:
|
|
result = subprocess.run([nvcc_path, "--version"], capture_output=True, text=True)
|
|
if result.returncode == 0:
|
|
match = re.search(r"release (\d+\.\d+)", result.stdout)
|
|
if match:
|
|
cuda_version = match.group(1)
|
|
print(f"[ComfyUI-RIFE-TensorRT] Detected CUDA via CUDA_HOME: {cuda_version}")
|
|
except Exception:
|
|
pass
|
|
|
|
if not cuda_version:
|
|
print("[ComfyUI-RIFE-TensorRT] WARNING: Could not detect CUDA version automatically.")
|
|
print("The NVIDIA CUDA Toolkit must be installed before TensorRT can work.")
|
|
print("Please run 'python install.py' manually after installing CUDA.")
|
|
failed_marker.touch()
|
|
return False
|
|
|
|
major_version = int(cuda_version.split('.')[0])
|
|
|
|
if major_version == 13:
|
|
print("[ComfyUI-RIFE-TensorRT] Installing CUDA 13 TensorRT packages (RTX 50 series)")
|
|
req_file = "requirements_cu13.txt"
|
|
elif major_version == 12:
|
|
print("[ComfyUI-RIFE-TensorRT] Installing CUDA 12 TensorRT packages (RTX 30/40 series)")
|
|
req_file = "requirements_cu12.txt"
|
|
else:
|
|
print(f"[ComfyUI-RIFE-TensorRT] Unsupported CUDA version: {cuda_version}")
|
|
failed_marker.touch()
|
|
return False
|
|
|
|
req_path = node_dir / req_file
|
|
if not req_path.exists():
|
|
print(f"[ComfyUI-RIFE-TensorRT] Missing requirements file: {req_path}")
|
|
failed_marker.touch()
|
|
return False
|
|
|
|
# Install base dependencies first, then the CUDA-specific TensorRT wheels.
|
|
for req_name in ["requirements.txt", req_file]:
|
|
req_file_path = node_dir / req_name
|
|
if not req_file_path.exists():
|
|
continue
|
|
print(f"[ComfyUI-RIFE-TensorRT] Installing from {req_name}...")
|
|
result = subprocess.run(
|
|
[sys.executable, "-m", "pip", "install", "--no-input", "--prefer-binary", "-r", str(req_file_path)],
|
|
capture_output=True
|
|
)
|
|
if result.returncode != 0:
|
|
print(f"[ComfyUI-RIFE-TensorRT] Failed to install {req_name}")
|
|
print(result.stderr.decode(errors="replace"))
|
|
failed_marker.touch()
|
|
return False
|
|
# Show pip stdout so the user can see progress / warnings.
|
|
stdout = result.stdout.decode(errors="replace").strip()
|
|
if stdout:
|
|
print(stdout)
|
|
|
|
installed_marker.touch()
|
|
print("[ComfyUI-RIFE-TensorRT] TensorRT installation completed successfully!")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"[ComfyUI-RIFE-TensorRT] Auto-installation failed: {e}")
|
|
print("Please run 'python install.py' manually to install TensorRT")
|
|
try:
|
|
failed_marker.touch()
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
# Run auto-install on module import
|
|
_auto_install_tensorrt()
|
|
|
|
# Auto-detect CUDA toolkit and add DLL path before importing polygraphy
|
|
def _setup_cuda_dll_path():
|
|
"""Auto-detect CUDA toolkit and add cudart64 DLL path on Windows."""
|
|
if not sys.platform.startswith("win"):
|
|
return
|
|
|
|
cuda_root = None
|
|
|
|
# Check for CUDA_PATH or CUDA_HOME environment variables
|
|
cuda_root = os.environ.get("CUDA_PATH") or os.environ.get("CUDA_HOME")
|
|
|
|
if not cuda_root:
|
|
# Try default Windows install location
|
|
program_files = os.environ.get("PROGRAMFILES")
|
|
if program_files:
|
|
cuda_base = Path(program_files) / "NVIDIA GPU Computing Toolkit" / "CUDA"
|
|
if cuda_base.exists():
|
|
# Find highest version directory
|
|
versions = sorted([d for d in cuda_base.iterdir() if d.is_dir()], reverse=True)
|
|
if versions:
|
|
cuda_root = str(versions[0])
|
|
|
|
if cuda_root:
|
|
cuda_path = Path(cuda_root)
|
|
# CUDA 13.0+ puts cudart64 in bin/x64 subdirectory
|
|
cuda_bin_x64 = cuda_path / "bin" / "x64"
|
|
if cuda_bin_x64.exists() and any(cuda_bin_x64.glob("cudart64*.dll")):
|
|
os.add_dll_directory(str(cuda_bin_x64))
|
|
return
|
|
# Fallback to regular bin directory for older CUDA versions
|
|
cuda_bin = cuda_path / "bin"
|
|
if cuda_bin.exists() and any(cuda_bin.glob("cudart64*.dll")):
|
|
os.add_dll_directory(str(cuda_bin))
|
|
return
|
|
|
|
# CUDA toolkit not found - print warning with download link
|
|
print("[ComfyUI-Rife-TensorRT] WARNING: CUDA toolkit not found.")
|
|
print(" Set CUDA_PATH environment variable or install CUDA toolkit.")
|
|
print(" Download: https://developer.nvidia.com/cuda-13-0-2-download-archive")
|
|
|
|
_setup_cuda_dll_path()
|
|
|
|
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")
|
|
|
|
# Default resolution profiles (fallback if config is missing)
|
|
DEFAULT_RESOLUTION_PROFILES = {
|
|
"small": {"min": 384, "opt": 720, "max": 1080},
|
|
"medium": {"min": 672, "opt": 1080, "max": 1312}
|
|
}
|
|
|
|
# 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": ["rife425_ensemble_False_scale_1_sim", "rife49_ensemble_True_scale_1_sim"],
|
|
"default": "rife425_ensemble_False_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 CustomResolutionConfig:
|
|
"""Node to configure custom resolution dimensions for TensorRT engine building."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"min_dim": ("INT", {"default": 384, "min": 64, "max": 4096, "step": 8, "tooltip": "Minimum resolution dimension"}),
|
|
"opt_dim": ("INT", {"default": 720, "min": 64, "max": 4096, "step": 8, "tooltip": "Optimal resolution dimension (most common)"}),
|
|
"max_dim": ("INT", {"default": 1312, "min": 64, "max": 4096, "step": 8, "tooltip": "Maximum resolution dimension"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("RIFE_RESOLUTION_CONFIG",)
|
|
RETURN_NAMES = ("resolution_config",)
|
|
FUNCTION = "configure"
|
|
CATEGORY = "⚡️ TensorRT/RIFE"
|
|
DESCRIPTION = "Configure custom resolution dimensions for RIFE TensorRT engine."
|
|
|
|
def configure(self, min_dim, opt_dim, max_dim):
|
|
config = {
|
|
"min": min_dim,
|
|
"opt": opt_dim,
|
|
"max": max_dim,
|
|
}
|
|
return (config,)
|
|
|
|
|
|
class AutoLoadRifeTensorrtModel:
|
|
@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.")
|
|
|
|
# Resolution profile configuration
|
|
profile_config = LOAD_RIFE_NODE_CONFIG.get("resolution_profile", {})
|
|
profile_options = profile_config.get("options", ["small", "medium"])
|
|
# Ensure 'custom' is always available
|
|
if "custom" not in profile_options:
|
|
profile_options = profile_options + ["custom"]
|
|
profile_default = profile_config.get("default", "small")
|
|
profile_tooltip = profile_config.get("tooltip", "Resolution range for TensorRT engine. Use 'custom' with the INT inputs below.")
|
|
|
|
return {
|
|
"required": {
|
|
"model": (model_options, {"default": model_default, "tooltip": model_tooltip}),
|
|
"precision": (precision_options, {"default": precision_default, "tooltip": precision_tooltip}),
|
|
"resolution_profile": (profile_options, {"default": profile_default, "tooltip": profile_tooltip}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 16, "step": 1,
|
|
"tooltip": "Max batch size for the engine profile. "
|
|
"Higher values allow batched inference (multiple frame pairs per GPU call) "
|
|
"for better throughput, but use more VRAM and require a rebuild. "
|
|
"Set to 1 for the original behaviour."}),
|
|
},
|
|
"optional": {
|
|
"custom_config": ("RIFE_RESOLUTION_CONFIG", {"tooltip": "Custom resolution config (used when profile='custom')"}),
|
|
}
|
|
}
|
|
|
|
RETURN_NAMES = ("rife_trt_model",)
|
|
RETURN_TYPES = ("RIFE_TRT_MODEL",)
|
|
CATEGORY = "⚡️ TensorRT/RIFE"
|
|
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, resolution_profile, batch_size=1, custom_config=None):
|
|
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")
|
|
|
|
# Get resolution dimensions based on profile
|
|
if resolution_profile == "custom":
|
|
if custom_config is None:
|
|
rife_logger.warning("Custom profile selected but no custom_config provided. Using defaults that cover both small and medium ranges.")
|
|
dim_min, dim_opt, dim_max = 384, 720, 1312
|
|
else:
|
|
dim_min = custom_config.get("min", 384)
|
|
dim_opt = custom_config.get("opt", 720)
|
|
dim_max = custom_config.get("max", 1312)
|
|
# Use dimensions in profile name for custom engines
|
|
profile_name = f"custom_{dim_min}_{dim_opt}_{dim_max}"
|
|
else:
|
|
profiles = LOAD_RIFE_NODE_CONFIG.get("resolution_profiles", DEFAULT_RESOLUTION_PROFILES)
|
|
profile = profiles.get(resolution_profile, DEFAULT_RESOLUTION_PROFILES["small"])
|
|
dim_min = profile.get("min", 384)
|
|
dim_opt = profile.get("opt", 720)
|
|
dim_max = profile.get("max", 1080)
|
|
profile_name = resolution_profile
|
|
rife_logger.info(f"Using resolution profile '{profile_name}': min={dim_min}, opt={dim_opt}, max={dim_max}")
|
|
|
|
# Build tensorrt model path with detailed naming (includes profile)
|
|
engine_channel = 3
|
|
engine_min_batch = 1
|
|
engine_opt_batch = max(1, batch_size)
|
|
engine_max_batch = max(1, batch_size)
|
|
engine_min_h, engine_opt_h, engine_max_h = dim_min, dim_opt, dim_max
|
|
engine_min_w, engine_opt_w, engine_max_w = dim_min, dim_opt, dim_max
|
|
tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{model}_{precision}_{profile_name}_{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):
|
|
# Try multiple download sources in case the primary one fails
|
|
onnx_download_urls = [
|
|
f"https://huggingface.co/huchukato/garage/resolve/main/onnx/rife/{model}.onnx",
|
|
f"https://huggingface.co/yuvraj108c/rife-onnx/resolve/main/{model}.onnx",
|
|
f"https://huggingface.co/marduk191/rife/resolve/main/{model}.onnx",
|
|
]
|
|
download_success = False
|
|
for url in onnx_download_urls:
|
|
try:
|
|
rife_logger.info(f"Downloading {url}")
|
|
download_file(url=url, save_path=onnx_model_path)
|
|
download_success = True
|
|
break
|
|
except Exception as e:
|
|
rife_logger.warning(f"Download failed from {url}: {e}")
|
|
continue
|
|
if not download_success:
|
|
raise RuntimeError(
|
|
f"Failed to download ONNX model '{model}.onnx' from all sources. "
|
|
f"You can manually export it using export_onnx.py and place it in: {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,)
|
|
|
|
|
|
class MakeInterpolationStateList:
|
|
"""Build an InterpolationStateList to skip or keep specific frame pairs.
|
|
Ported from Fannovel16/ComfyUI-Frame-Interpolation.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"frame_indices": ("STRING", {"multiline": True, "default": "1,2,3",
|
|
"tooltip": "Comma-separated list of frame-pair indices (0-based). "
|
|
"Pair i is formed by frame i and frame i+1."}),
|
|
"is_skip_list": ("BOOLEAN", {"default": True,
|
|
"tooltip": "True = skip interpolation for the listed pairs. "
|
|
"False = interpolate ONLY the listed pairs."}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INTERPOLATION_STATES",)
|
|
FUNCTION = "create_options"
|
|
CATEGORY = "⚡️ TensorRT/RIFE"
|
|
DESCRIPTION = "Build a skip/keep list controlling which frame pairs get interpolated by RIFE."
|
|
|
|
def create_options(self, frame_indices: str, is_skip_list: bool):
|
|
frame_indices_list = [int(item.strip()) for item in frame_indices.split(',') if item.strip() != '']
|
|
interpolation_state_list = InterpolationStateList(
|
|
frame_indices=frame_indices_list,
|
|
is_skip_list=is_skip_list,
|
|
)
|
|
return (interpolation_state_list,)
|
|
|
|
|
|
class AutoRifeTensorrt:
|
|
@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 (uniform). Ignored if multiplier_list is connected."}),
|
|
"keep_model_loaded": ("BOOLEAN", {"default": False, "tooltip": "Keep model loaded in memory after processing"}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 16,
|
|
"tooltip": "Number of interpolation tasks per GPU call. "
|
|
"Higher values improve throughput but use more VRAM. "
|
|
"Must be <= the batch_size used when building the engine. "
|
|
"Set to 1 for the safest behaviour."}),
|
|
},
|
|
"optional": {
|
|
"multiplier_list": ("STRING", {"multiline": True, "default": "",
|
|
"tooltip": "Per-pair multiplier schedule, comma-separated (e.g. '2,2,4,4,2'). "
|
|
"When connected, overrides 'multiplier'. Missing trailing pairs default to 2."}),
|
|
"interpolation_states": ("INTERPOLATION_STATES", {"tooltip": "Optional skip/keep list from MakeInterpolationStateList"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", )
|
|
FUNCTION = "vfi"
|
|
CATEGORY = "⚡️ TensorRT/RIFE"
|
|
OUTPUT_NODE=True
|
|
|
|
def vfi(
|
|
self,
|
|
frames,
|
|
rife_trt_model,
|
|
clear_cache_after_n_frames=100,
|
|
multiplier=2,
|
|
keep_model_loaded=False,
|
|
batch_size=1,
|
|
multiplier_list="",
|
|
interpolation_states=None,
|
|
):
|
|
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()
|
|
|
|
# Auto-detect CUDA graph usage based on CUDA version
|
|
use_cuda_graph = False
|
|
try:
|
|
import torch
|
|
cuda_version = torch.version.cuda
|
|
if cuda_version and cuda_version.startswith('13'):
|
|
use_cuda_graph = True
|
|
print(f"✅ CUDA {cuda_version} detected - enabling CUDA graph for better performance")
|
|
else:
|
|
use_cuda_graph = False
|
|
print(f"⚠️ CUDA {cuda_version} detected - disabling CUDA graph for stability")
|
|
except:
|
|
use_cuda_graph = False
|
|
print("⚠️ Could not detect CUDA version - disabling CUDA graph for safety")
|
|
|
|
# Use the provided model directly
|
|
engine = rife_trt_model
|
|
logger(f"Using loaded TensorRT engine")
|
|
|
|
# Activate engine and retrieve profile bounds
|
|
engine.activate()
|
|
bounds = engine.get_input_profile_bounds()
|
|
|
|
# Determine if tiling/padding is needed
|
|
needs_tiling = False
|
|
min_hw = (1, 1)
|
|
max_hw = (4096, 4096)
|
|
if bounds and 'img0' in bounds:
|
|
min_shape = bounds['img0'][0]
|
|
max_shape = bounds['img0'][2]
|
|
min_hw = (min_shape[2], min_shape[3])
|
|
max_hw = (max_shape[2], max_shape[3])
|
|
needs_tiling = H > max_hw[0] or W > max_hw[1] or H < min_hw[0] or W < min_hw[1]
|
|
if needs_tiling:
|
|
logger(f"Frame {H}x{W} outside profile range "
|
|
f"[{min_hw[0]}-{max_hw[0]}]x[{min_hw[1]}-{max_hw[1]}] — enabling tiling/padding")
|
|
|
|
frames = preprocess_frames(frames)
|
|
|
|
if needs_tiling:
|
|
# Disable CUDA graph when tiling (buffers are reallocated per tile)
|
|
use_cuda_graph = False
|
|
logger("Tiling mode: CUDA graph disabled")
|
|
|
|
def _raw_infer(f0, f1, t):
|
|
"""Per-tile inference: reallocate buffers for the tile shape."""
|
|
_, _, th, tw = f0.shape
|
|
tile_shape_dict = {
|
|
"img0": {"shape": (1, 3, th, tw)},
|
|
"img1": {"shape": (1, 3, th, tw)},
|
|
"output": {"shape": (1, 3, th, tw)},
|
|
}
|
|
engine.allocate_buffers(shape_dict=tile_shape_dict)
|
|
timestep_t = torch.tensor([t], dtype=torch.float32).to(get_torch_device())
|
|
output = engine.infer({"img0": f0, "img1": f1, "timestep": timestep_t}, cudaStream, use_cuda_graph)
|
|
return output['output']
|
|
|
|
def return_middle_frame(frame_0, frame_1, timestep):
|
|
return infer_tiled(frame_0, frame_1, timestep, _raw_infer, min_hw, max_hw)
|
|
|
|
# Batch inference is not compatible with tiling (each tile has different shape)
|
|
batch_infer_fn = None
|
|
effective_batch_size = 1
|
|
else:
|
|
# Fast path: single allocation, original behaviour
|
|
engine.allocate_buffers(shape_dict=shape_dict)
|
|
|
|
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, use_cuda_graph)
|
|
result = output['output']
|
|
return result
|
|
|
|
# Batch inference: allocate buffers for the requested batch size
|
|
if batch_size > 1:
|
|
batch_shape_dict = {
|
|
"img0": {"shape": (batch_size, 3, H, W)},
|
|
"img1": {"shape": (batch_size, 3, H, W)},
|
|
"output": {"shape": (batch_size, 3, H, W)},
|
|
}
|
|
# Check if the engine profile supports this batch size
|
|
engine_max_batch = max_hw[0] # placeholder, real check below
|
|
if bounds and 'img0' in bounds:
|
|
engine_max_batch = bounds['img0'][2][0] # max batch dim
|
|
if engine_max_batch >= batch_size:
|
|
logger(f"Batch inference enabled: batch_size={batch_size}")
|
|
# Pre-allocate buffers for the full batch
|
|
engine.allocate_buffers(shape_dict=batch_shape_dict)
|
|
|
|
def batch_infer_fn(frame0_batch, frame1_batch, timestep_batch):
|
|
B_actual = frame0_batch.shape[0]
|
|
# If the actual batch is smaller than the allocated batch,
|
|
# we still pass the full buffer but only use the first B_actual rows.
|
|
# TensorRT requires the input shape to match what was set in allocate_buffers.
|
|
# So we pad to the full batch size if needed.
|
|
if B_actual < batch_size:
|
|
pad_shape = (batch_size - B_actual,) + frame0_batch.shape[1:]
|
|
pad0 = torch.zeros(pad_shape, dtype=frame0_batch.dtype, device=frame0_batch.device)
|
|
pad1 = torch.zeros(pad_shape, dtype=frame1_batch.dtype, device=frame1_batch.device)
|
|
frame0_batch = torch.cat([frame0_batch, pad0], dim=0)
|
|
frame1_batch = torch.cat([frame1_batch, pad1], dim=0)
|
|
timestep_batch = torch.cat([timestep_batch, torch.zeros(batch_size - B_actual, dtype=timestep_batch.dtype, device=timestep_batch.device)])
|
|
|
|
timestep_t = timestep_batch.to(get_torch_device())
|
|
output = engine.infer({"img0": frame0_batch, "img1": frame1_batch, "timestep": timestep_t}, cudaStream, use_cuda_graph)
|
|
return output['output'][:B_actual]
|
|
|
|
effective_batch_size = batch_size
|
|
else:
|
|
logger(f"⚠️ Engine max_batch={engine_max_batch} < requested batch_size={batch_size}. "
|
|
"Falling back to sequential inference. Rebuild engine with batch_size >= "
|
|
f"{batch_size} to enable batched inference.")
|
|
batch_infer_fn = None
|
|
effective_batch_size = 1
|
|
else:
|
|
batch_infer_fn = None
|
|
effective_batch_size = 1
|
|
|
|
# Resolve multiplier: per-pair schedule (list) overrides uniform int
|
|
effective_multiplier = multiplier
|
|
if multiplier_list and multiplier_list.strip():
|
|
try:
|
|
effective_multiplier = [int(x.strip()) for x in multiplier_list.split(',') if x.strip() != '']
|
|
logger(f"Using per-pair multiplier schedule: {effective_multiplier}")
|
|
except ValueError as e:
|
|
logger(f"Invalid multiplier_list '{multiplier_list}': {e}. Falling back to uniform multiplier={multiplier}")
|
|
effective_multiplier = multiplier
|
|
|
|
result = generate_frames_rife(
|
|
frames,
|
|
clear_cache_after_n_frames,
|
|
effective_multiplier,
|
|
return_middle_frame,
|
|
interpolation_states=interpolation_states,
|
|
batch_size=effective_batch_size,
|
|
batch_infer_function=batch_infer_fn,
|
|
)
|
|
out = postprocess_frames(result)
|
|
|
|
if not keep_model_loaded:
|
|
engine.reset()
|
|
|
|
return (out,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AutoRifeTensorrt": AutoRifeTensorrt,
|
|
"AutoLoadRifeTensorrtModel": AutoLoadRifeTensorrtModel,
|
|
"CustomResolutionConfig": CustomResolutionConfig,
|
|
"MakeInterpolationStateList": MakeInterpolationStateList,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AutoRifeTensorrt": "Auto RIFE TensorRT",
|
|
"AutoLoadRifeTensorrtModel": "(Down)load RIFE TensorRT Model",
|
|
"CustomResolutionConfig": "RIFE Custom Resolution Config",
|
|
"MakeInterpolationStateList": "RIFE Interpolation State List",
|
|
}
|
|
|
|
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
|
|