From 4952023d35e072ca5ade852a883bbb9d0382557b Mon Sep 17 00:00:00 2001 From: yuvraj108c Date: Thu, 3 Oct 2024 07:14:29 +0000 Subject: [PATCH] v1 --- .gitignore | 2 + __init__.py | 85 ++++++++ export_onnx.py | 142 +++++++++++++ export_trt.py | 41 ++++ requirements.txt | 4 + rife_arch.py | 537 +++++++++++++++++++++++++++++++++++++++++++++++ trt_utilities.py | 258 +++++++++++++++++++++++ vfi_utilities.py | 84 ++++++++ 8 files changed, 1153 insertions(+) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 export_onnx.py create mode 100644 export_trt.py create mode 100644 requirements.txt create mode 100644 rife_arch.py create mode 100644 trt_utilities.py create mode 100644 vfi_utilities.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..10a149b --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +models +__pycache__ \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..6434b90 --- /dev/null +++ b/__init__.py @@ -0,0 +1,85 @@ +import torch +import os +from comfy.model_management import get_torch_device +from comfy.utils import ProgressBar +from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife, logger +from .trt_utilities import Engine +import folder_paths +import time + +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": 10, "min": 1, "max": 1000}), + "multiplier": ("INT", {"default": 2, "min": 1}), + }, + } + + RETURN_TYPES = ("IMAGE", ) + FUNCTION = "vfi" + CATEGORY = "tensorrt" + + def float_to_tensor(self, value, height, width): + return torch.full((1, 1, height, width), value) + + def vfi( + self, + frames, + engine, + clear_cache_after_n_frames=10, + 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)}, + } + + # cache tensorrt engine in memory + cudaStream = torch.cuda.current_stream().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) + # 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) + return (out,) + + +NODE_CLASS_MAPPINGS = { + "RifeTensorrt": RifeTensorrt, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RifeTensorrt": "Rife Tensorrt", +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/export_onnx.py b/export_onnx.py new file mode 100644 index 0000000..8df1ffb --- /dev/null +++ b/export_onnx.py @@ -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="rife49.pth", ensemble=True, scale_factor=1) \ No newline at end of file diff --git a/export_trt.py b/export_trt.py new file mode 100644 index 0000000..3292e62 --- /dev/null +++ b/export_trt.py @@ -0,0 +1,41 @@ +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) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a0e8f28 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +einops +colored +polygraphy +tensorrt==10.4.0 \ No newline at end of file diff --git a/rife_arch.py b/rife_arch.py new file mode 100644 index 0000000..5b1d722 --- /dev/null +++ b/rife_arch.py @@ -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 \ No newline at end of file diff --git a/trt_utilities.py b/trt_utilities.py new file mode 100644 index 0000000..ebe5649 --- /dev/null +++ b/trt_utilities.py @@ -0,0 +1,258 @@ +import torch +from torch.cuda import nvtx +from collections import OrderedDict +import numpy as np +from polygraphy.backend.common import bytes_from_path +from polygraphy import util +from polygraphy.backend.trt import ModifyNetworkOutputs, Profile +from polygraphy.backend.trt import ( + engine_from_bytes, + engine_from_network, + network_from_onnx_path, + save_engine, +) +from polygraphy.logger import G_LOGGER +import tensorrt as trt +from logging import error, warning +from tqdm import tqdm +import copy +# from .vfi_utilities import logger + +TRT_LOGGER = trt.Logger(trt.Logger.ERROR) +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() +} + +class TQDMProgressMonitor(trt.IProgressMonitor): + 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): + leave = False + try: + if parent_phase is not None: + nbIndents = ( + self._active_phases.get(parent_phase, {}).get( + "nbIndents", self.max_indent + ) + + 1 + ) + if nbIndents >= self.max_indent: + return + else: + nbIndents = 0 + leave = True + self._active_phases[phase_name] = { + "tq": tqdm( + total=num_steps, desc=phase_name, leave=leave, position=nbIndents + ), + "nbIndents": nbIndents, + "parent_phase": parent_phase, + } + except KeyboardInterrupt: + # The phase_start callback cannot directly cancel the build, so request the cancellation from within step_complete. + _step_result = False + + def phase_finish(self, phase_name): + try: + if phase_name in self._active_phases.keys(): + self._active_phases[phase_name]["tq"].update( + self._active_phases[phase_name]["tq"].total + - self._active_phases[phase_name]["tq"].n + ) + + parent_phase = self._active_phases[phase_name].get("parent_phase", None) + while parent_phase is not None: + self._active_phases[parent_phase]["tq"].refresh() + parent_phase = self._active_phases[parent_phase].get( + "parent_phase", None + ) + if ( + self._active_phases[phase_name]["parent_phase"] + in self._active_phases.keys() + ): + self._active_phases[ + self._active_phases[phase_name]["parent_phase"] + ]["tq"].refresh() + del self._active_phases[phase_name] + pass + except KeyboardInterrupt: + _step_result = False + + def step_complete(self, phase_name, step): + try: + if phase_name in self._active_phases.keys(): + self._active_phases[phase_name]["tq"].update( + step - self._active_phases[phase_name]["tq"].n + ) + return self._step_result + except KeyboardInterrupt: + # There is no need to propagate this exception to TensorRT. We can simply cancel the build. + return False + + +class Engine: + def __init__( + self, + engine_path, + ): + self.engine_path = engine_path + self.engine = None + self.context = None + self.buffers = OrderedDict() + self.tensors = OrderedDict() + self.cuda_graph_instance = None # cuda graph + + def __del__(self): + del self.engine + del self.context + del self.buffers + del self.tensors + + def reset(self, engine_path=None): + del self.engine + del self.context + del self.buffers + del self.tensors + self.engine_path = engine_path + + self.buffers = OrderedDict() + self.tensors = OrderedDict() + self.inputs = {} + self.outputs = {} + + def build( + self, + onnx_path, + fp16, + input_profile=None, + enable_refit=False, + enable_preview=False, + enable_all_tactics=False, + timing_cache=None, + update_output_names=None, + ): + print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}") + 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"] = [] + + network = network_from_onnx_path( + onnx_path, 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 + + profiles = copy.deepcopy(p) + for profile in profiles: + # Last profile is used for set_calibration_profile. + calib_profile = profile.fill_defaults(network[1]).to_trt( + builder, network[1] + ) + config.add_optimization_profile(calib_profile) + + try: + engine = engine_from_network( + network, + config, + ) + except Exception as e: + error(f"Failed to build engine: {e}") + return 1 + try: + save_engine(engine, path=self.engine_path) + except Exception as e: + error(f"Failed to save engine: {e}") + return 1 + return 0 + + def load(self): + # logger(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): + 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"): + 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): + 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()) + 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): + for idx in range(self.engine.num_io_tensors): + name = self.engine.get_tensor_name(idx) + shape = self.context.get_tensor_shape(name) + print(name, shape) diff --git a/vfi_utilities.py b/vfi_utilities.py new file mode 100644 index 0000000..37aaa31 --- /dev/null +++ b/vfi_utilities.py @@ -0,0 +1,84 @@ +# https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/blob/main/vfi_utils.py + +import os +import torch +import typing +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 + +DEVICE = get_torch_device() + +def load_file_from_github_release(model_type, ckpt_name): + 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, get_ckpt_container_path(model_type)) + 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 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") + +def postprocess_frames(frames): + 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 + ): + + output_frames = torch.zeros(multiplier*frames.shape[0], *frames.shape[1:], device="cpu") + out_len = 0 + + number_of_frames_processed_since_last_cleared_cuda_cache = 0 + pbar = ProgressBar(len(frames)) + + for frame_itr in range(len(frames) - 1): # Skip the final frame since there are no frames after it + + frame_0 = frames[frame_itr:frame_itr+1] + frame_1 = frames[frame_itr+1:frame_itr+2] + output_frames[out_len] = frame_0 # Start with first frame + out_len += 1 + + 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 + + # Try to avoid a memory overflow by clearing cuda cache regularly + number_of_frames_processed_since_last_cleared_cuda_cache += 1 + 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...") + + pbar.update(1) + + + # Append final frame + output_frames[out_len] = frames[-1:] + out_len += 1 + + # clear cache for courtesy + soft_empty_cache() + logger("Final clearing cache done ...") +# + res = output_frames[:out_len] + return res \ No newline at end of file