v1
This commit is contained in:
yuvraj108c
@@ -0,0 +1,2 @@
|
||||
models
|
||||
__pycache__
|
||||
+85
@@ -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']
|
||||
+142
@@ -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)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,4 @@
|
||||
einops
|
||||
colored
|
||||
polygraphy
|
||||
tensorrt==10.4.0
|
||||
+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,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)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user