add files
This commit is contained in:
Threads_Of_Fate
@@ -0,0 +1,3 @@
|
||||
__pycache__
|
||||
RIFE_RTX_ENGINE.cache
|
||||
*.bak
|
||||
@@ -1,3 +1,4 @@
|
||||
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
@@ -198,4 +199,4 @@
|
||||
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.
|
||||
limitations under the License.
|
||||
@@ -1,2 +1,76 @@
|
||||
# ComfyUI-RIFE-TensorRT-RTX
|
||||
ComfyUI custom node that performs RIFE Interpolation using Tensorrt-RTX
|
||||
# ComfyUI RIFE TensorRT-RTX
|
||||
|
||||
[](https://www.python.org/downloads/release/python-31313/)
|
||||
[](https://developer.nvidia.com/cuda-downloads)
|
||||
[](https://developer.nvidia.com/tensorrt)
|
||||
[](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.en)
|
||||
|
||||
## Overview
|
||||
|
||||
This project is a fork of the original [ComfyUI_RIFE_TensorRT_Auto](https://github.com/silveroxides/ComfyUI_RIFE_TensorRT_Auto) by [silveroxides](https://github.com/silveroxides).
|
||||
This project uses TensorRT-RTX to do RIFE interpolation
|
||||
|
||||
## Features
|
||||
- **NVIDIA TensorRT Acceleration**: Leverages Tensor Cores for up to 30x faster frame interpolation compared to standard local methods.
|
||||
- **Dynamic Shape Support**: Configure minimum, optimal, and maximum resolutions to handle varying input sizes without rebuilding the engine for every image.
|
||||
- **Automatic Engine Building**: Automatically compiles specialized `.trt` files from your `.onnx` models if a matching engine isn't found.
|
||||
- **Persistent Timing Cache**: Saves hardware-specific optimization data to a local cache to speed up subsequent engine builds.
|
||||
- **Memory Efficient**: Optimized for VRAM management with automatic garbage collection and cache clearing after inference.
|
||||
|
||||
|
||||
## Examples
|
||||
|
||||
|
||||
|
||||
## Installation
|
||||
|
||||
### ComfyUI Manager
|
||||
|
||||
This project has not been submitted to the ComfyUI Manager registry yet. But you can still install it this way:
|
||||
|
||||
1. Open ComfyUI Manager.
|
||||
2. Click the `Custom Nodes Manager` to open the custom nodes manager page.
|
||||
3. On the bottom right corner, click the `Install via Git URL` button.
|
||||
4. Enter the URL of this repository: `https://github.com/ThreadsOfFate/ComfyUI-RIFE-TensorRT-RTX`
|
||||
5. Click "Confirm".
|
||||
6. Restart ComfyUI.
|
||||
|
||||
### Manual
|
||||
|
||||
1. On the github page, click on the green `<> Code` button and then "Download ZIP".
|
||||
2. Extract the root folder within the downloaded ZIP file to your ComfyUI `custom_nodes` directory.
|
||||
3. Using the same python environment that runs ComfyUI, install the required dependencies: `python -m pip install -r custom_nodes/ComfyUI-Upscaler-TensorRT-Advanced/requirements.txt`.
|
||||
4. Restart ComfyUI.
|
||||
|
||||
|
||||
### Requirements
|
||||
pip install -r requirements.txt
|
||||
|
||||
|
||||
### CUDA Toolkit Required
|
||||
Please download NVIDIA CUDA Toolkit [https://developer.nvidia.com/cuda/toolkit](https://developer.nvidia.com/cuda/toolkit)
|
||||
|
||||
###Nodes Included
|
||||
🟦 RIFE RTX Engine Builder (RIFERTXEngineLoader)
|
||||
Loads a TensorRT-RTX engine. If the engine file for your specific GPU and resolution settings does not exist, it will automatically build one from the source ONNX model.
|
||||
🟦 RIFE RTX Frame Interpolation (RIFERTXFrameInterpolation)
|
||||
The core execution node. It takes an image and a loaded engine to perform the upscale. It supports an optional resize input for custom final dimensions.
|
||||
🟦 RIFE RTX Engine Config (RIFERTXEngineConfig)
|
||||
Allows you to define the resolution bounds (min/opt/max) for the engine.
|
||||
Min/Max: The range of resolutions the engine can handle.
|
||||
Opt: The resolution the engine is most optimized for.
|
||||
|
||||
|
||||
GPU: NVIDIA RTX Series (TensorRT-RTX requires Tensor Cores).
|
||||
Software: NVIDIA Drivers and a compatible version of TensorRT (handled via requirements.txt).
|
||||
|
||||
###Troubleshooting
|
||||
First Build Time: The first time you load a model with new shape settings, it may take several minutes to compile the engine. Check the console for progress via the TQDM bar.
|
||||
VRAM Errors: If the build fails, try reducing the width_max or height_max in the Dynamic Shape Config.
|
||||
|
||||
|
||||
## 🛠️ Supported Models
|
||||
|
||||
rife47_ensemble_True_scale_1_sim
|
||||
rife48_ensemble_True_scale_1_sim
|
||||
rife49_ensemble_True_scale_1_sim
|
||||
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
"""ComfyUI-Rife-TensorRT Extension Initialization.
|
||||
|
||||
This module acts as the entry point for ComfyUI, exporting the node classes
|
||||
and display names required to register the TensorRT-RTX RIFE implementation
|
||||
into the ComfyUI node graph.
|
||||
|
||||
Attributes:
|
||||
NODE_CLASS_MAPPINGS (dict): Maps internal lookup strings to node classes.
|
||||
NODE_DISPLAY_NAME_MAPPINGS (dict): Maps internal lookup strings to
|
||||
human-readable UI names.
|
||||
WEB_DIRECTORY (str): Path to the directory containing frontend
|
||||
javascript extensions.
|
||||
"""
|
||||
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
WEB_DIRECTORY = "./js"
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
@@ -0,0 +1,309 @@
|
||||
"""TensorRT-RTX RIFE Implementation for ComfyUI.
|
||||
|
||||
This module provides custom nodes for video frame interpolation (VFI) using
|
||||
NVIDIA's TensorRT-RTX. It includes nodes for configuring engine resolutions,
|
||||
building/loading optimized .trt engines, and performing high-speed inference.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
import folder_paths
|
||||
import logging
|
||||
import tensorrt_rtx as trt
|
||||
import time
|
||||
import torch
|
||||
|
||||
import comfy.model_management as mm
|
||||
|
||||
from pathlib import Path
|
||||
from polygraphy import cuda
|
||||
from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife
|
||||
from .trt_utilities import Engine
|
||||
|
||||
def update_folder_names_and_paths(key, targets=[]):
|
||||
"""Registers specific folder paths within ComfyUI's path manager.
|
||||
|
||||
Ensures that custom subdirectories (like onnx_rife) are recognized by
|
||||
ComfyUI's internal file system and UI loaders.
|
||||
|
||||
Args:
|
||||
key (str): The primary key to register in folder_names_and_paths.
|
||||
targets (list): Fallback keys to inherit paths from if the primary
|
||||
key is missing.
|
||||
"""
|
||||
# check for existing key
|
||||
base = folder_paths.folder_names_and_paths.get(key, ([], {}))
|
||||
base = base[0] if isinstance(base[0], (list, set, tuple)) else []
|
||||
# find base key & add w/ fallback, sanity check + warning
|
||||
target = next((x for x in targets if x in folder_paths.folder_names_and_paths), targets[0])
|
||||
orig, _ = folder_paths.folder_names_and_paths.get(target, ([], {}))
|
||||
folder_paths.folder_names_and_paths[key] = (orig or base, {".onnx"})
|
||||
if base and base != orig:
|
||||
logging.warning(f"Unknown file list already present on key {key}: {base}")
|
||||
|
||||
|
||||
# Auto-detect CUDA toolkit and add DLL path before importing polygraphy
|
||||
def _setup_cuda_dll_path():
|
||||
"""Auto-detects the CUDA toolkit and adds cudart64 DLL path on Windows.
|
||||
|
||||
This is necessary for TensorRT-RTX to correctly locate runtime libraries
|
||||
when running in a portable environment or on systems where CUDA is not
|
||||
globally in the system PATH.
|
||||
"""
|
||||
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")
|
||||
|
||||
# Initialize Environment
|
||||
_setup_cuda_dll_path()
|
||||
os.environ["POLYGRAPHY_USE_TENSORRT_RTX"] = "1"
|
||||
update_folder_names_and_paths("onnx_rife", ["onnx"])
|
||||
|
||||
class RIFERTXEngineOptions:
|
||||
"""Data container for RIFE shape boundaries.
|
||||
|
||||
Attributes:
|
||||
min_dim (int): Minimum supported resolution.
|
||||
opt_dim (int): Targeted optimal resolution.
|
||||
max_dim (int): Maximum supported resolution.
|
||||
"""
|
||||
min_dim: int = 384
|
||||
opt_dim: int = 720
|
||||
max_dim: int = 1312
|
||||
|
||||
class RIFERTXEngineConfig:
|
||||
"""ComfyUI Node for configuring resolution profiles.
|
||||
|
||||
Provides a UI interface to set the dynamic shape profiles required
|
||||
during the TensorRT engine build process.
|
||||
"""
|
||||
@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_RTX_ENGINE_CONFIG",)
|
||||
RETURN_NAMES = ("RIFE_ENGINE_CONFIG",)
|
||||
FUNCTION = "configure"
|
||||
CATEGORY = "TensorRT_RTX/RIFE"
|
||||
DESCRIPTION = "Configure custom resolution dimensions for RIFE TensorRT engine."
|
||||
|
||||
def configure(self, min_dim, opt_dim, max_dim):
|
||||
"""Wraps dimensions into a configuration dictionary.
|
||||
|
||||
Args:
|
||||
min_dim (int): Minimum resolution.
|
||||
opt_dim (int): Optimal resolution.
|
||||
max_dim (int): Maximum resolution.
|
||||
|
||||
Returns:
|
||||
tuple: Contains the configuration dictionary.
|
||||
"""
|
||||
config = {"min_d": min_dim, "opt_d": opt_dim, "max_d": max_dim,}
|
||||
return (config,)
|
||||
|
||||
|
||||
class RIFERTXEngineLoader:
|
||||
"""ComfyUI Node for loading or building RIFE TensorRT engines.
|
||||
|
||||
Checks if a .trt engine matching the requested configuration exists.
|
||||
If not, it compiles the ONNX model into a TensorRT engine.
|
||||
"""
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
onnx_names = [x for x in folder_paths.get_filename_list("onnx_rife") if "rife" in x.lower()]
|
||||
return {
|
||||
"required": { "onnx_name": (onnx_names,)},
|
||||
"optional": { "config": ("RIFE_RTX_ENGINE_CONFIG", {"tooltip": "Options for building the TensorRT engine"}),}
|
||||
}
|
||||
|
||||
RETURN_NAMES = ("RIFE_ENGINE",)
|
||||
RETURN_TYPES = ("RIFE_RTX_ENGINE",)
|
||||
CATEGORY = "TensorRT_RTX/RIFE"
|
||||
DESCRIPTION = "Load RIFE tensorrt_rtx models, they will be built automatically if not found."
|
||||
FUNCTION = "load_rife_tensorrt_rtx_model"
|
||||
|
||||
def load_rife_tensorrt_rtx_model(self, onnx_name, config=None):
|
||||
"""Loads or builds the RIFE TensorRT engine.
|
||||
|
||||
Args:
|
||||
onnx_name (str): Name of the source ONNX file.
|
||||
config (dict, optional): Resolution profile. Defaults to RIFERTXEngineOptions.
|
||||
|
||||
Returns:
|
||||
tuple: The loaded Engine object.
|
||||
"""
|
||||
if config is None:
|
||||
config = RIFERTXEngineOptions()
|
||||
|
||||
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"{onnx_name}")
|
||||
|
||||
logging.info(f"Using resolution profile '{profile_name}': min={config.min_dim}, opt={config.opt_dim}, max={config.max_dim}")
|
||||
|
||||
# Build tensorrt model path with detailed naming (includes profile)
|
||||
engine_channel = 3
|
||||
e_min, e_opt, e_max = config.min_d, config.opt_d, config.max_d
|
||||
tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{onnx_name}_1x{engine_channel}x{e_min}x{min_}_1x{engine_channel}x{e_opt}x{e_opt}_1x{engine_channel}x{e_max}x{e_max}_{trt.__version__}.trt")
|
||||
|
||||
engine = Engine(tensorrt_model_path)
|
||||
|
||||
if not os.path.exists(tensorrt_model_path):
|
||||
logging.info(f"Building TensorRT_RTX engine for {onnx_model_path}: {tensorrt_model_path}")
|
||||
|
||||
mm.soft_empty_cache()
|
||||
s = time.time()
|
||||
|
||||
result = engine.build(
|
||||
onnx_path=onnx_model_path,
|
||||
input_profile=[
|
||||
{
|
||||
"img0": [(1, engine_channel, e_min, e_min), (1, engine_channel, e_opt, e_opt), (1, engine_channel, e_max, e_max)],
|
||||
"img1": [(1, engine_channel, e_min, e_min), (1, engine_channel, e_opt, e_opt), (1, engine_channel, e_max, e_max)],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
if result != 0:
|
||||
raise ValueError(f"Failed to build the engine")
|
||||
e = time.time()
|
||||
logging.info(f"Time taken to build: {(e-s)} seconds")
|
||||
|
||||
logging.info(f"Loading TensorRT_RTX engine: {tensorrt_model_path}")
|
||||
mm.soft_empty_cache()
|
||||
engine.load()
|
||||
|
||||
return (engine,)
|
||||
|
||||
class RIFERTXFrameInterpolation:
|
||||
"""ComfyUI Node for performing VFI inference.
|
||||
|
||||
Uses the loaded TensorRT engine to interpolate frames in a video sequence.
|
||||
"""
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"frames": ("IMAGE", {"tooltip": "Input frames for video frame interpolation"}),
|
||||
"engine": ("RIFE_RTX_ENGINE", {"tooltip": "Tensorrt_RTX 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"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", )
|
||||
FUNCTION = "vfi"
|
||||
CATEGORY = "TensorRT_RTX/RIFE"
|
||||
OUTPUT_NODE=True
|
||||
|
||||
def vfi(
|
||||
self,
|
||||
frames,
|
||||
engine,
|
||||
clear_cache_after_n_frames=100,
|
||||
multiplier=2,
|
||||
use_cuda_graph=True,
|
||||
keep_model_loaded=False,
|
||||
):
|
||||
"""Interpolates frames using the TensorRT engine.
|
||||
|
||||
Args:
|
||||
frames (torch.Tensor): Input image batch.
|
||||
engine (Engine): The active TensorRT engine.
|
||||
clear_cache_after_n_frames (int): Frequency of CUDA cache clearing.
|
||||
multiplier (int): Target frame rate multiplier.
|
||||
use_cuda_graph (bool): Whether to use CUDA graph optimization.
|
||||
keep_model_loaded (bool): If True, keeps engine in VRAM after use.
|
||||
|
||||
Returns:
|
||||
tuple: Interpolated image batch.
|
||||
"""
|
||||
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()
|
||||
|
||||
# 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(mm.get_torch_device())
|
||||
output = engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream, use_cuda_graph)
|
||||
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:
|
||||
engine.reset()
|
||||
|
||||
return (out,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"RIFERTXFrameInterpolation": RIFERTXFrameInterpolation,
|
||||
"RIFERTXEngineLoader": RIFERTXEngineLoader,
|
||||
"RIFERTXEngineConfig": RIFERTXEngineConfig,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"RIFERTXFrameInterpolation": "RIFE RTX Frame Interpolation",
|
||||
"RIFERTXEngineLoader": "RIFE RTX Engine Loader",
|
||||
"RIFERTXEngineConfig": "RIFE RTX Engine Config",
|
||||
}
|
||||
+537
@@ -0,0 +1,537 @@
|
||||
# https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/tree/main/vfi_models/rife
|
||||
|
||||
"""
|
||||
26-Dez-21
|
||||
https://github.com/hzwer/Practical-RIFE
|
||||
https://github.com/hzwer/Practical-RIFE/blob/main/model/warplayer.py
|
||||
https://github.com/HolyWu/vs-rife/blob/master/vsrife/__init__.py
|
||||
"""
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
from torch.optim import AdamW
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
import warnings
|
||||
|
||||
device = "cuda:0"
|
||||
backwarp_tenGrid = {}
|
||||
|
||||
|
||||
class ResConv(nn.Module):
|
||||
def __init__(self, c, dilation=1):
|
||||
super(ResConv, self).__init__()
|
||||
self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1)
|
||||
self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True)
|
||||
self.relu = nn.LeakyReLU(0.2, True)
|
||||
|
||||
def forward(self, x):
|
||||
return self.relu(self.conv(x) * self.beta + x)
|
||||
|
||||
|
||||
def warp(tenInput, tenFlow):
|
||||
k = (str(tenFlow.device), str(tenFlow.size()))
|
||||
if k not in backwarp_tenGrid:
|
||||
tenHorizontal = (
|
||||
torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device)
|
||||
.view(1, 1, 1, tenFlow.shape[3])
|
||||
.expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1)
|
||||
)
|
||||
tenVertical = (
|
||||
torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device)
|
||||
.view(1, 1, tenFlow.shape[2], 1)
|
||||
.expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3])
|
||||
)
|
||||
backwarp_tenGrid[k] = torch.cat([tenHorizontal, tenVertical], 1).to(device)
|
||||
|
||||
tenFlow = torch.cat(
|
||||
[
|
||||
tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0),
|
||||
tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0),
|
||||
],
|
||||
1,
|
||||
)
|
||||
|
||||
g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1)
|
||||
|
||||
if tenInput.type() == "torch.cuda.HalfTensor":
|
||||
g = g.half()
|
||||
|
||||
return torch.nn.functional.grid_sample(
|
||||
input=tenInput,
|
||||
grid=g,
|
||||
mode="bilinear",
|
||||
padding_mode="border",
|
||||
align_corners=True,
|
||||
)
|
||||
|
||||
|
||||
def conv(
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
dilation=1,
|
||||
arch_ver="4.0",
|
||||
):
|
||||
if arch_ver == "4.0":
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
bias=True,
|
||||
),
|
||||
nn.PReLU(out_planes),
|
||||
)
|
||||
if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]:
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
bias=True,
|
||||
),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
)
|
||||
|
||||
|
||||
def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
bias=True,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
bias=True,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1, arch_ver="4.0"):
|
||||
if arch_ver == "4.0":
|
||||
return nn.Sequential(
|
||||
torch.nn.ConvTranspose2d(
|
||||
in_channels=in_planes,
|
||||
out_channels=out_planes,
|
||||
kernel_size=4,
|
||||
stride=2,
|
||||
padding=1,
|
||||
bias=True,
|
||||
),
|
||||
nn.PReLU(out_planes),
|
||||
)
|
||||
if arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]:
|
||||
return nn.Sequential(
|
||||
torch.nn.ConvTranspose2d(
|
||||
in_channels=in_planes,
|
||||
out_channels=out_planes,
|
||||
kernel_size=4,
|
||||
stride=2,
|
||||
padding=1,
|
||||
bias=True,
|
||||
),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
)
|
||||
|
||||
|
||||
class Conv2(nn.Module):
|
||||
def __init__(self, in_planes, out_planes, stride=2, arch_ver="4.0"):
|
||||
super(Conv2, self).__init__()
|
||||
self.conv1 = conv(in_planes, out_planes, 3, stride, 1, arch_ver=arch_ver)
|
||||
self.conv2 = conv(out_planes, out_planes, 3, 1, 1, arch_ver=arch_ver)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.conv2(x)
|
||||
return x
|
||||
|
||||
|
||||
class IFBlock(nn.Module):
|
||||
def __init__(self, in_planes, c=64, arch_ver="4.0"):
|
||||
super(IFBlock, self).__init__()
|
||||
self.arch_ver = arch_ver
|
||||
self.conv0 = nn.Sequential(
|
||||
conv(in_planes, c // 2, 3, 2, 1, arch_ver=arch_ver),
|
||||
conv(c // 2, c, 3, 2, 1, arch_ver=arch_ver),
|
||||
)
|
||||
self.arch_ver = arch_ver
|
||||
|
||||
if arch_ver in ["4.0", "4.2", "4.3"]:
|
||||
self.convblock = nn.Sequential(
|
||||
conv(c, c, arch_ver=arch_ver),
|
||||
conv(c, c, arch_ver=arch_ver),
|
||||
conv(c, c, arch_ver=arch_ver),
|
||||
conv(c, c, arch_ver=arch_ver),
|
||||
conv(c, c, arch_ver=arch_ver),
|
||||
conv(c, c, arch_ver=arch_ver),
|
||||
conv(c, c, arch_ver=arch_ver),
|
||||
conv(c, c, arch_ver=arch_ver),
|
||||
)
|
||||
self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)
|
||||
|
||||
if arch_ver in ["4.5", "4.6", "4.7", "4.10"]:
|
||||
self.convblock = nn.Sequential(
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
ResConv(c),
|
||||
)
|
||||
if arch_ver == "4.5":
|
||||
self.lastconv = nn.Sequential(
|
||||
nn.ConvTranspose2d(c, 4 * 5, 4, 2, 1), nn.PixelShuffle(2)
|
||||
)
|
||||
if arch_ver in ["4.6", "4.7", "4.10"]:
|
||||
self.lastconv = nn.Sequential(
|
||||
nn.ConvTranspose2d(c, 4 * 6, 4, 2, 1), nn.PixelShuffle(2)
|
||||
)
|
||||
|
||||
def forward(self, x, flow=None, scale=1):
|
||||
x = F.interpolate(
|
||||
x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
|
||||
)
|
||||
if flow is not None:
|
||||
flow = (
|
||||
F.interpolate(
|
||||
flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
|
||||
)
|
||||
* 1.0
|
||||
/ scale
|
||||
)
|
||||
x = torch.cat((x, flow), 1)
|
||||
feat = self.conv0(x)
|
||||
if self.arch_ver == "4.0":
|
||||
feat = self.convblock(feat) + feat
|
||||
if self.arch_ver in ["4.2", "4.3", "4.5", "4.6", "4.7", "4.10"]:
|
||||
feat = self.convblock(feat)
|
||||
|
||||
tmp = self.lastconv(feat)
|
||||
if self.arch_ver in ["4.0", "4.2", "4.3"]:
|
||||
tmp = F.interpolate(
|
||||
tmp, scale_factor=scale * 2, mode="bilinear", align_corners=False
|
||||
)
|
||||
flow = tmp[:, :4] * scale * 2
|
||||
if self.arch_ver in ["4.5", "4.6", "4.7", "4.10"]:
|
||||
tmp = F.interpolate(
|
||||
tmp, scale_factor=scale, mode="bilinear", align_corners=False
|
||||
)
|
||||
flow = tmp[:, :4] * scale
|
||||
mask = tmp[:, 4:5]
|
||||
return flow, mask
|
||||
|
||||
|
||||
class Contextnet(nn.Module):
|
||||
def __init__(self, arch_ver="4.0"):
|
||||
super(Contextnet, self).__init__()
|
||||
c = 16
|
||||
self.conv1 = Conv2(3, c, arch_ver=arch_ver)
|
||||
self.conv2 = Conv2(c, 2 * c, arch_ver=arch_ver)
|
||||
self.conv3 = Conv2(2 * c, 4 * c, arch_ver=arch_ver)
|
||||
self.conv4 = Conv2(4 * c, 8 * c, arch_ver=arch_ver)
|
||||
|
||||
def forward(self, x, flow):
|
||||
x = self.conv1(x)
|
||||
flow = (
|
||||
F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False)
|
||||
* 0.5
|
||||
)
|
||||
f1 = warp(x, flow)
|
||||
x = self.conv2(x)
|
||||
flow = (
|
||||
F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False)
|
||||
* 0.5
|
||||
)
|
||||
f2 = warp(x, flow)
|
||||
x = self.conv3(x)
|
||||
flow = (
|
||||
F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False)
|
||||
* 0.5
|
||||
)
|
||||
f3 = warp(x, flow)
|
||||
x = self.conv4(x)
|
||||
flow = (
|
||||
F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False)
|
||||
* 0.5
|
||||
)
|
||||
f4 = warp(x, flow)
|
||||
return [f1, f2, f3, f4]
|
||||
|
||||
|
||||
class Unet(nn.Module):
|
||||
def __init__(self, arch_ver="4.0"):
|
||||
super(Unet, self).__init__()
|
||||
c = 16
|
||||
self.down0 = Conv2(17, 2 * c, arch_ver=arch_ver)
|
||||
self.down1 = Conv2(4 * c, 4 * c, arch_ver=arch_ver)
|
||||
self.down2 = Conv2(8 * c, 8 * c, arch_ver=arch_ver)
|
||||
self.down3 = Conv2(16 * c, 16 * c, arch_ver=arch_ver)
|
||||
self.up0 = deconv(32 * c, 8 * c, arch_ver=arch_ver)
|
||||
self.up1 = deconv(16 * c, 4 * c, arch_ver=arch_ver)
|
||||
self.up2 = deconv(8 * c, 2 * c, arch_ver=arch_ver)
|
||||
self.up3 = deconv(4 * c, c, arch_ver=arch_ver)
|
||||
self.conv = nn.Conv2d(c, 3, 3, 1, 1)
|
||||
|
||||
def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1):
|
||||
s0 = self.down0(
|
||||
torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1)
|
||||
)
|
||||
s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
|
||||
s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
|
||||
s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
|
||||
x = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
|
||||
x = self.up1(torch.cat((x, s2), 1))
|
||||
x = self.up2(torch.cat((x, s1), 1))
|
||||
x = self.up3(torch.cat((x, s0), 1))
|
||||
x = self.conv(x)
|
||||
return torch.sigmoid(x)
|
||||
|
||||
|
||||
"""
|
||||
currently supports 4.0-4.12
|
||||
|
||||
4.0: 4.0, 4.1
|
||||
4.2: 4.2
|
||||
4.3: 4.3, 4.4
|
||||
4.5: 4.5
|
||||
4.6: 4.6
|
||||
4.7: 4.7, 4.8, 4.9
|
||||
4.10: 4.10 4.11 4.12
|
||||
"""
|
||||
|
||||
|
||||
class IFNet(nn.Module):
|
||||
def __init__(self, arch_ver="4.0"):
|
||||
super(IFNet, self).__init__()
|
||||
self.arch_ver = arch_ver
|
||||
if arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]:
|
||||
self.block0 = IFBlock(7, c=192, arch_ver=arch_ver)
|
||||
self.block1 = IFBlock(8 + 4, c=128, arch_ver=arch_ver)
|
||||
self.block2 = IFBlock(8 + 4, c=96, arch_ver=arch_ver)
|
||||
self.block3 = IFBlock(8 + 4, c=64, arch_ver=arch_ver)
|
||||
if arch_ver in ["4.7"]:
|
||||
self.block0 = IFBlock(7 + 8, c=192, arch_ver=arch_ver)
|
||||
self.block1 = IFBlock(8 + 4 + 8, c=128, arch_ver=arch_ver)
|
||||
self.block2 = IFBlock(8 + 4 + 8, c=96, arch_ver=arch_ver)
|
||||
self.block3 = IFBlock(8 + 4 + 8, c=64, arch_ver=arch_ver)
|
||||
self.encode = nn.Sequential(
|
||||
nn.Conv2d(3, 16, 3, 2, 1), nn.ConvTranspose2d(16, 4, 4, 2, 1)
|
||||
)
|
||||
if arch_ver in ["4.10"]:
|
||||
self.block0 = IFBlock(7 + 16, c=192)
|
||||
self.block1 = IFBlock(8 + 4 + 16, c=128)
|
||||
self.block2 = IFBlock(8 + 4 + 16, c=96)
|
||||
self.block3 = IFBlock(8 + 4 + 16, c=64)
|
||||
self.encode = nn.Sequential(
|
||||
nn.Conv2d(3, 32, 3, 2, 1),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
nn.Conv2d(32, 32, 3, 1, 1),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
nn.Conv2d(32, 32, 3, 1, 1),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
nn.ConvTranspose2d(32, 8, 4, 2, 1),
|
||||
)
|
||||
|
||||
if arch_ver in ["4.0", "4.2", "4.3"]:
|
||||
self.contextnet = Contextnet(arch_ver=arch_ver)
|
||||
self.unet = Unet(arch_ver=arch_ver)
|
||||
self.arch_ver = arch_ver
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img0,
|
||||
img1,
|
||||
timestep,
|
||||
scale_list=[8,4,2,1],
|
||||
training=True,
|
||||
fastmode=True,
|
||||
ensemble=False,
|
||||
return_flow=False,
|
||||
):
|
||||
img0 = torch.clamp(img0, 0, 1)
|
||||
img1 = torch.clamp(img1, 0, 1)
|
||||
|
||||
scale_list = [8,4,2,1]
|
||||
|
||||
training=True
|
||||
fastmode=True
|
||||
ensemble=True
|
||||
return_flow=False
|
||||
|
||||
n, c, h, w = img0.shape
|
||||
ph = ((h - 1) // 64 + 1) * 64
|
||||
pw = ((w - 1) // 64 + 1) * 64
|
||||
padding = (0, pw - w, 0, ph - h)
|
||||
img0 = F.pad(img0, padding)
|
||||
img1 = F.pad(img1, padding)
|
||||
x = torch.cat((img0, img1), 1)
|
||||
|
||||
if training == False:
|
||||
channel = x.shape[1] // 2
|
||||
img0 = x[:, :channel]
|
||||
img1 = x[:, channel:]
|
||||
if not torch.is_tensor(timestep):
|
||||
timestep = (x[:, :1].clone() * 0 + 1) * timestep
|
||||
else:
|
||||
timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3])
|
||||
|
||||
timestep = timestep.to(device)
|
||||
flow_list = []
|
||||
merged = []
|
||||
mask_list = []
|
||||
|
||||
if self.arch_ver in ["4.7", "4.10"]:
|
||||
f0 = self.encode(img0[:, :3])
|
||||
f1 = self.encode(img1[:, :3])
|
||||
|
||||
|
||||
warped_img0 = img0
|
||||
warped_img1 = img1
|
||||
flow = None
|
||||
mask = None
|
||||
block = [self.block0, self.block1, self.block2, self.block3]
|
||||
|
||||
for i in range(4):
|
||||
if flow is None:
|
||||
# 4.7+
|
||||
if self.arch_ver in ["4.7", "4.10"]:
|
||||
flow, mask = block[i](
|
||||
torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1),
|
||||
None,
|
||||
scale=scale_list[i],
|
||||
)
|
||||
|
||||
if ensemble:
|
||||
f_, m_ = block[i](
|
||||
torch.cat(
|
||||
(img1[:, :3], img0[:, :3], f1, f0, 1 - timestep), 1
|
||||
),
|
||||
None,
|
||||
scale=scale_list[i],
|
||||
)
|
||||
flow = (flow + torch.cat((f_[:, 2:4], f_[:, :2]), 1)) / 2
|
||||
mask = (mask + (-m_)) / 2
|
||||
|
||||
else:
|
||||
# 4.7+
|
||||
if self.arch_ver in ["4.7", "4.10"]:
|
||||
fd, m0 = block[i](
|
||||
torch.cat(
|
||||
(
|
||||
warped_img0[:, :3],
|
||||
warped_img1[:, :3],
|
||||
warp(f0, flow[:, :2]),
|
||||
warp(f1, flow[:, 2:4]),
|
||||
timestep,
|
||||
mask,
|
||||
),
|
||||
1,
|
||||
),
|
||||
flow,
|
||||
scale=scale_list[i],
|
||||
)
|
||||
flow = flow + fd
|
||||
|
||||
# 4.0-4.6 ensemble
|
||||
if ensemble and self.arch_ver in [
|
||||
"4.0",
|
||||
"4.2",
|
||||
"4.3",
|
||||
"4.5",
|
||||
"4.6",
|
||||
]:
|
||||
f1, m1 = block[i](
|
||||
torch.cat(
|
||||
(
|
||||
warped_img1[:, :3],
|
||||
warped_img0[:, :3],
|
||||
1 - timestep,
|
||||
-mask,
|
||||
),
|
||||
1,
|
||||
),
|
||||
torch.cat((flow[:, 2:4], flow[:, :2]), 1),
|
||||
scale=scale_list[i],
|
||||
)
|
||||
f0 = (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2
|
||||
m0 = (m0 + (-m1)) / 2
|
||||
|
||||
# 4.7+ ensemble
|
||||
if ensemble and self.arch_ver in ["4.7", "4.10"]:
|
||||
wf0 = warp(f0, flow[:, :2])
|
||||
wf1 = warp(f1, flow[:, 2:4])
|
||||
|
||||
f_, m_ = block[i](
|
||||
torch.cat(
|
||||
(
|
||||
warped_img1[:, :3],
|
||||
warped_img0[:, :3],
|
||||
wf1,
|
||||
wf0,
|
||||
1 - timestep,
|
||||
-mask,
|
||||
),
|
||||
1,
|
||||
),
|
||||
torch.cat((flow[:, 2:4], flow[:, :2]), 1),
|
||||
scale=scale_list[i],
|
||||
)
|
||||
fd = (fd + torch.cat((f_[:, 2:4], f_[:, :2]), 1)) / 2
|
||||
mask = (m0 + (-m_)) / 2
|
||||
|
||||
if self.arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]:
|
||||
flow = flow + f0
|
||||
mask = mask + m0
|
||||
|
||||
if not ensemble and self.arch_ver in ["4.7", "4.10"]:
|
||||
mask = m0
|
||||
|
||||
mask_list.append(mask)
|
||||
flow_list.append(flow)
|
||||
warped_img0 = warp(img0, flow[:, :2])
|
||||
warped_img1 = warp(img1, flow[:, 2:4])
|
||||
merged.append((warped_img0, warped_img1))
|
||||
|
||||
if self.arch_ver in ["4.0", "4.1", "4.2", "4.3", "4.4", "4.5", "4.6"]:
|
||||
mask_list[3] = torch.sigmoid(mask_list[3])
|
||||
merged[3] = merged[3][0] * mask_list[3] + merged[3][1] * (1 - mask_list[3])
|
||||
|
||||
if self.arch_ver in ["4.7", "4.10"]:
|
||||
mask = torch.sigmoid(mask)
|
||||
merged[3] = warped_img0 * mask + warped_img1 * (1 - mask)
|
||||
|
||||
if not fastmode and self.arch_ver in ["4.0", "4.2", "4.3"]:
|
||||
c0 = self.contextnet(img0, flow[:, :2])
|
||||
c1 = self.contextnet(img1, flow[:, 2:4])
|
||||
tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
|
||||
res = tmp[:, :3] * 2 - 1
|
||||
merged[3] = torch.clamp(merged[3] + res, 0, 1)
|
||||
|
||||
result = merged[3][:, :, :h, :w]
|
||||
return result
|
||||
@@ -0,0 +1,142 @@
|
||||
# https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/tree/main/vfi_utils.py
|
||||
# https://github.com/yester31/TensorRT_Examples/blob/main/timm_to_trt_python1/onnx_export.py
|
||||
|
||||
import torch
|
||||
import pathlib
|
||||
import traceback
|
||||
import os
|
||||
from urllib.parse import urlparse
|
||||
from torch.hub import download_url_to_file, get_dir
|
||||
from rife_arch import IFNet
|
||||
import onnx
|
||||
from onnxsim import simplify
|
||||
|
||||
MODEL_TYPE = pathlib.Path(__file__).parent.name
|
||||
CKPT_NAME_VER_DICT = {
|
||||
"rife40.pth": "4.0",
|
||||
"rife41.pth": "4.0",
|
||||
"rife42.pth": "4.2",
|
||||
"rife43.pth": "4.3",
|
||||
"rife44.pth": "4.3",
|
||||
"rife45.pth": "4.5",
|
||||
"rife46.pth": "4.6",
|
||||
"rife47.pth": "4.7",
|
||||
"rife48.pth": "4.7",
|
||||
"rife49.pth": "4.7",
|
||||
"sudo_rife4_269.662_testV1_scale1.pth": "4.0"
|
||||
}
|
||||
BASE_MODEL_DOWNLOAD_URLS = [
|
||||
"https://github.com/styler00dollar/VSGAN-tensorrt-docker/releases/download/models/",
|
||||
"https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/releases/download/models/",
|
||||
"https://github.com/dajes/frame-interpolation-pytorch/releases/download/v1.0.0/"
|
||||
]
|
||||
TORCH_DEVICE = "cuda:0"
|
||||
|
||||
def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
|
||||
"""Load file form http url, will download models if necessary.
|
||||
|
||||
Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
|
||||
|
||||
Args:
|
||||
url (str): URL to be downloaded.
|
||||
model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
|
||||
Default: None.
|
||||
progress (bool): Whether to show the download progress. Default: True.
|
||||
file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
|
||||
|
||||
Returns:
|
||||
str: The path to the downloaded file.
|
||||
"""
|
||||
if model_dir is None: # use the pytorch hub_dir
|
||||
hub_dir = get_dir()
|
||||
model_dir = os.path.join(hub_dir, 'checkpoints')
|
||||
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
|
||||
parts = urlparse(url)
|
||||
file_name = os.path.basename(parts.path)
|
||||
if file_name is not None:
|
||||
file_name = file_name
|
||||
cached_file = os.path.abspath(os.path.join(model_dir, file_name))
|
||||
if not os.path.exists(cached_file):
|
||||
print(f'Downloading: "{url}" to {cached_file}\n')
|
||||
download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
|
||||
return cached_file
|
||||
|
||||
def load_file_from_github_release(model_type, ckpt_name):
|
||||
os.makedirs("models", exist_ok=True)
|
||||
error_strs = []
|
||||
for i, base_model_download_url in enumerate(BASE_MODEL_DOWNLOAD_URLS):
|
||||
try:
|
||||
return load_file_from_url(base_model_download_url + ckpt_name, "models")
|
||||
except Exception:
|
||||
traceback_str = traceback.format_exc()
|
||||
if i < len(BASE_MODEL_DOWNLOAD_URLS) - 1:
|
||||
print("Failed! Trying another endpoint.")
|
||||
error_strs.append(f"Error when downloading from: {base_model_download_url + ckpt_name}\n\n{traceback_str}")
|
||||
|
||||
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 export_onnx(ckpt_name, ensemble, scale_factor):
|
||||
print(f"PyTorch version: {torch.__version__}")
|
||||
print(f"ONNX version: {onnx.__version__}")
|
||||
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
print(f"Using device: {device}")
|
||||
|
||||
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.eval().to(TORCH_DEVICE)
|
||||
|
||||
# # dummy data
|
||||
img0 = torch.randn(1, 3, 512, 512).to(TORCH_DEVICE)
|
||||
img1 = torch.randn(1, 3, 512, 512).to(TORCH_DEVICE)
|
||||
timestep = torch.tensor([0.5], dtype=torch.float32).to(TORCH_DEVICE)
|
||||
|
||||
# result = (interpolation_model(img0, img1, timestep))
|
||||
# print(result)
|
||||
|
||||
onnx_save_name = f"{ckpt_name.split('.')[0]}_ensemble_{ensemble}_scale_{scale_factor}.onnx"
|
||||
onnx_save_path = os.path.join("models",onnx_save_name)
|
||||
|
||||
dynamic_axes = {
|
||||
'img0': {2: 'height', 3: 'width'},
|
||||
'img1': {2: 'height', 3: 'width'},
|
||||
'output': {2: 'height', 3: 'width'},
|
||||
}
|
||||
|
||||
with torch.no_grad(): # Disable gradients for efficiency
|
||||
torch.onnx.export(interpolation_model,
|
||||
(img0, img1, timestep), # ensemble + scale_factor set in forward fn
|
||||
onnx_save_path,
|
||||
export_params=True,
|
||||
opset_version=19,
|
||||
do_constant_folding=True,
|
||||
input_names=['img0', 'img1', 'timestep'],
|
||||
output_names=['output'],
|
||||
dynamic_axes=dynamic_axes
|
||||
)
|
||||
print(f"ONNX model exported to: {onnx_save_path}")
|
||||
|
||||
# Verify the exported ONNX model
|
||||
onnx_model = onnx.load(onnx_save_path)
|
||||
onnx.checker.check_model(onnx_model) # Perform a validity check
|
||||
print("ONNX model validation successful!")
|
||||
|
||||
# print(onnx.helper.printable_graph(onnx_model.graph))
|
||||
|
||||
sim_model_path = os.path.join("models", f"{ckpt_name.split('.')[0]}_ensemble_{ensemble}_scale_{scale_factor}_sim.onnx")
|
||||
print("=> ONNX simplify start!")
|
||||
sim_onnx_model, check = simplify(onnx_model) # convert(simplify)
|
||||
onnx.save(sim_onnx_model, sim_model_path)
|
||||
print("=> ONNX simplify done!")
|
||||
|
||||
sim_model = onnx.load(sim_model_path)
|
||||
onnx.checker.check_model(sim_model)
|
||||
print("=> ONNX Model exported at ", sim_model_path)
|
||||
print("=> sim ONNX Model check done!")
|
||||
|
||||
|
||||
export_onnx(ckpt_name="rife47.pth", ensemble=True, scale_factor=1)
|
||||
@@ -0,0 +1,44 @@
|
||||
import sys
|
||||
import os
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
||||
import torch
|
||||
import time
|
||||
from trt_utilities import Engine
|
||||
import tensorrt
|
||||
|
||||
print("Tensorrt version: ", tensorrt.__version__)
|
||||
|
||||
def export_trt(trt_path=None, onnx_path=None, use_fp16=True):
|
||||
if trt_path is None:
|
||||
trt_path = input(
|
||||
"Enter the path to save the TensorRT engine (e.g ./realesrgan.engine): ")
|
||||
if onnx_path is None:
|
||||
onnx_path = input(
|
||||
"Enter the path to the ONNX model (e.g ./realesrgan.onnx): ")
|
||||
|
||||
engine = Engine(trt_path)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
s = time.time()
|
||||
ret = engine.build(
|
||||
onnx_path,
|
||||
use_fp16,
|
||||
enable_preview=True,
|
||||
input_profile=[
|
||||
# any sizes from 256x256 to 3840x3840, batch size 1
|
||||
{
|
||||
"img0": [(1, 3, 256, 256), (1, 3, 512, 512), (1, 3, 3840, 3840)],
|
||||
"img1": [(1, 3, 256, 256), (1, 3, 512, 512), (1, 3, 3840, 3840)],
|
||||
},
|
||||
],
|
||||
)
|
||||
e = time.time()
|
||||
print(f"Time taken to build: {(e-s)} seconds")
|
||||
print(f"Tensorrt engine saved at: {trt_path}")
|
||||
|
||||
return ret
|
||||
|
||||
|
||||
export_trt(trt_path="./models/rife49_ensemble_True_scale_1_sim.engine",
|
||||
onnx_path="./models/rife49_ensemble_True_scale_1_sim.onnx", use_fp16=True)
|
||||
@@ -0,0 +1,455 @@
|
||||
"""TensorRT Acceleration Utilities for RIFE.
|
||||
|
||||
This module provides the low-level interface for NVIDIA TensorRT, handling
|
||||
engine compilation, timing caches, memory allocation, and optimized inference
|
||||
using CUDA Graphs.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import io
|
||||
import os
|
||||
|
||||
import cuda.bindings.runtime as cudart
|
||||
import onnx
|
||||
import numpy as np
|
||||
import tensorrt_rtx as trt
|
||||
import torch
|
||||
|
||||
from collections import OrderedDict
|
||||
from logging import error, warning
|
||||
from onnxconverter_common import float16
|
||||
from onnxsim import simplify
|
||||
|
||||
# Use polygraphy TensorRTX backend
|
||||
os.environ["POLYGRAPHY_USE_TENSORRT_RTX"] = "1"
|
||||
|
||||
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,
|
||||
network_from_onnx_bytes,
|
||||
save_engine,
|
||||
)
|
||||
from polygraphy.logger import G_LOGGER
|
||||
|
||||
from torch.cuda import nvtx
|
||||
from tqdm import tqdm
|
||||
|
||||
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()
|
||||
}
|
||||
|
||||
# Path for your persistent hardware profile
|
||||
node_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
MASTER_CACHE_FILE = f"{node_dir}\\RIFE_RTX_ENGINE.cache"
|
||||
|
||||
def get_timing_cache(config):
|
||||
"""Retrieves or creates a TensorRT timing cache.
|
||||
|
||||
Args:
|
||||
config (trt.IBuilderConfig): The builder configuration object.
|
||||
|
||||
Returns:
|
||||
trt.ITimingCache: The loaded or new timing cache.
|
||||
"""
|
||||
if os.path.exists(MASTER_CACHE_FILE):
|
||||
with open(MASTER_CACHE_FILE, "rb") as f:
|
||||
cache_data = f.read()
|
||||
return config.create_timing_cache(cache_data)
|
||||
return config.create_timing_cache(b"")
|
||||
|
||||
def save_timing_cache(timing_cache):
|
||||
"""Saves the current timing cache to disk for future builds.
|
||||
|
||||
Args:
|
||||
timing_cache (trt.ITimingCache): The cache to serialize.
|
||||
"""
|
||||
with open(MASTER_CACHE_FILE, "wb") as f:
|
||||
f.write(timing_cache.serialize())
|
||||
|
||||
def CUASSERT(cuda_ret):
|
||||
"""Asserts that a CUDA call was successful.
|
||||
https://github.com/Jeff-LiangF/streamv2v/blob/18c1a3bd56ff348d54a3300605936980bb13b03c/src/streamv2v/acceleration/tensorrt/utilities.py
|
||||
|
||||
Args:
|
||||
cuda_ret (tuple): The return value from a cuda-python call.
|
||||
|
||||
Returns:
|
||||
Any: The result of the CUDA call if successful.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the CUDA call returned an error.
|
||||
"""
|
||||
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):
|
||||
"""Progress monitor for TensorRT engine builds using TQDM bars."""
|
||||
|
||||
def __init__(self):
|
||||
trt.IProgressMonitor.__init__(self)
|
||||
self._active_phases = {}
|
||||
self._step_result = True
|
||||
self.max_indent = 5
|
||||
|
||||
def phase_start(self, phase_name, parent_phase, num_steps):
|
||||
"""Called when a build phase starts."""
|
||||
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):
|
||||
"""Called when a build phase finishes."""
|
||||
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):
|
||||
"""Called when a step within a phase completes."""
|
||||
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:
|
||||
"""Wrapper for TensorRT engine management.
|
||||
|
||||
Handles loading, building, and executing inference with optional
|
||||
CUDA Graph support.
|
||||
|
||||
Attributes:
|
||||
engine_path (str): File path to the .trt engine.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
engine_path,
|
||||
):
|
||||
self.engine_path = engine_path
|
||||
self.engine = None
|
||||
self.context = None
|
||||
self.buffers = OrderedDict()
|
||||
self.tensors = OrderedDict()
|
||||
self.cuda_graph_instance = None # cuda graph
|
||||
self.graph = None
|
||||
|
||||
def __del__(self):
|
||||
"""Ensures CUDA resources are freed on deletion."""
|
||||
if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
|
||||
try:
|
||||
cudart.cudaGraphDestroy(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):
|
||||
"""Cleans up engine and CUDA resources.
|
||||
|
||||
Args:
|
||||
engine_path (str, optional): New path to set for the engine.
|
||||
"""
|
||||
if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
|
||||
try:
|
||||
cudart.cudaGraphDestroy(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,
|
||||
input_profile=None,
|
||||
enable_refit=False,
|
||||
enable_preview=False,
|
||||
enable_all_tactics=False,
|
||||
timing_cache=None,
|
||||
update_output_names=None,
|
||||
):
|
||||
"""Compiles an ONNX model into a TensorRT engine.
|
||||
|
||||
Args:
|
||||
onnx_path (str): Path to source ONNX file.
|
||||
input_profile (list): List of dicts defining min/opt/max shapes.
|
||||
**kwargs: Configuration flags (enable_refit, etc).
|
||||
|
||||
Returns:
|
||||
int: 0 if success, 1 if failure.
|
||||
"""
|
||||
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"] = []
|
||||
|
||||
# Load ONNX file then convert the model to fp16
|
||||
# Probably might cause issues if onnx file given is quantized or whatever
|
||||
onnx_bytes = onnx.load(onnx_path)
|
||||
onnx_precision = onnx_bytes.graph.input[0].type.tensor_type.elem_type
|
||||
onnx_simp, check = simplify(float16.convert_float_to_float16(onnx_bytes) if onnx_precision == onnx.TensorProto.FLOAT else onnx_bytes)
|
||||
|
||||
if not check:
|
||||
error(f"Simplified ONNX model: {onnx_path} could not be validated")
|
||||
return 1
|
||||
|
||||
# Keep simplified model in memory to not overwrite the onnx file in models directory
|
||||
# We dont want to be destructive
|
||||
onnx_buffer = io.BytesIO()
|
||||
onnx_buffer.write(onnx_simp.SerializeToString())
|
||||
onnx_buffer.seek(0)
|
||||
onnx_buffer.name = "onnx_simp_mem_rife.onnx"
|
||||
|
||||
#network = network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM])
|
||||
network = network_from_onnx_bytes(onnx_buffer.getvalue(), flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM])
|
||||
|
||||
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()
|
||||
config.progress_monitor = TQDMProgressMonitor()
|
||||
|
||||
#config.set_flag(trt.BuilderFlag.FP16) if fp16 else None
|
||||
config.set_flag(trt.BuilderFlag.REFIT) if enable_refit else None
|
||||
|
||||
# Limit workspace (e.g., to 2GB or 4GB depending on your GPU)
|
||||
# 1 << 30 is 1GB
|
||||
#config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * (1 << 30))
|
||||
config.builder_optimization_level = 5
|
||||
|
||||
# Use Master Cache
|
||||
timing_cache = get_timing_cache(config)
|
||||
config.set_timing_cache(timing_cache, False)
|
||||
|
||||
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)
|
||||
|
||||
result = 0
|
||||
|
||||
try:
|
||||
engine = engine_from_network(network, config)
|
||||
save_engine(engine, path=self.engine_path)
|
||||
except Exception as e:
|
||||
error(f"Failed to build engine: {e}")
|
||||
result = 1
|
||||
|
||||
# Update Master Cache with new findings
|
||||
save_timing_cache(timing_cache)
|
||||
print(f"Updated Master Cache.")
|
||||
|
||||
return result
|
||||
|
||||
def load(self):
|
||||
"""Loads the serialized engine from disk."""
|
||||
self.engine = engine_from_bytes(bytes_from_path(self.engine_path))
|
||||
|
||||
def activate(self, reuse_device_memory=None):
|
||||
"""Initializes the execution context."""
|
||||
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 allocate_buffers(self, shape_dict=None, device="cuda"):
|
||||
"""Allocates GPU memory for input and output tensors.
|
||||
|
||||
Args:
|
||||
shape_dict (dict, optional): Specific shapes for tensors.
|
||||
device (str): Device to allocate on (default "cuda").
|
||||
"""
|
||||
if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:
|
||||
try:
|
||||
cudart.cudaGraphDestroy(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)
|
||||
|
||||
dtype = trt.nptype(self.engine.get_tensor_dtype(name))
|
||||
if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
|
||||
self.context.set_input_shape(name, shape)
|
||||
tensor = torch.empty(
|
||||
tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]
|
||||
).to(device=device)
|
||||
self.tensors[binding] = tensor
|
||||
nvtx.range_pop()
|
||||
|
||||
def infer(self, feed_dict, stream, use_cuda_graph=False):
|
||||
"""Performs inference on the provided data.
|
||||
|
||||
Args:
|
||||
feed_dict (dict): Input tensors.
|
||||
stream (cuda.Stream): CUDA stream for execution.
|
||||
use_cuda_graph (bool): If True, uses/captures CUDA graphs.
|
||||
|
||||
Returns:
|
||||
OrderedDict: The output 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.")
|
||||
|
||||
return self.tensors
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Utility functions for RIFE Video Frame Interpolation (VFI).
|
||||
|
||||
This module contains preprocessing, postprocessing, and the core interpolation
|
||||
loop logic used to generate intermediate frames between a sequence of images.
|
||||
|
||||
https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/blob/main/vfi_utils.py
|
||||
"""
|
||||
|
||||
import einops
|
||||
import os
|
||||
import torch
|
||||
import typing
|
||||
|
||||
import numpy as np
|
||||
|
||||
from comfy.model_management import soft_empty_cache, get_torch_device
|
||||
from comfy.utils import ProgressBar
|
||||
|
||||
DEVICE = get_torch_device()
|
||||
|
||||
def preprocess_frames(frames: torch.Tensor) -> torch.Tensor:
|
||||
"""Prepares ComfyUI image batches for model inference.
|
||||
|
||||
Converts images from (Batch, Height, Width, Channels) to
|
||||
(Batch, Channels, Height, Width) and ensures only RGB channels are present.
|
||||
|
||||
Args:
|
||||
frames (torch.Tensor): Input tensor in BHWC format.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Preprocessed tensor in BCHW format.
|
||||
"""
|
||||
return einops.rearrange(frames[..., :3], "n h w c -> n c h w")
|
||||
|
||||
def postprocess_frames(frames: torch.Tensor) -> torch.Tensor:
|
||||
"""Converts model output back to ComfyUI standard format.
|
||||
|
||||
Converts images from (Batch, Channels, Height, Width) back to
|
||||
(Batch, Height, Width, Channels) and moves them to CPU memory.
|
||||
|
||||
Args:
|
||||
frames (torch.Tensor): Output tensor from the model in BCHW format.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Postprocessed tensor in BHWC format on CPU.
|
||||
"""
|
||||
return einops.rearrange(frames, "n c h w -> n h w c")[..., :3].cpu()
|
||||
|
||||
def generate_frames_rife(
|
||||
frames,
|
||||
clear_cache_after_n_frames,
|
||||
multiplier,
|
||||
return_middle_frame_function
|
||||
):
|
||||
"""Core interpolation loop for generating intermediate frames.
|
||||
|
||||
Iterates through a sequence of frames and calls the provided inference
|
||||
function to generate 'n' intermediate frames based on the multiplier.
|
||||
|
||||
Args:
|
||||
frames (torch.Tensor): The preprocessed frame sequence (BCHW).
|
||||
clear_cache_after_n_frames (int): Interval for VRAM cache clearing.
|
||||
multiplier (int): The factor by which to increase the frame count.
|
||||
return_middle_frame_function (Callable): Function that takes two
|
||||
frames and a timestep and returns the interpolated frame.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The complete sequence including original and
|
||||
interpolated frames.
|
||||
"""
|
||||
output_frames = torch.zeros(multiplier*frames.shape[0], *frames.shape[1:], device="cpu")
|
||||
out_len = 0
|
||||
cache_counter = 0
|
||||
|
||||
pbar = ProgressBar(len(frames))
|
||||
|
||||
for frame_itr in range(len(frames) - 1):
|
||||
|
||||
frame_0 = frames[frame_itr:frame_itr+1]
|
||||
frame_1 = frames[frame_itr+1:frame_itr+2]
|
||||
|
||||
# Store the current original frame
|
||||
output_frames[out_len] = frame_0
|
||||
out_len += 1
|
||||
|
||||
# Generate intermediate frames
|
||||
for middle_i in range(1, multiplier):
|
||||
timestep = middle_i/multiplier
|
||||
middle_frame = return_middle_frame_function(frame_0, frame_1, timestep).detach().cpu()
|
||||
|
||||
# Copy middle frames to output
|
||||
output_frames[out_len] = middle_frame
|
||||
out_len +=1
|
||||
|
||||
# VRAM Management
|
||||
cache_counter += 1
|
||||
if cache_counter >= clear_cache_after_n_frames:
|
||||
soft_empty_cache()
|
||||
cache_counter = 0
|
||||
|
||||
pbar.update(1)
|
||||
|
||||
# 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}")
|
||||
out_len += 1
|
||||
|
||||
# clear cache for courtesy
|
||||
soft_empty_cache()
|
||||
|
||||
# Slice to actual length in case of rounding or logic offsets
|
||||
return output_frames[:out_len]
|
||||
Reference in New Issue
Block a user