From 8370a2f1efd299a586e7c208b663eec7e541eb70 Mon Sep 17 00:00:00 2001 From: peteromallet Date: Sat, 25 May 2024 20:05:19 +0200 Subject: [PATCH] FIxing classes --- SteerableMotion.py | 5 +- .../ComfyUI_Frame_Interpolation/__init__.py | 38 - .../vfi_models/amt/__init__.py | 87 - .../vfi_models/amt/amt_arch.py | 1590 --------- .../vfi_models/cain/__init__.py | 64 - .../vfi_models/cain/cain_arch.py | 74 - .../vfi_models/cain/cain_encdec_arch.py | 95 - .../vfi_models/cain/cain_noca_arch.py | 73 - .../vfi_models/cain/common.py | 361 --- .../vfi_models/eisai/__init__.py | 84 - .../vfi_models/eisai/eisai_arch.py | 2586 --------------- .../vfi_models/film/__init__.py | 6 +- .../vfi_models/flavr/__init__.py | 115 - .../vfi_models/flavr/flavr_arch.py | 217 -- .../vfi_models/flavr/resnet_3D.py | 288 -- .../vfi_models/gmfss_fortuna/GMFSS_Fortuna.py | 24 - .../gmfss_fortuna/GMFSS_Fortuna_arch.py | 1850 ----------- .../gmfss_fortuna/GMFSS_Fortuna_union.py | 23 - .../gmfss_fortuna/GMFSS_Fortuna_union_arch.py | 1857 ----------- .../vfi_models/gmfss_fortuna/__init__.py | 143 - .../vfi_models/ifrnet/IFRNet_L_arch.py | 293 -- .../vfi_models/ifrnet/IFRNet_S_arch.py | 293 -- .../vfi_models/ifrnet/__init__.py | 57 - .../vfi_models/ifunet/IFUNet_arch.py | 766 ----- .../vfi_models/ifunet/__init__.py | 59 - .../vfi_models/m2m/M2M_arch.py | 1037 ------ .../vfi_models/m2m/__init__.py | 60 - .../vfi_models/ops/__init__.py | 22 - .../vfi_models/ops/cupy_ops/__init__.py | 11 - .../vfi_models/ops/cupy_ops/adacof.py | 491 --- .../vfi_models/ops/cupy_ops/batch_edt.py | 119 - .../vfi_models/ops/cupy_ops/correlation.py | 413 --- .../vfi_models/ops/cupy_ops/costvol.py | 317 -- .../vfi_models/ops/cupy_ops/sepconv.py | 332 -- .../vfi_models/ops/cupy_ops/softsplat.py | 440 --- .../vfi_models/ops/cupy_ops/utils.py | 242 -- .../vfi_models/ops/taichi_ops/__init__.py | 150 - .../vfi_models/ops/taichi_ops/adacof.py | 6 - .../vfi_models/ops/taichi_ops/batch_edt.py | 2 - .../vfi_models/ops/taichi_ops/correlation.py | 15 - .../vfi_models/ops/taichi_ops/costvol.py | 26 - .../ops/taichi_ops/raw_softsplat.py | 126 - .../vfi_models/ops/taichi_ops/sepconv.py | 39 - .../vfi_models/ops/taichi_ops/utils.py | 11 - .../ops/taichi_ops/worker_process.py | 26 - .../vfi_models/rife/__init__.py | 107 - .../vfi_models/rife/rife_arch.py | 581 ---- .../vfi_models/sepconv/__init__.py | 56 - .../vfi_models/sepconv/sepconv_enhanced.py | 748 ----- .../vfi_models/stmfnet/__init__.py | 100 - .../vfi_models/stmfnet/stmfnet_arch.py | 2858 ----------------- .../vfi_models/xvfi/__init__.py | 115 - .../vfi_models/xvfi/xvfi_arch.py | 506 --- .../ComfyUI_Frame_Interpolation/vfi_utils.py | 12 +- 54 files changed, 12 insertions(+), 20004 deletions(-) delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/amt/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/amt/amt_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/cain/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_encdec_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_noca_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/cain/common.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/eisai/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/eisai/eisai_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/flavr_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/resnet_3D.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_union.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_union_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/IFRNet_L_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/IFRNet_S_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ifunet/IFUNet_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ifunet/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/m2m/M2M_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/m2m/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/adacof.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/batch_edt.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/correlation.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/costvol.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/sepconv.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/softsplat.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/utils.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/adacof.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/batch_edt.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/correlation.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/costvol.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/raw_softsplat.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/sepconv.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/utils.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/worker_process.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/rife/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/rife/rife_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/sepconv/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/sepconv/sepconv_enhanced.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/stmfnet/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/stmfnet/stmfnet_arch.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/xvfi/__init__.py delete mode 100644 imports/ComfyUI_Frame_Interpolation/vfi_models/xvfi/xvfi_arch.py diff --git a/SteerableMotion.py b/SteerableMotion.py index 91edfdc..e2d823f 100644 --- a/SteerableMotion.py +++ b/SteerableMotion.py @@ -10,6 +10,7 @@ import matplotlib.pyplot as plt # Local application/library specific imports from .imports.ComfyUI_IPAdapter_plus.IPAdapterPlus import IPAdapterBatchImport, IPAdapterTiledBatchImport, IPAdapterTiledImport, PrepImageForClipVisionImport, IPAdapterAdvancedImport, IPAdapterNoiseImport from .imports.AdvancedControlNet.nodes_sparsectrl import SparseIndexMethodNodeImport +from .imports.ComfyUI_Frame_Interpolation.vfi_models.film import FILM_VFIImport import matplotlib import gc @@ -565,7 +566,7 @@ class BatchCreativeInterpolationNode: # import the class FILM_VFI from ComfyUI-Frame-Interpolation/vfi_models/film/__init__.py -from .imports.ComfyUI_Frame_Interpolation.vfi_models.film import FILM_VFI + # from .imports.AdvancedControlNet.nodes_sparsectrl import SparseIndexMethodNodeImport @@ -592,7 +593,7 @@ class DropFramesByIndex: frames_to_drop = sorted(frames_to_drop, reverse=True) # Create instance of FILM_VFI within the function - film_vfi = FILM_VFI() # Assuming FILM_VFI does not require any special setup + film_vfi = FILM_VFIImport() # Assuming FILM_VFI does not require any special setup for index in frames_to_drop: if 0 < index < images.shape[0] - 1: diff --git a/imports/ComfyUI_Frame_Interpolation/__init__.py b/imports/ComfyUI_Frame_Interpolation/__init__.py index b362e2c..8e7fb3f 100644 --- a/imports/ComfyUI_Frame_Interpolation/__init__.py +++ b/imports/ComfyUI_Frame_Interpolation/__init__.py @@ -2,41 +2,3 @@ import os import sys sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) -from .other_nodes import Gradually_More_Denoise_KSampler - -#Some models are commented out because the code is not completed -#from vfi_models.eisai import EISAI_VFI -from vfi_models.gmfss_fortuna import GMFSS_Fortuna_VFI -from vfi_models.ifrnet import IFRNet_VFI -from vfi_models.ifunet import IFUnet_VFI -from vfi_models.m2m import M2M_VFI -from vfi_models.rife import RIFE_VFI -from vfi_models.sepconv import SepconvVFI -from vfi_models.amt import AMT_VFI -from vfi_models.film import FILM_VFI -from vfi_models.stmfnet import STMFNet_VFI -from vfi_models.flavr import FLAVR_VFI -from vfi_models.cain import CAIN_VFI -from vfi_utils import MakeInterpolationStateList, FloatToInt - -NODE_CLASS_MAPPINGS = { - "KSampler Gradually Adding More Denoise (efficient)": Gradually_More_Denoise_KSampler, -# "EISAI VFI": EISAI_VFI, - "GMFSS Fortuna VFI": GMFSS_Fortuna_VFI, - "IFRNet VFI": IFRNet_VFI, - "IFUnet VFI": IFUnet_VFI, - "M2M VFI": M2M_VFI, - "RIFE VFI": RIFE_VFI, - "Sepconv VFI": SepconvVFI, - "AMT VFI": AMT_VFI, - "FILM VFI": FILM_VFI, - "Make Interpolation State List": MakeInterpolationStateList, - "STMFNet VFI": STMFNet_VFI, - "FLAVR VFI": FLAVR_VFI, - "CAIN VFI": CAIN_VFI, - "VFI FloatToInt": FloatToInt -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "RIFE VFI": "RIFE VFI (recommend rife47 and rife49)" -} \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/amt/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/amt/__init__.py deleted file mode 100644 index 20740c2..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/amt/__init__.py +++ /dev/null @@ -1,87 +0,0 @@ -import pathlib -import torch -from torch.utils.data import DataLoader -import pathlib -from vfi_utils import load_file_from_direct_url, preprocess_frames, postprocess_frames, generic_frame_loop, InterpolationStateList -import typing -from comfy.model_management import get_torch_device -from .amt_arch import AMT_S, AMT_L, AMT_G, InputPadder - -#https://github.com/MCG-NKU/AMT/tree/main/cfgs -CKPT_CONFIGS = { - "amt-s.pth": { - "network": AMT_S, - "params": { "corr_radius": 3, "corr_lvls": 4, "num_flows": 3 } - }, - "amt-l.pth": { - "network": AMT_L, - "params": { "corr_radius": 3, "corr_lvls": 4, "num_flows": 5 } - }, - "amt-g.pth": { - "network": AMT_G, - "params": { "corr_radius": 3, "corr_lvls": 4, "num_flows": 5 } - }, - "gopro_amt-s.pth": { - "network": AMT_S, - "params": { "corr_radius": 3, "corr_lvls": 4, "num_flows": 3 } - } -} - - -MODEL_TYPE = pathlib.Path(__file__).parent.name - -class AMT_VFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (list(CKPT_CONFIGS.keys()), ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 1, "min": 1, "max": 100}), - "multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}) - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames: typing.SupportsInt = 1, - multiplier: typing.SupportsInt = 2, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - model_path = load_file_from_direct_url(MODEL_TYPE, f"https://huggingface.co/lalala125/AMT/resolve/main/{ckpt_name}") - ckpt_config = CKPT_CONFIGS[ckpt_name] - - interpolation_model = ckpt_config["network"](**ckpt_config["params"]) - interpolation_model.load_state_dict(torch.load(model_path)["state_dict"]) - interpolation_model.eval().to(get_torch_device()) - - frames = preprocess_frames(frames) - padder = InputPadder(frames.shape, 16) - frames = padder.pad(frames) - - def return_middle_frame(frame_0, frame_1, timestep, model): - return model( - frame_0, - frame_1, - embt=torch.FloatTensor([timestep] * frame_0.shape[0]).view(frame_0.shape[0], 1, 1, 1).to(get_torch_device()), - scale_factor=1.0, - eval=True - )["imgt_pred"] - - args = [interpolation_model] - out = generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, - interpolation_states=optional_interpolation_states, dtype=torch.float32) - out = padder.unpad(out) - out = postprocess_frames(out) - return (out,) - diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/amt/amt_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/amt/amt_arch.py deleted file mode 100644 index 448c4d5..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/amt/amt_arch.py +++ /dev/null @@ -1,1590 +0,0 @@ -""" -https://github.com/MCG-NKU/AMT/blob/main/utils/dist_utils.py -https://github.com/MCG-NKU/AMT/blob/main/utils/flow_utils.py -https://github.com/MCG-NKU/AMT/blob/main/utils/utils.py -https://github.com/MCG-NKU/AMT/blob/main/networks/blocks/feat_enc.py -https://github.com/MCG-NKU/AMT/blob/main/networks/blocks/ifrnet.py -https://github.com/MCG-NKU/AMT/blob/main/networks/blocks/multi_flow.py -https://github.com/MCG-NKU/AMT/blob/main/networks/blocks/raft.py -https://github.com/MCG-NKU/AMT/blob/main/networks/AMT-S.py -https://github.com/MCG-NKU/AMT/blob/main/networks/AMT-L.py -https://github.com/MCG-NKU/AMT/blob/main/networks/AMT-G.py -""" -#Removed imageio by removing readImage, writeImage -#The model will receive image tensors from other ComfyUI's nodes so they are unneccessary - -import torch -import torch.nn as nn -import numpy as np -from PIL import ImageFile -import torch.nn.functional as F -ImageFile.LOAD_TRUNCATED_IMAGES = True -import re -import sys -import random - -def warp(img, flow): - B, _, H, W = flow.shape - xx = torch.linspace(-1.0, 1.0, W).view(1, 1, 1, W).expand(B, -1, H, -1) - yy = torch.linspace(-1.0, 1.0, H).view(1, 1, H, 1).expand(B, -1, -1, W) - grid = torch.cat([xx, yy], 1).to(img) - flow_ = torch.cat([flow[:, 0:1, :, :] / ((W - 1.0) / 2.0), flow[:, 1:2, :, :] / ((H - 1.0) / 2.0)], 1) - grid_ = (grid + flow_).permute(0, 2, 3, 1) - output = F.grid_sample(input=img, grid=grid_, mode='bilinear', padding_mode='border', align_corners=True) - return output - - -def make_colorwheel(): - """ - Generates a color wheel for optical flow visualization as presented in: - Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007) - URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf - Code follows the original C++ source code of Daniel Scharstein. - Code follows the the Matlab source code of Deqing Sun. - Returns: - np.ndarray: Color wheel - """ - - RY = 15 - YG = 6 - GC = 4 - CB = 11 - BM = 13 - MR = 6 - - ncols = RY + YG + GC + CB + BM + MR - colorwheel = np.zeros((ncols, 3)) - col = 0 - - # RY - colorwheel[0:RY, 0] = 255 - colorwheel[0:RY, 1] = np.floor(255*np.arange(0,RY)/RY) - col = col+RY - # YG - colorwheel[col:col+YG, 0] = 255 - np.floor(255*np.arange(0,YG)/YG) - colorwheel[col:col+YG, 1] = 255 - col = col+YG - # GC - colorwheel[col:col+GC, 1] = 255 - colorwheel[col:col+GC, 2] = np.floor(255*np.arange(0,GC)/GC) - col = col+GC - # CB - colorwheel[col:col+CB, 1] = 255 - np.floor(255*np.arange(CB)/CB) - colorwheel[col:col+CB, 2] = 255 - col = col+CB - # BM - colorwheel[col:col+BM, 2] = 255 - colorwheel[col:col+BM, 0] = np.floor(255*np.arange(0,BM)/BM) - col = col+BM - # MR - colorwheel[col:col+MR, 2] = 255 - np.floor(255*np.arange(MR)/MR) - colorwheel[col:col+MR, 0] = 255 - return colorwheel - -def flow_uv_to_colors(u, v, convert_to_bgr=False): - """ - Applies the flow color wheel to (possibly clipped) flow components u and v. - According to the C++ source code of Daniel Scharstein - According to the Matlab source code of Deqing Sun - Args: - u (np.ndarray): Input horizontal flow of shape [H,W] - v (np.ndarray): Input vertical flow of shape [H,W] - convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False. - Returns: - np.ndarray: Flow visualization image of shape [H,W,3] - """ - flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8) - colorwheel = make_colorwheel() # shape [55x3] - ncols = colorwheel.shape[0] - rad = np.sqrt(np.square(u) + np.square(v)) - a = np.arctan2(-v, -u)/np.pi - fk = (a+1) / 2*(ncols-1) - k0 = np.floor(fk).astype(np.int32) - k1 = k0 + 1 - k1[k1 == ncols] = 0 - f = fk - k0 - for i in range(colorwheel.shape[1]): - tmp = colorwheel[:,i] - col0 = tmp[k0] / 255.0 - col1 = tmp[k1] / 255.0 - col = (1-f)*col0 + f*col1 - idx = (rad <= 1) - col[idx] = 1 - rad[idx] * (1-col[idx]) - col[~idx] = col[~idx] * 0.75 # out of range - # Note the 2-i => BGR instead of RGB - ch_idx = 2-i if convert_to_bgr else i - flow_image[:,:,ch_idx] = np.floor(255 * col) - return flow_image - -def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False): - """ - Expects a two dimensional flow image of shape. - Args: - flow_uv (np.ndarray): Flow UV image of shape [H,W,2] - clip_flow (float, optional): Clip maximum of flow values. Defaults to None. - convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False. - Returns: - np.ndarray: Flow visualization image of shape [H,W,3] - """ - assert flow_uv.ndim == 3, 'input flow must have three dimensions' - assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]' - if clip_flow is not None: - flow_uv = np.clip(flow_uv, 0, clip_flow) - u = flow_uv[:,:,0] - v = flow_uv[:,:,1] - rad = np.sqrt(np.square(u) + np.square(v)) - rad_max = np.max(rad) - epsilon = 1e-5 - u = u / (rad_max + epsilon) - v = v / (rad_max + epsilon) - return flow_uv_to_colors(u, v, convert_to_bgr) - - - - - - - - - - - -class AverageMeter(): - def __init__(self): - self.reset() - - def reset(self): - self.val = 0. - self.avg = 0. - self.sum = 0. - self.count = 0 - - def update(self, val, n=1): - self.val = val - self.sum += val * n - self.count += n - self.avg = self.sum / self.count - - -class AverageMeterGroups: - def __init__(self) -> None: - self.meter_dict = dict() - - def update(self, dict, n=1): - for name, val in dict.items(): - if self.meter_dict.get(name) is None: - self.meter_dict[name] = AverageMeter() - self.meter_dict[name].update(val, n) - - def reset(self, name=None): - if name is None: - for v in self.meter_dict.values(): - v.reset() - else: - meter = self.meter_dict.get(name) - if meter is not None: - meter.reset() - - def avg(self, name): - meter = self.meter_dict.get(name) - if meter is not None: - return meter.avg - - -class InputPadder: - """ Pads images such that dimensions are divisible by divisor """ - def __init__(self, dims, divisor=16): - self.ht, self.wd = dims[-2:] - pad_ht = (((self.ht // divisor) + 1) * divisor - self.ht) % divisor - pad_wd = (((self.wd // divisor) + 1) * divisor - self.wd) % divisor - self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2] - - def pad(self, input_tensor): - return F.pad(input_tensor, self._pad, mode='replicate') - - def unpad(self, input_tensor): - return self._unpad(input_tensor) - - def _unpad(self, x): - ht, wd = x.shape[-2:] - c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]] - return x[..., c[0]:c[1], c[2]:c[3]] - - -def img2tensor(img): - if img.shape[-1] > 3: - img = img[:,:,:3] - return torch.tensor(img).permute(2, 0, 1).unsqueeze(0) / 255.0 - - -def tensor2img(img_t): - return (img_t * 255.).detach( - ).squeeze(0).permute(1, 2, 0).cpu().numpy( - ).clip(0, 255).astype(np.uint8) - -def seed_all(seed): - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -def readPFM(file): - file = open(file, 'rb') - - color = None - width = None - height = None - scale = None - endian = None - - header = file.readline().rstrip() - if header.decode("ascii") == 'PF': - color = True - elif header.decode("ascii") == 'Pf': - color = False - else: - raise Exception('Not a PFM file.') - - dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode("ascii")) - if dim_match: - width, height = list(map(int, dim_match.groups())) - else: - raise Exception('Malformed PFM header.') - - scale = float(file.readline().decode("ascii").rstrip()) - if scale < 0: - endian = '<' - scale = -scale - else: - endian = '>' - - data = np.fromfile(file, endian + 'f') - shape = (height, width, 3) if color else (height, width) - - data = np.reshape(data, shape) - data = np.flipud(data) - return data, scale - - -def writePFM(file, image, scale=1): - file = open(file, 'wb') - - color = None - - if image.dtype.name != 'float32': - raise Exception('Image dtype must be float32.') - - image = np.flipud(image) - - if len(image.shape) == 3 and image.shape[2] == 3: - color = True - elif len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1: - color = False - else: - raise Exception('Image must have H x W x 3, H x W x 1 or H x W dimensions.') - - file.write('PF\n' if color else 'Pf\n'.encode()) - file.write('%d %d\n'.encode() % (image.shape[1], image.shape[0])) - - endian = image.dtype.byteorder - - if endian == '<' or endian == '=' and sys.byteorder == 'little': - scale = -scale - - file.write('%f\n'.encode() % scale) - - image.tofile(file) - - -def readFlow(name): - if name.endswith('.pfm') or name.endswith('.PFM'): - return readPFM(name)[0][:,:,0:2] - - f = open(name, 'rb') - - header = f.read(4) - if header.decode("utf-8") != 'PIEH': - raise Exception('Flow file header does not contain PIEH') - - width = np.fromfile(f, np.int32, 1).squeeze() - height = np.fromfile(f, np.int32, 1).squeeze() - - flow = np.fromfile(f, np.float32, width * height * 2).reshape((height, width, 2)) - - return flow.astype(np.float32) - -def writeFlow(name, flow): - f = open(name, 'wb') - f.write('PIEH'.encode('utf-8')) - np.array([flow.shape[1], flow.shape[0]], dtype=np.int32).tofile(f) - flow = flow.astype(np.float32) - flow.tofile(f) - - -def readFloat(name): - f = open(name, 'rb') - - if(f.readline().decode("utf-8")) != 'float\n': - raise Exception('float file %s did not contain keyword' % name) - - dim = int(f.readline()) - - dims = [] - count = 1 - for i in range(0, dim): - d = int(f.readline()) - dims.append(d) - count *= d - - dims = list(reversed(dims)) - - data = np.fromfile(f, np.float32, count).reshape(dims) - if dim > 2: - data = np.transpose(data, (2, 1, 0)) - data = np.transpose(data, (1, 0, 2)) - - return data - - -def writeFloat(name, data): - f = open(name, 'wb') - - dim=len(data.shape) - if dim>3: - raise Exception('bad float file dimension: %d' % dim) - - f.write(('float\n').encode('ascii')) - f.write(('%d\n' % dim).encode('ascii')) - - if dim == 1: - f.write(('%d\n' % data.shape[0]).encode('ascii')) - else: - f.write(('%d\n' % data.shape[1]).encode('ascii')) - f.write(('%d\n' % data.shape[0]).encode('ascii')) - for i in range(2, dim): - f.write(('%d\n' % data.shape[i]).encode('ascii')) - - data = data.astype(np.float32) - if dim==2: - data.tofile(f) - - else: - np.transpose(data, (2, 0, 1)).tofile(f) - - -def check_dim_and_resize(tensor_list): - shape_list = [] - for t in tensor_list: - shape_list.append(t.shape[2:]) - - if len(set(shape_list)) > 1: - desired_shape = shape_list[0] - print(f'Inconsistent size of input video frames. All frames will be resized to {desired_shape}') - - resize_tensor_list = [] - for t in tensor_list: - resize_tensor_list.append(torch.nn.functional.interpolate(t, size=tuple(desired_shape), mode='bilinear')) - - tensor_list = resize_tensor_list - - return tensor_list - - - - - - - - - - - -class BottleneckBlock(nn.Module): - def __init__(self, in_planes, planes, norm_fn='group', stride=1): - super(BottleneckBlock, self).__init__() - - self.conv1 = nn.Conv2d(in_planes, planes//4, kernel_size=1, padding=0) - self.conv2 = nn.Conv2d(planes//4, planes//4, kernel_size=3, padding=1, stride=stride) - self.conv3 = nn.Conv2d(planes//4, planes, kernel_size=1, padding=0) - self.relu = nn.ReLU(inplace=True) - - num_groups = planes // 8 - - if norm_fn == 'group': - self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4) - self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4) - self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) - if not stride == 1: - self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) - - elif norm_fn == 'batch': - self.norm1 = nn.BatchNorm2d(planes//4) - self.norm2 = nn.BatchNorm2d(planes//4) - self.norm3 = nn.BatchNorm2d(planes) - if not stride == 1: - self.norm4 = nn.BatchNorm2d(planes) - - elif norm_fn == 'instance': - self.norm1 = nn.InstanceNorm2d(planes//4) - self.norm2 = nn.InstanceNorm2d(planes//4) - self.norm3 = nn.InstanceNorm2d(planes) - if not stride == 1: - self.norm4 = nn.InstanceNorm2d(planes) - - elif norm_fn == 'none': - self.norm1 = nn.Sequential() - self.norm2 = nn.Sequential() - self.norm3 = nn.Sequential() - if not stride == 1: - self.norm4 = nn.Sequential() - - if stride == 1: - self.downsample = None - - else: - self.downsample = nn.Sequential( - nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4) - - - def forward(self, x): - y = x - y = self.relu(self.norm1(self.conv1(y))) - y = self.relu(self.norm2(self.conv2(y))) - y = self.relu(self.norm3(self.conv3(y))) - - if self.downsample is not None: - x = self.downsample(x) - - return self.relu(x+y) - - -class ResidualBlock(nn.Module): - def __init__(self, in_planes, planes, norm_fn='group', stride=1): - super(ResidualBlock, self).__init__() - - self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride) - self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1) - self.relu = nn.ReLU(inplace=True) - - num_groups = planes // 8 - - if norm_fn == 'group': - self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) - self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) - if not stride == 1: - self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) - - elif norm_fn == 'batch': - self.norm1 = nn.BatchNorm2d(planes) - self.norm2 = nn.BatchNorm2d(planes) - if not stride == 1: - self.norm3 = nn.BatchNorm2d(planes) - - elif norm_fn == 'instance': - self.norm1 = nn.InstanceNorm2d(planes) - self.norm2 = nn.InstanceNorm2d(planes) - if not stride == 1: - self.norm3 = nn.InstanceNorm2d(planes) - - elif norm_fn == 'none': - self.norm1 = nn.Sequential() - self.norm2 = nn.Sequential() - if not stride == 1: - self.norm3 = nn.Sequential() - - if stride == 1: - self.downsample = None - - else: - self.downsample = nn.Sequential( - nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3) - - - def forward(self, x): - y = x - y = self.relu(self.norm1(self.conv1(y))) - y = self.relu(self.norm2(self.conv2(y))) - - if self.downsample is not None: - x = self.downsample(x) - - return self.relu(x+y) - - -class SmallEncoder(nn.Module): - def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0): - super(SmallEncoder, self).__init__() - self.norm_fn = norm_fn - - if self.norm_fn == 'group': - self.norm1 = nn.GroupNorm(num_groups=8, num_channels=32) - - elif self.norm_fn == 'batch': - self.norm1 = nn.BatchNorm2d(32) - - elif self.norm_fn == 'instance': - self.norm1 = nn.InstanceNorm2d(32) - - elif self.norm_fn == 'none': - self.norm1 = nn.Sequential() - - self.conv1 = nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3) - self.relu1 = nn.ReLU(inplace=True) - - self.in_planes = 32 - self.layer1 = self._make_layer(32, stride=1) - self.layer2 = self._make_layer(64, stride=2) - self.layer3 = self._make_layer(96, stride=2) - - self.dropout = None - if dropout > 0: - self.dropout = nn.Dropout2d(p=dropout) - - self.conv2 = nn.Conv2d(96, output_dim, kernel_size=1) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): - if m.weight is not None: - nn.init.constant_(m.weight, 1) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - def _make_layer(self, dim, stride=1): - layer1 = BottleneckBlock(self.in_planes, dim, self.norm_fn, stride=stride) - layer2 = BottleneckBlock(dim, dim, self.norm_fn, stride=1) - layers = (layer1, layer2) - - self.in_planes = dim - return nn.Sequential(*layers) - - - def forward(self, x): - - # if input is list, combine batch dimension - is_list = isinstance(x, tuple) or isinstance(x, list) - if is_list: - batch_dim = x[0].shape[0] - x = torch.cat(x, dim=0) - - x = self.conv1(x) - x = self.norm1(x) - x = self.relu1(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - x = self.conv2(x) - - if self.training and self.dropout is not None: - x = self.dropout(x) - - if is_list: - x = torch.split(x, [batch_dim, batch_dim], dim=0) - - return x - -class BasicEncoder(nn.Module): - def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0): - super(BasicEncoder, self).__init__() - self.norm_fn = norm_fn - - if self.norm_fn == 'group': - self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) - - elif self.norm_fn == 'batch': - self.norm1 = nn.BatchNorm2d(64) - - elif self.norm_fn == 'instance': - self.norm1 = nn.InstanceNorm2d(64) - - elif self.norm_fn == 'none': - self.norm1 = nn.Sequential() - - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) - self.relu1 = nn.ReLU(inplace=True) - - self.in_planes = 64 - self.layer1 = self._make_layer(64, stride=1) - self.layer2 = self._make_layer(72, stride=2) - self.layer3 = self._make_layer(128, stride=2) - - # output convolution - self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1) - - self.dropout = None - if dropout > 0: - self.dropout = nn.Dropout2d(p=dropout) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): - if m.weight is not None: - nn.init.constant_(m.weight, 1) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - def _make_layer(self, dim, stride=1): - layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) - layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) - layers = (layer1, layer2) - - self.in_planes = dim - return nn.Sequential(*layers) - - - def forward(self, x): - - # if input is list, combine batch dimension - is_list = isinstance(x, tuple) or isinstance(x, list) - if is_list: - batch_dim = x[0].shape[0] - x = torch.cat(x, dim=0) - - x = self.conv1(x) - x = self.norm1(x) - x = self.relu1(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - - x = self.conv2(x) - - if self.training and self.dropout is not None: - x = self.dropout(x) - - if is_list: - x = torch.split(x, [batch_dim, batch_dim], dim=0) - - return x - -class LargeEncoder(nn.Module): - def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0): - super(LargeEncoder, self).__init__() - self.norm_fn = norm_fn - - if self.norm_fn == 'group': - self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) - - elif self.norm_fn == 'batch': - self.norm1 = nn.BatchNorm2d(64) - - elif self.norm_fn == 'instance': - self.norm1 = nn.InstanceNorm2d(64) - - elif self.norm_fn == 'none': - self.norm1 = nn.Sequential() - - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) - self.relu1 = nn.ReLU(inplace=True) - - self.in_planes = 64 - self.layer1 = self._make_layer(64, stride=1) - self.layer2 = self._make_layer(112, stride=2) - self.layer3 = self._make_layer(160, stride=2) - self.layer3_2 = self._make_layer(160, stride=1) - - # output convolution - self.conv2 = nn.Conv2d(self.in_planes, output_dim, kernel_size=1) - - self.dropout = None - if dropout > 0: - self.dropout = nn.Dropout2d(p=dropout) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): - if m.weight is not None: - nn.init.constant_(m.weight, 1) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - def _make_layer(self, dim, stride=1): - layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) - layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) - layers = (layer1, layer2) - - self.in_planes = dim - return nn.Sequential(*layers) - - - def forward(self, x): - - # if input is list, combine batch dimension - is_list = isinstance(x, tuple) or isinstance(x, list) - if is_list: - batch_dim = x[0].shape[0] - x = torch.cat(x, dim=0) - - x = self.conv1(x) - x = self.norm1(x) - x = self.relu1(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - x = self.layer3_2(x) - - x = self.conv2(x) - - if self.training and self.dropout is not None: - x = self.dropout(x) - - if is_list: - x = torch.split(x, [batch_dim, batch_dim], dim=0) - - return x - - - - - - - - - - - -def resize(x, scale_factor): - return F.interpolate(x, scale_factor=scale_factor, mode="bilinear", align_corners=False) - -def convrelu(in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, groups=1, bias=True): - return nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=bias), - nn.PReLU(out_channels) - ) - -class ResBlock(nn.Module): - def __init__(self, in_channels, side_channels, bias=True): - super(ResBlock, self).__init__() - self.side_channels = side_channels - self.conv1 = nn.Sequential( - nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias), - nn.PReLU(in_channels) - ) - self.conv2 = nn.Sequential( - nn.Conv2d(side_channels, side_channels, kernel_size=3, stride=1, padding=1, bias=bias), - nn.PReLU(side_channels) - ) - self.conv3 = nn.Sequential( - nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias), - nn.PReLU(in_channels) - ) - self.conv4 = nn.Sequential( - nn.Conv2d(side_channels, side_channels, kernel_size=3, stride=1, padding=1, bias=bias), - nn.PReLU(side_channels) - ) - self.conv5 = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias) - self.prelu = nn.PReLU(in_channels) - - def forward(self, x): - out = self.conv1(x) - - res_feat = out[:, :-self.side_channels, ...] - side_feat = out[:, -self.side_channels:, :, :] - side_feat = self.conv2(side_feat) - out = self.conv3(torch.cat([res_feat, side_feat], 1)) - - res_feat = out[:, :-self.side_channels, ...] - side_feat = out[:, -self.side_channels:, :, :] - side_feat = self.conv4(side_feat) - out = self.conv5(torch.cat([res_feat, side_feat], 1)) - - out = self.prelu(x + out) - return out - -class Encoder(nn.Module): - def __init__(self, channels, large=False): - super(Encoder, self).__init__() - self.channels = channels - prev_ch = 3 - for idx, ch in enumerate(channels, 1): - k = 7 if large and idx == 1 else 3 - p = 3 if k ==7 else 1 - self.register_module(f'pyramid{idx}', - nn.Sequential( - convrelu(prev_ch, ch, k, 2, p), - convrelu(ch, ch, 3, 1, 1) - )) - prev_ch = ch - - def forward(self, in_x): - fs = [] - for idx in range(len(self.channels)): - out_x = getattr(self, f'pyramid{idx+1}')(in_x) - fs.append(out_x) - in_x = out_x - return fs - -class InitDecoder(nn.Module): - def __init__(self, in_ch, out_ch, skip_ch) -> None: - super().__init__() - self.convblock = nn.Sequential( - convrelu(in_ch*2+1, in_ch*2), - ResBlock(in_ch*2, skip_ch), - nn.ConvTranspose2d(in_ch*2, out_ch+4, 4, 2, 1, bias=True) - ) - def forward(self, f0, f1, embt): - h, w = f0.shape[2:] - embt = embt.repeat(1, 1, h, w) - out = self.convblock(torch.cat([f0, f1, embt], 1)) - flow0, flow1 = torch.chunk(out[:, :4, ...], 2, 1) - ft_ = out[:, 4:, ...] - return flow0, flow1, ft_ - -class IntermediateDecoder(nn.Module): - def __init__(self, in_ch, out_ch, skip_ch) -> None: - super().__init__() - self.convblock = nn.Sequential( - convrelu(in_ch*3+4, in_ch*3), - ResBlock(in_ch*3, skip_ch), - nn.ConvTranspose2d(in_ch*3, out_ch+4, 4, 2, 1, bias=True) - ) - def forward(self, ft_, f0, f1, flow0_in, flow1_in): - f0_warp = warp(f0, flow0_in) - f1_warp = warp(f1, flow1_in) - f_in = torch.cat([ft_, f0_warp, f1_warp, flow0_in, flow1_in], 1) - out = self.convblock(f_in) - flow0, flow1 = torch.chunk(out[:, :4, ...], 2, 1) - ft_ = out[:, 4:, ...] - flow0 = flow0 + 2.0 * resize(flow0_in, scale_factor=2.0) - flow1 = flow1 + 2.0 * resize(flow1_in, scale_factor=2.0) - return flow0, flow1, ft_ - - - - - - - - - - - -def multi_flow_combine(comb_block, img0, img1, flow0, flow1, - mask=None, img_res=None, mean=None): - ''' - A parallel implementation of multiple flow field warping - comb_block: An nn.Seqential object. - img shape: [b, c, h, w] - flow shape: [b, 2*num_flows, h, w] - mask (opt): - If 'mask' is None, the function conduct a simple average. - img_res (opt): - If 'img_res' is None, the function adds zero instead. - mean (opt): - If 'mean' is None, the function adds zero instead. - ''' - b, c, h, w = flow0.shape - num_flows = c // 2 - flow0 = flow0.reshape(b, num_flows, 2, h, w).reshape(-1, 2, h, w) - flow1 = flow1.reshape(b, num_flows, 2, h, w).reshape(-1, 2, h, w) - - mask = mask.reshape(b, num_flows, 1, h, w - ).reshape(-1, 1, h, w) if mask is not None else None - img_res = img_res.reshape(b, num_flows, 3, h, w - ).reshape(-1, 3, h, w) if img_res is not None else 0 - img0 = torch.stack([img0] * num_flows, 1).reshape(-1, 3, h, w) - img1 = torch.stack([img1] * num_flows, 1).reshape(-1, 3, h, w) - mean = torch.stack([mean] * num_flows, 1).reshape(-1, 1, 1, 1 - ) if mean is not None else 0 - - img0_warp = warp(img0, flow0) - img1_warp = warp(img1, flow1) - img_warps = mask * img0_warp + (1 - mask) * img1_warp + mean + img_res - img_warps = img_warps.reshape(b, num_flows, 3, h, w) - imgt_pred = img_warps.mean(1) + comb_block(img_warps.view(b, -1, h, w)) - return imgt_pred - - -class MultiFlowDecoder(nn.Module): - def __init__(self, in_ch, skip_ch, num_flows=3): - super(MultiFlowDecoder, self).__init__() - self.num_flows = num_flows - self.convblock = nn.Sequential( - convrelu(in_ch*3+4, in_ch*3), - ResBlock(in_ch*3, skip_ch), - nn.ConvTranspose2d(in_ch*3, 8*num_flows, 4, 2, 1, bias=True) - ) - - def forward(self, ft_, f0, f1, flow0, flow1): - n = self.num_flows - f0_warp = warp(f0, flow0) - f1_warp = warp(f1, flow1) - out = self.convblock(torch.cat([ft_, f0_warp, f1_warp, flow0, flow1], 1)) - delta_flow0, delta_flow1, mask, img_res = torch.split(out, [2*n, 2*n, n, 3*n], 1) - mask = torch.sigmoid(mask) - - flow0 = delta_flow0 + 2.0 * resize(flow0, scale_factor=2.0 - ).repeat(1, self.num_flows, 1, 1) - flow1 = delta_flow1 + 2.0 * resize(flow1, scale_factor=2.0 - ).repeat(1, self.num_flows, 1, 1) - - return flow0, flow1, mask, img_res - - - - - - - - - - - -def resize(x, scale_factor): - return F.interpolate(x, scale_factor=scale_factor, mode="bilinear", align_corners=False) - - -def bilinear_sampler(img, coords, mask=False): - """ Wrapper for grid_sample, uses pixel coordinates """ - H, W = img.shape[-2:] - xgrid, ygrid = coords.split([1,1], dim=-1) - xgrid = 2*xgrid/(W-1) - 1 - ygrid = 2*ygrid/(H-1) - 1 - - grid = torch.cat([xgrid, ygrid], dim=-1) - img = F.grid_sample(img, grid, align_corners=True) - - if mask: - mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) - return img, mask.float() - - return img - - -def coords_grid(batch, ht, wd, device): - coords = torch.meshgrid(torch.arange(ht, device=device), - torch.arange(wd, device=device), - indexing='ij') - coords = torch.stack(coords[::-1], dim=0).float() - return coords[None].repeat(batch, 1, 1, 1) - - -class SmallUpdateBlock(nn.Module): - def __init__(self, cdim, hidden_dim, flow_dim, corr_dim, fc_dim, - corr_levels=4, radius=3, scale_factor=None): - super(SmallUpdateBlock, self).__init__() - cor_planes = corr_levels * (2 * radius + 1) **2 - self.scale_factor = scale_factor - - self.convc1 = nn.Conv2d(2 * cor_planes, corr_dim, 1, padding=0) - self.convf1 = nn.Conv2d(4, flow_dim*2, 7, padding=3) - self.convf2 = nn.Conv2d(flow_dim*2, flow_dim, 3, padding=1) - self.conv = nn.Conv2d(corr_dim+flow_dim, fc_dim, 3, padding=1) - - self.gru = nn.Sequential( - nn.Conv2d(fc_dim+4+cdim, hidden_dim, 3, padding=1), - nn.LeakyReLU(negative_slope=0.1, inplace=True), - nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), - ) - - self.feat_head = nn.Sequential( - nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), - nn.LeakyReLU(negative_slope=0.1, inplace=True), - nn.Conv2d(hidden_dim, cdim, 3, padding=1), - ) - - self.flow_head = nn.Sequential( - nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), - nn.LeakyReLU(negative_slope=0.1, inplace=True), - nn.Conv2d(hidden_dim, 4, 3, padding=1), - ) - - self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True) - - def forward(self, net, flow, corr): - net = resize(net, 1 / self.scale_factor - ) if self.scale_factor is not None else net - cor = self.lrelu(self.convc1(corr)) - flo = self.lrelu(self.convf1(flow)) - flo = self.lrelu(self.convf2(flo)) - cor_flo = torch.cat([cor, flo], dim=1) - inp = self.lrelu(self.conv(cor_flo)) - inp = torch.cat([inp, flow, net], dim=1) - - out = self.gru(inp) - delta_net = self.feat_head(out) - delta_flow = self.flow_head(out) - - if self.scale_factor is not None: - delta_net = resize(delta_net, scale_factor=self.scale_factor) - delta_flow = self.scale_factor * resize(delta_flow, scale_factor=self.scale_factor) - - return delta_net, delta_flow - - -class BasicUpdateBlock(nn.Module): - def __init__(self, cdim, hidden_dim, flow_dim, corr_dim, corr_dim2, - fc_dim, corr_levels=4, radius=3, scale_factor=None, out_num=1): - super(BasicUpdateBlock, self).__init__() - cor_planes = corr_levels * (2 * radius + 1) **2 - - self.scale_factor = scale_factor - self.convc1 = nn.Conv2d(2 * cor_planes, corr_dim, 1, padding=0) - self.convc2 = nn.Conv2d(corr_dim, corr_dim2, 3, padding=1) - self.convf1 = nn.Conv2d(4, flow_dim*2, 7, padding=3) - self.convf2 = nn.Conv2d(flow_dim*2, flow_dim, 3, padding=1) - self.conv = nn.Conv2d(flow_dim+corr_dim2, fc_dim, 3, padding=1) - - self.gru = nn.Sequential( - nn.Conv2d(fc_dim+4+cdim, hidden_dim, 3, padding=1), - nn.LeakyReLU(negative_slope=0.1, inplace=True), - nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), - ) - - self.feat_head = nn.Sequential( - nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), - nn.LeakyReLU(negative_slope=0.1, inplace=True), - nn.Conv2d(hidden_dim, cdim, 3, padding=1), - ) - - self.flow_head = nn.Sequential( - nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1), - nn.LeakyReLU(negative_slope=0.1, inplace=True), - nn.Conv2d(hidden_dim, 4*out_num, 3, padding=1), - ) - - self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True) - - def forward(self, net, flow, corr): - net = resize(net, 1 / self.scale_factor - ) if self.scale_factor is not None else net - cor = self.lrelu(self.convc1(corr)) - cor = self.lrelu(self.convc2(cor)) - flo = self.lrelu(self.convf1(flow)) - flo = self.lrelu(self.convf2(flo)) - cor_flo = torch.cat([cor, flo], dim=1) - inp = self.lrelu(self.conv(cor_flo)) - inp = torch.cat([inp, flow, net], dim=1) - - out = self.gru(inp) - delta_net = self.feat_head(out) - delta_flow = self.flow_head(out) - - if self.scale_factor is not None: - delta_net = resize(delta_net, scale_factor=self.scale_factor) - delta_flow = self.scale_factor * resize(delta_flow, scale_factor=self.scale_factor) - return delta_net, delta_flow - - -class BidirCorrBlock: - def __init__(self, fmap1, fmap2, num_levels=4, radius=4): - self.num_levels = num_levels - self.radius = radius - self.corr_pyramid = [] - self.corr_pyramid_T = [] - - corr = BidirCorrBlock.corr(fmap1, fmap2) - batch, h1, w1, dim, h2, w2 = corr.shape - corr_T = corr.clone().permute(0, 4, 5, 3, 1, 2) - - corr = corr.reshape(batch*h1*w1, dim, h2, w2) - corr_T = corr_T.reshape(batch*h2*w2, dim, h1, w1) - - self.corr_pyramid.append(corr) - self.corr_pyramid_T.append(corr_T) - - for _ in range(self.num_levels-1): - corr = F.avg_pool2d(corr, 2, stride=2) - corr_T = F.avg_pool2d(corr_T, 2, stride=2) - self.corr_pyramid.append(corr) - self.corr_pyramid_T.append(corr_T) - - def __call__(self, coords0, coords1): - r = self.radius - coords0 = coords0.permute(0, 2, 3, 1) - coords1 = coords1.permute(0, 2, 3, 1) - assert coords0.shape == coords1.shape, f"coords0 shape: [{coords0.shape}] is not equal to [{coords1.shape}]" - batch, h1, w1, _ = coords0.shape - - out_pyramid = [] - out_pyramid_T = [] - for i in range(self.num_levels): - corr = self.corr_pyramid[i] - corr_T = self.corr_pyramid_T[i] - - dx = torch.linspace(-r, r, 2*r+1, device=coords0.device) - dy = torch.linspace(-r, r, 2*r+1, device=coords0.device) - delta = torch.stack(torch.meshgrid(dy, dx, indexing='ij'), axis=-1) - delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2) - - centroid_lvl_0 = coords0.reshape(batch*h1*w1, 1, 1, 2) / 2**i - centroid_lvl_1 = coords1.reshape(batch*h1*w1, 1, 1, 2) / 2**i - coords_lvl_0 = centroid_lvl_0 + delta_lvl - coords_lvl_1 = centroid_lvl_1 + delta_lvl - - corr = bilinear_sampler(corr, coords_lvl_0) - corr_T = bilinear_sampler(corr_T, coords_lvl_1) - corr = corr.view(batch, h1, w1, -1) - corr_T = corr_T.view(batch, h1, w1, -1) - out_pyramid.append(corr) - out_pyramid_T.append(corr_T) - - out = torch.cat(out_pyramid, dim=-1) - out_T = torch.cat(out_pyramid_T, dim=-1) - return out.permute(0, 3, 1, 2).contiguous().float(), out_T.permute(0, 3, 1, 2).contiguous().float() - - @staticmethod - def corr(fmap1, fmap2): - batch, dim, ht, wd = fmap1.shape - fmap1 = fmap1.view(batch, dim, ht*wd) - fmap2 = fmap2.view(batch, dim, ht*wd) - - corr = torch.matmul(fmap1.transpose(1,2), fmap2) - corr = corr.view(batch, ht, wd, 1, ht, wd) - return corr / torch.sqrt(torch.tensor(dim).float()) - - - - - - - - - - - -class AMT_S(nn.Module): - def __init__(self, - corr_radius=3, - corr_lvls=4, - num_flows=3, - channels=[20, 32, 44, 56], - skip_channels=20): - super(AMT_S, self).__init__() - self.radius = corr_radius - self.corr_levels = corr_lvls - self.num_flows = num_flows - self.channels = channels - self.skip_channels = skip_channels - - self.feat_encoder = SmallEncoder(output_dim=84, norm_fn='instance', dropout=0.) - self.encoder = Encoder(channels) - - self.decoder4 = InitDecoder(channels[3], channels[2], skip_channels) - self.decoder3 = IntermediateDecoder(channels[2], channels[1], skip_channels) - self.decoder2 = IntermediateDecoder(channels[1], channels[0], skip_channels) - self.decoder1 = MultiFlowDecoder(channels[0], skip_channels, num_flows) - - self.update4 = self._get_updateblock(44) - self.update3 = self._get_updateblock(32, 2) - self.update2 = self._get_updateblock(20, 4) - - self.comb_block = nn.Sequential( - nn.Conv2d(3*num_flows, 6*num_flows, 3, 1, 1), - nn.PReLU(6*num_flows), - nn.Conv2d(6*num_flows, 3, 3, 1, 1), - ) - - def _get_updateblock(self, cdim, scale_factor=None): - return SmallUpdateBlock(cdim=cdim, hidden_dim=76, flow_dim=20, corr_dim=64, - fc_dim=68, scale_factor=scale_factor, - corr_levels=self.corr_levels, radius=self.radius) - - def _corr_scale_lookup(self, corr_fn, coord, flow0, flow1, embt, downsample=1): - # convert t -> 0 to 0 -> 1 | convert t -> 1 to 1 -> 0 - # based on linear assumption - t1_scale = 1. / embt - t0_scale = 1. / (1. - embt) - if downsample != 1: - inv = 1 / downsample - flow0 = inv * resize(flow0, scale_factor=inv) - flow1 = inv * resize(flow1, scale_factor=inv) - - corr0, corr1 = corr_fn(coord + flow1 * t1_scale, coord + flow0 * t0_scale) - corr = torch.cat([corr0, corr1], dim=1) - flow = torch.cat([flow0, flow1], dim=1) - return corr, flow - - def forward(self, img0, img1, embt, scale_factor=1.0, eval=False, **kwargs): - mean_ = torch.cat([img0, img1], 2).mean(1, keepdim=True).mean(2, keepdim=True).mean(3, keepdim=True) - img0 = img0 - mean_ - img1 = img1 - mean_ - img0_ = resize(img0, scale_factor) if scale_factor != 1.0 else img0 - img1_ = resize(img1, scale_factor) if scale_factor != 1.0 else img1 - b, _, h, w = img0_.shape - coord = coords_grid(b, h // 8, w // 8, img0.device) - - fmap0, fmap1 = self.feat_encoder([img0_, img1_]) # [1, 128, H//8, W//8] - corr_fn = BidirCorrBlock(fmap0, fmap1, radius=self.radius, num_levels=self.corr_levels) - - # f0_1: [1, c0, H//2, W//2] | f0_2: [1, c1, H//4, W//4] - # f0_3: [1, c2, H//8, W//8] | f0_4: [1, c3, H//16, W//16] - f0_1, f0_2, f0_3, f0_4 = self.encoder(img0_) - f1_1, f1_2, f1_3, f1_4 = self.encoder(img1_) - - ######################################### the 4th decoder ######################################### - up_flow0_4, up_flow1_4, ft_3_ = self.decoder4(f0_4, f1_4, embt) - corr_4, flow_4 = self._corr_scale_lookup(corr_fn, coord, - up_flow0_4, up_flow1_4, - embt, downsample=1) - - # residue update with lookup corr - delta_ft_3_, delta_flow_4 = self.update4(ft_3_, flow_4, corr_4) - delta_flow0_4, delta_flow1_4 = torch.chunk(delta_flow_4, 2, 1) - up_flow0_4 = up_flow0_4 + delta_flow0_4 - up_flow1_4 = up_flow1_4 + delta_flow1_4 - ft_3_ = ft_3_ + delta_ft_3_ - - ######################################### the 3rd decoder ######################################### - up_flow0_3, up_flow1_3, ft_2_ = self.decoder3(ft_3_, f0_3, f1_3, up_flow0_4, up_flow1_4) - corr_3, flow_3 = self._corr_scale_lookup(corr_fn, - coord, up_flow0_3, up_flow1_3, - embt, downsample=2) - - # residue update with lookup corr - delta_ft_2_, delta_flow_3 = self.update3(ft_2_, flow_3, corr_3) - delta_flow0_3, delta_flow1_3 = torch.chunk(delta_flow_3, 2, 1) - up_flow0_3 = up_flow0_3 + delta_flow0_3 - up_flow1_3 = up_flow1_3 + delta_flow1_3 - ft_2_ = ft_2_ + delta_ft_2_ - - ######################################### the 2nd decoder ######################################### - up_flow0_2, up_flow1_2, ft_1_ = self.decoder2(ft_2_, f0_2, f1_2, up_flow0_3, up_flow1_3) - corr_2, flow_2 = self._corr_scale_lookup(corr_fn, - coord, up_flow0_2, up_flow1_2, - embt, downsample=4) - - # residue update with lookup corr - delta_ft_1_, delta_flow_2 = self.update2(ft_1_, flow_2, corr_2) - delta_flow0_2, delta_flow1_2 = torch.chunk(delta_flow_2, 2, 1) - up_flow0_2 = up_flow0_2 + delta_flow0_2 - up_flow1_2 = up_flow1_2 + delta_flow1_2 - ft_1_ = ft_1_ + delta_ft_1_ - - ######################################### the 1st decoder ######################################### - up_flow0_1, up_flow1_1, mask, img_res = self.decoder1(ft_1_, f0_1, f1_1, up_flow0_2, up_flow1_2) - - if scale_factor != 1.0: - up_flow0_1 = resize(up_flow0_1, scale_factor=(1.0/scale_factor)) * (1.0/scale_factor) - up_flow1_1 = resize(up_flow1_1, scale_factor=(1.0/scale_factor)) * (1.0/scale_factor) - mask = resize(mask, scale_factor=(1.0/scale_factor)) - img_res = resize(img_res, scale_factor=(1.0/scale_factor)) - - # Merge multiple predictions - imgt_pred = multi_flow_combine(self.comb_block, img0, img1, up_flow0_1, up_flow1_1, - mask, img_res, mean_) - imgt_pred = torch.clamp(imgt_pred, 0, 1) - - if eval: - return { 'imgt_pred': imgt_pred, } - else: - up_flow0_1 = up_flow0_1.reshape(b, self.num_flows, 2, h, w) - up_flow1_1 = up_flow1_1.reshape(b, self.num_flows, 2, h, w) - return { - 'imgt_pred': imgt_pred, - 'flow0_pred': [up_flow0_1, up_flow0_2, up_flow0_3, up_flow0_4], - 'flow1_pred': [up_flow1_1, up_flow1_2, up_flow1_3, up_flow1_4], - 'ft_pred': [ft_1_, ft_2_, ft_3_], - } - - - - - - - - - - - -class AMT_L(nn.Module): - def __init__(self, - corr_radius=3, - corr_lvls=4, - num_flows=5, - channels=[48, 64, 72, 128], - skip_channels=48 - ): - super(AMT_L, self).__init__() - self.radius = corr_radius - self.corr_levels = corr_lvls - self.num_flows = num_flows - - self.feat_encoder = BasicEncoder(output_dim=128, norm_fn='instance', dropout=0.) - self.encoder = Encoder([48, 64, 72, 128], large=True) - - self.decoder4 = InitDecoder(channels[3], channels[2], skip_channels) - self.decoder3 = IntermediateDecoder(channels[2], channels[1], skip_channels) - self.decoder2 = IntermediateDecoder(channels[1], channels[0], skip_channels) - self.decoder1 = MultiFlowDecoder(channels[0], skip_channels, num_flows) - - self.update4 = self._get_updateblock(72, None) - self.update3 = self._get_updateblock(64, 2.0) - self.update2 = self._get_updateblock(48, 4.0) - - self.comb_block = nn.Sequential( - nn.Conv2d(3*self.num_flows, 6*self.num_flows, 7, 1, 3), - nn.PReLU(6*self.num_flows), - nn.Conv2d(6*self.num_flows, 3, 7, 1, 3), - ) - - def _get_updateblock(self, cdim, scale_factor=None): - return BasicUpdateBlock(cdim=cdim, hidden_dim=128, flow_dim=48, - corr_dim=256, corr_dim2=160, fc_dim=124, - scale_factor=scale_factor, corr_levels=self.corr_levels, - radius=self.radius) - - def _corr_scale_lookup(self, corr_fn, coord, flow0, flow1, embt, downsample=1): - # convert t -> 0 to 0 -> 1 | convert t -> 1 to 1 -> 0 - # based on linear assumption - t1_scale = 1. / embt - t0_scale = 1. / (1. - embt) - if downsample != 1: - inv = 1 / downsample - flow0 = inv * resize(flow0, scale_factor=inv) - flow1 = inv * resize(flow1, scale_factor=inv) - - corr0, corr1 = corr_fn(coord + flow1 * t1_scale, coord + flow0 * t0_scale) - corr = torch.cat([corr0, corr1], dim=1) - flow = torch.cat([flow0, flow1], dim=1) - return corr, flow - - def forward(self, img0, img1, embt, scale_factor=1.0, eval=False, **kwargs): - mean_ = torch.cat([img0, img1], 2).mean(1, keepdim=True).mean(2, keepdim=True).mean(3, keepdim=True) - img0 = img0 - mean_ - img1 = img1 - mean_ - img0_ = resize(img0, scale_factor) if scale_factor != 1.0 else img0 - img1_ = resize(img1, scale_factor) if scale_factor != 1.0 else img1 - b, _, h, w = img0_.shape - coord = coords_grid(b, h // 8, w // 8, img0.device) - - fmap0, fmap1 = self.feat_encoder([img0_, img1_]) # [1, 128, H//8, W//8] - corr_fn = BidirCorrBlock(fmap0, fmap1, radius=self.radius, num_levels=self.corr_levels) - - # f0_1: [1, c0, H//2, W//2] | f0_2: [1, c1, H//4, W//4] - # f0_3: [1, c2, H//8, W//8] | f0_4: [1, c3, H//16, W//16] - f0_1, f0_2, f0_3, f0_4 = self.encoder(img0_) - f1_1, f1_2, f1_3, f1_4 = self.encoder(img1_) - - ######################################### the 4th decoder ######################################### - up_flow0_4, up_flow1_4, ft_3_ = self.decoder4(f0_4, f1_4, embt) - corr_4, flow_4 = self._corr_scale_lookup(corr_fn, coord, - up_flow0_4, up_flow1_4, - embt, downsample=1) - - # residue update with lookup corr - delta_ft_3_, delta_flow_4 = self.update4(ft_3_, flow_4, corr_4) - delta_flow0_4, delta_flow1_4 = torch.chunk(delta_flow_4, 2, 1) - up_flow0_4 = up_flow0_4 + delta_flow0_4 - up_flow1_4 = up_flow1_4 + delta_flow1_4 - ft_3_ = ft_3_ + delta_ft_3_ - - ######################################### the 3rd decoder ######################################### - up_flow0_3, up_flow1_3, ft_2_ = self.decoder3(ft_3_, f0_3, f1_3, up_flow0_4, up_flow1_4) - corr_3, flow_3 = self._corr_scale_lookup(corr_fn, - coord, up_flow0_3, up_flow1_3, - embt, downsample=2) - - # residue update with lookup corr - delta_ft_2_, delta_flow_3 = self.update3(ft_2_, flow_3, corr_3) - delta_flow0_3, delta_flow1_3 = torch.chunk(delta_flow_3, 2, 1) - up_flow0_3 = up_flow0_3 + delta_flow0_3 - up_flow1_3 = up_flow1_3 + delta_flow1_3 - ft_2_ = ft_2_ + delta_ft_2_ - - ######################################### the 2nd decoder ######################################### - up_flow0_2, up_flow1_2, ft_1_ = self.decoder2(ft_2_, f0_2, f1_2, up_flow0_3, up_flow1_3) - corr_2, flow_2 = self._corr_scale_lookup(corr_fn, - coord, up_flow0_2, up_flow1_2, - embt, downsample=4) - - # residue update with lookup corr - delta_ft_1_, delta_flow_2 = self.update2(ft_1_, flow_2, corr_2) - delta_flow0_2, delta_flow1_2 = torch.chunk(delta_flow_2, 2, 1) - up_flow0_2 = up_flow0_2 + delta_flow0_2 - up_flow1_2 = up_flow1_2 + delta_flow1_2 - ft_1_ = ft_1_ + delta_ft_1_ - - ######################################### the 1st decoder ######################################### - up_flow0_1, up_flow1_1, mask, img_res = self.decoder1(ft_1_, f0_1, f1_1, up_flow0_2, up_flow1_2) - - if scale_factor != 1.0: - up_flow0_1 = resize(up_flow0_1, scale_factor=(1.0/scale_factor)) * (1.0/scale_factor) - up_flow1_1 = resize(up_flow1_1, scale_factor=(1.0/scale_factor)) * (1.0/scale_factor) - mask = resize(mask, scale_factor=(1.0/scale_factor)) - img_res = resize(img_res, scale_factor=(1.0/scale_factor)) - - # Merge multiple predictions - imgt_pred = multi_flow_combine(self.comb_block, img0, img1, up_flow0_1, up_flow1_1, - mask, img_res, mean_) - imgt_pred = torch.clamp(imgt_pred, 0, 1) - - if eval: - return { 'imgt_pred': imgt_pred, } - else: - up_flow0_1 = up_flow0_1.reshape(b, self.num_flows, 2, h, w) - up_flow1_1 = up_flow1_1.reshape(b, self.num_flows, 2, h, w) - return { - 'imgt_pred': imgt_pred, - 'flow0_pred': [up_flow0_1, up_flow0_2, up_flow0_3, up_flow0_4], - 'flow1_pred': [up_flow1_1, up_flow1_2, up_flow1_3, up_flow1_4], - 'ft_pred': [ft_1_, ft_2_, ft_3_], - } - - - - - - - - - - - -class AMT_G(nn.Module): - def __init__(self, - corr_radius=3, - corr_lvls=4, - num_flows=5, - channels=[84, 96, 112, 128], - skip_channels=84): - super(AMT_G, self).__init__() - self.radius = corr_radius - self.corr_levels = corr_lvls - self.num_flows = num_flows - - self.feat_encoder = LargeEncoder(output_dim=128, norm_fn='instance', dropout=0.) - self.encoder = Encoder(channels, large=True) - self.decoder4 = InitDecoder(channels[3], channels[2], skip_channels) - self.decoder3 = IntermediateDecoder(channels[2], channels[1], skip_channels) - self.decoder2 = IntermediateDecoder(channels[1], channels[0], skip_channels) - self.decoder1 = MultiFlowDecoder(channels[0], skip_channels, num_flows) - - self.update4 = self._get_updateblock(112, None) - self.update3_low = self._get_updateblock(96, 2.0) - self.update2_low = self._get_updateblock(84, 4.0) - - self.update3_high = self._get_updateblock(96, None) - self.update2_high = self._get_updateblock(84, None) - - self.comb_block = nn.Sequential( - nn.Conv2d(3*self.num_flows, 6*self.num_flows, 7, 1, 3), - nn.PReLU(6*self.num_flows), - nn.Conv2d(6*self.num_flows, 3, 7, 1, 3), - ) - - def _get_updateblock(self, cdim, scale_factor=None): - return BasicUpdateBlock(cdim=cdim, hidden_dim=192, flow_dim=64, - corr_dim=256, corr_dim2=192, fc_dim=188, - scale_factor=scale_factor, corr_levels=self.corr_levels, - radius=self.radius) - - def _corr_scale_lookup(self, corr_fn, coord, flow0, flow1, embt, downsample=1): - # convert t -> 0 to 0 -> 1 | convert t -> 1 to 1 -> 0 - # based on linear assumption - t1_scale = 1. / embt - t0_scale = 1. / (1. - embt) - if downsample != 1: - inv = 1 / downsample - flow0 = inv * resize(flow0, scale_factor=inv) - flow1 = inv * resize(flow1, scale_factor=inv) - - corr0, corr1 = corr_fn(coord + flow1 * t1_scale, coord + flow0 * t0_scale) - corr = torch.cat([corr0, corr1], dim=1) - flow = torch.cat([flow0, flow1], dim=1) - return corr, flow - - def forward(self, img0, img1, embt, scale_factor=1.0, eval=False, **kwargs): - mean_ = torch.cat([img0, img1], 2).mean(1, keepdim=True).mean(2, keepdim=True).mean(3, keepdim=True) - img0 = img0 - mean_ - img1 = img1 - mean_ - img0_ = resize(img0, scale_factor) if scale_factor != 1.0 else img0 - img1_ = resize(img1, scale_factor) if scale_factor != 1.0 else img1 - b, _, h, w = img0_.shape - coord = coords_grid(b, h // 8, w // 8, img0.device) - - fmap0, fmap1 = self.feat_encoder([img0_, img1_]) # [1, 128, H//8, W//8] - corr_fn = BidirCorrBlock(fmap0, fmap1, radius=self.radius, num_levels=self.corr_levels) - - # f0_1: [1, c0, H//2, W//2] | f0_2: [1, c1, H//4, W//4] - # f0_3: [1, c2, H//8, W//8] | f0_4: [1, c3, H//16, W//16] - f0_1, f0_2, f0_3, f0_4 = self.encoder(img0_) - f1_1, f1_2, f1_3, f1_4 = self.encoder(img1_) - - ######################################### the 4th decoder ######################################### - up_flow0_4, up_flow1_4, ft_3_ = self.decoder4(f0_4, f1_4, embt) - corr_4, flow_4 = self._corr_scale_lookup(corr_fn, coord, - up_flow0_4, up_flow1_4, - embt, downsample=1) - - # residue update with lookup corr - delta_ft_3_, delta_flow_4 = self.update4(ft_3_, flow_4, corr_4) - delta_flow0_4, delta_flow1_4 = torch.chunk(delta_flow_4, 2, 1) - up_flow0_4 = up_flow0_4 + delta_flow0_4 - up_flow1_4 = up_flow1_4 + delta_flow1_4 - ft_3_ = ft_3_ + delta_ft_3_ - - ######################################### the 3rd decoder ######################################### - up_flow0_3, up_flow1_3, ft_2_ = self.decoder3(ft_3_, f0_3, f1_3, up_flow0_4, up_flow1_4) - corr_3, flow_3 = self._corr_scale_lookup(corr_fn, - coord, up_flow0_3, up_flow1_3, - embt, downsample=2) - - # residue update with lookup corr - delta_ft_2_, delta_flow_3 = self.update3_low(ft_2_, flow_3, corr_3) - delta_flow0_3, delta_flow1_3 = torch.chunk(delta_flow_3, 2, 1) - up_flow0_3 = up_flow0_3 + delta_flow0_3 - up_flow1_3 = up_flow1_3 + delta_flow1_3 - ft_2_ = ft_2_ + delta_ft_2_ - - # residue update with lookup corr (hr) - corr_3 = resize(corr_3, scale_factor=2.0) - up_flow_3 = torch.cat([up_flow0_3, up_flow1_3], dim=1) - delta_ft_2_, delta_up_flow_3 = self.update3_high(ft_2_, up_flow_3, corr_3) - ft_2_ += delta_ft_2_ - up_flow0_3 += delta_up_flow_3[:, 0:2] - up_flow1_3 += delta_up_flow_3[:, 2:4] - - ######################################### the 2nd decoder ######################################### - up_flow0_2, up_flow1_2, ft_1_ = self.decoder2(ft_2_, f0_2, f1_2, up_flow0_3, up_flow1_3) - corr_2, flow_2 = self._corr_scale_lookup(corr_fn, - coord, up_flow0_2, up_flow1_2, - embt, downsample=4) - - # residue update with lookup corr - delta_ft_1_, delta_flow_2 = self.update2_low(ft_1_, flow_2, corr_2) - delta_flow0_2, delta_flow1_2 = torch.chunk(delta_flow_2, 2, 1) - up_flow0_2 = up_flow0_2 + delta_flow0_2 - up_flow1_2 = up_flow1_2 + delta_flow1_2 - ft_1_ = ft_1_ + delta_ft_1_ - - # residue update with lookup corr (hr) - corr_2 = resize(corr_2, scale_factor=4.0) - up_flow_2 = torch.cat([up_flow0_2, up_flow1_2], dim=1) - delta_ft_1_, delta_up_flow_2 = self.update2_high(ft_1_, up_flow_2, corr_2) - ft_1_ += delta_ft_1_ - up_flow0_2 += delta_up_flow_2[:, 0:2] - up_flow1_2 += delta_up_flow_2[:, 2:4] - - ######################################### the 1st decoder ######################################### - up_flow0_1, up_flow1_1, mask, img_res = self.decoder1(ft_1_, f0_1, f1_1, up_flow0_2, up_flow1_2) - - if scale_factor != 1.0: - up_flow0_1 = resize(up_flow0_1, scale_factor=(1.0/scale_factor)) * (1.0/scale_factor) - up_flow1_1 = resize(up_flow1_1, scale_factor=(1.0/scale_factor)) * (1.0/scale_factor) - mask = resize(mask, scale_factor=(1.0/scale_factor)) - img_res = resize(img_res, scale_factor=(1.0/scale_factor)) - - # Merge multiple predictions - imgt_pred = multi_flow_combine(self.comb_block, img0, img1, up_flow0_1, up_flow1_1, - mask, img_res, mean_) - imgt_pred = torch.clamp(imgt_pred, 0, 1) - - if eval: - return { 'imgt_pred': imgt_pred, } - else: - up_flow0_1 = up_flow0_1.reshape(b, self.num_flows, 2, h, w) - up_flow1_1 = up_flow1_1.reshape(b, self.num_flows, 2, h, w) - return { - 'imgt_pred': imgt_pred, - 'flow0_pred': [up_flow0_1, up_flow0_2, up_flow0_3, up_flow0_4], - 'flow1_pred': [up_flow1_1, up_flow1_2, up_flow1_3, up_flow1_4], - 'ft_pred': [ft_1_, ft_2_, ft_3_], - } \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/__init__.py deleted file mode 100644 index 8d1166d..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/__init__.py +++ /dev/null @@ -1,64 +0,0 @@ -import torch -from torch.utils.data import DataLoader -import pathlib -from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames, generic_frame_loop, InterpolationStateList -import typing -from comfy.model_management import get_torch_device - -MODEL_TYPE = pathlib.Path(__file__).parent.name -CKPT_NAMES = ["pretrained_cain.pth"] - - -class CAIN_VFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (CKPT_NAMES, ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}), - "multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}) - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames: typing.SupportsInt = 1, - multiplier: typing.SupportsInt = 2, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - from .cain_arch import CAIN - model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) - sd = torch.load(model_path)["state_dict"] - sd = {key.replace('module.', ''): value for key, value in sd.items()} - - - global interpolation_model - interpolation_model = CAIN(depth=3) - interpolation_model.load_state_dict(sd) - interpolation_model.eval().to(get_torch_device()) - del sd - - frames = preprocess_frames(frames) - - - def return_middle_frame(frame_0, frame_1, timestep, model): - #CAIN does some direct modifications to input frame tensors so we need to clone them - return model(frame_0.detach().clone(), frame_1.detach().clone())[0] - - args = [interpolation_model] - out = postprocess_frames( - generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, - interpolation_states=optional_interpolation_states, use_timestep=False, dtype=torch.float32) - ) - return (out,) \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_arch.py deleted file mode 100644 index abd8db3..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_arch.py +++ /dev/null @@ -1,74 +0,0 @@ -import math -import numpy as np - -import torch -import torch.nn as nn - -from .common import * - - -class Encoder(nn.Module): - def __init__(self, in_channels=3, depth=3): - super(Encoder, self).__init__() - - # Shuffle pixels to expand in channel dimension - # shuffler_list = [PixelShuffle(0.5) for i in range(depth)] - # self.shuffler = nn.Sequential(*shuffler_list) - self.shuffler = PixelShuffle(1 / 2**depth) - - relu = nn.LeakyReLU(0.2, True) - - # FF_RCAN or FF_Resblocks - self.interpolate = Interpolation(5, 12, in_channels * (4**depth), act=relu) - - def forward(self, x1, x2): - """ - Encoder: Shuffle-spread --> Feature Fusion --> Return fused features - """ - feats1 = self.shuffler(x1) - feats2 = self.shuffler(x2) - - feats = self.interpolate(feats1, feats2) - - return feats - - -class Decoder(nn.Module): - def __init__(self, depth=3): - super(Decoder, self).__init__() - - # shuffler_list = [PixelShuffle(2) for i in range(depth)] - # self.shuffler = nn.Sequential(*shuffler_list) - self.shuffler = PixelShuffle(2**depth) - - def forward(self, feats): - out = self.shuffler(feats) - return out - - -class CAIN(nn.Module): - def __init__(self, depth=3): - super(CAIN, self).__init__() - - self.encoder = Encoder(in_channels=3, depth=depth) - self.decoder = Decoder(depth=depth) - - def forward(self, x1, x2): - x1, m1 = sub_mean(x1) - x2, m2 = sub_mean(x2) - - if not self.training: - paddingInput, paddingOutput = InOutPaddings(x1) - x1 = paddingInput(x1) - x2 = paddingInput(x2) - - feats = self.encoder(x1, x2) - out = self.decoder(feats) - - if not self.training: - out = paddingOutput(out) - - mi = (m1 + m2) / 2 - out += mi - - return out, feats \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_encdec_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_encdec_arch.py deleted file mode 100644 index ea10ed8..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_encdec_arch.py +++ /dev/null @@ -1,95 +0,0 @@ -import math -import numpy as np - -import torch -import torch.nn as nn - -from .common import * -from comfy.model_management import get_torch_device - -class Encoder(nn.Module): - def __init__(self, in_channels=3, depth=3, nf_start=32, norm=False): - super(Encoder, self).__init__() - self.device = get_torch_device() - - nf = nf_start - relu = nn.LeakyReLU(negative_slope=0.2, inplace=True) - - self.body = nn.Sequential( - ConvNorm(in_channels, nf * 1, 7, stride=1, norm=norm), - relu, - ConvNorm(nf * 1, nf * 2, 5, stride=2, norm=norm), - relu, - ConvNorm(nf * 2, nf * 4, 5, stride=2, norm=norm), - relu, - ConvNorm(nf * 4, nf * 6, 5, stride=2, norm=norm) - ) - - self.interpolate = Interpolation(5, 12, nf * 6, reduction=16, act=relu) - - def forward(self, x1, x2): - """ - Encoder: Feature Extraction --> Feature Fusion --> Return - """ - feats1 = self.body(x1) - feats2 = self.body(x2) - - feats = self.interpolate(feats1, feats2) - - return feats - - -class Decoder(nn.Module): - def __init__(self, in_channels=192, out_channels=3, depth=3, norm=False, up_mode='shuffle'): - super(Decoder, self).__init__() - self.device = get_torch_device() - - relu = nn.LeakyReLU(negative_slope=0.2, inplace=True) - - nf = [in_channels, (in_channels*2)//3, in_channels//3, in_channels//6] - #nf = [192, 128, 64, 32] - #nf = [186, 124, 62, 31] - self.body = nn.Sequential( - UpConvNorm(nf[0], nf[1], mode=up_mode, norm=norm), - ResBlock(nf[1], nf[1], norm=norm, act=relu), - UpConvNorm(nf[1], nf[2], mode=up_mode, norm=norm), - ResBlock(nf[2], nf[2], norm=norm, act=relu), - UpConvNorm(nf[2], nf[3], mode=up_mode, norm=norm), - ResBlock(nf[3], nf[3], norm=norm, act=relu), - conv7x7(nf[3], out_channels) - ) - - def forward(self, feats): - out = self.body(feats) - #out = self.conv_final(out) - - return out - - -class CAIN_EncDec(nn.Module): - def __init__(self, depth=3, n_resblocks=3, start_filts=32, up_mode='shuffle'): - super(CAIN_EncDec, self).__init__() - self.depth = depth - - self.encoder = Encoder(in_channels=3, depth=depth, norm=False) - self.decoder = Decoder(in_channels=start_filts*6, depth=depth, norm=False, up_mode=up_mode) - - def forward(self, x1, x2): - x1, m1 = sub_mean(x1) - x2, m2 = sub_mean(x2) - - if not self.training: - paddingInput, paddingOutput = InOutPaddings(x1) - x1 = paddingInput(x1) - x2 = paddingInput(x2) - - feats = self.encoder(x1, x2) - out = self.decoder(feats) - - if not self.training: - out = paddingOutput(out) - - mi = (m1 + m2)/2 - out += mi - - return out, feats \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_noca_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_noca_arch.py deleted file mode 100644 index 08fbb11..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/cain_noca_arch.py +++ /dev/null @@ -1,73 +0,0 @@ -import math -import numpy as np - -import torch -import torch.nn as nn - -from .common import * -from comfy.model_management import get_torch_device - -class Encoder(nn.Module): - def __init__(self, in_channels=3, depth=3): - super(Encoder, self).__init__() - self.device = get_torch_device() - - self.shuffler = PixelShuffle(1/2**depth) - # self.shuffler = nn.Sequential( - # PixelShuffle(1/2), - # PixelShuffle(1/2), - # PixelShuffle(1/2)) - self.interpolate = Interpolation_res(5, 12, in_channels * (4**depth)) - - def forward(self, x1, x2): - feats1 = self.shuffler(x1) - feats2 = self.shuffler(x2) - - feats = self.interpolate(feats1, feats2) - - return feats - - -class Decoder(nn.Module): - def __init__(self, depth=3): - super(Decoder, self).__init__() - self.device = get_torch_device() - - self.shuffler = PixelShuffle(2**depth) - # self.shuffler = nn.Sequential( - # PixelShuffle(2), - # PixelShuffle(2), - # PixelShuffle(2)) - - def forward(self, feats): - out = self.shuffler(feats) - return out - - -class CAIN_NoCA(nn.Module): - def __init__(self, depth=3): - super(CAIN_NoCA, self).__init__() - self.depth = depth - - self.encoder = Encoder(in_channels=3, depth=depth) - self.decoder = Decoder(depth=depth) - - def forward(self, x1, x2): - x1, m1 = sub_mean(x1) - x2, m2 = sub_mean(x2) - - if not self.training: - paddingInput, paddingOutput = InOutPaddings(x1) - x1 = paddingInput(x1) - x2 = paddingInput(x2) - - feats = self.encoder(x1, x2) - out = self.decoder(feats) - - if not self.training: - out = paddingOutput(out) - - mi = (m1 + m2) / 2 - out += mi - - return out, feats \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/common.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/common.py deleted file mode 100644 index aa1600a..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/cain/common.py +++ /dev/null @@ -1,361 +0,0 @@ -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - -def sub_mean(x): - mean = x.mean(2, keepdim=True).mean(3, keepdim=True) - x -= mean - return x, mean - -def InOutPaddings(x): - w, h = x.size(3), x.size(2) - padding_width, padding_height = 0, 0 - if w != ((w >> 7) << 7): - padding_width = (((w >> 7) + 1) << 7) - w - if h != ((h >> 7) << 7): - padding_height = (((h >> 7) + 1) << 7) - h - paddingInput = nn.ReflectionPad2d(padding=[padding_width // 2, padding_width - padding_width // 2, - padding_height // 2, padding_height - padding_height // 2]) - paddingOutput = nn.ReflectionPad2d(padding=[0 - padding_width // 2, padding_width // 2 - padding_width, - 0 - padding_height // 2, padding_height // 2 - padding_height]) - return paddingInput, paddingOutput - - -class ConvNorm(nn.Module): - def __init__(self, in_feat, out_feat, kernel_size, stride=1, norm=False): - super(ConvNorm, self).__init__() - - reflection_padding = kernel_size // 2 - self.reflection_pad = nn.ReflectionPad2d(reflection_padding) - self.conv = nn.Conv2d(in_feat, out_feat, stride=stride, kernel_size=kernel_size, bias=True) - - self.norm = norm - if norm == 'IN': - self.norm = nn.InstanceNorm2d(out_feat, track_running_stats=True) - elif norm == 'BN': - self.norm = nn.BatchNorm2d(out_feat) - - def forward(self, x): - out = self.reflection_pad(x) - out = self.conv(out) - if self.norm: - out = self.norm(out) - return out - - -class UpConvNorm(nn.Module): - def __init__(self, in_channels, out_channels, mode='transpose', norm=False): - super(UpConvNorm, self).__init__() - - if mode == 'transpose': - self.upconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1) - elif mode == 'shuffle': - self.upconv = nn.Sequential( - ConvNorm(in_channels, 4*out_channels, kernel_size=3, stride=1, norm=norm), - PixelShuffle(2)) - else: - # out_channels is always going to be the same as in_channels - self.upconv = nn.Sequential( - nn.Upsample(mode='bilinear', scale_factor=2, align_corners=False), - ConvNorm(in_channels, out_channels, kernel_size=1, stride=1, norm=norm)) - - def forward(self, x): - out = self.upconv(x) - return out - - - -class meanShift(nn.Module): - def __init__(self, rgbRange, rgbMean, sign, nChannel=3): - super(meanShift, self).__init__() - if nChannel == 1: - l = rgbMean[0] * rgbRange * float(sign) - - self.shifter = nn.Conv2d(1, 1, kernel_size=1, stride=1, padding=0) - self.shifter.weight.data = torch.eye(1).view(1, 1, 1, 1) - self.shifter.bias.data = torch.Tensor([l]) - elif nChannel == 3: - r = rgbMean[0] * rgbRange * float(sign) - g = rgbMean[1] * rgbRange * float(sign) - b = rgbMean[2] * rgbRange * float(sign) - - self.shifter = nn.Conv2d(3, 3, kernel_size=1, stride=1, padding=0) - self.shifter.weight.data = torch.eye(3).view(3, 3, 1, 1) - self.shifter.bias.data = torch.Tensor([r, g, b]) - else: - r = rgbMean[0] * rgbRange * float(sign) - g = rgbMean[1] * rgbRange * float(sign) - b = rgbMean[2] * rgbRange * float(sign) - self.shifter = nn.Conv2d(6, 6, kernel_size=1, stride=1, padding=0) - self.shifter.weight.data = torch.eye(6).view(6, 6, 1, 1) - self.shifter.bias.data = torch.Tensor([r, g, b, r, g, b]) - - # Freeze the meanShift layer - for params in self.shifter.parameters(): - params.requires_grad = False - - def forward(self, x): - x = self.shifter(x) - - return x - - -""" CONV - (BN) - RELU - CONV - (BN) """ -class ResBlock(nn.Module): - def __init__(self, in_feat, out_feat, kernel_size=3, reduction=False, bias=True, # 'reduction' is just for placeholder - norm=False, act=nn.ReLU(True), downscale=False): - super(ResBlock, self).__init__() - - self.body = nn.Sequential( - ConvNorm(in_feat, out_feat, kernel_size=kernel_size, stride=2 if downscale else 1), - act, - ConvNorm(out_feat, out_feat, kernel_size=kernel_size, stride=1) - ) - - self.downscale = None - if downscale: - self.downscale = nn.Conv2d(in_feat, out_feat, kernel_size=1, stride=2) - - def forward(self, x): - res = x - out = self.body(x) - if self.downscale is not None: - res = self.downscale(res) - out += res - - return out - - -## Channel Attention (CA) Layer -class CALayer(nn.Module): - def __init__(self, channel, reduction=16): - super(CALayer, self).__init__() - # global average pooling: feature --> point - self.avg_pool = nn.AdaptiveAvgPool2d(1) - # feature channel downscale and upscale --> channel weight - self.conv_du = nn.Sequential( - nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True), - nn.ReLU(inplace=True), - nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True), - nn.Sigmoid() - ) - - def forward(self, x): - y = self.avg_pool(x) - y = self.conv_du(y) - return x * y, y - - -## Residual Channel Attention Block (RCAB) -class RCAB(nn.Module): - def __init__(self, in_feat, out_feat, kernel_size, reduction, bias=True, - norm=False, act=nn.ReLU(True), downscale=False, return_ca=False): - super(RCAB, self).__init__() - - self.body = nn.Sequential( - ConvNorm(in_feat, out_feat, kernel_size, stride=2 if downscale else 1, norm=norm), - act, - ConvNorm(out_feat, out_feat, kernel_size, stride=1, norm=norm), - CALayer(out_feat, reduction) - ) - self.downscale = downscale - if downscale: - self.downConv = nn.Conv2d(in_feat, out_feat, kernel_size=3, stride=2, padding=1) - self.return_ca = return_ca - - def forward(self, x): - res = x - out, ca = self.body(x) - if self.downscale: - res = self.downConv(res) - out += res - - if self.return_ca: - return out, ca - else: - return out - - -## Residual Group (RG) -class ResidualGroup(nn.Module): - def __init__(self, Block, n_resblocks, n_feat, kernel_size, reduction, act, norm=False): - super(ResidualGroup, self).__init__() - - modules_body = [Block(n_feat, n_feat, kernel_size, reduction, bias=True, norm=norm, act=act) - for _ in range(n_resblocks)] - modules_body.append(ConvNorm(n_feat, n_feat, kernel_size, stride=1, norm=norm)) - self.body = nn.Sequential(*modules_body) - - def forward(self, x): - res = self.body(x) - res += x - return res - - -def pixel_shuffle(input, scale_factor): - batch_size, channels, in_height, in_width = input.size() - - out_channels = int(int(channels / scale_factor) / scale_factor) - out_height = int(in_height * scale_factor) - out_width = int(in_width * scale_factor) - - if scale_factor >= 1: - input_view = input.contiguous().view(batch_size, out_channels, scale_factor, scale_factor, in_height, in_width) - shuffle_out = input_view.permute(0, 1, 4, 2, 5, 3).contiguous() - else: - block_size = int(1 / scale_factor) - input_view = input.contiguous().view(batch_size, channels, out_height, block_size, out_width, block_size) - shuffle_out = input_view.permute(0, 1, 3, 5, 2, 4).contiguous() - - return shuffle_out.view(batch_size, out_channels, out_height, out_width) - - -class PixelShuffle(nn.Module): - def __init__(self, scale_factor): - super(PixelShuffle, self).__init__() - self.scale_factor = scale_factor - - def forward(self, x): - return pixel_shuffle(x, self.scale_factor) - def extra_repr(self): - return 'scale_factor={}'.format(self.scale_factor) - - -def conv(in_channels, out_channels, kernel_size, - stride=1, bias=True, groups=1): - return nn.Conv2d( - in_channels, - out_channels, - kernel_size=kernel_size, - padding=kernel_size//2, - stride=1, - bias=bias, - groups=groups) - - -def conv1x1(in_channels, out_channels, stride=1, bias=True, groups=1): - return nn.Conv2d( - in_channels, - out_channels, - kernel_size=1, - stride=stride, - bias=bias, - groups=groups) - -def conv3x3(in_channels, out_channels, stride=1, - padding=1, bias=True, groups=1): - return nn.Conv2d( - in_channels, - out_channels, - kernel_size=3, - stride=stride, - padding=padding, - bias=bias, - groups=groups) - -def conv5x5(in_channels, out_channels, stride=1, - padding=2, bias=True, groups=1): - return nn.Conv2d( - in_channels, - out_channels, - kernel_size=5, - stride=stride, - padding=padding, - bias=bias, - groups=groups) - -def conv7x7(in_channels, out_channels, stride=1, - padding=3, bias=True, groups=1): - return nn.Conv2d( - in_channels, - out_channels, - kernel_size=7, - stride=stride, - padding=padding, - bias=bias, - groups=groups) - -def upconv2x2(in_channels, out_channels, mode='shuffle'): - if mode == 'transpose': - return nn.ConvTranspose2d( - in_channels, - out_channels, - kernel_size=4, - stride=2, - padding=1) - elif mode == 'shuffle': - return nn.Sequential( - conv3x3(in_channels, 4*out_channels), - PixelShuffle(2)) - else: - # out_channels is always going to be the same as in_channels - return nn.Sequential( - nn.Upsample(mode='bilinear', scale_factor=2, align_corners=False), - conv1x1(in_channels, out_channels)) - - - -class Interpolation(nn.Module): - def __init__(self, n_resgroups, n_resblocks, n_feats, - reduction=16, act=nn.LeakyReLU(0.2, True), norm=False): - super(Interpolation, self).__init__() - - # define modules: head, body, tail - self.headConv = conv3x3(n_feats * 2, n_feats) - - modules_body = [ - ResidualGroup( - RCAB, - n_resblocks=n_resblocks, - n_feat=n_feats, - kernel_size=3, - reduction=reduction, - act=act, - norm=norm) - for _ in range(n_resgroups)] - self.body = nn.Sequential(*modules_body) - - self.tailConv = conv3x3(n_feats, n_feats) - - def forward(self, x0, x1): - # Build input tensor - x = torch.cat([x0, x1], dim=1) - x = self.headConv(x) - - res = self.body(x) - res += x - - out = self.tailConv(res) - return out - - -class Interpolation_res(nn.Module): - def __init__(self, n_resgroups, n_resblocks, n_feats, - act=nn.LeakyReLU(0.2, True), norm=False): - super(Interpolation_res, self).__init__() - - # define modules: head, body, tail (reduces concatenated inputs to n_feat) - self.headConv = conv3x3(n_feats * 2, n_feats) - - modules_body = [ResidualGroup(ResBlock, n_resblocks=n_resblocks, n_feat=n_feats, kernel_size=3, - reduction=0, act=act, norm=norm) - for _ in range(n_resgroups)] - self.body = nn.Sequential(*modules_body) - - self.tailConv = conv3x3(n_feats, n_feats) - - def forward(self, x0, x1): - # Build input tensor - x = torch.cat([x0, x1], dim=1) - x = self.headConv(x) - - res = x - for m in self.body: - res = m(res) - res += x - - x = self.tailConv(res) - - return x \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/eisai/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/eisai/__init__.py deleted file mode 100644 index 3602505..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/eisai/__init__.py +++ /dev/null @@ -1,84 +0,0 @@ -import pathlib -from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames, generic_frame_loop, InterpolationStateList -import typing -import torch -import torch.nn as nn -from comfy.model_management import soft_empty_cache, get_torch_device - -MODEL_TYPE = pathlib.Path(__file__).parent.name -MODEL_FILE_NAMES = { - "ssl": "eisai_ssl.pt", - "dtm": "eisai_dtm.pt", - "raft": "eisai_anime_interp_full.ckpt" -} - -class EISAI(nn.Module): - def __init__(self, model_file_names) -> None: - from .eisai_arch import SoftsplatLite, DTM, RAFT - super(EISAI, self).__init__() - self.raft = RAFT(load_file_from_github_release(MODEL_TYPE, model_file_names["raft"])) - self.raft.to(get_torch_device()).eval() - - self.ssl = SoftsplatLite() - self.ssl.load_state_dict(torch.load(load_file_from_github_release(MODEL_TYPE, model_file_names["ssl"]))) - self.ssl.to(get_torch_device()).eval() - - self.dtm = DTM() - self.dtm.load_state_dict(torch.load(load_file_from_github_release(MODEL_TYPE, model_file_names["dtm"]))) - self.dtm.to(get_torch_device()).eval() - - def forward(self, img0, img1, t): - with torch.no_grad(): - flow0, _ = self.raft(img0, img1) - flow1, _ = self.raft(img1, img0) - x = { - "images": torch.stack([img0, img1], dim=1), - "flows": torch.stack([flow0, flow1], dim=1), - } - out_ssl, _ = self.ssl(x, t=t, return_more=True) - out_dtm, _ = self.dtm(x, out_ssl, _, return_more=False) - return out_dtm[:, :3] - -class EISAI_VFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (["eisai"], ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}), - "multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}), - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames = 10, - multiplier: typing.SupportsInt = 2, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - interpolation_model = EISAI(MODEL_FILE_NAMES) - interpolation_model.eval().to(get_torch_device()) - frames = preprocess_frames(frames) - - def return_middle_frame(frame_0, frame_1, timestep, model): - return model(frame_0, frame_1, t=timestep) - - scale = 1 - - args = [interpolation_model, scale] - out = postprocess_frames( - generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, - interpolation_states=optional_interpolation_states, dtype=torch.float32) - ) - return (out,) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/eisai/eisai_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/eisai/eisai_arch.py deleted file mode 100644 index 4a3abae..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/eisai/eisai_arch.py +++ /dev/null @@ -1,2586 +0,0 @@ -""" -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_scripts/interpolate.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/models/ssldtm.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/util_v0.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/twodee_v0.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/pytorch_v0.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/distance_transform_v0.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/sketchers_v1.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/interpolator_v0.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/gridnet_v1.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/flow_v0.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/softsplat_v0.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/raft_v1/rfr_new.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/raft_v1/extractor.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/raft_v1/update.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/raft_v1/corr.py -https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/raft_v1/utils.py -""" - -import copy -import cv2 -import torch.nn.functional as F -import torchvision.transforms.functional as F -import gc -from PIL import Image, ImageFile, ImageFont, ImageDraw -import inspect -from scipy import interpolate -import kornia -import math -from argparse import Namespace -import torch.nn as nn -import numpy as np -import os -from functools import partial -import pathlib -import PIL -import re -import requests -from scipy.spatial.transform import Rotation -import scipy -import shutil -import torchvision.transforms as T -import time -import torch -import torchvision as tv -import zlib -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from tqdm.auto import tqdm as std_tqdm -from tqdm.auto import trange as std_trange -from vfi_models.ops import FunctionSoftsplat, batch_edt -from comfy.model_management import get_torch_device - -device = get_torch_device() -autocast = torch.autocast -tqdm = partial(std_tqdm, dynamic_ncols=True) -trange = partial(std_trange, dynamic_ncols=True) -ImageFile.LOAD_TRUNCATED_IMAGES = True - - -def pixel_ij(x, rounding=True): - if isinstance(x, np.ndarray): - x = x.tolist() - return tuple( - pixel_rounder(i, rounding) - for i in (x if isinstance(x, tuple) or isinstance(x, list) else (x, x)) - ) - - -def rescale_dry(x, factor): - h, w = x[-2:] if isinstance(x, tuple) or isinstance(x, list) else I(x).size - return (h * factor, w * factor) - - -def pixel_rounder(n, mode): - if mode == True or mode == "round": - return round(n) - elif mode == "ceil": - return math.ceil(n) - elif mode == "floor": - return math.floor(n) - else: - return n - - -def diam(x): - if isinstance(x, tuple) or isinstance(x, list): - h, w = x[-2:] - elif isinstance(x, I): - h, w = x.size - else: - h, w = x.shape[-2:] - return np.sqrt(h**2 + w**2) - - -def pixel_logit(x, pixel_margin=1): - x = (x * (255 - 2 * pixel_margin) + pixel_margin) / 255 - return torch.log(x / (1 - x)) - - -class InputPadder: - """Pads images such that dimensions are divisible by 8""" - - def __init__(self, dims): - self.ht, self.wd = dims[-2:] - pad_ht = (((self.ht // 8) + 1) * 8 - self.ht) % 8 - pad_wd = (((self.wd // 8) + 1) * 8 - self.wd) % 8 - self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, 0, pad_ht] - - def pad(self, *inputs): - return [F.pad(x, self._pad, mode="replicate") for x in inputs] - - def unpad(self, x): - ht, wd = x.shape[-2:] - c = [self._pad[2], ht - self._pad[3], self._pad[0], wd - self._pad[1]] - return x[..., c[0] : c[1], c[2] : c[3]] - - -def forward_interpolate(flow): - flow = flow.detach().cpu().numpy() - dx, dy = flow[0], flow[1] - - ht, wd = dx.shape - x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht)) - - x1 = x0 + dx - y1 = y0 + dy - - x1 = x1.reshape(-1) - y1 = y1.reshape(-1) - dx = dx.reshape(-1) - dy = dy.reshape(-1) - - valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht) - x1 = x1[valid] - y1 = y1[valid] - dx = dx[valid] - dy = dy[valid] - - flow_x = interpolate.griddata((x1, y1), dx, (x0, y0), method="cubic", fill_value=0) - - flow_y = interpolate.griddata((x1, y1), dy, (x0, y0), method="cubic", fill_value=0) - - flow = np.stack([flow_x, flow_y], axis=0) - return torch.from_numpy(flow).float() - - -def bilinear_sampler(img, coords, mode="bilinear", mask=False): - """Wrapper for grid_sample, uses pixel coordinates""" - H, W = img.shape[-2:] - xgrid, ygrid = coords.split([1, 1], dim=-1) - xgrid = 2 * xgrid / (W - 1) - 1 - ygrid = 2 * ygrid / (H - 1) - 1 - - grid = torch.cat([xgrid, ygrid], dim=-1) - # print(img.size()) - img = F.grid_sample(img, grid, align_corners=True) - - if mask: - mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) - return img, mask.float() - - return img - - -def coords_grid(batch, ht, wd): - coords = torch.meshgrid(torch.arange(ht), torch.arange(wd)) - coords = torch.stack(coords[::-1], dim=0).float() - return coords[None].repeat(batch, 1, 1, 1) - - -def upflow8(flow, mode="bilinear"): - new_size = (8 * flow.shape[2], 8 * flow.shape[3]) - return 8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True) - - -class CorrBlock: - def __init__(self, fmap1, fmap2, num_levels=4, radius=4): - self.num_levels = num_levels - self.radius = radius - self.corr_pyramid = [] - - # all pairs correlation - corr = CorrBlock.corr(fmap1, fmap2) - - batch, h1, w1, dim, h2, w2 = corr.shape - corr = corr.reshape(batch * h1 * w1, dim, h2, w2) - - self.corr_pyramid.append(corr) - for i in range(self.num_levels - 1): - corr = F.avg_pool2d(corr, 2, stride=2) - self.corr_pyramid.append(corr) - - def __call__(self, coords): - r = self.radius - coords = coords.permute(0, 2, 3, 1) - batch, h1, w1, _ = coords.shape - - out_pyramid = [] - for i in range(self.num_levels): - corr = self.corr_pyramid[i] - dx = torch.linspace(-r, r, 2 * r + 1) - dy = torch.linspace(-r, r, 2 * r + 1) - delta = torch.stack(torch.meshgrid(dy, dx), dim=-1).to(coords.device) - - centroid_lvl = coords.reshape(batch * h1 * w1, 1, 1, 2) / 2**i - delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2) - coords_lvl = centroid_lvl + delta_lvl - - corr = bilinear_sampler(corr, coords_lvl) - corr = corr.view(batch, h1, w1, -1) - out_pyramid.append(corr) - - out = torch.cat(out_pyramid, dim=-1) - return out.permute(0, 3, 1, 2).contiguous().float() - - @staticmethod - def corr(fmap1, fmap2): - batch, dim, ht, wd = fmap1.shape - fmap1 = fmap1.view(batch, dim, ht * wd) - fmap2 = fmap2.view(batch, dim, ht * wd) - - corr = torch.matmul(fmap1.transpose(1, 2), fmap2) - corr = corr.view(batch, ht, wd, 1, ht, wd) - return corr / torch.sqrt(torch.tensor(dim).float()) - - -class FlowHead(nn.Module): - def __init__(self, input_dim=128, hidden_dim=256): - super(FlowHead, self).__init__() - self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) - self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1) - self.relu = nn.ReLU(inplace=True) - - def forward(self, x): - return self.conv2(self.relu(self.conv1(x))) - - -class ConvGRU(nn.Module): - def __init__(self, hidden_dim=128, input_dim=192 + 128): - super(ConvGRU, self).__init__() - self.convz = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1) - self.convr = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1) - self.convq = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1) - - def forward(self, h, x): - hx = torch.cat([h, x], dim=1) - - z = torch.sigmoid(self.convz(hx)) - r = torch.sigmoid(self.convr(hx)) - q = torch.tanh(self.convq(torch.cat([r * h, x], dim=1))) - - h = (1 - z) * h + z * q - return h - - -class SepConvGRU(nn.Module): - def __init__(self, hidden_dim=128, input_dim=192 + 128): - super(SepConvGRU, self).__init__() - self.convz1 = nn.Conv2d( - hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2) - ) - self.convr1 = nn.Conv2d( - hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2) - ) - self.convq1 = nn.Conv2d( - hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2) - ) - - self.convz2 = nn.Conv2d( - hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0) - ) - self.convr2 = nn.Conv2d( - hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0) - ) - self.convq2 = nn.Conv2d( - hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0) - ) - - def forward(self, h, x): - # horizontal - hx = torch.cat([h, x], dim=1) - z = torch.sigmoid(self.convz1(hx)) - r = torch.sigmoid(self.convr1(hx)) - q = torch.tanh(self.convq1(torch.cat([r * h, x], dim=1))) - h = (1 - z) * h + z * q - - # vertical - hx = torch.cat([h, x], dim=1) - z = torch.sigmoid(self.convz2(hx)) - r = torch.sigmoid(self.convr2(hx)) - q = torch.tanh(self.convq2(torch.cat([r * h, x], dim=1))) - h = (1 - z) * h + z * q - - return h - - -class SmallMotionEncoder(nn.Module): - def __init__(self, args): - super(SmallMotionEncoder, self).__init__() - cor_planes = args.corr_levels * (2 * args.corr_radius + 1) ** 2 - self.convc1 = nn.Conv2d(cor_planes, 96, 1, padding=0) - self.convf1 = nn.Conv2d(2, 64, 7, padding=3) - self.convf2 = nn.Conv2d(64, 32, 3, padding=1) - self.conv = nn.Conv2d(128, 80, 3, padding=1) - - def forward(self, flow, corr): - cor = F.relu(self.convc1(corr)) - flo = F.relu(self.convf1(flow)) - flo = F.relu(self.convf2(flo)) - cor_flo = torch.cat([cor, flo], dim=1) - out = F.relu(self.conv(cor_flo)) - return torch.cat([out, flow], dim=1) - - -class BasicMotionEncoder(nn.Module): - def __init__(self, args): - super(BasicMotionEncoder, self).__init__() - cor_planes = args.corr_levels * (2 * args.corr_radius + 1) ** 2 - self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0) - self.convc2 = nn.Conv2d(256, 192, 3, padding=1) - self.convf1 = nn.Conv2d(2, 128, 7, padding=3) - self.convf2 = nn.Conv2d(128, 64, 3, padding=1) - self.conv = nn.Conv2d(64 + 192, 128 - 2, 3, padding=1) - - def forward(self, flow, corr): - cor = F.relu(self.convc1(corr)) - cor = F.relu(self.convc2(cor)) - flo = F.relu(self.convf1(flow)) - flo = F.relu(self.convf2(flo)) - - cor_flo = torch.cat([cor, flo], dim=1) - out = F.relu(self.conv(cor_flo)) - return torch.cat([out, flow], dim=1) - - -class SmallUpdateBlock(nn.Module): - def __init__(self, args, hidden_dim=96): - super(SmallUpdateBlock, self).__init__() - self.encoder = SmallMotionEncoder(args) - self.gru = ConvGRU(hidden_dim=hidden_dim, input_dim=82 + 64) - self.flow_head = FlowHead(hidden_dim, hidden_dim=128) - - def forward(self, net, inp, corr, flow): - motion_features = self.encoder(flow, corr) - inp = torch.cat([inp, motion_features], dim=1) - net = self.gru(net, inp) - delta_flow = self.flow_head(net) - - return net, None, delta_flow - - -class BasicUpdateBlock(nn.Module): - def __init__(self, args, hidden_dim=128, input_dim=128): - super(BasicUpdateBlock, self).__init__() - self.args = args - self.encoder = BasicMotionEncoder(args) - self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128 + hidden_dim) - self.flow_head = FlowHead(hidden_dim, hidden_dim=256) - - self.mask = nn.Sequential( - nn.Conv2d(128, 256, 3, padding=1), - nn.ReLU(inplace=True), - nn.Conv2d(256, 64 * 9, 1, padding=0), - ) - - def forward(self, net, inp, corr, flow, upsample=True): - motion_features = self.encoder(flow, corr) - inp = torch.cat([inp, motion_features], dim=1) - - net = self.gru(net, inp) - delta_flow = self.flow_head(net) - - # scale mask to balence gradients - mask = 0.25 * self.mask(net) - return net, mask, delta_flow - - -class ResidualBlock(nn.Module): - def __init__(self, in_planes, planes, norm_fn="group", stride=1): - super(ResidualBlock, self).__init__() - - self.conv1 = nn.Conv2d( - in_planes, planes, kernel_size=3, padding=1, stride=stride - ) - self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1) - self.relu = nn.ReLU(inplace=True) - - num_groups = planes // 8 - - if norm_fn == "group": - self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) - self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) - if not stride == 1: - self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) - - elif norm_fn == "batch": - self.norm1 = nn.BatchNorm2d(planes) - self.norm2 = nn.BatchNorm2d(planes) - if not stride == 1: - self.norm3 = nn.BatchNorm2d(planes) - - elif norm_fn == "instance": - self.norm1 = nn.InstanceNorm2d(planes) - self.norm2 = nn.InstanceNorm2d(planes) - if not stride == 1: - self.norm3 = nn.InstanceNorm2d(planes) - - elif norm_fn == "none": - self.norm1 = nn.Sequential() - self.norm2 = nn.Sequential() - if not stride == 1: - self.norm3 = nn.Sequential() - - if stride == 1: - self.downsample = None - - else: - self.downsample = nn.Sequential( - nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3 - ) - - def forward(self, x): - y = x - y = self.relu(self.norm1(self.conv1(y))) - y = self.relu(self.norm2(self.conv2(y))) - - if self.downsample is not None: - x = self.downsample(x) - - return self.relu(x + y) - - -class BottleneckBlock(nn.Module): - def __init__(self, in_planes, planes, norm_fn="group", stride=1): - super(BottleneckBlock, self).__init__() - - self.conv1 = nn.Conv2d(in_planes, planes // 4, kernel_size=1, padding=0) - self.conv2 = nn.Conv2d( - planes // 4, planes // 4, kernel_size=3, padding=1, stride=stride - ) - self.conv3 = nn.Conv2d(planes // 4, planes, kernel_size=1, padding=0) - self.relu = nn.ReLU(inplace=True) - - num_groups = planes // 8 - - if norm_fn == "group": - self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // 4) - self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // 4) - self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) - if not stride == 1: - self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) - - elif norm_fn == "batch": - self.norm1 = nn.BatchNorm2d(planes // 4) - self.norm2 = nn.BatchNorm2d(planes // 4) - self.norm3 = nn.BatchNorm2d(planes) - if not stride == 1: - self.norm4 = nn.BatchNorm2d(planes) - - elif norm_fn == "instance": - self.norm1 = nn.InstanceNorm2d(planes // 4) - self.norm2 = nn.InstanceNorm2d(planes // 4) - self.norm3 = nn.InstanceNorm2d(planes) - if not stride == 1: - self.norm4 = nn.InstanceNorm2d(planes) - - elif norm_fn == "none": - self.norm1 = nn.Sequential() - self.norm2 = nn.Sequential() - self.norm3 = nn.Sequential() - if not stride == 1: - self.norm4 = nn.Sequential() - - if stride == 1: - self.downsample = None - - else: - self.downsample = nn.Sequential( - nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4 - ) - - def forward(self, x): - y = x - y = self.relu(self.norm1(self.conv1(y))) - y = self.relu(self.norm2(self.conv2(y))) - y = self.relu(self.norm3(self.conv3(y))) - - if self.downsample is not None: - x = self.downsample(x) - - return self.relu(x + y) - - -class BasicEncoder(nn.Module): - def __init__(self, output_dim=128, norm_fn="batch", dropout=0.0): - super(BasicEncoder, self).__init__() - self.norm_fn = norm_fn - - if self.norm_fn == "group": - self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) - - elif self.norm_fn == "batch": - self.norm1 = nn.BatchNorm2d(64) - - elif self.norm_fn == "instance": - self.norm1 = nn.InstanceNorm2d(64) - - elif self.norm_fn == "none": - self.norm1 = nn.Sequential() - - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) - self.relu1 = nn.ReLU(inplace=True) - - self.in_planes = 64 - self.layer1 = self._make_layer(64, stride=1) - self.layer2 = self._make_layer(96, stride=2) - self.layer3 = self._make_layer(128, stride=2) - - # output convolution - self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1) - - self.dropout = None - if dropout > 0: - self.dropout = nn.Dropout2d(p=dropout) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") - elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): - if m.weight is not None: - nn.init.constant_(m.weight, 1) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - def _make_layer(self, dim, stride=1): - layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) - layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) - layers = (layer1, layer2) - - self.in_planes = dim - return nn.Sequential(*layers) - - def forward(self, x): - # if input is list, combine batch dimension - is_list = isinstance(x, tuple) or isinstance(x, list) - if is_list: - batch_dim = x[0].shape[0] - x = torch.cat(x, dim=0) - - x = self.conv1(x) - x = self.norm1(x) - x = self.relu1(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - - x = self.conv2(x) - - if self.training and self.dropout is not None: - x = self.dropout(x) - - if is_list: - x = torch.split(x, [batch_dim, batch_dim], dim=0) - - return x - - -class BasicEncoder1(nn.Module): - def __init__(self, output_dim=128, norm_fn="batch", dropout=0.0): - super(BasicEncoder1, self).__init__() - self.norm_fn = norm_fn - - if self.norm_fn == "group": - self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) - - elif self.norm_fn == "batch": - self.norm1 = nn.BatchNorm2d(64) - - elif self.norm_fn == "instance": - self.norm1 = nn.InstanceNorm2d(64) - - elif self.norm_fn == "none": - self.norm1 = nn.Sequential() - - self.conv1 = nn.Conv2d(2, 64, kernel_size=7, stride=2, padding=3) - self.relu1 = nn.ReLU(inplace=True) - - self.in_planes = 64 - self.layer1 = self._make_layer(64, stride=1) - self.layer2 = self._make_layer(96, stride=2) - self.layer3 = self._make_layer(128, stride=2) - - # output convolution - self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1) - - self.dropout = None - if dropout > 0: - self.dropout = nn.Dropout2d(p=dropout) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") - elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): - if m.weight is not None: - nn.init.constant_(m.weight, 1) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - def _make_layer(self, dim, stride=1): - layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) - layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) - layers = (layer1, layer2) - - self.in_planes = dim - return nn.Sequential(*layers) - - def forward(self, x): - # if input is list, combine batch dimension - is_list = isinstance(x, tuple) or isinstance(x, list) - if is_list: - batch_dim = x[0].shape[0] - x = torch.cat(x, dim=0) - - x = self.conv1(x) - x = self.norm1(x) - x = self.relu1(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - - x = self.conv2(x) - - if self.training and self.dropout is not None: - x = self.dropout(x) - - if is_list: - x = torch.split(x, [batch_dim, batch_dim], dim=0) - - return x - - -class SmallEncoder(nn.Module): - def __init__(self, output_dim=128, norm_fn="batch", dropout=0.0): - super(SmallEncoder, self).__init__() - self.norm_fn = norm_fn - - if self.norm_fn == "group": - self.norm1 = nn.GroupNorm(num_groups=8, num_channels=32) - - elif self.norm_fn == "batch": - self.norm1 = nn.BatchNorm2d(32) - - elif self.norm_fn == "instance": - self.norm1 = nn.InstanceNorm2d(32) - - elif self.norm_fn == "none": - self.norm1 = nn.Sequential() - - self.conv1 = nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3) - self.relu1 = nn.ReLU(inplace=True) - - self.in_planes = 32 - self.layer1 = self._make_layer(32, stride=1) - self.layer2 = self._make_layer(64, stride=2) - self.layer3 = self._make_layer(96, stride=2) - - self.dropout = None - if dropout > 0: - self.dropout = nn.Dropout2d(p=dropout) - - self.conv2 = nn.Conv2d(96, output_dim, kernel_size=1) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") - elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): - if m.weight is not None: - nn.init.constant_(m.weight, 1) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - def _make_layer(self, dim, stride=1): - layer1 = BottleneckBlock(self.in_planes, dim, self.norm_fn, stride=stride) - layer2 = BottleneckBlock(dim, dim, self.norm_fn, stride=1) - layers = (layer1, layer2) - - self.in_planes = dim - return nn.Sequential(*layers) - - def forward(self, x): - # if input is list, combine batch dimension - is_list = isinstance(x, tuple) or isinstance(x, list) - if is_list: - batch_dim = x[0].shape[0] - x = torch.cat(x, dim=0) - - x = self.conv1(x) - x = self.norm1(x) - x = self.relu1(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - x = self.conv2(x) - - if self.training and self.dropout is not None: - x = self.dropout(x) - - if is_list: - x = torch.split(x, [batch_dim, batch_dim], dim=0) - - return x - - -################################################## -# RFR is implemented based on RAFT optical flow # -################################################## - - -def backwarp(img, flow): - _, _, H, W = img.size() - - u = flow[:, 0, :, :] - v = flow[:, 1, :, :] - - gridX, gridY = np.meshgrid(np.arange(W), np.arange(H)) - - gridX = torch.tensor( - gridX, - requires_grad=False, - ).cuda() - gridY = torch.tensor( - gridY, - requires_grad=False, - ).cuda() - x = gridX.unsqueeze(0).expand_as(u).float() + u - y = gridY.unsqueeze(0).expand_as(v).float() + v - # range -1 to 1 - x = 2 * (x / (W - 1) - 0.5) - y = 2 * (y / (H - 1) - 0.5) - # stacking X and Y - grid = torch.stack((x, y), dim=3) - # Sample pixels using bilinear interpolation. - imgOut = torch.nn.functional.grid_sample(img, grid, align_corners=True) - - return imgOut - - -class ErrorAttention(nn.Module): - """A three-layer network for predicting mask""" - - def __init__(self, input, output): - super(ErrorAttention, self).__init__() - self.conv1 = nn.Conv2d(input, 32, 5, padding=2) - self.conv2 = nn.Conv2d(32, 32, 3, padding=1) - self.conv3 = nn.Conv2d(38, output, 3, padding=1) - self.prelu1 = nn.PReLU() - self.prelu2 = nn.PReLU() - - def forward(self, x1): - x = self.prelu1(self.conv1(x1)) - x = self.prelu2(torch.cat([self.conv2(x), x1], dim=1)) - x = self.conv3(x) - return x - - -class RFR(nn.Module): - def __init__(self, args): - super(RFR, self).__init__() - self.attention2 = ErrorAttention(6, 1) - self.hidden_dim = hdim = 128 - self.context_dim = cdim = 128 - args.corr_levels = 4 - args.corr_radius = 4 - args.dropout = 0 - self.args = args - - # feature network, context network, and update block - self.fnet = BasicEncoder(output_dim=256, norm_fn="none", dropout=args.dropout) - # self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='none', dropout=args.dropout) - self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim) - - def freeze_bn(self): - for m in self.modules(): - if isinstance(m, nn.BatchNorm2d): - m.eval() - - def initialize_flow(self, img): - """Flow is represented as difference between two coordinate grids flow = coords1 - coords0""" - N, C, H, W = img.shape - coords0 = coords_grid(N, H // 8, W // 8).to(img.device) - coords1 = coords_grid(N, H // 8, W // 8).to(img.device) - - # optical flow computed as difference: flow = coords1 - coords0 - return coords0, coords1 - - def upsample_flow(self, flow, mask): - """Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination""" - N, _, H, W = flow.shape - mask = mask.view(N, 1, 9, 8, 8, H, W) - mask = torch.softmax(mask, dim=2) - - up_flow = F.unfold(8 * flow, [3, 3], padding=1) - up_flow = up_flow.view(N, 2, 9, 1, 1, H, W) - - up_flow = torch.sum(mask * up_flow, dim=2) - up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) - return up_flow.reshape(N, 2, 8 * H, 8 * W) - - def forward( - self, image1, image2, iters=12, flow_init=None, upsample=True, test_mode=False - ): - H, W = image1.size()[2:4] - H8 = H // 8 * 8 - W8 = W // 8 * 8 - - if flow_init is not None: - flow_init_resize = F.interpolate( - flow_init, size=(H8 // 8, W8 // 8), mode="nearest" - ) - - flow_init_resize[:, :1] = ( - flow_init_resize[:, :1].clone() * (W8 // 8 * 1.0) / flow_init.size()[3] - ) - flow_init_resize[:, 1:] = ( - flow_init_resize[:, 1:].clone() * (H8 // 8 * 1.0) / flow_init.size()[2] - ) - - if not hasattr(self.args, "not_use_rfr_mask") or ( - hasattr(self.args, "not_use_rfr_mask") - and (not self.args.not_use_rfr_mask) - ): - im18 = F.interpolate(image1, size=(H8 // 8, W8 // 8), mode="bilinear") - im28 = F.interpolate(image2, size=(H8 // 8, W8 // 8), mode="bilinear") - - warp21 = backwarp(im28, flow_init_resize) - error21 = torch.sum(torch.abs(warp21 - im18), dim=1, keepdim=True) - # print('errormin', error21.min(), error21.max()) - f12init = ( - torch.exp( - -self.attention2( - torch.cat([im18, error21, flow_init_resize], dim=1) - ) - ** 2 - ) - * flow_init_resize - ) - else: - flow_init_resize = None - flow_init = torch.zeros( - image1.size()[0], 2, image1.size()[2] // 8, image1.size()[3] // 8 - ).cuda() - error21 = torch.zeros( - image1.size()[0], 1, image1.size()[2] // 8, image1.size()[3] // 8 - ).cuda() - - f12_init = flow_init - # print('None inital flow!') - - image1 = F.interpolate(image1, size=(H8, W8), mode="bilinear") - image2 = F.interpolate(image2, size=(H8, W8), mode="bilinear") - - f12s, f12, f12_init = self.forward_pred( - image1, image2, iters, flow_init_resize, upsample, test_mode - ) - - if hasattr(self.args, "requires_sq_flow") and self.args.requires_sq_flow: - for ii in range(len(f12s)): - f12s[ii] = F.interpolate(f12s[ii], size=(H, W), mode="bilinear") - f12s[ii][:, :1] = f12s[ii][:, :1].clone() / (1.0 * W8) * W - f12s[ii][:, 1:] = f12s[ii][:, 1:].clone() / (1.0 * H8) * H - if self.training: - return f12s - else: - return [f12s[-1]], f12_init - else: - f12[:, :1] = f12[:, :1].clone() / (1.0 * W8) * W - f12[:, 1:] = f12[:, 1:].clone() / (1.0 * H8) * H - - f12 = F.interpolate(f12, size=(H, W), mode="bilinear") - # print('wo!!') - return ( - f12, - f12_init, - error21, - ) - - def forward_pred( - self, image1, image2, iters=12, flow_init=None, upsample=True, test_mode=False - ): - """Estimate optical flow between pair of frames""" - - image1 = image1.contiguous() - image2 = image2.contiguous() - - hdim = self.hidden_dim - cdim = self.context_dim - - # run the feature network - with autocast(device.type, enabled=self.args.mixed_precision): - fmap1, fmap2 = self.fnet([image1, image2]) - fmap1 = fmap1.float() - fmap2 = fmap2.float() - corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius) - - # run the context network - with autocast(device.type, enabled=self.args.mixed_precision): - cnet = self.fnet(image1) - net, inp = torch.split(cnet, [hdim, cdim], dim=1) - net = torch.tanh(net) - inp = torch.relu(inp) - - coords0, coords1 = self.initialize_flow(image1) - - if flow_init is not None: - coords1 = coords1 + flow_init - - flow_predictions = [] - for itr in range(iters): - coords1 = coords1.detach() - if itr == 0: - if flow_init is not None: - coords1 = coords1 + flow_init - corr = corr_fn(coords1) # index correlation volume - - flow = coords1 - coords0 - with autocast(device.type, enabled=self.args.mixed_precision): - net, up_mask, delta_flow = self.update_block(net, inp, corr, flow) - - # F(t+1) = F(t) + \Delta(t) - coords1 = coords1 + delta_flow - - # upsample predictions - if up_mask is None: - flow_up = upflow8(coords1 - coords0) - else: - flow_up = self.upsample_flow(coords1 - coords0, up_mask) - - flow_predictions.append(flow_up) - - return flow_predictions, flow_up, flow_init - -####################### WARPING ####################### - - -# expects batched tensors, considered low-level operation -# img: bs, ch, h, w -# flow: bs, xy (pix displace), h, w -def flow_backwarp( - img, flow, resample="bilinear", padding_mode="border", align_corners=False -): - if len(img.shape) != 4: - img = img[None,] - if len(flow.shape) != 4: - flow = flow[None,] - q = ( - 2 - * flow - / torch.tensor( - [ - flow.shape[-2], - flow.shape[-1], - ], - device=flow.device, - dtype=torch.float, - )[None, :, None, None] - ) - q = q + torch.stack( - torch.meshgrid( - torch.linspace(-1, 1, flow.shape[-2]), - torch.linspace(-1, 1, flow.shape[-1]), - ) - )[ - None, - ].to( - flow.device - ) - if img.dtype != q.dtype: - img = img.type(q.dtype) - - return nn.functional.grid_sample( - img, - q.flip(dims=(1,)).permute(0, 2, 3, 1), - mode=resample, # nearest, bicubic, bilinear - padding_mode=padding_mode, # border, zeros, reflection - align_corners=align_corners, - ) - - -backwarp = flow_warp = flow_backwarp - - -# mode: sum, avg, lin, softmax -# lin/softmax w/out metric defaults to avg -# must use gpu, move back to cpu if retain_device -# typical metric: -20 * | img0 - backwarp(img1,flow) | -# From Fannovel16: Changed mode params for common ops. -def flow_forewarp( - img, flow, mode="average", metric=None, mask=False, retain_device=True -): - # setup - #if mode == "sum": - # mode = "summation" - #elif mode == "avg": - # mode = "average" - if mode in ["lin", "linear"]: - #mode = "linear" if metric is not None else "average" - mode = "linear" if metric is not None else "avg" - elif mode in ["sm", "softmax"]: - #mode = "softmax" if metric is not None else "average" - mode = "soft" if metric is not None else "avg" - if len(img.shape) != 4: - img = img[None,] - if len(flow.shape) != 4: - flow = flow[None,] - if metric is not None and len(metric.shape) != 4: - metric = metric[None,] - flow = flow.flip(dims=(1,)) - if img.dtype != torch.float32: - img = img.type(torch.float32) - if flow.dtype != torch.float32: - flow = flow.type(torch.float32) - if metric is not None and metric.dtype != torch.float32: - metric = metric.type(torch.float32) - - # move to gpu if necessary - assert img.device == flow.device - if metric is not None: - assert img.device == metric.device - was_cpu = img.device.type == "cpu" - if was_cpu: - img = img.to("cuda") - flow = flow.to("cuda") - if metric is not None: - metric = metric.to("cuda") - - # add mask - if mask: - bs, ch, h, w = img.shape - img = torch.cat( - [img, torch.ones(bs, 1, h, w, dtype=img.dtype, device=img.device)], dim=1 - ) - - # forward, move back to cpu if desired - ans = FunctionSoftsplat(img, flow, metric, mode) - if was_cpu and retain_device: - ans = ans.cpu() - return ans - - -forewarp = flow_forewarp - - -# resizing utility -def flow_resize(flow, size, mode="nearest", align_corners=False): - # flow: bs,xy,h,w - size = pixel_ij(size, rounding=True) - if flow.dtype != torch.float: - flow = flow.float() - if len(flow.shape) == 3: - flow = flow[None,] - if flow.shape[-2:] == size: - return flow - return ( - nn.functional.interpolate( - flow, - size=size, - mode=mode, - align_corners=align_corners if mode != "nearest" else None, - ) - * torch.tensor( - [b / a for a, b in zip(flow.shape[-2:], size)], - device=flow.device, - )[None, :, None, None] - ) - - -####################### TRADITIONAL ####################### - -# dense -_lucaskanade = lambda a, b: np.moveaxis( - cv2.optflow.calcOpticalFlowSparseToDense( - a, - b, # grid_step=5, sigma=0.5, - ), - 2, - 0, -)[ - None, -] -_farneback = lambda a, b: np.moveaxis( - cv2.calcOpticalFlowFarneback( - a, - b, - None, - 0.6, - 3, - 25, - 7, - 5, - 1.2, - cv2.OPTFLOW_FARNEBACK_GAUSSIAN, - ), - 2, - 0, -)[ - None, -] -_dtvl1_ = cv2.optflow.createOptFlow_DualTVL1() -_dtvl1 = lambda a, b: np.moveaxis( - _dtvl1_.calc( - a, - b, - None, - ), - 2, - 0, -)[ - None, -] -_simple = lambda a, b: np.moveaxis( - cv2.optflow.calcOpticalFlowSF( - a, - b, - 3, - 5, - 5, - ), - 2, - 0, -)[ - None, -] -_pca_ = cv2.optflow.createOptFlow_PCAFlow() -_pca = lambda a, b: np.moveaxis( - _pca_.calc( - a, - b, - None, - ), - 2, - 0, -)[ - None, -] -_drlof = lambda a, b: np.moveaxis( - cv2.optflow.calcOpticalFlowDenseRLOF( - a, - b, - None, - ), - 2, - 0, -)[ - None, -] -_deepflow_ = cv2.optflow.createOptFlow_DeepFlow() -_deepflow = lambda a, b: np.moveaxis( - _deepflow_.calc( - a, - b, - None, - ), - 2, - 0, -)[ - None, -] - - -def cv2flow(a, b, method="lucaskanade", back=False): - if method == "lucaskanade": - f = _lucaskanade - a = a.convert("L").cv2() - b = b.convert("L").cv2() - elif method == "farneback": - f = _farneback - a = a.convert("L").cv2() - b = b.convert("L").cv2() - elif method == "dtvl1": - f = _dtvl1 - a = a.convert("L").cv2() - b = b.convert("L").cv2() - elif method == "simple": - f = _simple - a = a.convert("RGB").cv2() - b = b.convert("RGB").cv2() - elif method == "pca": - f = _pca - a = a.convert("L").cv2() - b = b.convert("L").cv2() - elif method == "drlof": - f = _drlof - a = a.convert("RGB").cv2() - b = b.convert("RGB").cv2() - elif method == "deepflow": - f = _deepflow - a = a.convert("L").cv2() - b = b.convert("L").cv2() - else: - assert 0 - ans = f(b, a) - if back: - ans = np.concatenate( - [ - ans, - f(a, b), - ] - ) - return torch.tensor(ans).flip(dims=(1,)) - - -####################### FLOWNET2 ####################### - - -def flownet2(img_a, img_b, mode="shm", back=False): - # package - url = f"http://localhost:8109/get-flow" - if mode == "shm": - t = time.time() - fn_a = img_a.save(mkfile(f"/dev/shm/_flownet2/{t}/img_a.png")) - fn_b = img_b.save(mkfile(f"/dev/shm/_flownet2/{t}/img_b.png")) - elif mode == "net": - assert False, "not impl" - q = u2d.img2uri(img.pil("RGB")) - q.decode() - resp = requests.get( - url, - params={ - "img_a": fn_a, - "img_b": fn_b, - "mode": mode, - "back": back, - # 'vis': vis, - }, - ) - - # return - ans = {"response": resp} - if resp.status_code == 200: - j = resp.json() - ans["time"] = j["time"] - ans["output"] = { - "flow": torch.tensor(load(j["fn_flow"])), - } - # if vis: - # ans['output']['vis'] = I(j['fn_vis']) - if mode == "shm": - shutil.rmtree(f"/dev/shm/_flownet2/{t}") - return ans - - -####################### VISUALIZATION ####################### - - -class Gridnet(nn.Module): - def __init__(self, channels_0, channels_1, channels_2, total_dropout_p, depth): - super().__init__() - self.channels_0 = ch0 = channels_0 - self.channels_1 = ch1 = channels_1 - self.channels_2 = ch2 = channels_2 - self.total_dropout_p = p = total_dropout_p - self.depth = depth - self.encoders = nn.ModuleList( - [GridnetEncoder(ch0, ch1, ch2) for i in range(self.depth)] - ) - self.decoders = nn.ModuleList( - [GridnetDecoder(ch0, ch1, ch2) for i in range(self.depth)] - ) - self.total_dropout = GridnetTotalDropout(p) - return - - def forward(self, x): - for e, enc in enumerate(self.encoders): - t = [self.total_dropout(i) for i in t] if e != 0 else x - t = enc(t) - for d, dec in enumerate(self.decoders): - t = [self.total_dropout(i) for i in t] - t = dec(t) - return t - - -class GridnetEncoder(nn.Module): - def __init__(self, channels_0, channels_1, channels_2): - super().__init__() - self.channels_0 = ch0 = channels_0 - self.channels_1 = ch1 = channels_1 - self.channels_2 = ch2 = channels_2 - self.resnet_0 = GridnetResnet(ch0) - self.resnet_1 = GridnetResnet(ch1) - self.resnet_2 = GridnetResnet(ch2) - self.downsample_01 = GridnetDownsample(ch0, ch1) - self.downsample_12 = GridnetDownsample(ch1, ch2) - return - - def forward(self, x): - out = [ - None, - ] * 3 - out[0] = self.resnet_0(x[0]) - out[1] = self.resnet_1(x[1]) + self.downsample_01(out[0]) - out[2] = self.resnet_2(x[2]) + self.downsample_12(out[1]) - return out - - -class GridnetDecoder(nn.Module): - def __init__(self, channels_0, channels_1, channels_2): - super().__init__() - self.channels_0 = ch0 = channels_0 - self.channels_1 = ch1 = channels_1 - self.channels_2 = ch2 = channels_2 - self.resnet_0 = GridnetResnet(ch0) - self.resnet_1 = GridnetResnet(ch1) - self.resnet_2 = GridnetResnet(ch2) - self.upsample_10 = GridnetUpsample(ch1, ch0) - self.upsample_21 = GridnetUpsample(ch2, ch1) - return - - def forward(self, x): - out = [ - None, - ] * 3 - out[2] = self.resnet_2(x[2]) - out[1] = self.resnet_1(x[1]) + self.upsample_21(out[2]) - out[0] = self.resnet_0(x[0]) + self.upsample_10(out[1]) - return out - - -class GridnetConverter(nn.Module): - def __init__(self, channels_in, channels_out): - super().__init__() - self.channels_in = cin = channels_in - self.channels_out = cout = channels_out - self.nets = nn.ModuleList( - [ - nn.Sequential( - nn.PReLU(a), - nn.Conv2d(a, b, kernel_size=1, padding=0), - nn.BatchNorm2d(b), - ) - for a, b in zip(cin, cout) - ] - ) - return - - def forward(self, x): - return [m(q) for m, q in zip(self.nets, x)] - - -class GridnetResnet(nn.Module): - def __init__(self, channels): - super().__init__() - self.channels = ch = channels - self.net = nn.Sequential( - nn.PReLU(ch), - nn.Conv2d(ch, ch, kernel_size=3, padding=1), - nn.BatchNorm2d(ch), - nn.PReLU(ch), - nn.Conv2d(ch, ch, kernel_size=3, padding=1), - nn.BatchNorm2d(ch), - ) - return - - def forward(self, x): - return x + self.net(x) - - -class GridnetDownsample(nn.Module): - def __init__(self, channels_in, channels_out): - super().__init__() - self.channels_in = chin = channels_in - self.channels_out = chout = channels_out - self.net = nn.Sequential( - nn.PReLU(chin), - nn.Conv2d(chin, chin, kernel_size=3, padding=1, stride=2), - nn.BatchNorm2d(chin), - nn.PReLU(chin), - nn.Conv2d(chin, chout, kernel_size=3, padding=1), - nn.BatchNorm2d(chout), - ) - return - - def forward(self, x): - return self.net(x) - - -class GridnetUpsample(nn.Module): - def __init__(self, channels_in, channels_out): - super().__init__() - self.channels_in = chin = channels_in - self.channels_out = chout = channels_out - self.net = nn.Sequential( - nn.Upsample(scale_factor=2, mode="nearest"), - nn.PReLU(chin), - nn.Conv2d(chin, chout, kernel_size=3, padding=1), - nn.BatchNorm2d(chout), - nn.PReLU(chout), - nn.Conv2d(chout, chout, kernel_size=3, padding=1), - nn.BatchNorm2d(chout), - ) - return - - def forward(self, x): - return self.net(x) - - -class GridnetTotalDropout(nn.Module): - def __init__(self, p): - super().__init__() - assert 0 <= p < 1 - self.p = p - self.weight = 1 / (1 - p) - return - - def get_drop(self, x): - d = torch.rand(len(x))[:, None, None, None] < self.p - d = (1 - d.float()).to(x.device) * self.weight - return d - - def forward(self, x, force_drop=None): - if force_drop is True: - ans = x * self.get_drop(x) - elif force_drop is False: - ans = x - else: - if self.training: - ans = x * self.get_drop(x) - else: - ans = x - return ans - - -class Interpolator(nn.Module): - def __init__(self, size, mode="bilinear"): - super().__init__() - self.size = size - self.mode = mode - return - - def forward(self, x, is_flow=False): - if x.shape[-2] == self.size: - return x - if len(x.shape) == 4: - # bs,ch,h,w - bs, ch, h, w = x.shape - ans = nn.functional.interpolate( - x, - size=self.size, - mode=self.mode, - align_corners=(False, None)[self.mode == "nearest"], - ) - if is_flow: - ans = ( - ans - * torch.tensor( - [b / a for a, b in zip((h, w), self.size)], - device=ans.device, - )[None, :, None, None] - ) - return ans - elif len(x.shape) == 5: - # bs,k,ch,h,w (merge bs and k) - bs, k, ch, h, w = x.shape - return self.forward( - x.view(bs * k, ch, h, w), - is_flow=is_flow, - ).view(bs, k, ch, *self.size) - else: - assert 0 - - -###################### CANNY ###################### - - -def canny(img, a=100, b=200): - img = I(img).convert("L") - return I(cv2.Canny(img.cv2(), a, b)) - - -# https://www.pyimagesearch.com/2015/04/06/zero-parameter-automatic-canny-edge-detection-with-python-and-opencv/ -def canny_pis(img, sigma=0.33): - # compute the median of the single channel pixel intensities - img = I(img).convert("L").uint8(ch_last=False) - v = np.median(img) - # apply automatic Canny edge detection using the computed median - lower = int(max(0, (1.0 - sigma) * v)) - upper = int(min(255, (1.0 + sigma) * v)) - edged = cv2.Canny(img[0], lower, upper) - # return the edged image - return I(edged) - - -# https://en.wikipedia.org/wiki/Otsu%27s_method -def canny_otsu(img): - img = I(img).convert("L").uint8(ch_last=False) - high, _ = cv2.threshold(img[0], 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) - low = 0.5 * high - return I(cv2.Canny(img[0], low, high)) - - -def xdog(img, t=1.0, epsilon=0.04, phi=100, sigma=3, k=1.6): - img = I(img).convert("L").uint8(ch_last=False) - grey = np.asarray(img, dtype=np.float32) - g0 = scipy.ndimage.gaussian_filter(grey, sigma) - g1 = scipy.ndimage.gaussian_filter(grey, sigma * k) - - # ans = ((1+p) * g0 - p * g1) / 255 - ans = (g0 - t * g1) / 255 - ans = 1 + np.tanh(phi * (ans - epsilon)) * (ans < epsilon) - return ans - - -def dog(img, t=1.0, sigma=1.0, k=1.6, epsilon=0.01, kernel_factor=4, clip=True): - img = I(img).convert("L").tensor()[None] - kern0 = max(2 * int(sigma * kernel_factor) + 1, 3) - kern1 = max(2 * int(sigma * k * kernel_factor) + 1, 3) - g0 = kornia.filters.gaussian_blur2d( - img, - (kern0, kern0), - (sigma, sigma), - border_type="replicate", - ) - g1 = kornia.filters.gaussian_blur2d( - img, - (kern1, kern1), - (sigma * k, sigma * k), - border_type="replicate", - ) - ans = 0.5 + t * (g1 - g0) - epsilon - ans = ans.clip(0, 1) if clip else ans - return ans[0].numpy() - - -# input: (bs,rgb(a),h,w) or (bs,1,h,w) -# returns: (bs,1,h,w) -def batch_dog(img, t=1.0, sigma=1.0, k=1.6, epsilon=0.01, kernel_factor=4, clip=True): - # to grayscale if needed - bs, ch, h, w = img.shape - if ch in [3, 4]: - img = kornia.color.rgb_to_grayscale(img[:, :3]) - else: - assert ch == 1 - - # calculate dog - kern0 = max(2 * int(sigma * kernel_factor) + 1, 3) - kern1 = max(2 * int(sigma * k * kernel_factor) + 1, 3) - g0 = kornia.filters.gaussian_blur2d( - img, - (kern0, kern0), - (sigma, sigma), - border_type="replicate", - ) - g1 = kornia.filters.gaussian_blur2d( - img, - (kern1, kern1), - (sigma * k, sigma * k), - border_type="replicate", - ) - ans = 0.5 + t * (g1 - g0) - epsilon - ans = ans.clip(0, 1) if clip else ans - return ans - - -############### DERIVED DISTANCES ############### - -# input: (bs,h,w) or (bs,1,h,w) -# returns: (bs,) -# normalized s.t. metric is same across proportional image scales - - -# average of two asymmetric distances -# normalized by diameter and area -def batch_chamfer_distance(gt, pred, block=1024, return_more=False): - t = batch_chamfer_distance_t(gt, pred, block=block) - p = batch_chamfer_distance_p(gt, pred, block=block) - cd = (t + p) / 2 - return cd - - -def batch_chamfer_distance_t(gt, pred, block=1024, return_more=False): - assert gt.device == pred.device and gt.shape == pred.shape - bs, h, w = gt.shape[0], gt.shape[-2], gt.shape[-1] - dpred = batch_edt(pred, block=block) - cd = (gt * dpred).float().mean((-2, -1)) / np.sqrt(h**2 + w**2) - if len(cd.shape) == 2: - assert cd.shape[1] == 1 - cd = cd.squeeze(1) - return cd - - -def batch_chamfer_distance_p(gt, pred, block=1024, return_more=False): - assert gt.device == pred.device and gt.shape == pred.shape - bs, h, w = gt.shape[0], gt.shape[-2], gt.shape[-1] - dgt = batch_edt(gt, block=block) - cd = (pred * dgt).float().mean((-2, -1)) / np.sqrt(h**2 + w**2) - if len(cd.shape) == 2: - assert cd.shape[1] == 1 - cd = cd.squeeze(1) - return cd - - -# normalized by diameter -# always between [0,1] -def batch_hausdorff_distance(gt, pred, block=1024, return_more=False): - assert gt.device == pred.device and gt.shape == pred.shape - bs, h, w = gt.shape[0], gt.shape[-2], gt.shape[-1] - dgt = batch_edt(gt, block=block) - dpred = batch_edt(pred, block=block) - hd = torch.stack( - [ - (dgt * pred).amax(dim=(-2, -1)), - (dpred * gt).amax(dim=(-2, -1)), - ] - ).amax(dim=0).float() / np.sqrt(h**2 + w**2) - if len(hd.shape) == 2: - assert hd.shape[1] == 1 - hd = hd.squeeze(1) - return hd - - -#################### UTILITIES #################### - - -def reset_parameters(model): - for layer in model.children(): - if hasattr(layer, "reset_parameters"): - layer.reset_parameters() - return model - - -def channel_squeeze(x, dim=1): - a = x.shape[:dim] - b = x.shape[dim + 2 :] - return x.reshape(*a, -1, *b) - - -def channel_unsqueeze(x, shape, dim=1): - a = x.shape[:dim] - b = x.shape[dim + 1 :] - return x.reshape(*a, *shape, *b) - - -def default_collate(items, device=None): - return to(dict(torch.utils.data.dataloader.default_collate(items)), device) - - -def to(x, device): - if device is None: - return x - if issubclass(x.__class__, dict): - return dict( - { - k: v.to(device) if isinstance(v, torch.Tensor) else v - for k, v in x.items() - } - ) - if isinstance(x, torch.Tensor): - return x.to(device) - if isinstance(x, np.ndarray): - return torch.tensor(x).to(device) - assert 0, "data not understood" - - -################ PARSING ################ - -from argparse import Namespace - -# args: all args -# bargs: base args -# pargs: data processing args -# largs: data loading args -# margs: model args -# targs: training args - - -# typically used to read dataset filters -def read_filter(fn, cast=None, sort=True, sort_key=None): - if cast is None: - cast = lambda x: x - ans = [cast(line) for line in read(fn).split("\n") if line != ""] - if sort: - return sorted(ans, key=sort_key) - else: - return ans - - -################ FILE MANAGEMENT ################ - - -def mkfile(fn, parents=True, exist_ok=True): - dn = "/".join(fn.split("/")[:-1]) - mkdir(dn, parents=parents, exist_ok=exist_ok) - return fn - - -def mkdir(dn, parents=True, exist_ok=True): - pathlib.Path(dn).mkdir(parents=parents, exist_ok=exist_ok) - return dn if (not dn[-1] == "/" or dn == "/") else dn[:-1] - - -def fstrip(fn, return_more=False): - dspl = fn.split("/") - dn = "/".join(dspl[:-1]) if len(dspl) > 1 else "." - fn = dspl[-1] - fspl = fn.split(".") - if len(fspl) == 1: - bn = fspl[0] - ext = "" - else: - bn = ".".join(fspl[:-1]) - ext = fspl[-1] - if return_more: - return Namespace( - dn=dn, - fn=fn, - path=f"{dn}/{fn}", - bn_path=f"{dn}/{bn}", - bn=bn, - ext=ext, - ) - else: - return bn - - -def read(fn, mode="r"): - with open(fn, mode) as handle: - return handle.read() - - -def write(text, fn, mode="w"): - mkfile(fn, parents=True, exist_ok=True) - with open(fn, mode) as handle: - return handle.write(text) - - -import pickle - - -def dump(obj, fn, mode="wb"): - mkfile(fn, parents=True, exist_ok=True) - with open(fn, mode) as handle: - return pickle.dump(obj, handle) - - -def load(fn, mode="rb"): - with open(fn, mode) as handle: - return pickle.load(handle) - - -import json - - -def jwrite(x, fn, mode="w", indent="\t", sort_keys=False): - mkfile(fn, parents=True, exist_ok=True) - with open(fn, mode) as handle: - return json.dump(x, handle, indent=indent, sort_keys=sort_keys) - - -def jread(fn, mode="r"): - with open(fn, mode) as handle: - return json.load(handle) - - -try: - import yaml - - def ywrite(x, fn, mode="w", default_flow_style=False): - mkfile(fn, parents=True, exist_ok=True) - with open(fn, mode) as handle: - return yaml.dump(x, handle, default_flow_style=default_flow_style) - - def yread(fn, mode="r"): - with open(fn, mode) as handle: - return yaml.safe_load(handle) - -except: - pass - -try: - import pyunpack -except: - pass - -try: - import mysql - import mysql.connector -except: - pass - - -################ MISC ################ - -hakase = "./env/__hakase__.jpg" -if not os.path.isfile(hakase): - hakase = "./__env__/__hakase__.jpg" - - -def mem(units="m"): - return ( - psProcess(os.getpid()).memory_info().rss - / { - "b": 1, - "k": 1e3, - "m": 1e6, - "g": 1e9, - "t": 1e12, - }[units[0].lower()] - ) - - -def chunk(array, length, colwise=True): - if colwise: - return [array[i : i + length] for i in range(0, len(array), length)] - else: - return chunk(array, int(math.ceil(len(array) / length)), colwise=True) - - -def classtree(x): - return inspect.getclasstree(inspect.getmro(x)) - - -################ AESTHETIC ################ - - -class Table: - def __init__( - self, - table, - delimiter=" ", - orientation="br", - double_colon=True, - ): - self.delimiter = delimiter - self.orientation = orientation - self.t = Table.parse(table, delimiter, orientation, double_colon) - return - - # rendering - def __str__(self): - return self.render() - - def __repr__(self): - return self.render() - - def render(self): - # set up empty entry - empty = ("", Table._spec(self.orientation, transpose=False)) - - # calculate table size - t = copy.deepcopy(self.t) - totalrows = len(t) - totalcols = [len(r) for r in t] - assert min(totalcols) == max(totalcols) - totalcols = totalcols[0] - - # string-ify - for i in range(totalrows): - for j in range(totalcols): - x, s = t[i][j] - sp = s[11] - if sp: - x = eval(f'f"{{{x}{sp}}}"') - Table._put((str(x), s), t, (i, j), empty) - - # expand delimiters - _repl = ( - lambda s: s[:2] + (1, 0, 0, 0, 0) + s[7:10] + (1,) + s[11:] - if s[2] - else s[:2] + (0, 0, 0, 0, 0) + s[7:10] + (1,) + s[11:] - ) - for i, row in enumerate(t): - for j, (x, s_own) in enumerate(row): - # expand delim_up(^) - if s_own[3]: - u, v = i, j - while 0 <= u: - _, s = t[u][v] - if (i, j) != (u, v) and (s[2] and not s[10]): - break - Table._put((x, _repl(s)), t, (u, v), empty) - u -= 1 - - # expand delim_down(v) - if s_own[4]: - u, v = i, j - while u < totalrows: - _, s = t[u][v] - if (i, j) != (u, v) and (s[2] and not s[10]): - break - Table._put((x, _repl(s)), t, (u, v), empty) - u += 1 - - # expand delim_right(>) - if s_own[5]: - u, v = i, j - while v < totalcols: - _, s = t[u][v] - if (i, j) != (u, v) and (s[2] and not s[10]): - break - Table._put((x, _repl(s)), t, (u, v), empty) - v += 1 - - # expand delim_left(<) - if s_own[6]: - u, v = i, j - while 0 <= v: - _, s = t[u][v] - if (i, j) != (u, v) and (s[2] and not s[10]): - break - Table._put((x, _repl(s)), t, (u, v), empty) - v -= 1 - - # justification calculation - widths = [ - 0, - ] * totalcols # j - heights = [ - 0, - ] * totalrows # i - for i, row in enumerate(t): - for j, (x, s) in enumerate(row): - # height caclulation - heights[i] = max(heights[i], x.count("\n")) - - # width calculation; non-delim fillers no contribution - if s[2] or not s[10]: - w = max(len(q) for q in x.split("\n")) - widths[j] = max(widths[j], w) - # no newline ==> height=1 - heights = [h + 1 for h in heights] - - # render table - rend = [] - roff = 0 - for i, row in enumerate(t): - for j, (x, s) in enumerate(row): - w, h = widths[j], heights[i] - - # expand fillers and delimiters - if s[2] or s[10]: - xs = x.split("\n") - xw0 = min(len(l) for l in xs) - xw1 = max(len(l) for l in xs) - xh = len(xs) - if (xw0 == xw1 == w) and (xh == h): - pass - elif xw0 == xw1 == w: - x = "\n".join( - [ - xs[0], - ] - * h - ) - elif xh == h: - x = "\n".join([(l[0] if l else "") * w for l in xs]) - else: - x = x[0] if x else " " - x = "\n".join( - [ - x * w, - ] - * h - ) - - # justify horizontally - x = [l.rjust(w) if s[0] else l.ljust(w) for l in x.split("\n")] - - # justify vertically - plus = [ - " " * w, - ] * (h - len(x)) - x = plus + x if not s[1] else x + plus - - # input to table - for r, xline in enumerate(x): - Table._put(xline, rend, (roff + r, j), None) - roff += h - - # return rendered string - return "\n".join(["".join(r) for r in rend]) - - # parsing - def _spec(s, transpose=False): - if ":" in s: - i = s.index(":") - sp = s[i:] - s = s[:i] - else: - sp = "" - s = s.lower() - return ( - int("r" in s), # 0:: 0:left(l) 1:right(r) - int("t" in s), # 1:: 0:bottom(b) 1:top(t) - int(any([i in s for i in [".", "<", ">", "^", "v"]])), # 2:: delim_here(.) - int("^" in s if not transpose else "<" in s), # 3:: delim_up(^) - int("v" in s if not transpose else ">" in s), # 4:: delim_down(v) - int(">" in s if not transpose else "v" in s), # 5:: delim_right(>) - int("<" in s if not transpose else "^" in s), # 6:: delim_left(<) - int("+" in s), # 7:: subtable(+) - int("-" in s if not transpose else "|" in s), # 8:: subtable_horiz(-) - int("|" in s if not transpose else "-" in s), # 9:: subtable_vert(|) - int("_" in s), # 10:: fill(_); if delim, overwrite; else fit - sp, # 11:: special(:) f-string for numbers - ) - - def _put(obj, t, ij, empty): - i, j = ij - while i >= len(t): - t.append([]) - while j >= len(t[i]): - t[i].append(empty) - t[i][j] = obj - return - - def parse( - table, - delimiter=" ", - orientation="br", - double_colon=True, - ): - # disabling transpose - transpose = False - - # set up empty entry - empty = ("", Table._spec(orientation, transpose)) - - # transpose - t = [] - for i, row in enumerate(table): - for j, item in enumerate(row): - ij = (i, j) if not transpose else (j, i) - if type(item) == tuple and len(item) == 2 and type(item[1]) == str: - item = (item[0], Table._spec(item[1], transpose)) - elif double_colon and type(item) == str and "::" in item: - x, s = item.split("::") - item = (x, Table._spec(s, transpose)) - else: - item = (item, Table._spec(orientation, transpose)) - Table._put(item, t, ij, empty) - - # normalization - maxcol = 0 - maxrow = len(t) - for i, row in enumerate(t): - # take element number into account - maxcol = max(maxcol, len([i for i in row if not i[1][2]])) - - # take subtables into account - for j, (x, s) in enumerate(row): - if s[7]: - r = len(x) - maxrow = max(maxrow, i + r) - c = max(len(q) for q in x) - maxcol = max(maxcol, j + c) - elif s[8]: - c = len(x) - maxcol = max(maxcol, j + c) - elif s[9]: - r = len(x) - maxrow = max(maxrow, i + r) - totalcols = 2 * maxcol + 1 - totalrows = maxrow - t += [[]] * (totalrows - len(t)) - newt = [] - delim = (delimiter, Table._spec("._" + orientation, transpose)) - for i, row in enumerate(t): - wasd = False - tcount = 0 - for j in range(totalcols): - item = t[i][tcount] if tcount < len(t[i]) else empty - isd = item[1][2] - if wasd and isd: - Table._put(empty, newt, (i, j), empty) - wasd = False - elif wasd and not isd: - Table._put(item, newt, (i, j), empty) - tcount += 1 - wasd = False - elif not wasd and isd: - Table._put(item, newt, (i, j), empty) - tcount += 1 - wasd = True - elif not wasd and not isd: - Table._put(delim, newt, (i, j), empty) - wasd = True - t = newt - - # normalization: add dummy last column for delimiter - for row in t: - row.append(empty) - - # expand subtables - delim_cols = [i for i in range(totalcols) if i % 2 == 0] - while True: - # find a table - ij = None - for i, row in enumerate(t): - for j, item in enumerate(row): - st, s = item - if s[7]: - ij = i, j, 7, st, s - break - elif s[8]: - ij = i, j, 8, st, s - break - elif s[9]: - ij = i, j, 9, st, s - break - if ij is not None: - break - if ij is None: - break - - # replace its specs - i, j, k, st, s = ij - s = list(s) - s[7] = s[8] = s[9] = 0 - s = tuple(s) - - # expand it - if k == 7: # 2d table - for x, row in enumerate(st): - for y, obj in enumerate(row): - a = i + x if not transpose else i + y - b = j + 2 * y if not transpose else j + 2 * x - Table._put((obj, s), t, (a, b), None) - if k == 8: # subtable_horiz - for y, obj in enumerate(st): - Table._put((obj, s), t, (i, j + 2 * y), None) - if k == 9: # subtable_vert - for x, obj in enumerate(st): - Table._put((obj, s), t, (i + x, j), None) - - # return, finally - return t - - -class Resnet(nn.Module): - def __init__(self, channels): - super().__init__() - self.channels = ch = channels - self.net = nn.Sequential( - nn.PReLU(ch), - nn.Conv2d(ch, ch, kernel_size=3, padding=1), - nn.BatchNorm2d(ch), - nn.PReLU(ch), - nn.Conv2d(ch, ch, kernel_size=3, padding=1), - nn.BatchNorm2d(ch), - ) - return - - def forward(self, x): - return x + self.net(x) - - -class Synthesizer(nn.Module): - def __init__( - self, size, channels_image, channels_flow, channels_mask, channels_feature - ): - super().__init__() - self.size = size - self.diam = diam(self.size) - self.channels_image = cimg = channels_image - self.channels_flow = cflow = channels_flow - self.channels_mask = cmask = channels_mask - self.channels_feature = cfeat = channels_feature - self.channels = ch = cimg + cflow // 2 + cmask + cfeat - self.interpolator = Interpolator(self.size, mode="bilinear") - self.net = nn.Sequential( - nn.Conv2d(ch + 3, 64, kernel_size=1, padding=0), - Resnet(64), - nn.Sequential( - nn.PReLU(64), - nn.Conv2d(64, 32, kernel_size=3, padding=1), - nn.BatchNorm2d(32), - ), - Resnet(32), - nn.Sequential( - nn.PReLU(32), - nn.Conv2d(32, 16, kernel_size=3, padding=1), - nn.BatchNorm2d(16), - ), - Resnet(16), - nn.Sequential( - nn.PReLU(16), - nn.Conv2d(16, 3, kernel_size=3, padding=1), - ), - ) - return - - def forward(self, images, flows, masks, features, return_more=False): - itp = self.interpolator - images = [ - (images[0] + images[1]) / 2, - ] + images - logimgs = [itp(pixel_logit(i[:, :3])) for i in images] - cat = torch.cat( - [ - *logimgs, - *[itp(f).norm(dim=1, keepdim=True) / self.diam for f in flows], - *[itp(m) for m in masks], - *[itp(f) for f in features], - ], - dim=1, - ) - residual = self.net(cat) - return torch.sigmoid(logimgs[0] + 0.5 * residual), ( - locals() if return_more else None - ) - - -class FlowZMetric(nn.Module): - def __init__(self): - super().__init__() - return - - def forward(self, img0, img1, flow0, flow1, return_more=False): - # B(i0,f0) = i1 - # B(i1,f1) = i0 - # F(x,f0,z0) - # F(x,f1,z1) - img0 = kornia.color.rgb_to_lab(img0[:, :3]) - img1 = kornia.color.rgb_to_lab(img1[:, :3]) - return [ - -0.1 * (img1 - flow_backwarp(img0, flow0)).norm(dim=1, keepdim=True), # z0 - -0.1 * (img0 - flow_backwarp(img1, flow1)).norm(dim=1, keepdim=True), # z1 - ], (locals() if return_more else None) - - -class NEDT(nn.Module): - def __init__(self): - super().__init__() - return - - def forward( - self, - img, - t=2.0, - sigma_factor=1 / 540, - k=1.6, - epsilon=0.01, - kernel_factor=4, - exp_factor=540 / 15, - return_more=False, - ): - with torch.no_grad(): - dog = batch_dog( - img, - t=t, - sigma=img.shape[-2] * sigma_factor, - k=k, - epsilon=epsilon, - kernel_factor=kernel_factor, - clip=False, - ) - edt = batch_edt((dog > 0.5).float()) - ans = 1 - (-edt * exp_factor / max(edt.shape[-2:])).exp() - return ans, (locals() if return_more else None) - - -class HalfWarper(nn.Module): - def __init__(self): - super().__init__() - self.channels_image = 4 * 3 - self.channels_flow = 2 * 2 - self.channels_mask = 2 * 1 - self.channels = self.channels_image + self.channels_flow + self.channels_mask - - def morph_open(self, x, k): - if k == 0: - return x - else: - with torch.no_grad(): - return kornia.morphology.opening(x, torch.ones(k, k, device=x.device)) - - def forward(self, img0, img1, flow0, flow1, z0, z1, k, t=0.5, return_more=False): - # forewarps - flow0_ = (1 - t) * flow0 - flow1_ = t * flow1 - f01 = forewarp(img0, flow1_, mode="sm", metric=z1, mask=True) - f10 = forewarp(img1, flow0_, mode="sm", metric=z0, mask=True) - f01i, f01m = f01[:, :-1], self.morph_open(f01[:, -1:], k=k) - f10i, f10m = f10[:, :-1], self.morph_open(f10[:, -1:], k=k) - - # base guess - base0 = f01m * f01i + (1 - f01m) * f10i - base1 = f10m * f10i + (1 - f10m) * f01i - ans = [ - [ # images - base0, - base1, - f01i, - f10i, - ], - [ # flows - flow0_, - flow1_, - ], - [ # masks - f01m, - f10m, - ], - ] - return ans, (locals() if return_more else None) - - -class ResnetFeatureExtractor(nn.Module): - def __init__(self, inferserve_query, size_in=None): - super().__init__() - self.inferserve_query = iq = inferserve_query - self.size_in = si = size_in - if iq[0] == "torchvision": - # use pytorch pretrained resnet50 - self.base_hparams = None - resnet = tv.models.resnet50(pretrained=True) - - self.resize = T.Resize(256) - self.resnet_preprocess = T.Normalize( - mean=[0.485, 0.456, 0.406], - std=[0.229, 0.224, 0.225], - ) - self.conv1 = resnet.conv1 - self.bn1 = resnet.bn1 - self.relu = resnet.relu # 64ch, 128p (assuming 256p input) - self.maxpool = resnet.maxpool - self.layer1 = resnet.layer1 # 256ch, 64p - self.layer2 = resnet.layer2 # 512ch, 32p - else: - base = userving.infer_model_load(*iq).eval() - self.base_hparams = base.hparams - - self.resize = T.Resize(base.hparams.largs.size) - self.resnet_preprocess = base.resnet_preprocess - self.conv1 = base.resnet.conv1 - self.bn1 = base.resnet.bn1 - self.relu = base.resnet.relu # 64ch, 128p (assuming 256p input) - self.maxpool = base.resnet.maxpool - self.layer1 = base.resnet.layer1 # 256ch, 64p - self.layer2 = base.resnet.layer2 # 512ch, 32p - if self.size_in is None: - self.sizes_out = None - else: - s = self.resize.size - self.sizes_out = [ - pixel_ij( - rescale_dry(si, (s // 2) / si[0]), rounding="ceil" - ), # conv1, 128p - pixel_ij( - rescale_dry(si, (s // 4) / si[0]), rounding="ceil" - ), # layer1, 64p - pixel_ij( - rescale_dry(si, (s // 8) / si[0]), rounding="ceil" - ), # layer2, 32p - ] - self.channels = [ - 64, - 256, - 512, - ] - return - - def forward(self, x, force_sizes_out=False, return_more=False): - ans = [] - x = x[:, :3] - x = self.resize(x) - x = self.resnet_preprocess(x) - x = self.conv1(x) - x = self.bn1(x) - x = self.relu(x) - ans.append(x) # conv1 - x = self.maxpool(x) - x = self.layer1(x) - ans.append(x) # layer1 - x = self.layer2(x) - ans.append(x) # layer2 - if force_sizes_out or (self.sizes_out is None): - self.sizes_out = [tuple(q.shape[-2:]) for q in ans] - return ans, (locals() if return_more else None) - - -class NetNedt(nn.Module): - def __init__(self): - super().__init__() - chin = 3 + 1 + 4 + 4 + 1 + 1 - ch = 16 - chout = 1 - self.net = nn.Sequential( - nn.PReLU(chin), - nn.Conv2d(chin, ch, kernel_size=3, padding=1), - nn.BatchNorm2d(ch), - nn.PReLU(ch), - nn.Conv2d(ch, ch, kernel_size=3, padding=1), - nn.BatchNorm2d(ch), - nn.PReLU(ch), - nn.Conv2d(ch, chout, kernel_size=3, padding=1), - ) - return - - def forward(self, out_base, out_base_nedt, hw_imgs, hw_masks, return_more=False): - cat = torch.cat( - [ - out_base, # 3 - out_base_nedt, # 1 - hw_imgs[0], # 4 - hw_imgs[1], # 4 - hw_masks[0], # 1 - hw_masks[1], # 1 - ], - dim=1, - ) - log = pixel_logit(cat.clip(0, 1)) - ans = torch.sigmoid(self.net(log)) - return ans, (locals() if return_more else None) - - -class NetTail(nn.Module): - def __init__(self): - super().__init__() - chin = 3 + 1 + 1 - ch = 16 - chout = 3 - self.net = nn.Sequential( - nn.PReLU(chin), - nn.Conv2d(chin, ch, kernel_size=3, padding=1), - nn.BatchNorm2d(ch), - nn.PReLU(ch), - nn.Conv2d(ch, ch, kernel_size=3, padding=1), - nn.BatchNorm2d(ch), - nn.PReLU(ch), - nn.Conv2d(ch, ch, kernel_size=3, padding=1), - nn.BatchNorm2d(ch), - nn.PReLU(ch), - nn.Conv2d(ch, chout, kernel_size=3, padding=1), - ) - return - - def forward(self, out_base, out_base_nedt, pred_nedt, return_more=False): - cat = torch.cat( - [ - out_base, # 3 - out_base_nedt, # 1 - pred_nedt, # 1 - ], - dim=1, - ) - log = pixel_logit(cat.clip(0, 1)) - ans = torch.sigmoid(log[:, :3] + self.net(log)) - return ans, (locals() if return_more else None) - - -class SoftsplatLite(nn.Module): - def __init__(self): - super().__init__() - self.feature_extractor = ResnetFeatureExtractor( - ("torchvision", "resnet50"), - (540, 960), - ) - self.z_metric = FlowZMetric() - self.flow_downsamplers = [ - Interpolator(s, mode="bilinear") for s in self.feature_extractor.sizes_out - ] - self.gridnet_converter = GridnetConverter( - self.feature_extractor.channels, - [32, 64, 128], - ) - self.gridnet = Gridnet( - *[32, 64, 128], - total_dropout_p=0.0, - depth=1, # equivalent to u-net - ) - self.nedt = NEDT() - self.half_warper = HalfWarper() - self.synthesizer = Synthesizer( - (540, 960), - self.half_warper.channels_image, - self.half_warper.channels_flow, - self.half_warper.channels_mask, - self.gridnet.channels_0, - ) - return - - def forward(self, x, t=0.5, k=5, return_more=False): - rm = return_more - flow0, flow1 = x["flows"].swapaxes(0, 1) - img0, img1 = x["images"][:, 0], x["images"][:, -1] - (z0, z1), locs_z = self.z_metric(img0, img1, flow0, flow1, return_more=rm) - img0 = torch.cat([img0, self.nedt(img0)[0]], dim=1) - img1 = torch.cat([img1, self.nedt(img1)[0]], dim=1) - - # images and flows - (hw_imgs, hw_flows, hw_masks), locs_hw = self.half_warper( - img0, - img1, - flow0, - flow1, - z0, - z1, - k, - t=t, - return_more=rm, - ) - - # features - feats0, locs_fe0 = self.feature_extractor(img0, return_more=rm) - feats1, locs_fe1 = self.feature_extractor(img1, return_more=rm) - warps = [] - for ft0, ft1, ds in zip(feats0, feats1, self.flow_downsamplers): - (w, _, _), _ = self.half_warper( - ft0, - ft1, - ds(flow0, 1), - ds(flow1, 1), - ds(z0), - ds(z1), - k, - t=t, - ) - warps.append((w[0] + w[1]) / 2) - feats = self.gridnet(self.gridnet_converter(warps)) - - # synthesis - pred, locs_synth = self.synthesizer( - hw_imgs, - hw_flows, - hw_masks, - [ - feats[0], - ], - return_more=rm, - ) - return pred, (locals() if rm else None) - - -class DTM(nn.Module): - def __init__(self): - super().__init__() - self.net_nedt = NetNedt() - self.net_tail = NetTail() - self.nedt = NEDT() - return - - def forward(self, x, out_base, locs_base, return_more=False): - rm = return_more - with torch.no_grad(): - out_base_nedt, locs_base_nedt = self.nedt(out_base, return_more=rm) - hw_imgs, hw_masks = locs_base["hw_imgs"], locs_base["hw_masks"] - pred_nedt, locs_nedt = self.net_nedt( - out_base, out_base_nedt, hw_imgs, hw_masks, return_more=rm - ) - pred, locs_tail = self.net_tail( - out_base, out_base_nedt, pred_nedt.clone().detach(), return_more=rm - ) - return torch.cat([pred, pred_nedt], dim=1), (locals() if rm else None) - - -class RAFT(nn.Module): - def __init__(self, path="/workspace/tensorrt/models/anime_interp_full.ckpt"): - super().__init__() - self.raft = RFR( - Namespace( - small=False, - mixed_precision=False, - ) - ) - if path is not None: - sd = torch.load(path)["model_state_dict"] - self.raft.load_state_dict( - { - k[len("module.flownet.") :]: v - for k, v in sd.items() - if k.startswith("module.flownet.") - }, - strict=False, - ) - return - - def forward(self, img0, img1, flow0=None, iters=12, return_more=False): - if flow0 is not None: - flow0 = flow0.flip(dims=(1,)) - out = self.raft(img1, img0, iters=iters, flow_init=flow0) - return out[0].flip(dims=(1,)), (locals() if return_more else None) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/film/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/film/__init__.py index 1f89f1b..0ded7a4 100644 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/film/__init__.py +++ b/imports/ComfyUI_Frame_Interpolation/vfi_models/film/__init__.py @@ -3,7 +3,7 @@ from comfy.model_management import get_torch_device, soft_empty_cache import bisect import numpy as np import typing -from vfi_utils import InterpolationStateList, load_file_from_github_release, preprocess_frames, postprocess_frames +from vfi_utils import InterpolationStateListImport, load_file_from_github_release, preprocess_frames, postprocess_frames import pathlib import gc @@ -41,7 +41,7 @@ def inference(model, img_batch_1, img_batch_2, inter_frames): return [tensor.flip(0) for tensor in results] -class FILM_VFI: +class FILM_VFIImport: @classmethod def INPUT_TYPES(s): return { @@ -66,7 +66,7 @@ class FILM_VFI: frames: torch.Tensor, clear_cache_after_n_frames = 10, multiplier: typing.SupportsInt = 2, - optional_interpolation_states: InterpolationStateList = None, + optional_interpolation_states: InterpolationStateListImport = None, **kwargs ): interpolation_states = optional_interpolation_states diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/__init__.py deleted file mode 100644 index 64a4b6b..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/__init__.py +++ /dev/null @@ -1,115 +0,0 @@ -import torch -from comfy.model_management import get_torch_device, soft_empty_cache -import numpy as np -import typing -from vfi_utils import InterpolationStateList, load_file_from_github_release, preprocess_frames, postprocess_frames, assert_batch_size -import pathlib -import warnings -from .flavr_arch import UNet_3D_3D, InputPadder -import gc - -device = get_torch_device() -NBR_FRAME = 4 - -def build_flavr(model_path): - sd = torch.load(model_path)['state_dict'] - sd = {k.partition("module.")[-1]:v for k,v in sd.items()} - - #Ref: Class UNet_3D_3D - model = UNet_3D_3D("unet_18", n_inputs=NBR_FRAME, n_outputs=sd["outconv.1.weight"].shape[0] // 3, joinType="concat" , upmode="transpose") - model.load_state_dict(sd) - model.to(device).eval() - del sd - return model - -MODEL_TYPE = pathlib.Path(__file__).parent.name -CKPT_NAMES = ["FLAVR_2x.pth", "FLAVR_4x.pth", "FLAVR_8x.pth"] - -class FLAVR_VFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (CKPT_NAMES, ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}), - "multiplier": ("INT", {"default": 2, "min": 2, "max": 2}), #TODO: Implement recursively invoking interpolator for multi-frame interpolation - "duplicate_first_last_frames": ("BOOLEAN", {"default": False}) - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - #Reference: https://github.com/danier97/ST-MFNet/blob/main/interpolate_yuv.py#L93 - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames = 10, - multiplier: typing.SupportsInt = 2, - duplicate_first_last_frames: bool = False, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - if multiplier != 2: - warnings.warn("Currently, FLAVR only supports 2x interpolation. The process will continue but please set multiplier=2 afterward") - - assert_batch_size(frames, batch_size=4, vfi_name="ST-MFNet") - interpolation_states = optional_interpolation_states - model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) - model = build_flavr(model_path) - frames = preprocess_frames(frames) - padder = InputPadder(frames.shape, 16) - frames = padder.pad(frames) - - number_of_frames_processed_since_last_cleared_cuda_cache = 0 - output_frames = [] - for frame_itr in range(len(frames) - 3): - #Does skipping frame i+1 make sanse in this case? - if interpolation_states is not None and interpolation_states.is_frame_skipped(frame_itr) and interpolation_states.is_frame_skipped(frame_itr + 1): - continue - - #Ensure that input frames are in fp32 - the same dtype as model - frame0, frame1, frame2, frame3 = ( - frames[frame_itr:frame_itr+1].float(), - frames[frame_itr+1:frame_itr+2].float(), - frames[frame_itr+2:frame_itr+3].float(), - frames[frame_itr+3:frame_itr+4].float() - ) - new_frame = model([frame0.to(device), frame1.to(device), frame2.to(device), frame3.to(device)])[0].detach().cpu() - number_of_frames_processed_since_last_cleared_cuda_cache += 2 - - if frame_itr == 0: - output_frames.append(frame0) - if duplicate_first_last_frames: - output_frames.append(frame0) # repeat the first frame - output_frames.append(frame1) - output_frames.append(new_frame) - output_frames.append(frame2) - if frame_itr == len(frames) - 4: - output_frames.append(frame3) - if duplicate_first_last_frames: - output_frames.append(frame3) # repeat the last frame - - # Try to avoid a memory overflow by clearing cuda cache regularly - if number_of_frames_processed_since_last_cleared_cuda_cache >= clear_cache_after_n_frames: - print("Comfy-VFI: Clearing cache...", end = ' ') - soft_empty_cache() - number_of_frames_processed_since_last_cleared_cuda_cache = 0 - print("Done cache clearing") - gc.collect() - - dtype = torch.float32 - output_frames = [frame.cpu().to(dtype=dtype) for frame in output_frames] #Ensure all frames are in cpu - out = torch.cat(output_frames, dim=0) - out = padder.unpad(out) - # clear cache for courtesy - print("Comfy-VFI: Final clearing cache...", end=' ') - soft_empty_cache() - print("Done cache clearing") - return (postprocess_frames(out), ) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/flavr_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/flavr_arch.py deleted file mode 100644 index 60c27d3..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/flavr_arch.py +++ /dev/null @@ -1,217 +0,0 @@ -""" -https://github.com/tarun005/FLAVR/blob/main/model/FLAVR_arch.py -https://github.com/tarun005/FLAVR/blob/main/model/resnet_3D.py (only SEGating) -""" -import math -import numpy as np -import importlib - -import torch -import torch.nn as nn -import torch.nn.functional as F - -class SEGating(nn.Module): - - def __init__(self , inplanes , reduction=16): - - super().__init__() - - self.pool = nn.AdaptiveAvgPool3d(1) - self.attn_layer = nn.Sequential( - nn.Conv3d(inplanes , inplanes , kernel_size=1 , stride=1 , bias=True), - nn.Sigmoid() - ) - - def forward(self , x): - - out = self.pool(x) - y = self.attn_layer(out) - return x * y - -def joinTensors(X1 , X2 , type="concat"): - - if type == "concat": - return torch.cat([X1 , X2] , dim=1) - elif type == "add": - return X1 + X2 - else: - return X1 - - -class Conv_2d(nn.Module): - - def __init__(self, in_ch, out_ch, kernel_size, stride=1, padding=0, bias=False, batchnorm=False): - - super().__init__() - self.conv = [nn.Conv2d(in_ch, out_ch, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)] - - if batchnorm: - self.conv += [nn.BatchNorm2d(out_ch)] - - self.conv = nn.Sequential(*self.conv) - - def forward(self, x): - - return self.conv(x) - -class upConv3D(nn.Module): - - def __init__(self, in_ch, out_ch, kernel_size, stride, padding, upmode="transpose" , batchnorm=False): - - super().__init__() - - self.upmode = upmode - - if self.upmode=="transpose": - self.upconv = nn.ModuleList( - [nn.ConvTranspose3d(in_ch, out_ch, kernel_size=kernel_size, stride=stride, padding=padding), - SEGating(out_ch) - ] - ) - - else: - self.upconv = nn.ModuleList( - [nn.Upsample(mode='trilinear', scale_factor=(1,2,2), align_corners=False), - nn.Conv3d(in_ch, out_ch , kernel_size=1 , stride=1), - SEGating(out_ch) - ] - ) - - if batchnorm: - self.upconv += [nn.BatchNorm3d(out_ch)] - - self.upconv = nn.Sequential(*self.upconv) - - def forward(self, x): - - return self.upconv(x) - -class Conv_3d(nn.Module): - - def __init__(self, in_ch, out_ch, kernel_size, stride=1, padding=0, bias=True, batchnorm=False): - - super().__init__() - self.conv = [nn.Conv3d(in_ch, out_ch, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias), - SEGating(out_ch) - ] - - if batchnorm: - self.conv += [nn.BatchNorm3d(out_ch)] - - self.conv = nn.Sequential(*self.conv) - - def forward(self, x): - - return self.conv(x) - -class upConv2D(nn.Module): - - def __init__(self, in_ch, out_ch, kernel_size, stride, padding, upmode="transpose" , batchnorm=False): - - super().__init__() - - self.upmode = upmode - - if self.upmode=="transpose": - self.upconv = [nn.ConvTranspose2d(in_ch, out_ch, kernel_size=kernel_size, stride=stride, padding=padding)] - - else: - self.upconv = [ - nn.Upsample(mode='bilinear', scale_factor=2, align_corners=False), - nn.Conv2d(in_ch, out_ch , kernel_size=1 , stride=1) - ] - - if batchnorm: - self.upconv += [nn.BatchNorm2d(out_ch)] - - self.upconv = nn.Sequential(*self.upconv) - - def forward(self, x): - - return self.upconv(x) - - -class UNet_3D_3D(nn.Module): - def __init__(self, block , n_inputs, n_outputs, batchnorm=False , joinType="concat" , upmode="transpose"): - super().__init__() - - nf = [512 , 256 , 128 , 64] - out_channels = 3*n_outputs - self.joinType = joinType - self.n_outputs = n_outputs - - growth = 2 if joinType == "concat" else 1 - self.lrelu = nn.LeakyReLU(0.2, True) - - unet_3D = importlib.import_module(".resnet_3D", "models.flavr") - if n_outputs > 1: - unet_3D.useBias = True - self.encoder = getattr(unet_3D , block)(pretrained=False , bn=batchnorm) - - self.decoder = nn.Sequential( - Conv_3d(nf[0], nf[1] , kernel_size=3, padding=1, bias=True, batchnorm=batchnorm), - upConv3D(nf[1]*growth, nf[2], kernel_size=(3,4,4), stride=(1,2,2), padding=(1,1,1) , upmode=upmode, batchnorm=batchnorm), - upConv3D(nf[2]*growth, nf[3], kernel_size=(3,4,4), stride=(1,2,2), padding=(1,1,1) , upmode=upmode, batchnorm=batchnorm), - Conv_3d(nf[3]*growth, nf[3] , kernel_size=3, padding=1, bias=True, batchnorm=batchnorm), - upConv3D(nf[3]*growth , nf[3], kernel_size=(3,4,4), stride=(1,2,2), padding=(1,1,1) , upmode=upmode, batchnorm=batchnorm) - ) - - self.feature_fuse = Conv_2d(nf[3]*n_inputs , nf[3] , kernel_size=1 , stride=1, batchnorm=batchnorm) - - self.outconv = nn.Sequential( - nn.ReflectionPad2d(3), - nn.Conv2d(nf[3], out_channels , kernel_size=7 , stride=1, padding=0) - ) - - def forward(self, images): - - images = torch.stack(images , dim=2) - - ## Batch mean normalization works slightly better than global mean normalization, thanks to https://github.com/myungsub/CAIN - mean_ = images.mean(2, keepdim=True).mean(3, keepdim=True).mean(4,keepdim=True) - images = images-mean_ - - x_0 , x_1 , x_2 , x_3 , x_4 = self.encoder(images) - - dx_3 = self.lrelu(self.decoder[0](x_4)) - dx_3 = joinTensors(dx_3 , x_3 , type=self.joinType) - - dx_2 = self.lrelu(self.decoder[1](dx_3)) - dx_2 = joinTensors(dx_2 , x_2 , type=self.joinType) - - dx_1 = self.lrelu(self.decoder[2](dx_2)) - dx_1 = joinTensors(dx_1 , x_1 , type=self.joinType) - - dx_0 = self.lrelu(self.decoder[3](dx_1)) - dx_0 = joinTensors(dx_0 , x_0 , type=self.joinType) - - dx_out = self.lrelu(self.decoder[4](dx_0)) - dx_out = torch.cat(torch.unbind(dx_out , 2) , 1) - - out = self.lrelu(self.feature_fuse(dx_out)) - out = self.outconv(out) - - out = torch.split(out, dim=1, split_size_or_sections=3) - mean_ = mean_.squeeze(2) - out = [o+mean_ for o in out] - - return out - -class InputPadder: - """ Pads images such that dimensions are divisible by divisor """ - def __init__(self, dims, divisor=16): - self.ht, self.wd = dims[-2:] - pad_ht = (((self.ht // divisor) + 1) * divisor - self.ht) % divisor - pad_wd = (((self.wd // divisor) + 1) * divisor - self.wd) % divisor - self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2] - - def pad(self, input_tensor): - return F.pad(input_tensor, self._pad, mode='replicate') - - def unpad(self, input_tensor): - return self._unpad(input_tensor) - - def _unpad(self, x): - ht, wd = x.shape[-2:] - c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]] - return x[..., c[0]:c[1], c[2]:c[3]] \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/resnet_3D.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/resnet_3D.py deleted file mode 100644 index ba8b9a8..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/flavr/resnet_3D.py +++ /dev/null @@ -1,288 +0,0 @@ -# Modified from https://github.com/pytorch/vision/tree/master/torchvision/models/video - -import torch -import torch.nn as nn - -__all__ = ['unet_18', 'unet_34'] - -useBias = False - -class identity(nn.Module): - - def __init__(self , *args , **kwargs): - - super().__init__() - - def forward(self , x): - return x - -class Conv3DSimple(nn.Conv3d): - def __init__(self, - in_planes, - out_planes, - midplanes=None, - stride=1, - padding=1): - - super(Conv3DSimple, self).__init__( - in_channels=in_planes, - out_channels=out_planes, - kernel_size=(3, 3, 3), - stride=stride, - padding=padding, - bias=useBias) - - @staticmethod - def get_downsample_stride(stride , temporal_stride): - if temporal_stride: - return (temporal_stride, stride, stride) - else: - return (stride , stride , stride) - -class BasicStem(nn.Sequential): - """The default conv-batchnorm-relu stem - """ - def __init__(self): - super().__init__( - nn.Conv3d(3, 64, kernel_size=(3, 7, 7), stride=(1, 2, 2), - padding=(1, 3, 3), bias=useBias), - batchnorm(64), - nn.ReLU(inplace=False)) - - -class Conv2Plus1D(nn.Sequential): - - def __init__(self, - in_planes, - out_planes, - midplanes, - stride=1, - padding=1): - if not isinstance(stride , int): - temporal_stride , stride , stride = stride - else: - temporal_stride = stride - - super(Conv2Plus1D, self).__init__( - nn.Conv3d(in_planes, midplanes, kernel_size=(1, 3, 3), - stride=(1, stride, stride), padding=(0, padding, padding), - bias=False), - # batchnorm(midplanes), - nn.ReLU(inplace=True), - nn.Conv3d(midplanes, out_planes, kernel_size=(3, 1, 1), - stride=(temporal_stride, 1, 1), padding=(padding, 0, 0), - bias=False)) - - @staticmethod - def get_downsample_stride(stride , temporal_stride): - if temporal_stride: - return (temporal_stride, stride, stride) - else: - return (stride , stride , stride) - -class R2Plus1dStem(nn.Sequential): - """R(2+1)D stem is different than the default one as it uses separated 3D convolution - """ - def __init__(self): - super().__init__( - nn.Conv3d(3, 45, kernel_size=(1, 7, 7), - stride=(1, 2, 2), padding=(0, 3, 3), - bias=False), - batchnorm(45), - nn.ReLU(inplace=True), - nn.Conv3d(45, 64, kernel_size=(3, 1, 1), - stride=(1, 1, 1), padding=(1, 0, 0), - bias=False), - batchnorm(64), - nn.ReLU(inplace=True)) - - -class SEGating(nn.Module): - - def __init__(self , inplanes , reduction=16): - - super().__init__() - - self.pool = nn.AdaptiveAvgPool3d(1) - self.attn_layer = nn.Sequential( - nn.Conv3d(inplanes , inplanes , kernel_size=1 , stride=1 , bias=True), - nn.Sigmoid() - ) - - def forward(self , x): - - out = self.pool(x) - y = self.attn_layer(out) - return x * y - -class BasicBlock(nn.Module): - - expansion = 1 - - def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None): - midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes) - - super(BasicBlock, self).__init__() - self.conv1 = nn.Sequential( - conv_builder(inplanes, planes, midplanes, stride), - batchnorm(planes), - nn.ReLU(inplace=True) - ) - self.conv2 = nn.Sequential( - conv_builder(planes, planes, midplanes), - batchnorm(planes) - ) - self.fg = SEGating(planes) ## Feature Gating - self.relu = nn.ReLU(inplace=True) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - out = self.conv1(x) - out = self.conv2(out) - out = self.fg(out) - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - -class VideoResNet(nn.Module): - - def __init__(self, block, conv_makers, layers, - stem, zero_init_residual=False): - """Generic resnet video generator. - - Args: - block (nn.Module): resnet building block - conv_makers (list(functions)): generator function for each layer - layers (List[int]): number of blocks per layer - stem (nn.Module, optional): Resnet stem, if None, defaults to conv-bn-relu. Defaults to None. - """ - super(VideoResNet, self).__init__() - self.inplanes = 64 - - self.stem = stem() - - self.layer1 = self._make_layer(block, conv_makers[0], 64, layers[0], stride=1 ) - self.layer2 = self._make_layer(block, conv_makers[1], 128, layers[1], stride=2 , temporal_stride=1) - self.layer3 = self._make_layer(block, conv_makers[2], 256, layers[2], stride=2 , temporal_stride=1) - self.layer4 = self._make_layer(block, conv_makers[3], 512, layers[3], stride=1, temporal_stride=1) - - # init weights - self._initialize_weights() - - if zero_init_residual: - for m in self.modules(): - if isinstance(m, Bottleneck): - nn.init.constant_(m.bn3.weight, 0) - - def forward(self, x): - x_0 = self.stem(x) - x_1 = self.layer1(x_0) - x_2 = self.layer2(x_1) - x_3 = self.layer3(x_2) - x_4 = self.layer4(x_3) - return x_0 , x_1 , x_2 , x_3 , x_4 - - def _make_layer(self, block, conv_builder, planes, blocks, stride=1, temporal_stride=None): - downsample = None - - if stride != 1 or self.inplanes != planes * block.expansion: - ds_stride = conv_builder.get_downsample_stride(stride , temporal_stride) - downsample = nn.Sequential( - nn.Conv3d(self.inplanes, planes * block.expansion, - kernel_size=1, stride=ds_stride, bias=False), - batchnorm(planes * block.expansion) - ) - stride = ds_stride - - layers = [] - layers.append(block(self.inplanes, planes, conv_builder, stride, downsample )) - - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes, conv_builder )) - - return nn.Sequential(*layers) - - def _initialize_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv3d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', - nonlinearity='relu') - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm3d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.normal_(m.weight, 0, 0.01) - nn.init.constant_(m.bias, 0) - - -def _video_resnet(arch, pretrained=False, progress=True, **kwargs): - model = VideoResNet(**kwargs) - ## TODO: Other 3D resnet models, like S3D, r(2+1)D. - - if pretrained: - state_dict = load_state_dict_from_url(model_urls[arch], - progress=progress) - model.load_state_dict(state_dict) - return model - - -def unet_18(pretrained=False, bn=False, progress=True, **kwargs): - """ - Construct 18 layer Unet3D model as in - https://arxiv.org/abs/1711.11248 - - Args: - pretrained (bool): If True, returns a model pre-trained on Kinetics-400 - progress (bool): If True, displays a progress bar of the download to stderr - - Returns: - nn.Module: R3D-18 encoder - """ - global batchnorm - if bn: - batchnorm = nn.BatchNorm3d - else: - batchnorm = identity - - return _video_resnet('r3d_18', - pretrained, progress, - block=BasicBlock, - conv_makers=[Conv3DSimple] * 4, - layers=[2, 2, 2, 2], - stem=BasicStem, **kwargs) - -def unet_34(pretrained=False, bn=False, progress=True, **kwargs): - """ - Construct 34 layer Unet3D model as in - https://arxiv.org/abs/1711.11248 - - Args: - pretrained (bool): If True, returns a model pre-trained on Kinetics-400 - progress (bool): If True, displays a progress bar of the download to stderr - - Returns: - nn.Module: R3D-18 encoder - """ - global batchnorm - # bn = False - if bn: - batchnorm = nn.BatchNorm3d - else: - batchnorm = identity - - - return _video_resnet('r3d_34', - pretrained, progress, - block=BasicBlock, - conv_makers=[Conv3DSimple] * 4, - layers=[3, 4, 6, 3], - stem=BasicStem, **kwargs) \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna.py deleted file mode 100644 index 949e513..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna.py +++ /dev/null @@ -1,24 +0,0 @@ -import itertools -import numpy as np -import vapoursynth as vs -from .GMFSS_Fortuna_arch import Model_inference -import torch -import traceback - - -class GMFSS_Fortuna: - def __init__(self): - self.cache = False - self.amount_input_img = 2 - - torch.set_grad_enabled(False) - torch.backends.cudnn.enabled = True - torch.backends.cudnn.benchmark = True - - self.model = Model_inference() - self.model.eval() - - def execute(self, I0, I1, timestep): - with torch.inference_mode(): - middle = self.model(I0, I1, timestep).cpu() - return middle diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_arch.py deleted file mode 100644 index 0624896..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_arch.py +++ /dev/null @@ -1,1850 +0,0 @@ -""" -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/GMFSS_infer_b.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/softsplat.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/FusionNet_b.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/FeatureNet.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/MetricNet.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/IFNet_HDv3.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/gmflow.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/utils.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/position.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/geometry.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/matching.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/transformer.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/backbone.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/trident_conv.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/warplayer.py -""" - -from torch import nn -from torch.nn import functional as F -from torch.nn.modules.utils import _pair -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch -import math -from vfi_models.rife.rife_arch import IFNet -from vfi_models.ops import softsplat -from comfy.model_management import get_torch_device - -device = get_torch_device() -backwarp_tenGrid = {} - - -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) - return torch.nn.functional.grid_sample( - input=tenInput, - grid=g, - mode="bilinear", - padding_mode="border", - align_corners=True, - ) - - -class MultiScaleTridentConv(nn.Module): - def __init__( - self, - in_channels, - out_channels, - kernel_size, - stride=1, - strides=1, - paddings=0, - dilations=1, - dilation=1, - groups=1, - num_branch=1, - test_branch_idx=-1, - bias=False, - norm=None, - activation=None, - ): - super(MultiScaleTridentConv, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = _pair(kernel_size) - self.num_branch = num_branch - self.stride = _pair(stride) - self.groups = groups - self.with_bias = bias - self.dilation = dilation - if isinstance(paddings, int): - paddings = [paddings] * self.num_branch - if isinstance(dilations, int): - dilations = [dilations] * self.num_branch - if isinstance(strides, int): - strides = [strides] * self.num_branch - self.paddings = [_pair(padding) for padding in paddings] - self.dilations = [_pair(dilation) for dilation in dilations] - self.strides = [_pair(stride) for stride in strides] - self.test_branch_idx = test_branch_idx - self.norm = norm - self.activation = activation - - assert len({self.num_branch, len(self.paddings), len(self.strides)}) == 1 - - self.weight = nn.Parameter( - torch.Tensor(out_channels, in_channels // groups, *self.kernel_size) - ) - if bias: - self.bias = nn.Parameter(torch.Tensor(out_channels)) - else: - self.bias = None - - nn.init.kaiming_uniform_(self.weight, nonlinearity="relu") - if self.bias is not None: - nn.init.constant_(self.bias, 0) - - def forward(self, inputs): - num_branch = ( - self.num_branch if self.training or self.test_branch_idx == -1 else 1 - ) - assert len(inputs) == num_branch - - if self.training or self.test_branch_idx == -1: - outputs = [ - F.conv2d( - input, - self.weight, - self.bias, - stride, - padding, - self.dilation, - self.groups, - ) - for input, stride, padding in zip(inputs, self.strides, self.paddings) - ] - else: - outputs = [ - F.conv2d( - inputs[0], - self.weight, - self.bias, - self.strides[self.test_branch_idx] - if self.test_branch_idx == -1 - else self.strides[-1], - self.paddings[self.test_branch_idx] - if self.test_branch_idx == -1 - else self.paddings[-1], - self.dilation, - self.groups, - ) - ] - - if self.norm is not None: - outputs = [self.norm(x) for x in outputs] - if self.activation is not None: - outputs = [self.activation(x) for x in outputs] - return outputs - - -class ResidualBlock_class(nn.Module): - def __init__( - self, - in_planes, - planes, - norm_layer=nn.InstanceNorm2d, - stride=1, - dilation=1, - ): - super(ResidualBlock_class, self).__init__() - - self.conv1 = nn.Conv2d( - in_planes, - planes, - kernel_size=3, - dilation=dilation, - padding=dilation, - stride=stride, - bias=False, - ) - self.conv2 = nn.Conv2d( - planes, - planes, - kernel_size=3, - dilation=dilation, - padding=dilation, - bias=False, - ) - self.relu = nn.ReLU(inplace=True) - - self.norm1 = norm_layer(planes) - self.norm2 = norm_layer(planes) - if not stride == 1 or in_planes != planes: - self.norm3 = norm_layer(planes) - - if stride == 1 and in_planes == planes: - self.downsample = None - else: - self.downsample = nn.Sequential( - nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3 - ) - - def forward(self, x): - y = x - y = self.relu(self.norm1(self.conv1(y))) - y = self.relu(self.norm2(self.conv2(y))) - - if self.downsample is not None: - x = self.downsample(x) - - return self.relu(x + y) - - -class CNNEncoder(nn.Module): - def __init__( - self, - output_dim=128, - norm_layer=nn.InstanceNorm2d, - num_output_scales=1, - **kwargs, - ): - super(CNNEncoder, self).__init__() - self.num_branch = num_output_scales - - feature_dims = [64, 96, 128] - - self.conv1 = nn.Conv2d( - 3, feature_dims[0], kernel_size=7, stride=2, padding=3, bias=False - ) # 1/2 - self.norm1 = norm_layer(feature_dims[0]) - self.relu1 = nn.ReLU(inplace=True) - - self.in_planes = feature_dims[0] - self.layer1 = self._make_layer( - feature_dims[0], stride=1, norm_layer=norm_layer - ) # 1/2 - self.layer2 = self._make_layer( - feature_dims[1], stride=2, norm_layer=norm_layer - ) # 1/4 - - # highest resolution 1/4 or 1/8 - stride = 2 if num_output_scales == 1 else 1 - self.layer3 = self._make_layer( - feature_dims[2], - stride=stride, - norm_layer=norm_layer, - ) # 1/4 or 1/8 - - self.conv2 = nn.Conv2d(feature_dims[2], output_dim, 1, 1, 0) - - if self.num_branch > 1: - if self.num_branch == 4: - strides = (1, 2, 4, 8) - elif self.num_branch == 3: - strides = (1, 2, 4) - elif self.num_branch == 2: - strides = (1, 2) - else: - raise ValueError - - self.trident_conv = MultiScaleTridentConv( - output_dim, - output_dim, - kernel_size=3, - strides=strides, - paddings=1, - num_branch=self.num_branch, - ) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") - elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): - if m.weight is not None: - nn.init.constant_(m.weight, 1) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - def _make_layer(self, dim, stride=1, dilation=1, norm_layer=nn.InstanceNorm2d): - layer1 = ResidualBlock_class( - self.in_planes, dim, norm_layer=norm_layer, stride=stride, dilation=dilation - ) - layer2 = ResidualBlock_class( - dim, dim, norm_layer=norm_layer, stride=1, dilation=dilation - ) - - layers = (layer1, layer2) - - self.in_planes = dim - return nn.Sequential(*layers) - - def forward(self, x): - x = self.conv1(x) - x = self.norm1(x) - x = self.relu1(x) - - x = self.layer1(x) # 1/2 - x = self.layer2(x) # 1/4 - x = self.layer3(x) # 1/8 or 1/4 - - x = self.conv2(x) - - if self.num_branch > 1: - out = self.trident_conv([x] * self.num_branch) # high to low res - else: - out = [x] - - return out - - -def single_head_full_attention(q, k, v): - # q, k, v: [B, L, C] - assert q.dim() == k.dim() == v.dim() == 3 - - scores = torch.matmul(q, k.permute(0, 2, 1)) / (q.size(2) ** 0.5) # [B, L, L] - attn = torch.softmax(scores, dim=2) # [B, L, L] - out = torch.matmul(attn, v) # [B, L, C] - - return out - - -def generate_shift_window_attn_mask( - input_resolution, - window_size_h, - window_size_w, - shift_size_h, - shift_size_w, - device=get_torch_device(), -): - # Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py - # calculate attention mask for SW-MSA - h, w = input_resolution - img_mask = torch.zeros((1, h, w, 1)).to(device) # 1 H W 1 - h_slices = ( - slice(0, -window_size_h), - slice(-window_size_h, -shift_size_h), - slice(-shift_size_h, None), - ) - w_slices = ( - slice(0, -window_size_w), - slice(-window_size_w, -shift_size_w), - slice(-shift_size_w, None), - ) - cnt = 0 - for h in h_slices: - for w in w_slices: - img_mask[:, h, w, :] = cnt - cnt += 1 - - mask_windows = split_feature( - img_mask, num_splits=input_resolution[-1] // window_size_w, channel_last=True - ) - - mask_windows = mask_windows.view(-1, window_size_h * window_size_w) - attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( - attn_mask == 0, float(0.0) - ) - - return attn_mask - - -def single_head_split_window_attention( - q, - k, - v, - num_splits=1, - with_shift=False, - h=None, - w=None, - attn_mask=None, -): - # Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py - # q, k, v: [B, L, C] - assert q.dim() == k.dim() == v.dim() == 3 - - assert h is not None and w is not None - assert q.size(1) == h * w - - b, _, c = q.size() - - b_new = b * num_splits * num_splits - - window_size_h = h // num_splits - window_size_w = w // num_splits - - q = q.view(b, h, w, c) # [B, H, W, C] - k = k.view(b, h, w, c) - v = v.view(b, h, w, c) - - scale_factor = c**0.5 - - if with_shift: - assert attn_mask is not None # compute once - shift_size_h = window_size_h // 2 - shift_size_w = window_size_w // 2 - - q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) - k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) - v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) - - q = split_feature( - q, num_splits=num_splits, channel_last=True - ) # [B*K*K, H/K, W/K, C] - k = split_feature(k, num_splits=num_splits, channel_last=True) - v = split_feature(v, num_splits=num_splits, channel_last=True) - - scores = ( - torch.matmul(q.view(b_new, -1, c), k.view(b_new, -1, c).permute(0, 2, 1)) - / scale_factor - ) # [B*K*K, H/K*W/K, H/K*W/K] - - if with_shift: - scores += attn_mask.repeat(b, 1, 1) - - attn = torch.softmax(scores, dim=-1) - - out = torch.matmul(attn, v.view(b_new, -1, c)) # [B*K*K, H/K*W/K, C] - - out = merge_splits( - out.view(b_new, h // num_splits, w // num_splits, c), - num_splits=num_splits, - channel_last=True, - ) # [B, H, W, C] - - # shift back - if with_shift: - out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2)) - - out = out.view(b, -1, c) - - return out - - -class TransformerLayer(nn.Module): - def __init__( - self, - d_model=256, - nhead=1, - attention_type="swin", - no_ffn=False, - ffn_dim_expansion=4, - with_shift=False, - **kwargs, - ): - super(TransformerLayer, self).__init__() - - self.dim = d_model - self.nhead = nhead - self.attention_type = attention_type - self.no_ffn = no_ffn - - self.with_shift = with_shift - - # multi-head attention - self.q_proj = nn.Linear(d_model, d_model, bias=False) - self.k_proj = nn.Linear(d_model, d_model, bias=False) - self.v_proj = nn.Linear(d_model, d_model, bias=False) - - self.merge = nn.Linear(d_model, d_model, bias=False) - - self.norm1 = nn.LayerNorm(d_model) - - # no ffn after self-attn, with ffn after cross-attn - if not self.no_ffn: - in_channels = d_model * 2 - self.mlp = nn.Sequential( - nn.Linear(in_channels, in_channels * ffn_dim_expansion, bias=False), - nn.GELU(), - nn.Linear(in_channels * ffn_dim_expansion, d_model, bias=False), - ) - - self.norm2 = nn.LayerNorm(d_model) - - def forward( - self, - source, - target, - height=None, - width=None, - shifted_window_attn_mask=None, - attn_num_splits=None, - **kwargs, - ): - # source, target: [B, L, C] - query, key, value = source, target, target - - # single-head attention - query = self.q_proj(query) # [B, L, C] - key = self.k_proj(key) # [B, L, C] - value = self.v_proj(value) # [B, L, C] - - if self.attention_type == "swin" and attn_num_splits > 1: - if self.nhead > 1: - # we observe that multihead attention slows down the speed and increases the memory consumption - # without bringing obvious performance gains and thus the implementation is removed - raise NotImplementedError - else: - message = single_head_split_window_attention( - query, - key, - value, - num_splits=attn_num_splits, - with_shift=self.with_shift, - h=height, - w=width, - attn_mask=shifted_window_attn_mask, - ) - else: - message = single_head_full_attention(query, key, value) # [B, L, C] - - message = self.merge(message) # [B, L, C] - message = self.norm1(message) - - if not self.no_ffn: - message = self.mlp(torch.cat([source, message], dim=-1)) - message = self.norm2(message) - - return source + message - - -class TransformerBlock(nn.Module): - """self attention + cross attention + FFN""" - - def __init__( - self, - d_model=256, - nhead=1, - attention_type="swin", - ffn_dim_expansion=4, - with_shift=False, - **kwargs, - ): - super(TransformerBlock, self).__init__() - - self.self_attn = TransformerLayer( - d_model=d_model, - nhead=nhead, - attention_type=attention_type, - no_ffn=True, - ffn_dim_expansion=ffn_dim_expansion, - with_shift=with_shift, - ) - - self.cross_attn_ffn = TransformerLayer( - d_model=d_model, - nhead=nhead, - attention_type=attention_type, - ffn_dim_expansion=ffn_dim_expansion, - with_shift=with_shift, - ) - - def forward( - self, - source, - target, - height=None, - width=None, - shifted_window_attn_mask=None, - attn_num_splits=None, - **kwargs, - ): - # source, target: [B, L, C] - - # self attention - source = self.self_attn( - source, - source, - height=height, - width=width, - shifted_window_attn_mask=shifted_window_attn_mask, - attn_num_splits=attn_num_splits, - ) - - # cross attention and ffn - source = self.cross_attn_ffn( - source, - target, - height=height, - width=width, - shifted_window_attn_mask=shifted_window_attn_mask, - attn_num_splits=attn_num_splits, - ) - - return source - - -class FeatureTransformer(nn.Module): - def __init__( - self, - num_layers=6, - d_model=128, - nhead=1, - attention_type="swin", - ffn_dim_expansion=4, - **kwargs, - ): - super(FeatureTransformer, self).__init__() - - self.attention_type = attention_type - - self.d_model = d_model - self.nhead = nhead - - self.layers = nn.ModuleList( - [ - TransformerBlock( - d_model=d_model, - nhead=nhead, - attention_type=attention_type, - ffn_dim_expansion=ffn_dim_expansion, - with_shift=True - if attention_type == "swin" and i % 2 == 1 - else False, - ) - for i in range(num_layers) - ] - ) - - for p in self.parameters(): - if p.dim() > 1: - nn.init.xavier_uniform_(p) - - def forward( - self, - feature0, - feature1, - attn_num_splits=None, - **kwargs, - ): - b, c, h, w = feature0.shape - assert self.d_model == c - - feature0 = feature0.flatten(-2).permute(0, 2, 1) # [B, H*W, C] - feature1 = feature1.flatten(-2).permute(0, 2, 1) # [B, H*W, C] - - if self.attention_type == "swin" and attn_num_splits > 1: - # global and refine use different number of splits - window_size_h = h // attn_num_splits - window_size_w = w // attn_num_splits - - # compute attn mask once - shifted_window_attn_mask = generate_shift_window_attn_mask( - input_resolution=(h, w), - window_size_h=window_size_h, - window_size_w=window_size_w, - shift_size_h=window_size_h // 2, - shift_size_w=window_size_w // 2, - device=feature0.device, - ) # [K*K, H/K*W/K, H/K*W/K] - else: - shifted_window_attn_mask = None - - # concat feature0 and feature1 in batch dimension to compute in parallel - concat0 = torch.cat((feature0, feature1), dim=0) # [2B, H*W, C] - concat1 = torch.cat((feature1, feature0), dim=0) # [2B, H*W, C] - - for layer in self.layers: - concat0 = layer( - concat0, - concat1, - height=h, - width=w, - shifted_window_attn_mask=shifted_window_attn_mask, - attn_num_splits=attn_num_splits, - ) - - # update feature1 - concat1 = torch.cat(concat0.chunk(chunks=2, dim=0)[::-1], dim=0) - - feature0, feature1 = concat0.chunk(chunks=2, dim=0) # [B, H*W, C] - - # reshape back - feature0 = ( - feature0.view(b, h, w, c).permute(0, 3, 1, 2).contiguous() - ) # [B, C, H, W] - feature1 = ( - feature1.view(b, h, w, c).permute(0, 3, 1, 2).contiguous() - ) # [B, C, H, W] - - return feature0, feature1 - - -class FeatureFlowAttention(nn.Module): - """ - flow propagation with self-attention on feature - query: feature0, key: feature0, value: flow - """ - - def __init__( - self, - in_channels, - **kwargs, - ): - super(FeatureFlowAttention, self).__init__() - - self.q_proj = nn.Linear(in_channels, in_channels) - self.k_proj = nn.Linear(in_channels, in_channels) - - for p in self.parameters(): - if p.dim() > 1: - nn.init.xavier_uniform_(p) - - def forward( - self, - feature0, - flow, - local_window_attn=False, - local_window_radius=1, - **kwargs, - ): - # q, k: feature [B, C, H, W], v: flow [B, 2, H, W] - if local_window_attn: - return self.forward_local_window_attn( - feature0, flow, local_window_radius=local_window_radius - ) - - b, c, h, w = feature0.size() - - query = feature0.view(b, c, h * w).permute(0, 2, 1) # [B, H*W, C] - - # a note: the ``correct'' implementation should be: - # ``query = self.q_proj(query), key = self.k_proj(query)'' - # this problem is observed while cleaning up the code - # however, this doesn't affect the performance since the projection is a linear operation, - # thus the two projection matrices for key can be merged - # so I just leave it as is in order to not re-train all models :) - query = self.q_proj(query) # [B, H*W, C] - key = self.k_proj(query) # [B, H*W, C] - - value = flow.view(b, flow.size(1), h * w).permute(0, 2, 1) # [B, H*W, 2] - - scores = torch.matmul(query, key.permute(0, 2, 1)) / (c**0.5) # [B, H*W, H*W] - prob = torch.softmax(scores, dim=-1) - - out = torch.matmul(prob, value) # [B, H*W, 2] - out = out.view(b, h, w, value.size(-1)).permute(0, 3, 1, 2) # [B, 2, H, W] - - return out - - def forward_local_window_attn( - self, - feature0, - flow, - local_window_radius=1, - ): - assert flow.size(1) == 2 - assert local_window_radius > 0 - - b, c, h, w = feature0.size() - - feature0_reshape = self.q_proj( - feature0.view(b, c, -1).permute(0, 2, 1) - ).reshape( - b * h * w, 1, c - ) # [B*H*W, 1, C] - - kernel_size = 2 * local_window_radius + 1 - - feature0_proj = ( - self.k_proj(feature0.view(b, c, -1).permute(0, 2, 1)) - .permute(0, 2, 1) - .reshape(b, c, h, w) - ) - - feature0_window = F.unfold( - feature0_proj, kernel_size=kernel_size, padding=local_window_radius - ) # [B, C*(2R+1)^2), H*W] - - feature0_window = ( - feature0_window.view(b, c, kernel_size**2, h, w) - .permute(0, 3, 4, 1, 2) - .reshape(b * h * w, c, kernel_size**2) - ) # [B*H*W, C, (2R+1)^2] - - flow_window = F.unfold( - flow, kernel_size=kernel_size, padding=local_window_radius - ) # [B, 2*(2R+1)^2), H*W] - - flow_window = ( - flow_window.view(b, 2, kernel_size**2, h, w) - .permute(0, 3, 4, 2, 1) - .reshape(b * h * w, kernel_size**2, 2) - ) # [B*H*W, (2R+1)^2, 2] - - scores = torch.matmul(feature0_reshape, feature0_window) / ( - c**0.5 - ) # [B*H*W, 1, (2R+1)^2] - - prob = torch.softmax(scores, dim=-1) - - out = ( - torch.matmul(prob, flow_window) - .view(b, h, w, 2) - .permute(0, 3, 1, 2) - .contiguous() - ) # [B, 2, H, W] - - return out - - -def global_correlation_softmax( - feature0, - feature1, - pred_bidir_flow=False, -): - # global correlation - b, c, h, w = feature0.shape - feature0 = feature0.view(b, c, -1).permute(0, 2, 1) # [B, H*W, C] - feature1 = feature1.view(b, c, -1) # [B, C, H*W] - - correlation = torch.matmul(feature0, feature1).view(b, h, w, h, w) / ( - c**0.5 - ) # [B, H, W, H, W] - - # flow from softmax - init_grid = coords_grid(b, h, w).to(correlation.device) # [B, 2, H, W] - grid = init_grid.view(b, 2, -1).permute(0, 2, 1) # [B, H*W, 2] - - correlation = correlation.view(b, h * w, h * w) # [B, H*W, H*W] - - if pred_bidir_flow: - correlation = torch.cat( - (correlation, correlation.permute(0, 2, 1)), dim=0 - ) # [2*B, H*W, H*W] - init_grid = init_grid.repeat(2, 1, 1, 1) # [2*B, 2, H, W] - grid = grid.repeat(2, 1, 1) # [2*B, H*W, 2] - b = b * 2 - - prob = F.softmax(correlation, dim=-1) # [B, H*W, H*W] - - correspondence = ( - torch.matmul(prob, grid).view(b, h, w, 2).permute(0, 3, 1, 2) - ) # [B, 2, H, W] - - # when predicting bidirectional flow, flow is the concatenation of forward flow and backward flow - flow = correspondence - init_grid - - return flow, prob - - -def local_correlation_softmax( - feature0, - feature1, - local_radius, - padding_mode="zeros", -): - b, c, h, w = feature0.size() - coords_init = coords_grid(b, h, w).to(feature0.device) # [B, 2, H, W] - coords = coords_init.view(b, 2, -1).permute(0, 2, 1) # [B, H*W, 2] - - local_h = 2 * local_radius + 1 - local_w = 2 * local_radius + 1 - - window_grid = generate_window_grid( - -local_radius, - local_radius, - -local_radius, - local_radius, - local_h, - local_w, - device=feature0.device, - ) # [2R+1, 2R+1, 2] - window_grid = window_grid.reshape(-1, 2).repeat(b, 1, 1, 1) # [B, 1, (2R+1)^2, 2] - sample_coords = coords.unsqueeze(-2) + window_grid # [B, H*W, (2R+1)^2, 2] - - sample_coords_softmax = sample_coords - - # exclude coords that are out of image space - valid_x = (sample_coords[:, :, :, 0] >= 0) & ( - sample_coords[:, :, :, 0] < w - ) # [B, H*W, (2R+1)^2] - valid_y = (sample_coords[:, :, :, 1] >= 0) & ( - sample_coords[:, :, :, 1] < h - ) # [B, H*W, (2R+1)^2] - - valid = ( - valid_x & valid_y - ) # [B, H*W, (2R+1)^2], used to mask out invalid values when softmax - - # normalize coordinates to [-1, 1] - sample_coords_norm = normalize_coords(sample_coords, h, w) # [-1, 1] - window_feature = F.grid_sample( - feature1, sample_coords_norm, padding_mode=padding_mode, align_corners=True - ).permute( - 0, 2, 1, 3 - ) # [B, H*W, C, (2R+1)^2] - feature0_view = feature0.permute(0, 2, 3, 1).view(b, h * w, 1, c) # [B, H*W, 1, C] - - corr = torch.matmul(feature0_view, window_feature).view(b, h * w, -1) / ( - c**0.5 - ) # [B, H*W, (2R+1)^2] - - # mask invalid locations - corr[~valid] = -1e9 - - prob = F.softmax(corr, -1) # [B, H*W, (2R+1)^2] - - correspondence = ( - torch.matmul(prob.unsqueeze(-2), sample_coords_softmax) - .squeeze(-2) - .view(b, h, w, 2) - .permute(0, 3, 1, 2) - ) # [B, 2, H, W] - - flow = correspondence - coords_init - match_prob = prob - - return flow, match_prob - - -def coords_grid(b, h, w, homogeneous=False, device=None): - y, x = torch.meshgrid(torch.arange(h), torch.arange(w)) # [H, W] - - stacks = [x, y] - - if homogeneous: - ones = torch.ones_like(x) # [H, W] - stacks.append(ones) - - grid = torch.stack(stacks, dim=0).float() # [2, H, W] or [3, H, W] - - grid = grid[None].repeat(b, 1, 1, 1) # [B, 2, H, W] or [B, 3, H, W] - - if device is not None: - grid = grid.to(device) - - return grid - - -def generate_window_grid(h_min, h_max, w_min, w_max, len_h, len_w, device=None): - assert device is not None - - x, y = torch.meshgrid( - [ - torch.linspace(w_min, w_max, len_w, device=device), - torch.linspace(h_min, h_max, len_h, device=device), - ], - ) - grid = torch.stack((x, y), -1).transpose(0, 1).float() # [H, W, 2] - - return grid - - -def normalize_coords(coords, h, w): - # coords: [B, H, W, 2] - c = torch.Tensor([(w - 1) / 2.0, (h - 1) / 2.0]).float().to(coords.device) - return (coords - c) / c # [-1, 1] - - -def bilinear_sample( - img, sample_coords, mode="bilinear", padding_mode="zeros", return_mask=False -): - # img: [B, C, H, W] - # sample_coords: [B, 2, H, W] in image scale - if sample_coords.size(1) != 2: # [B, H, W, 2] - sample_coords = sample_coords.permute(0, 3, 1, 2) - - b, _, h, w = sample_coords.shape - - # Normalize to [-1, 1] - x_grid = 2 * sample_coords[:, 0] / (w - 1) - 1 - y_grid = 2 * sample_coords[:, 1] / (h - 1) - 1 - - grid = torch.stack([x_grid, y_grid], dim=-1) # [B, H, W, 2] - - img = F.grid_sample( - img, grid, mode=mode, padding_mode=padding_mode, align_corners=True - ) - - if return_mask: - mask = ( - (x_grid >= -1) & (y_grid >= -1) & (x_grid <= 1) & (y_grid <= 1) - ) # [B, H, W] - - return img, mask - - return img - - -def flow_warp(feature, flow, mask=False, padding_mode="zeros"): - b, c, h, w = feature.size() - assert flow.size(1) == 2 - - grid = coords_grid(b, h, w).to(flow.device) + flow # [B, 2, H, W] - - return bilinear_sample(feature, grid, padding_mode=padding_mode, return_mask=mask) - - -def forward_backward_consistency_check(fwd_flow, bwd_flow, alpha=0.01, beta=0.5): - # fwd_flow, bwd_flow: [B, 2, H, W] - # alpha and beta values are following UnFlow (https://arxiv.org/abs/1711.07837) - assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4 - assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2 - flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1) # [B, H, W] - - warped_bwd_flow = flow_warp(bwd_flow, fwd_flow) # [B, 2, H, W] - warped_fwd_flow = flow_warp(fwd_flow, bwd_flow) # [B, 2, H, W] - - diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1) # [B, H, W] - diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1) - - threshold = alpha * flow_mag + beta - - fwd_occ = (diff_fwd > threshold).float() # [B, H, W] - bwd_occ = (diff_bwd > threshold).float() - - return fwd_occ, bwd_occ - - -class PositionEmbeddingSine(nn.Module): - """ - This is a more standard version of the position embedding, very similar to the one - used by the Attention is all you need paper, generalized to work on images. - """ - - def __init__(self, num_pos_feats=64, temperature=10000, normalize=True, scale=None): - super().__init__() - self.num_pos_feats = num_pos_feats - self.temperature = temperature - self.normalize = normalize - if scale is not None and normalize is False: - raise ValueError("normalize should be True if scale is passed") - if scale is None: - scale = 2 * math.pi - self.scale = scale - - def forward(self, x): - # x = tensor_list.tensors # [B, C, H, W] - # mask = tensor_list.mask # [B, H, W], input with padding, valid as 0 - b, c, h, w = x.size() - mask = torch.ones((b, h, w), device=x.device) # [B, H, W] - y_embed = mask.cumsum(1, dtype=torch.float32) - x_embed = mask.cumsum(2, dtype=torch.float32) - if self.normalize: - eps = 1e-6 - y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale - x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale - - dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) - dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) - - pos_x = x_embed[:, :, :, None] / dim_t - pos_y = y_embed[:, :, :, None] / dim_t - pos_x = torch.stack( - (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 - ).flatten(3) - pos_y = torch.stack( - (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 - ).flatten(3) - pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) - return pos - - -def split_feature( - feature, - num_splits=2, - channel_last=False, -): - if channel_last: # [B, H, W, C] - b, h, w, c = feature.size() - assert h % num_splits == 0 and w % num_splits == 0 - - b_new = b * num_splits * num_splits - h_new = h // num_splits - w_new = w // num_splits - - feature = ( - feature.view(b, num_splits, h // num_splits, num_splits, w // num_splits, c) - .permute(0, 1, 3, 2, 4, 5) - .reshape(b_new, h_new, w_new, c) - ) # [B*K*K, H/K, W/K, C] - else: # [B, C, H, W] - b, c, h, w = feature.size() - assert h % num_splits == 0 and w % num_splits == 0 - - b_new = b * num_splits * num_splits - h_new = h // num_splits - w_new = w // num_splits - - feature = ( - feature.view(b, c, num_splits, h // num_splits, num_splits, w // num_splits) - .permute(0, 2, 4, 1, 3, 5) - .reshape(b_new, c, h_new, w_new) - ) # [B*K*K, C, H/K, W/K] - - return feature - - -def merge_splits( - splits, - num_splits=2, - channel_last=False, -): - if channel_last: # [B*K*K, H/K, W/K, C] - b, h, w, c = splits.size() - new_b = b // num_splits // num_splits - - splits = splits.view(new_b, num_splits, num_splits, h, w, c) - merge = ( - splits.permute(0, 1, 3, 2, 4, 5) - .contiguous() - .view(new_b, num_splits * h, num_splits * w, c) - ) # [B, H, W, C] - else: # [B*K*K, C, H/K, W/K] - b, c, h, w = splits.size() - new_b = b // num_splits // num_splits - - splits = splits.view(new_b, num_splits, num_splits, c, h, w) - merge = ( - splits.permute(0, 3, 1, 4, 2, 5) - .contiguous() - .view(new_b, c, num_splits * h, num_splits * w) - ) # [B, C, H, W] - - return merge - - -def normalize_img(img0, img1): - # loaded images are in [0, 255] - # normalize by ImageNet mean and std - mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(img1.device) - std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(img1.device) - img0 = (img0 - mean) / std - img1 = (img1 - mean) / std - - return img0, img1 - - -def feature_add_position(feature0, feature1, attn_splits, feature_channels): - pos_enc = PositionEmbeddingSine(num_pos_feats=feature_channels // 2) - - if attn_splits > 1: # add position in splited window - feature0_splits = split_feature(feature0, num_splits=attn_splits) - feature1_splits = split_feature(feature1, num_splits=attn_splits) - - position = pos_enc(feature0_splits) - - feature0_splits = feature0_splits + position - feature1_splits = feature1_splits + position - - feature0 = merge_splits(feature0_splits, num_splits=attn_splits) - feature1 = merge_splits(feature1_splits, num_splits=attn_splits) - else: - position = pos_enc(feature0) - - feature0 = feature0 + position - feature1 = feature1 + position - - return feature0, feature1 - - -class GMFlow(nn.Module): - def __init__( - self, - num_scales=2, - upsample_factor=4, - feature_channels=128, - attention_type="swin", - num_transformer_layers=6, - ffn_dim_expansion=4, - num_head=1, - **kwargs, - ): - super(GMFlow, self).__init__() - - self.num_scales = num_scales - self.feature_channels = feature_channels - self.upsample_factor = upsample_factor - self.attention_type = attention_type - self.num_transformer_layers = num_transformer_layers - - # CNN backbone - self.backbone = CNNEncoder( - output_dim=feature_channels, num_output_scales=num_scales - ) - - # Transformer - self.transformer = FeatureTransformer( - num_layers=num_transformer_layers, - d_model=feature_channels, - nhead=num_head, - attention_type=attention_type, - ffn_dim_expansion=ffn_dim_expansion, - ) - - # flow propagation with self-attn - self.feature_flow_attn = FeatureFlowAttention(in_channels=feature_channels) - - # convex upsampling: concat feature0 and flow as input - self.upsampler = nn.Sequential( - nn.Conv2d(2 + feature_channels, 256, 3, 1, 1), - nn.ReLU(inplace=True), - nn.Conv2d(256, upsample_factor**2 * 9, 1, 1, 0), - ) - - def extract_feature(self, img0, img1): - concat = torch.cat((img0, img1), dim=0) # [2B, C, H, W] - features = self.backbone( - concat - ) # list of [2B, C, H, W], resolution from high to low - - # reverse: resolution from low to high - features = features[::-1] - - feature0, feature1 = [], [] - - for i in range(len(features)): - feature = features[i] - chunks = torch.chunk(feature, 2, 0) # tuple - feature0.append(chunks[0]) - feature1.append(chunks[1]) - - return feature0, feature1 - - def upsample_flow( - self, - flow, - feature, - bilinear=False, - upsample_factor=8, - ): - if bilinear: - up_flow = ( - F.interpolate( - flow, - scale_factor=upsample_factor, - mode="bilinear", - align_corners=True, - ) - * upsample_factor - ) - - else: - # convex upsampling - concat = torch.cat((flow, feature), dim=1) - - mask = self.upsampler(concat) - b, flow_channel, h, w = flow.shape - mask = mask.view( - b, 1, 9, self.upsample_factor, self.upsample_factor, h, w - ) # [B, 1, 9, K, K, H, W] - mask = torch.softmax(mask, dim=2) - - up_flow = F.unfold(self.upsample_factor * flow, [3, 3], padding=1) - up_flow = up_flow.view( - b, flow_channel, 9, 1, 1, h, w - ) # [B, 2, 9, 1, 1, H, W] - - up_flow = torch.sum(mask * up_flow, dim=2) # [B, 2, K, K, H, W] - up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) # [B, 2, K, H, K, W] - up_flow = up_flow.reshape( - b, flow_channel, self.upsample_factor * h, self.upsample_factor * w - ) # [B, 2, K*H, K*W] - - return up_flow - - def forward( - self, - img0, - img1, - attn_splits_list=[2, 8], - corr_radius_list=[-1, 4], - prop_radius_list=[-1, 1], - pred_bidir_flow=False, - **kwargs, - ): - img0, img1 = normalize_img(img0, img1) # [B, 3, H, W] - - # resolution low to high - feature0_list, feature1_list = self.extract_feature( - img0, img1 - ) # list of features - - flow = None - - assert ( - len(attn_splits_list) - == len(corr_radius_list) - == len(prop_radius_list) - == self.num_scales - ) - - for scale_idx in range(self.num_scales): - feature0, feature1 = feature0_list[scale_idx], feature1_list[scale_idx] - - if pred_bidir_flow and scale_idx > 0: - # predicting bidirectional flow with refinement - feature0, feature1 = torch.cat((feature0, feature1), dim=0), torch.cat( - (feature1, feature0), dim=0 - ) - - upsample_factor = self.upsample_factor * ( - 2 ** (self.num_scales - 1 - scale_idx) - ) - - if scale_idx > 0: - flow = ( - F.interpolate( - flow, scale_factor=2, mode="bilinear", align_corners=True - ) - * 2 - ) - - if flow is not None: - flow = flow.detach() - feature1 = flow_warp(feature1, flow) # [B, C, H, W] - - attn_splits = attn_splits_list[scale_idx] - corr_radius = corr_radius_list[scale_idx] - prop_radius = prop_radius_list[scale_idx] - - # add position to features - feature0, feature1 = feature_add_position( - feature0, feature1, attn_splits, self.feature_channels - ) - - # Transformer - feature0, feature1 = self.transformer( - feature0, feature1, attn_num_splits=attn_splits - ) - - # correlation and softmax - if corr_radius == -1: # global matching - flow_pred = global_correlation_softmax( - feature0, feature1, pred_bidir_flow - )[0] - else: # local matching - flow_pred = local_correlation_softmax(feature0, feature1, corr_radius)[ - 0 - ] - - # flow or residual flow - flow = flow + flow_pred if flow is not None else flow_pred - - # upsample to the original resolution for supervison - if ( - self.training - ): # only need to upsample intermediate flow predictions at training time - flow_bilinear = self.upsample_flow( - flow, None, bilinear=True, upsample_factor=upsample_factor - ) - - # flow propagation with self-attn - if pred_bidir_flow and scale_idx == 0: - feature0 = torch.cat( - (feature0, feature1), dim=0 - ) # [2*B, C, H, W] for propagation - flow = self.feature_flow_attn( - feature0, - flow.detach(), - local_window_attn=prop_radius > 0, - local_window_radius=prop_radius, - ) - - # bilinear upsampling at training time except the last one - if self.training and scale_idx < self.num_scales - 1: - flow_up = self.upsample_flow( - flow, feature0, bilinear=True, upsample_factor=upsample_factor - ) - - if scale_idx == self.num_scales - 1: - flow_up = self.upsample_flow(flow, feature0) - - return flow_up - - -backwarp_tenGrid = {} - - -def backwarp(tenIn, tenflow): - if str(tenflow.shape) not in backwarp_tenGrid: - tenHor = ( - torch.linspace( - start=-1.0, - end=1.0, - steps=tenflow.shape[3], - dtype=tenflow.dtype, - device=tenflow.device, - ) - .view(1, 1, 1, -1) - .repeat(1, 1, tenflow.shape[2], 1) - ) - tenVer = ( - torch.linspace( - start=-1.0, - end=1.0, - steps=tenflow.shape[2], - dtype=tenflow.dtype, - device=tenflow.device, - ) - .view(1, 1, -1, 1) - .repeat(1, 1, 1, tenflow.shape[3]) - ) - - backwarp_tenGrid[str(tenflow.shape)] = torch.cat([tenHor, tenVer], 1).to(get_torch_device()) - # end - - tenflow = torch.cat( - [ - tenflow[:, 0:1, :, :] / ((tenIn.shape[3] - 1.0) / 2.0), - tenflow[:, 1:2, :, :] / ((tenIn.shape[2] - 1.0) / 2.0), - ], - 1, - ) - - return torch.nn.functional.grid_sample( - input=tenIn, - grid=(backwarp_tenGrid[str(tenflow.shape)] + tenflow).permute(0, 2, 3, 1), - mode="bilinear", - padding_mode="zeros", - align_corners=True, - ) - - -class MetricNet(nn.Module): - def __init__(self): - super(MetricNet, self).__init__() - self.metric_in = nn.Conv2d(14, 64, 3, 1, 1) - self.metric_net1 = nn.Sequential(nn.PReLU(), nn.Conv2d(64, 64, 3, 1, 1)) - self.metric_net2 = nn.Sequential(nn.PReLU(), nn.Conv2d(64, 64, 3, 1, 1)) - self.metric_net3 = nn.Sequential(nn.PReLU(), nn.Conv2d(64, 64, 3, 1, 1)) - self.metric_out = nn.Sequential(nn.PReLU(), nn.Conv2d(64, 2, 3, 1, 1)) - - def forward(self, img0, img1, flow01, flow10): - metric0 = F.l1_loss(img0, backwarp(img1, flow01), reduction="none").mean( - [1], True - ) - metric1 = F.l1_loss(img1, backwarp(img0, flow10), reduction="none").mean( - [1], True - ) - - fwd_occ, bwd_occ = forward_backward_consistency_check(flow01, flow10) - - flow01 = torch.cat( - [ - flow01[:, 0:1, :, :] / ((flow01.shape[3] - 1.0) / 2.0), - flow01[:, 1:2, :, :] / ((flow01.shape[2] - 1.0) / 2.0), - ], - 1, - ) - flow10 = torch.cat( - [ - flow10[:, 0:1, :, :] / ((flow10.shape[3] - 1.0) / 2.0), - flow10[:, 1:2, :, :] / ((flow10.shape[2] - 1.0) / 2.0), - ], - 1, - ) - - img = torch.cat((img0, img1), 1) - metric = torch.cat((-metric0, -metric1), 1) - flow = torch.cat((flow01, flow10), 1) - occ = torch.cat((fwd_occ.unsqueeze(1), bwd_occ.unsqueeze(1)), 1) - - feat = self.metric_in(torch.cat((img, metric, flow, occ), 1)) - feat = self.metric_net1(feat) + feat - feat = self.metric_net2(feat) + feat - feat = self.metric_net3(feat) + feat - metric = self.metric_out(feat) - - metric = torch.tanh(metric) * 10 - - return metric[:, :1], metric[:, 1:2] - - -class FeatureNet(nn.Module): - """The quadratic model""" - - def __init__(self): - super(FeatureNet, self).__init__() - self.block1 = nn.Sequential( - nn.PReLU(), - nn.Conv2d(3, 64, 3, 2, 1), - nn.PReLU(), - nn.Conv2d(64, 64, 3, 1, 1), - ) - self.block2 = nn.Sequential( - nn.PReLU(), - nn.Conv2d(64, 128, 3, 2, 1), - nn.PReLU(), - nn.Conv2d(128, 128, 3, 1, 1), - ) - self.block3 = nn.Sequential( - nn.PReLU(), - nn.Conv2d(128, 192, 3, 2, 1), - nn.PReLU(), - nn.Conv2d(192, 192, 3, 1, 1), - ) - - def forward(self, x): - x1 = self.block1(x) - x2 = self.block2(x1) - x3 = self.block3(x2) - - return x1, x2, x3 - - -# Residual Block -def ResidualBlock(in_channels, out_channels, stride=1): - return torch.nn.Sequential( - nn.PReLU(), - nn.Conv2d( - in_channels, - out_channels, - kernel_size=3, - stride=stride, - padding=1, - bias=True, - ), - nn.PReLU(), - nn.Conv2d( - out_channels, - out_channels, - kernel_size=3, - stride=stride, - padding=1, - bias=True, - ), - ) - - -# downsample block -def DownsampleBlock(in_channels, out_channels, stride=2): - return torch.nn.Sequential( - nn.PReLU(), - nn.Conv2d( - in_channels, - out_channels, - kernel_size=3, - stride=stride, - padding=1, - bias=True, - ), - nn.PReLU(), - nn.Conv2d( - out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True - ), - ) - - -# upsample block -def UpsampleBlock(in_channels, out_channels, stride=2): - return torch.nn.Sequential( - nn.PReLU(), - nn.ConvTranspose2d( - in_channels, - out_channels, - kernel_size=4, - stride=stride, - padding=1, - bias=True, - ), - nn.PReLU(), - nn.Conv2d( - out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True - ), - ) - - -class PixelShuffleBlcok(nn.Module): - def __init__(self, in_feat, num_feat, num_out_ch): - super(PixelShuffleBlcok, self).__init__() - self.conv_before_upsample = nn.Sequential( - nn.Conv2d(in_feat, num_feat, 3, 1, 1), nn.PReLU() - ) - self.upsample = nn.Sequential( - nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1), nn.PixelShuffle(2) - ) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - - def forward(self, x): - x = self.conv_before_upsample(x) - x = self.conv_last(self.upsample(x)) - return x - - -# grid network -class GridNet(nn.Module): - def __init__( - self, - in_channels=12, - in_channels1=128, - in_channels2=256, - in_channels3=384, - out_channels=3, - ): - super(GridNet, self).__init__() - - self.residual_model_head = ResidualBlock(in_channels, 64) - self.residual_model_head1 = ResidualBlock(in_channels1, 64) - self.residual_model_head2 = ResidualBlock(in_channels2, 128) - self.residual_model_head3 = ResidualBlock(in_channels3, 192) - - self.residual_model_01 = ResidualBlock(64, 64) - # self.residual_model_02=ResidualBlock(64, 64) - # self.residual_model_03=ResidualBlock(64, 64) - self.residual_model_04 = ResidualBlock(64, 64) - self.residual_model_05 = ResidualBlock(64, 64) - self.residual_model_tail = PixelShuffleBlcok(64, 64, out_channels) - - self.residual_model_11 = ResidualBlock(128, 128) - # self.residual_model_12=ResidualBlock(128, 128) - # self.residual_model_13=ResidualBlock(128, 128) - self.residual_model_14 = ResidualBlock(128, 128) - self.residual_model_15 = ResidualBlock(128, 128) - - self.residual_model_21 = ResidualBlock(192, 192) - # self.residual_model_22=ResidualBlock(192, 192) - # self.residual_model_23=ResidualBlock(192, 192) - self.residual_model_24 = ResidualBlock(192, 192) - self.residual_model_25 = ResidualBlock(192, 192) - - # - - self.downsample_model_10 = DownsampleBlock(64, 128) - self.downsample_model_20 = DownsampleBlock(128, 192) - - self.downsample_model_11 = DownsampleBlock(64, 128) - self.downsample_model_21 = DownsampleBlock(128, 192) - - # self.downsample_model_12=DownsampleBlock(64, 128) - # self.downsample_model_22=DownsampleBlock(128, 192) - - # - - # self.upsample_model_03=UpsampleBlock(128, 64) - # self.upsample_model_13=UpsampleBlock(192, 128) - - self.upsample_model_04 = UpsampleBlock(128, 64) - self.upsample_model_14 = UpsampleBlock(192, 128) - - self.upsample_model_05 = UpsampleBlock(128, 64) - self.upsample_model_15 = UpsampleBlock(192, 128) - - def forward(self, x, x1, x2, x3): - X00 = self.residual_model_head(x) + self.residual_model_head1( - x1 - ) # --- 182 ~ 185 - # X10 = self.residual_model_head1(x1) - - X01 = self.residual_model_01(X00) + X00 # --- 208 ~ 211 ,AddBackward1213 - - X10 = self.downsample_model_10(X00) + self.residual_model_head2( - x2 - ) # --- 186 ~ 189 - X20 = self.downsample_model_20(X10) + self.residual_model_head3( - x3 - ) # --- 190 ~ 193 - - residual_11 = ( - self.residual_model_11(X10) + X10 - ) # 201 ~ 204 , sum AddBackward1206 - downsample_11 = self.downsample_model_11(X01) # 214 ~ 217 - X11 = residual_11 + downsample_11 # --- AddBackward1218 - - residual_21 = ( - self.residual_model_21(X20) + X20 - ) # 194 ~ 197 , sum AddBackward1199 - downsample_21 = self.downsample_model_21(X11) # 219 ~ 222 - X21 = residual_21 + downsample_21 # AddBackward1223 - - X24 = self.residual_model_24(X21) + X21 # --- 224 ~ 227 , AddBackward1229 - X25 = self.residual_model_25(X24) + X24 # --- 230 ~ 233 , AddBackward1235 - - upsample_14 = self.upsample_model_14(X24) # 242 ~ 246 - residual_14 = self.residual_model_14(X11) + X11 # 248 ~ 251, AddBackward1253 - X14 = upsample_14 + residual_14 # --- AddBackward1254 - - upsample_04 = self.upsample_model_04(X14) # 268 ~ 272 - residual_04 = self.residual_model_04(X01) + X01 # 274 ~ 277, AddBackward1279 - X04 = upsample_04 + residual_04 # --- AddBackward1280 - - upsample_15 = self.upsample_model_15(X25) # 236 ~ 240 - residual_15 = self.residual_model_15(X14) + X14 # 255 ~ 258, AddBackward1260 - X15 = upsample_15 + residual_15 # AddBackward1261 - - upsample_05 = self.upsample_model_05(X15) # 262 ~ 266 - residual_05 = self.residual_model_05(X04) + X04 # 281 ~ 284,AddBackward1286 - X05 = upsample_05 + residual_05 # AddBackward1287 - - X_tail = self.residual_model_tail(X05) # 288 ~ 291 - - return X_tail -# end - -class Model: - def __init__(self): - self.flownet = GMFlow() - self.metricnet = MetricNet() - self.feat_ext = FeatureNet() - self.fusionnet = GridNet() - self.version = 3.9 - - def eval(self): - self.flownet.eval() - self.metricnet.eval() - self.feat_ext.eval() - self.fusionnet.eval() - - def device(self): - self.flownet.to(device) - self.metricnet.to(device) - self.feat_ext.to(device) - self.fusionnet.to(device) - - def load_model(self, path_dict): - #models/GMFSS_fortuna_flownet.pkl - self.flownet.load_state_dict(torch.load(path_dict["flownet"])) - #models/GMFSS_fortuna_metric.pkl - self.metricnet.load_state_dict(torch.load(path_dict["metricnet"])) - #models/GMFSS_fortuna_feat.pkl - self.feat_ext.load_state_dict(torch.load(path_dict["feat_ext"])) - #models/GMFSS_fortuna_fusionnet.pkl - self.fusionnet.load_state_dict(torch.load(path_dict["fusionnet"])) - - def reuse(self, img0, img1, scale): - feat11, feat12, feat13 = self.feat_ext(img0) - feat21, feat22, feat23 = self.feat_ext(img1) - - img0 = F.interpolate( - img0, scale_factor=0.5, mode="bilinear", align_corners=False - ) - img1 = F.interpolate( - img1, scale_factor=0.5, mode="bilinear", align_corners=False - ) - - if scale != 1.0: - imgf0 = F.interpolate( - img0, scale_factor=scale, mode="bilinear", align_corners=False - ) - imgf1 = F.interpolate( - img1, scale_factor=scale, mode="bilinear", align_corners=False - ) - else: - imgf0 = img0 - imgf1 = img1 - flow01 = self.flownet(imgf0, imgf1, return_flow=True) - flow10 = self.flownet(imgf1, imgf0, return_flow=True) - if scale != 1.0: - flow01 = ( - F.interpolate( - flow01, - scale_factor=1.0 / scale, - mode="bilinear", - align_corners=False, - ) - / scale - ) - flow10 = ( - F.interpolate( - flow10, - scale_factor=1.0 / scale, - mode="bilinear", - align_corners=False, - ) - / scale - ) - - metric0, metric1 = self.metricnet(img0, img1, flow01, flow10) - - return ( - flow01, - flow10, - metric0, - metric1, - feat11, - feat12, - feat13, - feat21, - feat22, - feat23, - ) - - def inference( - self, - img0, - img1, - flow01, - flow10, - metric0, - metric1, - feat11, - feat12, - feat13, - feat21, - feat22, - feat23, - timestep, - ): - F1t = timestep * flow01 - F2t = (1 - timestep) * flow10 - - Z1t = timestep * metric0 - Z2t = (1 - timestep) * metric1 - - img0 = F.interpolate( - img0, scale_factor=0.5, mode="bilinear", align_corners=False - ) - I1t = softsplat(img0, F1t, Z1t, strMode="soft") - img1 = F.interpolate( - img1, scale_factor=0.5, mode="bilinear", align_corners=False - ) - I2t = softsplat(img1, F2t, Z2t, strMode="soft") - - feat1t1 = softsplat(feat11, F1t, Z1t, strMode="soft") - feat2t1 = softsplat(feat21, F2t, Z2t, strMode="soft") - - F1td = ( - F.interpolate(F1t, scale_factor=0.5, mode="bilinear", align_corners=False) - * 0.5 - ) - Z1d = F.interpolate(Z1t, scale_factor=0.5, mode="bilinear", align_corners=False) - feat1t2 = softsplat(feat12, F1td, Z1d, strMode="soft") - F2td = ( - F.interpolate(F2t, scale_factor=0.5, mode="bilinear", align_corners=False) - * 0.5 - ) - Z2d = F.interpolate(Z2t, scale_factor=0.5, mode="bilinear", align_corners=False) - feat2t2 = softsplat(feat22, F2td, Z2d, strMode="soft") - - F1tdd = ( - F.interpolate(F1t, scale_factor=0.25, mode="bilinear", align_corners=False) - * 0.25 - ) - Z1dd = F.interpolate( - Z1t, scale_factor=0.25, mode="bilinear", align_corners=False - ) - feat1t3 = softsplat(feat13, F1tdd, Z1dd, strMode="soft") - F2tdd = ( - F.interpolate(F2t, scale_factor=0.25, mode="bilinear", align_corners=False) - * 0.25 - ) - Z2dd = F.interpolate( - Z2t, scale_factor=0.25, mode="bilinear", align_corners=False - ) - feat2t3 = softsplat(feat23, F2tdd, Z2dd, strMode="soft") - - out = self.fusionnet( - torch.cat([img0, I1t, I2t, img1], dim=1), - torch.cat([feat1t1, feat2t1], dim=1), - torch.cat([feat1t2, feat2t2], dim=1), - torch.cat([feat1t3, feat2t3], dim=1), - ) - - return torch.clamp(out, 0, 1) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_union.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_union.py deleted file mode 100644 index 41e92dd..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_union.py +++ /dev/null @@ -1,23 +0,0 @@ -import itertools -import numpy as np -import vapoursynth as vs -from .GMFSS_Fortuna_union_arch import Model_inference -import torch - - -class GMFSS_Fortuna_union: - def __init__(self): - self.cache = False - self.amount_input_img = 2 - - torch.set_grad_enabled(False) - torch.backends.cudnn.enabled = True - torch.backends.cudnn.benchmark = True - - self.model = Model_inference() - self.model.eval() - - def execute(self, I0, I1, timestep): - with torch.inference_mode(): - middle = self.model(I0, I1, timestep).cpu() - return middle diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_union_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_union_arch.py deleted file mode 100644 index 665e565..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/GMFSS_Fortuna_union_arch.py +++ /dev/null @@ -1,1857 +0,0 @@ -""" -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/GMFSS_infer_u.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/softsplat.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/FusionNet_u.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/FeatureNet.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/MetricNet.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/IFNet_HDv3.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/gmflow.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/utils.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/position.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/geometry.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/matching.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/transformer.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/backbone.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/gmflow/trident_conv.py -https://github.com/98mxr/GMFSS_Fortuna/blob/b5d0bd544e3f1eee6a059e49c69bcd3124c8343c/model/warplayer.py -""" - -from torch import nn -from torch.nn import functional as F -from torch.nn.modules.utils import _pair -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch -import math -from vfi_models.rife.rife_arch import IFNet -from vfi_models.ops import softsplat -from comfy.model_management import get_torch_device - -device = get_torch_device() -backwarp_tenGrid = {} - - -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) - return torch.nn.functional.grid_sample( - input=tenInput, - grid=g, - mode="bilinear", - padding_mode="border", - align_corners=True, - ) - - -class MultiScaleTridentConv(nn.Module): - def __init__( - self, - in_channels, - out_channels, - kernel_size, - stride=1, - strides=1, - paddings=0, - dilations=1, - dilation=1, - groups=1, - num_branch=1, - test_branch_idx=-1, - bias=False, - norm=None, - activation=None, - ): - super(MultiScaleTridentConv, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = _pair(kernel_size) - self.num_branch = num_branch - self.stride = _pair(stride) - self.groups = groups - self.with_bias = bias - self.dilation = dilation - if isinstance(paddings, int): - paddings = [paddings] * self.num_branch - if isinstance(dilations, int): - dilations = [dilations] * self.num_branch - if isinstance(strides, int): - strides = [strides] * self.num_branch - self.paddings = [_pair(padding) for padding in paddings] - self.dilations = [_pair(dilation) for dilation in dilations] - self.strides = [_pair(stride) for stride in strides] - self.test_branch_idx = test_branch_idx - self.norm = norm - self.activation = activation - - assert len({self.num_branch, len(self.paddings), len(self.strides)}) == 1 - - self.weight = nn.Parameter( - torch.Tensor(out_channels, in_channels // groups, *self.kernel_size) - ) - if bias: - self.bias = nn.Parameter(torch.Tensor(out_channels)) - else: - self.bias = None - - nn.init.kaiming_uniform_(self.weight, nonlinearity="relu") - if self.bias is not None: - nn.init.constant_(self.bias, 0) - - def forward(self, inputs): - num_branch = ( - self.num_branch if self.training or self.test_branch_idx == -1 else 1 - ) - assert len(inputs) == num_branch - - if self.training or self.test_branch_idx == -1: - outputs = [ - F.conv2d( - input, - self.weight, - self.bias, - stride, - padding, - self.dilation, - self.groups, - ) - for input, stride, padding in zip(inputs, self.strides, self.paddings) - ] - else: - outputs = [ - F.conv2d( - inputs[0], - self.weight, - self.bias, - self.strides[self.test_branch_idx] - if self.test_branch_idx == -1 - else self.strides[-1], - self.paddings[self.test_branch_idx] - if self.test_branch_idx == -1 - else self.paddings[-1], - self.dilation, - self.groups, - ) - ] - - if self.norm is not None: - outputs = [self.norm(x) for x in outputs] - if self.activation is not None: - outputs = [self.activation(x) for x in outputs] - return outputs - - -class ResidualBlock_class(nn.Module): - def __init__( - self, - in_planes, - planes, - norm_layer=nn.InstanceNorm2d, - stride=1, - dilation=1, - ): - super(ResidualBlock_class, self).__init__() - - self.conv1 = nn.Conv2d( - in_planes, - planes, - kernel_size=3, - dilation=dilation, - padding=dilation, - stride=stride, - bias=False, - ) - self.conv2 = nn.Conv2d( - planes, - planes, - kernel_size=3, - dilation=dilation, - padding=dilation, - bias=False, - ) - self.relu = nn.ReLU(inplace=True) - - self.norm1 = norm_layer(planes) - self.norm2 = norm_layer(planes) - if not stride == 1 or in_planes != planes: - self.norm3 = norm_layer(planes) - - if stride == 1 and in_planes == planes: - self.downsample = None - else: - self.downsample = nn.Sequential( - nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3 - ) - - def forward(self, x): - y = x - y = self.relu(self.norm1(self.conv1(y))) - y = self.relu(self.norm2(self.conv2(y))) - - if self.downsample is not None: - x = self.downsample(x) - - return self.relu(x + y) - - -class CNNEncoder(nn.Module): - def __init__( - self, - output_dim=128, - norm_layer=nn.InstanceNorm2d, - num_output_scales=1, - **kwargs, - ): - super(CNNEncoder, self).__init__() - self.num_branch = num_output_scales - - feature_dims = [64, 96, 128] - - self.conv1 = nn.Conv2d( - 3, feature_dims[0], kernel_size=7, stride=2, padding=3, bias=False - ) # 1/2 - self.norm1 = norm_layer(feature_dims[0]) - self.relu1 = nn.ReLU(inplace=True) - - self.in_planes = feature_dims[0] - self.layer1 = self._make_layer( - feature_dims[0], stride=1, norm_layer=norm_layer - ) # 1/2 - self.layer2 = self._make_layer( - feature_dims[1], stride=2, norm_layer=norm_layer - ) # 1/4 - - # highest resolution 1/4 or 1/8 - stride = 2 if num_output_scales == 1 else 1 - self.layer3 = self._make_layer( - feature_dims[2], - stride=stride, - norm_layer=norm_layer, - ) # 1/4 or 1/8 - - self.conv2 = nn.Conv2d(feature_dims[2], output_dim, 1, 1, 0) - - if self.num_branch > 1: - if self.num_branch == 4: - strides = (1, 2, 4, 8) - elif self.num_branch == 3: - strides = (1, 2, 4) - elif self.num_branch == 2: - strides = (1, 2) - else: - raise ValueError - - self.trident_conv = MultiScaleTridentConv( - output_dim, - output_dim, - kernel_size=3, - strides=strides, - paddings=1, - num_branch=self.num_branch, - ) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") - elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): - if m.weight is not None: - nn.init.constant_(m.weight, 1) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - def _make_layer(self, dim, stride=1, dilation=1, norm_layer=nn.InstanceNorm2d): - layer1 = ResidualBlock_class( - self.in_planes, dim, norm_layer=norm_layer, stride=stride, dilation=dilation - ) - layer2 = ResidualBlock_class( - dim, dim, norm_layer=norm_layer, stride=1, dilation=dilation - ) - - layers = (layer1, layer2) - - self.in_planes = dim - return nn.Sequential(*layers) - - def forward(self, x): - x = self.conv1(x) - x = self.norm1(x) - x = self.relu1(x) - - x = self.layer1(x) # 1/2 - x = self.layer2(x) # 1/4 - x = self.layer3(x) # 1/8 or 1/4 - - x = self.conv2(x) - - if self.num_branch > 1: - out = self.trident_conv([x] * self.num_branch) # high to low res - else: - out = [x] - - return out - - -def single_head_full_attention(q, k, v): - # q, k, v: [B, L, C] - assert q.dim() == k.dim() == v.dim() == 3 - - scores = torch.matmul(q, k.permute(0, 2, 1)) / (q.size(2) ** 0.5) # [B, L, L] - attn = torch.softmax(scores, dim=2) # [B, L, L] - out = torch.matmul(attn, v) # [B, L, C] - - return out - - -def generate_shift_window_attn_mask( - input_resolution, - window_size_h, - window_size_w, - shift_size_h, - shift_size_w, - device=get_torch_device(), -): - # Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py - # calculate attention mask for SW-MSA - h, w = input_resolution - img_mask = torch.zeros((1, h, w, 1)).to(device) # 1 H W 1 - h_slices = ( - slice(0, -window_size_h), - slice(-window_size_h, -shift_size_h), - slice(-shift_size_h, None), - ) - w_slices = ( - slice(0, -window_size_w), - slice(-window_size_w, -shift_size_w), - slice(-shift_size_w, None), - ) - cnt = 0 - for h in h_slices: - for w in w_slices: - img_mask[:, h, w, :] = cnt - cnt += 1 - - mask_windows = split_feature( - img_mask, num_splits=input_resolution[-1] // window_size_w, channel_last=True - ) - - mask_windows = mask_windows.view(-1, window_size_h * window_size_w) - attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( - attn_mask == 0, float(0.0) - ) - - return attn_mask - - -def single_head_split_window_attention( - q, - k, - v, - num_splits=1, - with_shift=False, - h=None, - w=None, - attn_mask=None, -): - # Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py - # q, k, v: [B, L, C] - assert q.dim() == k.dim() == v.dim() == 3 - - assert h is not None and w is not None - assert q.size(1) == h * w - - b, _, c = q.size() - - b_new = b * num_splits * num_splits - - window_size_h = h // num_splits - window_size_w = w // num_splits - - q = q.view(b, h, w, c) # [B, H, W, C] - k = k.view(b, h, w, c) - v = v.view(b, h, w, c) - - scale_factor = c**0.5 - - if with_shift: - assert attn_mask is not None # compute once - shift_size_h = window_size_h // 2 - shift_size_w = window_size_w // 2 - - q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) - k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) - v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2)) - - q = split_feature( - q, num_splits=num_splits, channel_last=True - ) # [B*K*K, H/K, W/K, C] - k = split_feature(k, num_splits=num_splits, channel_last=True) - v = split_feature(v, num_splits=num_splits, channel_last=True) - - scores = ( - torch.matmul(q.view(b_new, -1, c), k.view(b_new, -1, c).permute(0, 2, 1)) - / scale_factor - ) # [B*K*K, H/K*W/K, H/K*W/K] - - if with_shift: - scores += attn_mask.repeat(b, 1, 1) - - attn = torch.softmax(scores, dim=-1) - - out = torch.matmul(attn, v.view(b_new, -1, c)) # [B*K*K, H/K*W/K, C] - - out = merge_splits( - out.view(b_new, h // num_splits, w // num_splits, c), - num_splits=num_splits, - channel_last=True, - ) # [B, H, W, C] - - # shift back - if with_shift: - out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2)) - - out = out.view(b, -1, c) - - return out - - -class TransformerLayer(nn.Module): - def __init__( - self, - d_model=256, - nhead=1, - attention_type="swin", - no_ffn=False, - ffn_dim_expansion=4, - with_shift=False, - **kwargs, - ): - super(TransformerLayer, self).__init__() - - self.dim = d_model - self.nhead = nhead - self.attention_type = attention_type - self.no_ffn = no_ffn - - self.with_shift = with_shift - - # multi-head attention - self.q_proj = nn.Linear(d_model, d_model, bias=False) - self.k_proj = nn.Linear(d_model, d_model, bias=False) - self.v_proj = nn.Linear(d_model, d_model, bias=False) - - self.merge = nn.Linear(d_model, d_model, bias=False) - - self.norm1 = nn.LayerNorm(d_model) - - # no ffn after self-attn, with ffn after cross-attn - if not self.no_ffn: - in_channels = d_model * 2 - self.mlp = nn.Sequential( - nn.Linear(in_channels, in_channels * ffn_dim_expansion, bias=False), - nn.GELU(), - nn.Linear(in_channels * ffn_dim_expansion, d_model, bias=False), - ) - - self.norm2 = nn.LayerNorm(d_model) - - def forward( - self, - source, - target, - height=None, - width=None, - shifted_window_attn_mask=None, - attn_num_splits=None, - **kwargs, - ): - # source, target: [B, L, C] - query, key, value = source, target, target - - # single-head attention - query = self.q_proj(query) # [B, L, C] - key = self.k_proj(key) # [B, L, C] - value = self.v_proj(value) # [B, L, C] - - if self.attention_type == "swin" and attn_num_splits > 1: - if self.nhead > 1: - # we observe that multihead attention slows down the speed and increases the memory consumption - # without bringing obvious performance gains and thus the implementation is removed - raise NotImplementedError - else: - message = single_head_split_window_attention( - query, - key, - value, - num_splits=attn_num_splits, - with_shift=self.with_shift, - h=height, - w=width, - attn_mask=shifted_window_attn_mask, - ) - else: - message = single_head_full_attention(query, key, value) # [B, L, C] - - message = self.merge(message) # [B, L, C] - message = self.norm1(message) - - if not self.no_ffn: - message = self.mlp(torch.cat([source, message], dim=-1)) - message = self.norm2(message) - - return source + message - - -class TransformerBlock(nn.Module): - """self attention + cross attention + FFN""" - - def __init__( - self, - d_model=256, - nhead=1, - attention_type="swin", - ffn_dim_expansion=4, - with_shift=False, - **kwargs, - ): - super(TransformerBlock, self).__init__() - - self.self_attn = TransformerLayer( - d_model=d_model, - nhead=nhead, - attention_type=attention_type, - no_ffn=True, - ffn_dim_expansion=ffn_dim_expansion, - with_shift=with_shift, - ) - - self.cross_attn_ffn = TransformerLayer( - d_model=d_model, - nhead=nhead, - attention_type=attention_type, - ffn_dim_expansion=ffn_dim_expansion, - with_shift=with_shift, - ) - - def forward( - self, - source, - target, - height=None, - width=None, - shifted_window_attn_mask=None, - attn_num_splits=None, - **kwargs, - ): - # source, target: [B, L, C] - - # self attention - source = self.self_attn( - source, - source, - height=height, - width=width, - shifted_window_attn_mask=shifted_window_attn_mask, - attn_num_splits=attn_num_splits, - ) - - # cross attention and ffn - source = self.cross_attn_ffn( - source, - target, - height=height, - width=width, - shifted_window_attn_mask=shifted_window_attn_mask, - attn_num_splits=attn_num_splits, - ) - - return source - - -class FeatureTransformer(nn.Module): - def __init__( - self, - num_layers=6, - d_model=128, - nhead=1, - attention_type="swin", - ffn_dim_expansion=4, - **kwargs, - ): - super(FeatureTransformer, self).__init__() - - self.attention_type = attention_type - - self.d_model = d_model - self.nhead = nhead - - self.layers = nn.ModuleList( - [ - TransformerBlock( - d_model=d_model, - nhead=nhead, - attention_type=attention_type, - ffn_dim_expansion=ffn_dim_expansion, - with_shift=True - if attention_type == "swin" and i % 2 == 1 - else False, - ) - for i in range(num_layers) - ] - ) - - for p in self.parameters(): - if p.dim() > 1: - nn.init.xavier_uniform_(p) - - def forward( - self, - feature0, - feature1, - attn_num_splits=None, - **kwargs, - ): - b, c, h, w = feature0.shape - assert self.d_model == c - - feature0 = feature0.flatten(-2).permute(0, 2, 1) # [B, H*W, C] - feature1 = feature1.flatten(-2).permute(0, 2, 1) # [B, H*W, C] - - if self.attention_type == "swin" and attn_num_splits > 1: - # global and refine use different number of splits - window_size_h = h // attn_num_splits - window_size_w = w // attn_num_splits - - # compute attn mask once - shifted_window_attn_mask = generate_shift_window_attn_mask( - input_resolution=(h, w), - window_size_h=window_size_h, - window_size_w=window_size_w, - shift_size_h=window_size_h // 2, - shift_size_w=window_size_w // 2, - device=feature0.device, - ) # [K*K, H/K*W/K, H/K*W/K] - else: - shifted_window_attn_mask = None - - # concat feature0 and feature1 in batch dimension to compute in parallel - concat0 = torch.cat((feature0, feature1), dim=0) # [2B, H*W, C] - concat1 = torch.cat((feature1, feature0), dim=0) # [2B, H*W, C] - - for layer in self.layers: - concat0 = layer( - concat0, - concat1, - height=h, - width=w, - shifted_window_attn_mask=shifted_window_attn_mask, - attn_num_splits=attn_num_splits, - ) - - # update feature1 - concat1 = torch.cat(concat0.chunk(chunks=2, dim=0)[::-1], dim=0) - - feature0, feature1 = concat0.chunk(chunks=2, dim=0) # [B, H*W, C] - - # reshape back - feature0 = ( - feature0.view(b, h, w, c).permute(0, 3, 1, 2).contiguous() - ) # [B, C, H, W] - feature1 = ( - feature1.view(b, h, w, c).permute(0, 3, 1, 2).contiguous() - ) # [B, C, H, W] - - return feature0, feature1 - - -class FeatureFlowAttention(nn.Module): - """ - flow propagation with self-attention on feature - query: feature0, key: feature0, value: flow - """ - - def __init__( - self, - in_channels, - **kwargs, - ): - super(FeatureFlowAttention, self).__init__() - - self.q_proj = nn.Linear(in_channels, in_channels) - self.k_proj = nn.Linear(in_channels, in_channels) - - for p in self.parameters(): - if p.dim() > 1: - nn.init.xavier_uniform_(p) - - def forward( - self, - feature0, - flow, - local_window_attn=False, - local_window_radius=1, - **kwargs, - ): - # q, k: feature [B, C, H, W], v: flow [B, 2, H, W] - if local_window_attn: - return self.forward_local_window_attn( - feature0, flow, local_window_radius=local_window_radius - ) - - b, c, h, w = feature0.size() - - query = feature0.view(b, c, h * w).permute(0, 2, 1) # [B, H*W, C] - - # a note: the ``correct'' implementation should be: - # ``query = self.q_proj(query), key = self.k_proj(query)'' - # this problem is observed while cleaning up the code - # however, this doesn't affect the performance since the projection is a linear operation, - # thus the two projection matrices for key can be merged - # so I just leave it as is in order to not re-train all models :) - query = self.q_proj(query) # [B, H*W, C] - key = self.k_proj(query) # [B, H*W, C] - - value = flow.view(b, flow.size(1), h * w).permute(0, 2, 1) # [B, H*W, 2] - - scores = torch.matmul(query, key.permute(0, 2, 1)) / (c**0.5) # [B, H*W, H*W] - prob = torch.softmax(scores, dim=-1) - - out = torch.matmul(prob, value) # [B, H*W, 2] - out = out.view(b, h, w, value.size(-1)).permute(0, 3, 1, 2) # [B, 2, H, W] - - return out - - def forward_local_window_attn( - self, - feature0, - flow, - local_window_radius=1, - ): - assert flow.size(1) == 2 - assert local_window_radius > 0 - - b, c, h, w = feature0.size() - - feature0_reshape = self.q_proj( - feature0.view(b, c, -1).permute(0, 2, 1) - ).reshape( - b * h * w, 1, c - ) # [B*H*W, 1, C] - - kernel_size = 2 * local_window_radius + 1 - - feature0_proj = ( - self.k_proj(feature0.view(b, c, -1).permute(0, 2, 1)) - .permute(0, 2, 1) - .reshape(b, c, h, w) - ) - - feature0_window = F.unfold( - feature0_proj, kernel_size=kernel_size, padding=local_window_radius - ) # [B, C*(2R+1)^2), H*W] - - feature0_window = ( - feature0_window.view(b, c, kernel_size**2, h, w) - .permute(0, 3, 4, 1, 2) - .reshape(b * h * w, c, kernel_size**2) - ) # [B*H*W, C, (2R+1)^2] - - flow_window = F.unfold( - flow, kernel_size=kernel_size, padding=local_window_radius - ) # [B, 2*(2R+1)^2), H*W] - - flow_window = ( - flow_window.view(b, 2, kernel_size**2, h, w) - .permute(0, 3, 4, 2, 1) - .reshape(b * h * w, kernel_size**2, 2) - ) # [B*H*W, (2R+1)^2, 2] - - scores = torch.matmul(feature0_reshape, feature0_window) / ( - c**0.5 - ) # [B*H*W, 1, (2R+1)^2] - - prob = torch.softmax(scores, dim=-1) - - out = ( - torch.matmul(prob, flow_window) - .view(b, h, w, 2) - .permute(0, 3, 1, 2) - .contiguous() - ) # [B, 2, H, W] - - return out - - -def global_correlation_softmax( - feature0, - feature1, - pred_bidir_flow=False, -): - # global correlation - b, c, h, w = feature0.shape - feature0 = feature0.view(b, c, -1).permute(0, 2, 1) # [B, H*W, C] - feature1 = feature1.view(b, c, -1) # [B, C, H*W] - - correlation = torch.matmul(feature0, feature1).view(b, h, w, h, w) / ( - c**0.5 - ) # [B, H, W, H, W] - - # flow from softmax - init_grid = coords_grid(b, h, w).to(correlation.device) # [B, 2, H, W] - grid = init_grid.view(b, 2, -1).permute(0, 2, 1) # [B, H*W, 2] - - correlation = correlation.view(b, h * w, h * w) # [B, H*W, H*W] - - if pred_bidir_flow: - correlation = torch.cat( - (correlation, correlation.permute(0, 2, 1)), dim=0 - ) # [2*B, H*W, H*W] - init_grid = init_grid.repeat(2, 1, 1, 1) # [2*B, 2, H, W] - grid = grid.repeat(2, 1, 1) # [2*B, H*W, 2] - b = b * 2 - - prob = F.softmax(correlation, dim=-1) # [B, H*W, H*W] - - correspondence = ( - torch.matmul(prob, grid).view(b, h, w, 2).permute(0, 3, 1, 2) - ) # [B, 2, H, W] - - # when predicting bidirectional flow, flow is the concatenation of forward flow and backward flow - flow = correspondence - init_grid - - return flow, prob - - -def local_correlation_softmax( - feature0, - feature1, - local_radius, - padding_mode="zeros", -): - b, c, h, w = feature0.size() - coords_init = coords_grid(b, h, w).to(feature0.device) # [B, 2, H, W] - coords = coords_init.view(b, 2, -1).permute(0, 2, 1) # [B, H*W, 2] - - local_h = 2 * local_radius + 1 - local_w = 2 * local_radius + 1 - - window_grid = generate_window_grid( - -local_radius, - local_radius, - -local_radius, - local_radius, - local_h, - local_w, - device=feature0.device, - ) # [2R+1, 2R+1, 2] - window_grid = window_grid.reshape(-1, 2).repeat(b, 1, 1, 1) # [B, 1, (2R+1)^2, 2] - sample_coords = coords.unsqueeze(-2) + window_grid # [B, H*W, (2R+1)^2, 2] - - sample_coords_softmax = sample_coords - - # exclude coords that are out of image space - valid_x = (sample_coords[:, :, :, 0] >= 0) & ( - sample_coords[:, :, :, 0] < w - ) # [B, H*W, (2R+1)^2] - valid_y = (sample_coords[:, :, :, 1] >= 0) & ( - sample_coords[:, :, :, 1] < h - ) # [B, H*W, (2R+1)^2] - - valid = ( - valid_x & valid_y - ) # [B, H*W, (2R+1)^2], used to mask out invalid values when softmax - - # normalize coordinates to [-1, 1] - sample_coords_norm = normalize_coords(sample_coords, h, w) # [-1, 1] - window_feature = F.grid_sample( - feature1, sample_coords_norm, padding_mode=padding_mode, align_corners=True - ).permute( - 0, 2, 1, 3 - ) # [B, H*W, C, (2R+1)^2] - feature0_view = feature0.permute(0, 2, 3, 1).view(b, h * w, 1, c) # [B, H*W, 1, C] - - corr = torch.matmul(feature0_view, window_feature).view(b, h * w, -1) / ( - c**0.5 - ) # [B, H*W, (2R+1)^2] - - # mask invalid locations - corr[~valid] = -1e9 - - prob = F.softmax(corr, -1) # [B, H*W, (2R+1)^2] - - correspondence = ( - torch.matmul(prob.unsqueeze(-2), sample_coords_softmax) - .squeeze(-2) - .view(b, h, w, 2) - .permute(0, 3, 1, 2) - ) # [B, 2, H, W] - - flow = correspondence - coords_init - match_prob = prob - - return flow, match_prob - - -def coords_grid(b, h, w, homogeneous=False, device=None): - y, x = torch.meshgrid(torch.arange(h), torch.arange(w)) # [H, W] - - stacks = [x, y] - - if homogeneous: - ones = torch.ones_like(x) # [H, W] - stacks.append(ones) - - grid = torch.stack(stacks, dim=0).float() # [2, H, W] or [3, H, W] - - grid = grid[None].repeat(b, 1, 1, 1) # [B, 2, H, W] or [B, 3, H, W] - - if device is not None: - grid = grid.to(device) - - return grid - - -def generate_window_grid(h_min, h_max, w_min, w_max, len_h, len_w, device=None): - assert device is not None - - x, y = torch.meshgrid( - [ - torch.linspace(w_min, w_max, len_w, device=device), - torch.linspace(h_min, h_max, len_h, device=device), - ], - ) - grid = torch.stack((x, y), -1).transpose(0, 1).float() # [H, W, 2] - - return grid - - -def normalize_coords(coords, h, w): - # coords: [B, H, W, 2] - c = torch.Tensor([(w - 1) / 2.0, (h - 1) / 2.0]).float().to(coords.device) - return (coords - c) / c # [-1, 1] - - -def bilinear_sample( - img, sample_coords, mode="bilinear", padding_mode="zeros", return_mask=False -): - # img: [B, C, H, W] - # sample_coords: [B, 2, H, W] in image scale - if sample_coords.size(1) != 2: # [B, H, W, 2] - sample_coords = sample_coords.permute(0, 3, 1, 2) - - b, _, h, w = sample_coords.shape - - # Normalize to [-1, 1] - x_grid = 2 * sample_coords[:, 0] / (w - 1) - 1 - y_grid = 2 * sample_coords[:, 1] / (h - 1) - 1 - - grid = torch.stack([x_grid, y_grid], dim=-1) # [B, H, W, 2] - - img = F.grid_sample( - img, grid, mode=mode, padding_mode=padding_mode, align_corners=True - ) - - if return_mask: - mask = ( - (x_grid >= -1) & (y_grid >= -1) & (x_grid <= 1) & (y_grid <= 1) - ) # [B, H, W] - - return img, mask - - return img - - -def flow_warp(feature, flow, mask=False, padding_mode="zeros"): - b, c, h, w = feature.size() - assert flow.size(1) == 2 - - grid = coords_grid(b, h, w).to(flow.device) + flow # [B, 2, H, W] - - return bilinear_sample(feature, grid, padding_mode=padding_mode, return_mask=mask) - - -def forward_backward_consistency_check(fwd_flow, bwd_flow, alpha=0.01, beta=0.5): - # fwd_flow, bwd_flow: [B, 2, H, W] - # alpha and beta values are following UnFlow (https://arxiv.org/abs/1711.07837) - assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4 - assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2 - flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1) # [B, H, W] - - warped_bwd_flow = flow_warp(bwd_flow, fwd_flow) # [B, 2, H, W] - warped_fwd_flow = flow_warp(fwd_flow, bwd_flow) # [B, 2, H, W] - - diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1) # [B, H, W] - diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1) - - threshold = alpha * flow_mag + beta - - fwd_occ = (diff_fwd > threshold).float() # [B, H, W] - bwd_occ = (diff_bwd > threshold).float() - - return fwd_occ, bwd_occ - - -class PositionEmbeddingSine(nn.Module): - """ - This is a more standard version of the position embedding, very similar to the one - used by the Attention is all you need paper, generalized to work on images. - """ - - def __init__(self, num_pos_feats=64, temperature=10000, normalize=True, scale=None): - super().__init__() - self.num_pos_feats = num_pos_feats - self.temperature = temperature - self.normalize = normalize - if scale is not None and normalize is False: - raise ValueError("normalize should be True if scale is passed") - if scale is None: - scale = 2 * math.pi - self.scale = scale - - def forward(self, x): - # x = tensor_list.tensors # [B, C, H, W] - # mask = tensor_list.mask # [B, H, W], input with padding, valid as 0 - b, c, h, w = x.size() - mask = torch.ones((b, h, w), device=x.device) # [B, H, W] - y_embed = mask.cumsum(1, dtype=torch.float32) - x_embed = mask.cumsum(2, dtype=torch.float32) - if self.normalize: - eps = 1e-6 - y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale - x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale - - dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) - dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) - - pos_x = x_embed[:, :, :, None] / dim_t - pos_y = y_embed[:, :, :, None] / dim_t - pos_x = torch.stack( - (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 - ).flatten(3) - pos_y = torch.stack( - (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 - ).flatten(3) - pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) - return pos - - -def split_feature( - feature, - num_splits=2, - channel_last=False, -): - if channel_last: # [B, H, W, C] - b, h, w, c = feature.size() - assert h % num_splits == 0 and w % num_splits == 0 - - b_new = b * num_splits * num_splits - h_new = h // num_splits - w_new = w // num_splits - - feature = ( - feature.view(b, num_splits, h // num_splits, num_splits, w // num_splits, c) - .permute(0, 1, 3, 2, 4, 5) - .reshape(b_new, h_new, w_new, c) - ) # [B*K*K, H/K, W/K, C] - else: # [B, C, H, W] - b, c, h, w = feature.size() - assert h % num_splits == 0 and w % num_splits == 0 - - b_new = b * num_splits * num_splits - h_new = h // num_splits - w_new = w // num_splits - - feature = ( - feature.view(b, c, num_splits, h // num_splits, num_splits, w // num_splits) - .permute(0, 2, 4, 1, 3, 5) - .reshape(b_new, c, h_new, w_new) - ) # [B*K*K, C, H/K, W/K] - - return feature - - -def merge_splits( - splits, - num_splits=2, - channel_last=False, -): - if channel_last: # [B*K*K, H/K, W/K, C] - b, h, w, c = splits.size() - new_b = b // num_splits // num_splits - - splits = splits.view(new_b, num_splits, num_splits, h, w, c) - merge = ( - splits.permute(0, 1, 3, 2, 4, 5) - .contiguous() - .view(new_b, num_splits * h, num_splits * w, c) - ) # [B, H, W, C] - else: # [B*K*K, C, H/K, W/K] - b, c, h, w = splits.size() - new_b = b // num_splits // num_splits - - splits = splits.view(new_b, num_splits, num_splits, c, h, w) - merge = ( - splits.permute(0, 3, 1, 4, 2, 5) - .contiguous() - .view(new_b, c, num_splits * h, num_splits * w) - ) # [B, C, H, W] - - return merge - - -def normalize_img(img0, img1): - # loaded images are in [0, 255] - # normalize by ImageNet mean and std - mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(img1.device) - std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(img1.device) - img0 = (img0 - mean) / std - img1 = (img1 - mean) / std - - return img0, img1 - - -def feature_add_position(feature0, feature1, attn_splits, feature_channels): - pos_enc = PositionEmbeddingSine(num_pos_feats=feature_channels // 2) - - if attn_splits > 1: # add position in splited window - feature0_splits = split_feature(feature0, num_splits=attn_splits) - feature1_splits = split_feature(feature1, num_splits=attn_splits) - - position = pos_enc(feature0_splits) - - feature0_splits = feature0_splits + position - feature1_splits = feature1_splits + position - - feature0 = merge_splits(feature0_splits, num_splits=attn_splits) - feature1 = merge_splits(feature1_splits, num_splits=attn_splits) - else: - position = pos_enc(feature0) - - feature0 = feature0 + position - feature1 = feature1 + position - - return feature0, feature1 - - -class GMFlow(nn.Module): - def __init__( - self, - num_scales=2, - upsample_factor=4, - feature_channels=128, - attention_type="swin", - num_transformer_layers=6, - ffn_dim_expansion=4, - num_head=1, - **kwargs, - ): - super(GMFlow, self).__init__() - - self.num_scales = num_scales - self.feature_channels = feature_channels - self.upsample_factor = upsample_factor - self.attention_type = attention_type - self.num_transformer_layers = num_transformer_layers - - # CNN backbone - self.backbone = CNNEncoder( - output_dim=feature_channels, num_output_scales=num_scales - ) - - # Transformer - self.transformer = FeatureTransformer( - num_layers=num_transformer_layers, - d_model=feature_channels, - nhead=num_head, - attention_type=attention_type, - ffn_dim_expansion=ffn_dim_expansion, - ) - - # flow propagation with self-attn - self.feature_flow_attn = FeatureFlowAttention(in_channels=feature_channels) - - # convex upsampling: concat feature0 and flow as input - self.upsampler = nn.Sequential( - nn.Conv2d(2 + feature_channels, 256, 3, 1, 1), - nn.ReLU(inplace=True), - nn.Conv2d(256, upsample_factor**2 * 9, 1, 1, 0), - ) - - def extract_feature(self, img0, img1): - concat = torch.cat((img0, img1), dim=0) # [2B, C, H, W] - features = self.backbone( - concat - ) # list of [2B, C, H, W], resolution from high to low - - # reverse: resolution from low to high - features = features[::-1] - - feature0, feature1 = [], [] - - for i in range(len(features)): - feature = features[i] - chunks = torch.chunk(feature, 2, 0) # tuple - feature0.append(chunks[0]) - feature1.append(chunks[1]) - - return feature0, feature1 - - def upsample_flow( - self, - flow, - feature, - bilinear=False, - upsample_factor=8, - ): - if bilinear: - up_flow = ( - F.interpolate( - flow, - scale_factor=upsample_factor, - mode="bilinear", - align_corners=True, - ) - * upsample_factor - ) - - else: - # convex upsampling - concat = torch.cat((flow, feature), dim=1) - - mask = self.upsampler(concat) - b, flow_channel, h, w = flow.shape - mask = mask.view( - b, 1, 9, self.upsample_factor, self.upsample_factor, h, w - ) # [B, 1, 9, K, K, H, W] - mask = torch.softmax(mask, dim=2) - - up_flow = F.unfold(self.upsample_factor * flow, [3, 3], padding=1) - up_flow = up_flow.view( - b, flow_channel, 9, 1, 1, h, w - ) # [B, 2, 9, 1, 1, H, W] - - up_flow = torch.sum(mask * up_flow, dim=2) # [B, 2, K, K, H, W] - up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) # [B, 2, K, H, K, W] - up_flow = up_flow.reshape( - b, flow_channel, self.upsample_factor * h, self.upsample_factor * w - ) # [B, 2, K*H, K*W] - - return up_flow - - def forward( - self, - img0, - img1, - attn_splits_list=[2, 8], - corr_radius_list=[-1, 4], - prop_radius_list=[-1, 1], - pred_bidir_flow=False, - **kwargs, - ): - img0, img1 = normalize_img(img0, img1) # [B, 3, H, W] - - # resolution low to high - feature0_list, feature1_list = self.extract_feature( - img0, img1 - ) # list of features - - flow = None - - assert ( - len(attn_splits_list) - == len(corr_radius_list) - == len(prop_radius_list) - == self.num_scales - ) - - for scale_idx in range(self.num_scales): - feature0, feature1 = feature0_list[scale_idx], feature1_list[scale_idx] - - if pred_bidir_flow and scale_idx > 0: - # predicting bidirectional flow with refinement - feature0, feature1 = torch.cat((feature0, feature1), dim=0), torch.cat( - (feature1, feature0), dim=0 - ) - - upsample_factor = self.upsample_factor * ( - 2 ** (self.num_scales - 1 - scale_idx) - ) - - if scale_idx > 0: - flow = ( - F.interpolate( - flow, scale_factor=2, mode="bilinear", align_corners=True - ) - * 2 - ) - - if flow is not None: - flow = flow.detach() - feature1 = flow_warp(feature1, flow) # [B, C, H, W] - - attn_splits = attn_splits_list[scale_idx] - corr_radius = corr_radius_list[scale_idx] - prop_radius = prop_radius_list[scale_idx] - - # add position to features - feature0, feature1 = feature_add_position( - feature0, feature1, attn_splits, self.feature_channels - ) - - # Transformer - feature0, feature1 = self.transformer( - feature0, feature1, attn_num_splits=attn_splits - ) - - # correlation and softmax - if corr_radius == -1: # global matching - flow_pred = global_correlation_softmax( - feature0, feature1, pred_bidir_flow - )[0] - else: # local matching - flow_pred = local_correlation_softmax(feature0, feature1, corr_radius)[ - 0 - ] - - # flow or residual flow - flow = flow + flow_pred if flow is not None else flow_pred - - # upsample to the original resolution for supervison - if ( - self.training - ): # only need to upsample intermediate flow predictions at training time - flow_bilinear = self.upsample_flow( - flow, None, bilinear=True, upsample_factor=upsample_factor - ) - - # flow propagation with self-attn - if pred_bidir_flow and scale_idx == 0: - feature0 = torch.cat( - (feature0, feature1), dim=0 - ) # [2*B, C, H, W] for propagation - flow = self.feature_flow_attn( - feature0, - flow.detach(), - local_window_attn=prop_radius > 0, - local_window_radius=prop_radius, - ) - - # bilinear upsampling at training time except the last one - if self.training and scale_idx < self.num_scales - 1: - flow_up = self.upsample_flow( - flow, feature0, bilinear=True, upsample_factor=upsample_factor - ) - - if scale_idx == self.num_scales - 1: - flow_up = self.upsample_flow(flow, feature0) - - return flow_up - - -backwarp_tenGrid = {} - - -def backwarp(tenIn, tenflow): - if str(tenflow.shape) not in backwarp_tenGrid: - tenHor = ( - torch.linspace( - start=-1.0, - end=1.0, - steps=tenflow.shape[3], - dtype=tenflow.dtype, - device=tenflow.device, - ) - .view(1, 1, 1, -1) - .repeat(1, 1, tenflow.shape[2], 1) - ) - tenVer = ( - torch.linspace( - start=-1.0, - end=1.0, - steps=tenflow.shape[2], - dtype=tenflow.dtype, - device=tenflow.device, - ) - .view(1, 1, -1, 1) - .repeat(1, 1, 1, tenflow.shape[3]) - ) - - backwarp_tenGrid[str(tenflow.shape)] = torch.cat([tenHor, tenVer], 1).to(get_torch_device()) - # end - - tenflow = torch.cat( - [ - tenflow[:, 0:1, :, :] / ((tenIn.shape[3] - 1.0) / 2.0), - tenflow[:, 1:2, :, :] / ((tenIn.shape[2] - 1.0) / 2.0), - ], - 1, - ) - - return torch.nn.functional.grid_sample( - input=tenIn, - grid=(backwarp_tenGrid[str(tenflow.shape)] + tenflow).permute(0, 2, 3, 1), - mode="bilinear", - padding_mode="zeros", - align_corners=True, - ) - - -class MetricNet(nn.Module): - def __init__(self): - super(MetricNet, self).__init__() - self.metric_in = nn.Conv2d(14, 64, 3, 1, 1) - self.metric_net1 = nn.Sequential(nn.PReLU(), nn.Conv2d(64, 64, 3, 1, 1)) - self.metric_net2 = nn.Sequential(nn.PReLU(), nn.Conv2d(64, 64, 3, 1, 1)) - self.metric_net3 = nn.Sequential(nn.PReLU(), nn.Conv2d(64, 64, 3, 1, 1)) - self.metric_out = nn.Sequential(nn.PReLU(), nn.Conv2d(64, 2, 3, 1, 1)) - - def forward(self, img0, img1, flow01, flow10): - metric0 = F.l1_loss(img0, backwarp(img1, flow01), reduction="none").mean( - [1], True - ) - metric1 = F.l1_loss(img1, backwarp(img0, flow10), reduction="none").mean( - [1], True - ) - - fwd_occ, bwd_occ = forward_backward_consistency_check(flow01, flow10) - - flow01 = torch.cat( - [ - flow01[:, 0:1, :, :] / ((flow01.shape[3] - 1.0) / 2.0), - flow01[:, 1:2, :, :] / ((flow01.shape[2] - 1.0) / 2.0), - ], - 1, - ) - flow10 = torch.cat( - [ - flow10[:, 0:1, :, :] / ((flow10.shape[3] - 1.0) / 2.0), - flow10[:, 1:2, :, :] / ((flow10.shape[2] - 1.0) / 2.0), - ], - 1, - ) - - img = torch.cat((img0, img1), 1) - metric = torch.cat((-metric0, -metric1), 1) - flow = torch.cat((flow01, flow10), 1) - occ = torch.cat((fwd_occ.unsqueeze(1), bwd_occ.unsqueeze(1)), 1) - - feat = self.metric_in(torch.cat((img, metric, flow, occ), 1)) - feat = self.metric_net1(feat) + feat - feat = self.metric_net2(feat) + feat - feat = self.metric_net3(feat) + feat - metric = self.metric_out(feat) - - metric = torch.tanh(metric) * 10 - - return metric[:, :1], metric[:, 1:2] - - -class FeatureNet(nn.Module): - """The quadratic model""" - - def __init__(self): - super(FeatureNet, self).__init__() - self.block1 = nn.Sequential( - nn.PReLU(), - nn.Conv2d(3, 64, 3, 2, 1), - nn.PReLU(), - nn.Conv2d(64, 64, 3, 1, 1), - ) - self.block2 = nn.Sequential( - nn.PReLU(), - nn.Conv2d(64, 128, 3, 2, 1), - nn.PReLU(), - nn.Conv2d(128, 128, 3, 1, 1), - ) - self.block3 = nn.Sequential( - nn.PReLU(), - nn.Conv2d(128, 192, 3, 2, 1), - nn.PReLU(), - nn.Conv2d(192, 192, 3, 1, 1), - ) - - def forward(self, x): - x1 = self.block1(x) - x2 = self.block2(x1) - x3 = self.block3(x2) - - return x1, x2, x3 - - -# Residual Block -def ResidualBlock(in_channels, out_channels, stride=1): - return torch.nn.Sequential( - nn.PReLU(), - nn.Conv2d( - in_channels, - out_channels, - kernel_size=3, - stride=stride, - padding=1, - bias=True, - ), - nn.PReLU(), - nn.Conv2d( - out_channels, - out_channels, - kernel_size=3, - stride=stride, - padding=1, - bias=True, - ), - ) - - -# downsample block -def DownsampleBlock(in_channels, out_channels, stride=2): - return torch.nn.Sequential( - nn.PReLU(), - nn.Conv2d( - in_channels, - out_channels, - kernel_size=3, - stride=stride, - padding=1, - bias=True, - ), - nn.PReLU(), - nn.Conv2d( - out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True - ), - ) - - -# upsample block -def UpsampleBlock(in_channels, out_channels, stride=2): - return torch.nn.Sequential( - nn.PReLU(), - nn.ConvTranspose2d( - in_channels, - out_channels, - kernel_size=4, - stride=stride, - padding=1, - bias=True, - ), - nn.PReLU(), - nn.Conv2d( - out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True - ), - ) - - -class PixelShuffleBlcok(nn.Module): - def __init__(self, in_feat, num_feat, num_out_ch): - super(PixelShuffleBlcok, self).__init__() - self.conv_before_upsample = nn.Sequential( - nn.Conv2d(in_feat, num_feat, 3, 1, 1), nn.PReLU() - ) - self.upsample = nn.Sequential( - nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1), nn.PixelShuffle(2) - ) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - - def forward(self, x): - x = self.conv_before_upsample(x) - x = self.conv_last(self.upsample(x)) - return x - - -# grid network -class GridNet(nn.Module): - def __init__( - self, - in_channels=9, - in_channels1=128, - in_channels2=256, - in_channels3=384, - out_channels=3, - ): - super(GridNet, self).__init__() - - self.residual_model_head0 = ResidualBlock(in_channels, 64) - self.residual_model_head1 = ResidualBlock(in_channels1, 64) - self.residual_model_head2 = ResidualBlock(in_channels2, 128) - self.residual_model_head3 = ResidualBlock(in_channels3, 192) - - self.residual_model_01 = ResidualBlock(64, 64) - # self.residual_model_02=ResidualBlock(64, 64) - # self.residual_model_03=ResidualBlock(64, 64) - self.residual_model_04 = ResidualBlock(64, 64) - self.residual_model_05 = ResidualBlock(64, 64) - self.residual_model_tail = PixelShuffleBlcok(64, 64, out_channels) - - self.residual_model_11 = ResidualBlock(128, 128) - # self.residual_model_12=ResidualBlock(128, 128) - # self.residual_model_13=ResidualBlock(128, 128) - self.residual_model_14 = ResidualBlock(128, 128) - self.residual_model_15 = ResidualBlock(128, 128) - - self.residual_model_21 = ResidualBlock(192, 192) - # self.residual_model_22=ResidualBlock(192, 192) - # self.residual_model_23=ResidualBlock(192, 192) - self.residual_model_24 = ResidualBlock(192, 192) - self.residual_model_25 = ResidualBlock(192, 192) - - # - - self.downsample_model_10 = DownsampleBlock(64, 128) - self.downsample_model_20 = DownsampleBlock(128, 192) - - self.downsample_model_11 = DownsampleBlock(64, 128) - self.downsample_model_21 = DownsampleBlock(128, 192) - - # self.downsample_model_12=DownsampleBlock(64, 128) - # self.downsample_model_22=DownsampleBlock(128, 192) - - # - - # self.upsample_model_03=UpsampleBlock(128, 64) - # self.upsample_model_13=UpsampleBlock(192, 128) - - self.upsample_model_04 = UpsampleBlock(128, 64) - self.upsample_model_14 = UpsampleBlock(192, 128) - - self.upsample_model_05 = UpsampleBlock(128, 64) - self.upsample_model_15 = UpsampleBlock(192, 128) - - def forward(self, x, x1, x2, x3): - X00 = self.residual_model_head0(x) + self.residual_model_head1( - x1 - ) # --- 182 ~ 185 - # X10 = self.residual_model_head1(x1) - - X01 = self.residual_model_01(X00) + X00 # --- 208 ~ 211 ,AddBackward1213 - - X10 = self.downsample_model_10(X00) + self.residual_model_head2( - x2 - ) # --- 186 ~ 189 - X20 = self.downsample_model_20(X10) + self.residual_model_head3( - x3 - ) # --- 190 ~ 193 - - residual_11 = ( - self.residual_model_11(X10) + X10 - ) # 201 ~ 204 , sum AddBackward1206 - downsample_11 = self.downsample_model_11(X01) # 214 ~ 217 - X11 = residual_11 + downsample_11 # --- AddBackward1218 - - residual_21 = ( - self.residual_model_21(X20) + X20 - ) # 194 ~ 197 , sum AddBackward1199 - downsample_21 = self.downsample_model_21(X11) # 219 ~ 222 - X21 = residual_21 + downsample_21 # AddBackward1223 - - X24 = self.residual_model_24(X21) + X21 # --- 224 ~ 227 , AddBackward1229 - X25 = self.residual_model_25(X24) + X24 # --- 230 ~ 233 , AddBackward1235 - - upsample_14 = self.upsample_model_14(X24) # 242 ~ 246 - residual_14 = self.residual_model_14(X11) + X11 # 248 ~ 251, AddBackward1253 - X14 = upsample_14 + residual_14 # --- AddBackward1254 - - upsample_04 = self.upsample_model_04(X14) # 268 ~ 272 - residual_04 = self.residual_model_04(X01) + X01 # 274 ~ 277, AddBackward1279 - X04 = upsample_04 + residual_04 # --- AddBackward1280 - - upsample_15 = self.upsample_model_15(X25) # 236 ~ 240 - residual_15 = self.residual_model_15(X14) + X14 # 255 ~ 258, AddBackward1260 - X15 = upsample_15 + residual_15 # AddBackward1261 - - upsample_05 = self.upsample_model_05(X15) # 262 ~ 266 - residual_05 = self.residual_model_05(X04) + X04 # 281 ~ 284,AddBackward1286 - X05 = upsample_05 + residual_05 # AddBackward1287 - - X_tail = self.residual_model_tail(X05) # 288 ~ 291 - - return X_tail -# end - - -class Model: - def __init__(self): - self.flownet = GMFlow() - self.ifnet = IFNet(arch_ver="4.6") - self.metricnet = MetricNet() - self.feat_ext = FeatureNet() - self.fusionnet = GridNet() - self.version = 3.9 - - def eval(self): - self.flownet.eval() - self.ifnet.eval() - self.metricnet.eval() - self.feat_ext.eval() - self.fusionnet.eval() - - def device(self): - self.flownet.to(device) - self.ifnet.to(device) - self.metricnet.to(device) - self.feat_ext.to(device) - self.fusionnet.to(device) - - def load_model(self, path_dict): - #models/rife46.pth - self.ifnet.load_state_dict(torch.load(path_dict["ifnet"])) - #models/GMFSS_fortuna_flownet.pkl - self.flownet.load_state_dict(torch.load(path_dict["flownet"])) - #models/GMFSS_fortuna_union_metric.pkl - self.metricnet.load_state_dict(torch.load(path_dict["metricnet"])) - #models/GMFSS_fortuna_union_feat.pkl - self.feat_ext.load_state_dict(torch.load(path_dict["feat_ext"])) - #models/GMFSS_fortuna_union_fusionnet.pkl - self.fusionnet.load_state_dict(torch.load(path_dict["fusionnet"])) - - def reuse(self, img0, img1, scale): - feat11, feat12, feat13 = self.feat_ext(img0) - feat21, feat22, feat23 = self.feat_ext(img1) - - img0 = F.interpolate( - img0, scale_factor=0.5, mode="bilinear", align_corners=False - ) - img1 = F.interpolate( - img1, scale_factor=0.5, mode="bilinear", align_corners=False - ) - - if scale != 1.0: - imgf0 = F.interpolate( - img0, scale_factor=scale, mode="bilinear", align_corners=False - ) - imgf1 = F.interpolate( - img1, scale_factor=scale, mode="bilinear", align_corners=False - ) - else: - imgf0 = img0 - imgf1 = img1 - flow01 = self.flownet(imgf0, imgf1, return_flow=True) - flow10 = self.flownet(imgf1, imgf0, return_flow=True) - if scale != 1.0: - flow01 = ( - F.interpolate( - flow01, - scale_factor=1.0 / scale, - mode="bilinear", - align_corners=False, - ) - / scale - ) - flow10 = ( - F.interpolate( - flow10, - scale_factor=1.0 / scale, - mode="bilinear", - align_corners=False, - ) - / scale - ) - - metric0, metric1 = self.metricnet(img0, img1, flow01, flow10) - - return ( - flow01, - flow10, - metric0, - metric1, - feat11, - feat12, - feat13, - feat21, - feat22, - feat23, - ) - - def inference( - self, - img0, - img1, - flow01, - flow10, - metric0, - metric1, - feat11, - feat12, - feat13, - feat21, - feat22, - feat23, - timestep, - ): - F1t = timestep * flow01 - F2t = (1 - timestep) * flow10 - - Z1t = timestep * metric0 - Z2t = (1 - timestep) * metric1 - - img0 = F.interpolate( - img0, scale_factor=0.5, mode="bilinear", align_corners=False - ) - I1t = softsplat(img0, F1t, Z1t, strMode="soft") - img1 = F.interpolate( - img1, scale_factor=0.5, mode="bilinear", align_corners=False - ) - I2t = softsplat(img1, F2t, Z2t, strMode="soft") - - rife = self.ifnet(img0, img1, timestep, scale_list=[8, 4, 2, 1]) - - feat1t1 = softsplat(feat11, F1t, Z1t, strMode="soft") - feat2t1 = softsplat(feat21, F2t, Z2t, strMode="soft") - - F1td = ( - F.interpolate(F1t, scale_factor=0.5, mode="bilinear", align_corners=False) - * 0.5 - ) - Z1d = F.interpolate(Z1t, scale_factor=0.5, mode="bilinear", align_corners=False) - feat1t2 = softsplat(feat12, F1td, Z1d, strMode="soft") - F2td = ( - F.interpolate(F2t, scale_factor=0.5, mode="bilinear", align_corners=False) - * 0.5 - ) - Z2d = F.interpolate(Z2t, scale_factor=0.5, mode="bilinear", align_corners=False) - feat2t2 = softsplat(feat22, F2td, Z2d, strMode="soft") - - F1tdd = ( - F.interpolate(F1t, scale_factor=0.25, mode="bilinear", align_corners=False) - * 0.25 - ) - Z1dd = F.interpolate( - Z1t, scale_factor=0.25, mode="bilinear", align_corners=False - ) - feat1t3 = softsplat(feat13, F1tdd, Z1dd, strMode="soft") - F2tdd = ( - F.interpolate(F2t, scale_factor=0.25, mode="bilinear", align_corners=False) - * 0.25 - ) - Z2dd = F.interpolate( - Z2t, scale_factor=0.25, mode="bilinear", align_corners=False - ) - feat2t3 = softsplat(feat23, F2tdd, Z2dd, strMode="soft") - - out = self.fusionnet( - torch.cat([I1t, rife, I2t], dim=1), - torch.cat([feat1t1, feat2t1], dim=1), - torch.cat([feat1t2, feat2t2], dim=1), - torch.cat([feat1t3, feat2t3], dim=1), - ) - - return torch.clamp(out, 0, 1) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/__init__.py deleted file mode 100644 index 3e55bb6..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/gmfss_fortuna/__init__.py +++ /dev/null @@ -1,143 +0,0 @@ -import pathlib -from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames, generic_frame_loop, InterpolationStateList -import typing -import torch -import torch.nn as nn -import torch.nn.functional as F -from comfy.model_management import get_torch_device - - -GLOBAL_MODEL_TYPE = pathlib.Path(__file__).parent.name -CKPTS_PATH_CONFIG = { - "GMFSS_fortuna_union": { - "ifnet": ("rife", "rife46.pth"), - "flownet": (GLOBAL_MODEL_TYPE, "GMFSS_fortuna_flownet.pkl"), - "metricnet": (GLOBAL_MODEL_TYPE, "GMFSS_fortuna_union_metric.pkl"), - "feat_ext": (GLOBAL_MODEL_TYPE, "GMFSS_fortuna_union_feat.pkl"), - "fusionnet": (GLOBAL_MODEL_TYPE, "GMFSS_fortuna_union_fusionnet.pkl") - }, - "GMFSS_fortuna": { - "flownet": (GLOBAL_MODEL_TYPE, "GMFSS_fortuna_flownet.pkl"), - "metricnet": (GLOBAL_MODEL_TYPE, "GMFSS_fortuna_metric.pkl"), - "feat_ext": (GLOBAL_MODEL_TYPE, "GMFSS_fortuna_feat.pkl"), - "fusionnet": (GLOBAL_MODEL_TYPE, "GMFSS_fortuna_fusionnet.pkl") - } -} - -class CommonModelInference(nn.Module): - def __init__(self, model_type): - super(CommonModelInference, self).__init__() - from .GMFSS_Fortuna_arch import Model as GMFSS - from .GMFSS_Fortuna_union_arch import Model as GMFSS_Union - self.model = GMFSS_Union() if "union" in model_type else GMFSS() - self.model.eval() - self.model.device() - _model_path_config = CKPTS_PATH_CONFIG[model_type] - self.model.load_model({ - key: load_file_from_github_release(*_model_path_config[key]) - for key in _model_path_config - }) - - def forward(self, I0, I1, timestep, scale=1.0): - n, c, h, w = I0.shape - tmp = max(64, int(64 / scale)) - ph = ((h - 1) // tmp + 1) * tmp - pw = ((w - 1) // tmp + 1) * tmp - padding = (0, pw - w, 0, ph - h) - I0 = F.pad(I0, padding) - I1 = F.pad(I1, padding) - ( - flow01, - flow10, - metric0, - metric1, - feat11, - feat12, - feat13, - feat21, - feat22, - feat23, - ) = self.model.reuse(I0, I1, scale) - - output = self.model.inference( - I0, - I1, - flow01, - flow10, - metric0, - metric1, - feat11, - feat12, - feat13, - feat21, - feat22, - feat23, - timestep, - ) - return output[:, :, :h, :w] - -class GMFSS_Fortuna_VFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (list(CKPTS_PATH_CONFIG.keys()), ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}), - "multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}), - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames = 10, - multiplier: typing.SupportsInt = 2, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - """ - Perform video frame interpolation using a given checkpoint model. - - Args: - ckpt_name (str): The name of the checkpoint model to use. - frames (torch.Tensor): A tensor containing input video frames. - clear_cache_after_n_frames (int, optional): The number of frames to process before clearing CUDA cache - to prevent memory overflow. Defaults to 10. Lower numbers are safer but mean more processing time. - How high you should set it depends on how many input frames there are, input resolution (after upscaling), - how many times you want to multiply them, and how long you're willing to wait for the process to complete. - multiplier (int, optional): The multiplier for each input frame. 60 input frames * 2 = 120 output frames. Defaults to 2. - - Returns: - tuple: A tuple containing the output interpolated frames. - - Note: - This method interpolates frames in a video sequence using a specified checkpoint model. - It processes each frame sequentially, generating interpolated frames between them. - - To prevent memory overflow, it clears the CUDA cache after processing a specified number of frames. - """ - - interpolation_model = CommonModelInference(model_type=ckpt_name) - interpolation_model.eval().to(get_torch_device()) - frames = preprocess_frames(frames) - - def return_middle_frame(frame_0, frame_1, timestep, model, scale): - return model(frame_0, frame_1, timestep, scale) - - scale = 1 - - args = [interpolation_model, scale] - out = postprocess_frames( - generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, - interpolation_states=optional_interpolation_states, dtype=torch.float32) - ) - return (out,) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/IFRNet_L_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/IFRNet_L_arch.py deleted file mode 100644 index f9554a0..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/IFRNet_L_arch.py +++ /dev/null @@ -1,293 +0,0 @@ -# https://github.com/ltkong218/IFRNet/blob/main/models/IFRNet_L.py -# https://github.com/ltkong218/IFRNet/blob/main/utils.py -import torch -import torch.nn as nn -import torch.nn.functional as F -from comfy.model_management import get_torch_device - - -def warp(img, flow): - B, _, H, W = flow.shape - xx = torch.linspace(-1.0, 1.0, W).view(1, 1, 1, W).expand(B, -1, H, -1) - yy = torch.linspace(-1.0, 1.0, H).view(1, 1, H, 1).expand(B, -1, -1, W) - grid = torch.cat([xx, yy], 1).to(img) - flow_ = torch.cat( - [ - flow[:, 0:1, :, :] / ((W - 1.0) / 2.0), - flow[:, 1:2, :, :] / ((H - 1.0) / 2.0), - ], - 1, - ) - grid_ = (grid + flow_).permute(0, 2, 3, 1) - output = F.grid_sample( - input=img, - grid=grid_, - mode="bilinear", - padding_mode="border", - align_corners=True, - ) - return output - - -def get_robust_weight(flow_pred, flow_gt, beta): - epe = ((flow_pred.detach() - flow_gt) ** 2).sum(dim=1, keepdim=True) ** 0.5 - robust_weight = torch.exp(-beta * epe) - return robust_weight - - -def resize(x, scale_factor): - return F.interpolate( - x, scale_factor=scale_factor, mode="bilinear", align_corners=False - ) - - -def convrelu( - in_channels, - out_channels, - kernel_size=3, - stride=1, - padding=1, - dilation=1, - groups=1, - bias=True, -): - return nn.Sequential( - nn.Conv2d( - in_channels, - out_channels, - kernel_size, - stride, - padding, - dilation, - groups, - bias=bias, - ), - nn.PReLU(out_channels), - ) - - -class ResBlock(nn.Module): - def __init__(self, in_channels, side_channels, bias=True): - super(ResBlock, self).__init__() - self.side_channels = side_channels - self.conv1 = nn.Sequential( - nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias - ), - nn.PReLU(in_channels), - ) - self.conv2 = nn.Sequential( - nn.Conv2d( - side_channels, - side_channels, - kernel_size=3, - stride=1, - padding=1, - bias=bias, - ), - nn.PReLU(side_channels), - ) - self.conv3 = nn.Sequential( - nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias - ), - nn.PReLU(in_channels), - ) - self.conv4 = nn.Sequential( - nn.Conv2d( - side_channels, - side_channels, - kernel_size=3, - stride=1, - padding=1, - bias=bias, - ), - nn.PReLU(side_channels), - ) - self.conv5 = nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias - ) - self.prelu = nn.PReLU(in_channels) - - def forward(self, x): - out = self.conv1(x) - out[:, -self.side_channels :, :, :] = self.conv2( - out[:, -self.side_channels :, :, :] - ) - out = self.conv3(out) - out[:, -self.side_channels :, :, :] = self.conv4( - out[:, -self.side_channels :, :, :] - ) - out = self.prelu(x + self.conv5(out)) - return out - - -class Encoder(nn.Module): - def __init__(self): - super(Encoder, self).__init__() - self.pyramid1 = nn.Sequential( - convrelu(3, 64, 7, 2, 3), convrelu(64, 64, 3, 1, 1) - ) - self.pyramid2 = nn.Sequential( - convrelu(64, 96, 3, 2, 1), convrelu(96, 96, 3, 1, 1) - ) - self.pyramid3 = nn.Sequential( - convrelu(96, 144, 3, 2, 1), convrelu(144, 144, 3, 1, 1) - ) - self.pyramid4 = nn.Sequential( - convrelu(144, 192, 3, 2, 1), convrelu(192, 192, 3, 1, 1) - ) - - def forward(self, img): - f1 = self.pyramid1(img) - f2 = self.pyramid2(f1) - f3 = self.pyramid3(f2) - f4 = self.pyramid4(f3) - return f1, f2, f3, f4 - - -class Decoder4(nn.Module): - def __init__(self): - super(Decoder4, self).__init__() - self.convblock = nn.Sequential( - convrelu(384 + 1, 384), - ResBlock(384, 64), - nn.ConvTranspose2d(384, 148, 4, 2, 1, bias=True), - ) - - def forward(self, f0, f1, embt): - b, c, h, w = f0.shape - embt = embt.repeat(1, 1, h, w) - f_in = torch.cat([f0, f1, embt], 1) - f_out = self.convblock(f_in) - return f_out - - -class Decoder3(nn.Module): - def __init__(self): - super(Decoder3, self).__init__() - self.convblock = nn.Sequential( - convrelu(436, 432), - ResBlock(432, 64), - nn.ConvTranspose2d(432, 100, 4, 2, 1, bias=True), - ) - - def forward(self, ft_, f0, f1, up_flow0, up_flow1): - f0_warp = warp(f0, up_flow0) - f1_warp = warp(f1, up_flow1) - f_in = torch.cat([ft_, f0_warp, f1_warp, up_flow0, up_flow1], 1) - f_out = self.convblock(f_in) - return f_out - - -class Decoder2(nn.Module): - def __init__(self): - super(Decoder2, self).__init__() - self.convblock = nn.Sequential( - convrelu(292, 288), - ResBlock(288, 64), - nn.ConvTranspose2d(288, 68, 4, 2, 1, bias=True), - ) - - def forward(self, ft_, f0, f1, up_flow0, up_flow1): - f0_warp = warp(f0, up_flow0) - f1_warp = warp(f1, up_flow1) - f_in = torch.cat([ft_, f0_warp, f1_warp, up_flow0, up_flow1], 1) - f_out = self.convblock(f_in) - return f_out - - -class Decoder1(nn.Module): - def __init__(self): - super(Decoder1, self).__init__() - self.convblock = nn.Sequential( - convrelu(196, 192), - ResBlock(192, 64), - nn.ConvTranspose2d(192, 8, 4, 2, 1, bias=True), - ) - - def forward(self, ft_, f0, f1, up_flow0, up_flow1): - f0_warp = warp(f0, up_flow0) - f1_warp = warp(f1, up_flow1) - f_in = torch.cat([ft_, f0_warp, f1_warp, up_flow0, up_flow1], 1) - f_out = self.convblock(f_in) - return f_out - - -class IRFNet_L(nn.Module): - def __init__(self): - super(IRFNet_L, self).__init__() - self.encoder = Encoder() - self.decoder4 = Decoder4() - self.decoder3 = Decoder3() - self.decoder2 = Decoder2() - self.decoder1 = Decoder1() - - def forward(self, img0, img1, scale_factor=1.0, timestep=0.5): - # emb1 = torch.tensor(1/2).view(1, 1, 1, 1).float() - # emb2 = torch.tensor(2/2).view(1, 1, 1, 1).float() - # embt = torch.cat([emb1, emb2], 0) - 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) - - #Support multiple batches - embt = torch.tensor([timestep] * n).view(n, 1, 1, 1).float().to(get_torch_device()) - if "HalfTensor" in str(img0.type()): - embt = embt.half() - - mean_ = ( - torch.cat([img0, img1], 2) - .mean(1, keepdim=True) - .mean(2, keepdim=True) - .mean(3, keepdim=True) - ) - img0 = img0 - mean_ - img1 = img1 - mean_ - - img0_ = resize(img0, scale_factor=scale_factor) - img1_ = resize(img1, scale_factor=scale_factor) - - f0_1, f0_2, f0_3, f0_4 = self.encoder(img0_) - f1_1, f1_2, f1_3, f1_4 = self.encoder(img1_) - - out4 = self.decoder4(f0_4, f1_4, embt) - up_flow0_4 = out4[:, 0:2] - up_flow1_4 = out4[:, 2:4] - ft_3_ = out4[:, 4:] - - out3 = self.decoder3(ft_3_, f0_3, f1_3, up_flow0_4, up_flow1_4) - up_flow0_3 = out3[:, 0:2] + 2.0 * resize(up_flow0_4, scale_factor=2.0) - up_flow1_3 = out3[:, 2:4] + 2.0 * resize(up_flow1_4, scale_factor=2.0) - ft_2_ = out3[:, 4:] - - out2 = self.decoder2(ft_2_, f0_2, f1_2, up_flow0_3, up_flow1_3) - up_flow0_2 = out2[:, 0:2] + 2.0 * resize(up_flow0_3, scale_factor=2.0) - up_flow1_2 = out2[:, 2:4] + 2.0 * resize(up_flow1_3, scale_factor=2.0) - ft_1_ = out2[:, 4:] - - out1 = self.decoder1(ft_1_, f0_1, f1_1, up_flow0_2, up_flow1_2) - up_flow0_1 = out1[:, 0:2] + 2.0 * resize(up_flow0_2, scale_factor=2.0) - up_flow1_1 = out1[:, 2:4] + 2.0 * resize(up_flow1_2, scale_factor=2.0) - up_mask_1 = torch.sigmoid(out1[:, 4:5]) - up_res_1 = out1[:, 5:] - - up_flow0_1 = resize(up_flow0_1, scale_factor=(1.0 / scale_factor)) * ( - 1.0 / scale_factor - ) - up_flow1_1 = resize(up_flow1_1, scale_factor=(1.0 / scale_factor)) * ( - 1.0 / scale_factor - ) - up_mask_1 = resize(up_mask_1, scale_factor=(1.0 / scale_factor)) - up_res_1 = resize(up_res_1, scale_factor=(1.0 / scale_factor)) - - img0_warp = warp(img0, up_flow0_1) - img1_warp = warp(img1, up_flow1_1) - imgt_merge = up_mask_1 * img0_warp + (1 - up_mask_1) * img1_warp + mean_ - imgt_pred = imgt_merge + up_res_1 - imgt_pred = torch.clamp(imgt_pred, 0, 1) - return imgt_pred[:, :, :h, :w] diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/IFRNet_S_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/IFRNet_S_arch.py deleted file mode 100644 index 4a2282b..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/IFRNet_S_arch.py +++ /dev/null @@ -1,293 +0,0 @@ -# https://github.com/ltkong218/IFRNet/blob/main/models/IFRNet_S.py -# https://github.com/ltkong218/IFRNet/blob/main/utils.py -import torch -import torch.nn as nn -import torch.nn.functional as F -from comfy.model_management import get_torch_device - - -def warp(img, flow): - B, _, H, W = flow.shape - xx = torch.linspace(-1.0, 1.0, W).view(1, 1, 1, W).expand(B, -1, H, -1) - yy = torch.linspace(-1.0, 1.0, H).view(1, 1, H, 1).expand(B, -1, -1, W) - grid = torch.cat([xx, yy], 1).to(img) - flow_ = torch.cat( - [ - flow[:, 0:1, :, :] / ((W - 1.0) / 2.0), - flow[:, 1:2, :, :] / ((H - 1.0) / 2.0), - ], - 1, - ) - grid_ = (grid + flow_).permute(0, 2, 3, 1) - output = F.grid_sample( - input=img, - grid=grid_, - mode="bilinear", - padding_mode="border", - align_corners=True, - ) - return output - - -def get_robust_weight(flow_pred, flow_gt, beta): - epe = ((flow_pred.detach() - flow_gt) ** 2).sum(dim=1, keepdim=True) ** 0.5 - robust_weight = torch.exp(-beta * epe) - return robust_weight - - -def resize(x, scale_factor): - return F.interpolate( - x, scale_factor=scale_factor, mode="bilinear", align_corners=False - ) - - -def convrelu( - in_channels, - out_channels, - kernel_size=3, - stride=1, - padding=1, - dilation=1, - groups=1, - bias=True, -): - return nn.Sequential( - nn.Conv2d( - in_channels, - out_channels, - kernel_size, - stride, - padding, - dilation, - groups, - bias=bias, - ), - nn.PReLU(out_channels), - ) - - -class ResBlock(nn.Module): - def __init__(self, in_channels, side_channels, bias=True): - super(ResBlock, self).__init__() - self.side_channels = side_channels - self.conv1 = nn.Sequential( - nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias - ), - nn.PReLU(in_channels), - ) - self.conv2 = nn.Sequential( - nn.Conv2d( - side_channels, - side_channels, - kernel_size=3, - stride=1, - padding=1, - bias=bias, - ), - nn.PReLU(side_channels), - ) - self.conv3 = nn.Sequential( - nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias - ), - nn.PReLU(in_channels), - ) - self.conv4 = nn.Sequential( - nn.Conv2d( - side_channels, - side_channels, - kernel_size=3, - stride=1, - padding=1, - bias=bias, - ), - nn.PReLU(side_channels), - ) - self.conv5 = nn.Conv2d( - in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias - ) - self.prelu = nn.PReLU(in_channels) - - def forward(self, x): - out = self.conv1(x) - out[:, -self.side_channels :, :, :] = self.conv2( - out[:, -self.side_channels :, :, :] - ) - out = self.conv3(out) - out[:, -self.side_channels :, :, :] = self.conv4( - out[:, -self.side_channels :, :, :] - ) - out = self.prelu(x + self.conv5(out)) - return out - - -class Encoder(nn.Module): - def __init__(self): - super(Encoder, self).__init__() - self.pyramid1 = nn.Sequential( - convrelu(3, 24, 3, 2, 1), convrelu(24, 24, 3, 1, 1) - ) - self.pyramid2 = nn.Sequential( - convrelu(24, 36, 3, 2, 1), convrelu(36, 36, 3, 1, 1) - ) - self.pyramid3 = nn.Sequential( - convrelu(36, 54, 3, 2, 1), convrelu(54, 54, 3, 1, 1) - ) - self.pyramid4 = nn.Sequential( - convrelu(54, 72, 3, 2, 1), convrelu(72, 72, 3, 1, 1) - ) - - def forward(self, img): - f1 = self.pyramid1(img) - f2 = self.pyramid2(f1) - f3 = self.pyramid3(f2) - f4 = self.pyramid4(f3) - return f1, f2, f3, f4 - - -class Decoder4(nn.Module): - def __init__(self): - super(Decoder4, self).__init__() - self.convblock = nn.Sequential( - convrelu(144 + 1, 144), - ResBlock(144, 24), - nn.ConvTranspose2d(144, 58, 4, 2, 1, bias=True), - ) - - def forward(self, f0, f1, embt): - b, c, h, w = f0.shape - embt = embt.repeat(1, 1, h, w) - f_in = torch.cat([f0, f1, embt], 1) - f_out = self.convblock(f_in) - return f_out - - -class Decoder3(nn.Module): - def __init__(self): - super(Decoder3, self).__init__() - self.convblock = nn.Sequential( - convrelu(166, 162), - ResBlock(162, 24), - nn.ConvTranspose2d(162, 40, 4, 2, 1, bias=True), - ) - - def forward(self, ft_, f0, f1, up_flow0, up_flow1): - f0_warp = warp(f0, up_flow0) - f1_warp = warp(f1, up_flow1) - f_in = torch.cat([ft_, f0_warp, f1_warp, up_flow0, up_flow1], 1) - f_out = self.convblock(f_in) - return f_out - - -class Decoder2(nn.Module): - def __init__(self): - super(Decoder2, self).__init__() - self.convblock = nn.Sequential( - convrelu(112, 108), - ResBlock(108, 24), - nn.ConvTranspose2d(108, 28, 4, 2, 1, bias=True), - ) - - def forward(self, ft_, f0, f1, up_flow0, up_flow1): - f0_warp = warp(f0, up_flow0) - f1_warp = warp(f1, up_flow1) - f_in = torch.cat([ft_, f0_warp, f1_warp, up_flow0, up_flow1], 1) - f_out = self.convblock(f_in) - return f_out - - -class Decoder1(nn.Module): - def __init__(self): - super(Decoder1, self).__init__() - self.convblock = nn.Sequential( - convrelu(76, 72), - ResBlock(72, 24), - nn.ConvTranspose2d(72, 8, 4, 2, 1, bias=True), - ) - - def forward(self, ft_, f0, f1, up_flow0, up_flow1): - f0_warp = warp(f0, up_flow0) - f1_warp = warp(f1, up_flow1) - f_in = torch.cat([ft_, f0_warp, f1_warp, up_flow0, up_flow1], 1) - f_out = self.convblock(f_in) - return f_out - - -class IRFNet_S(nn.Module): - def __init__(self): - super(IRFNet_S, self).__init__() - self.encoder = Encoder() - self.decoder4 = Decoder4() - self.decoder3 = Decoder3() - self.decoder2 = Decoder2() - self.decoder1 = Decoder1() - - def forward(self, img0, img1, scale_factor=1.0, timestep=0.5): - # emb1 = torch.tensor(1/2).view(1, 1, 1, 1).float() - # emb2 = torch.tensor(2/2).view(1, 1, 1, 1).float() - # embt = torch.cat([emb1, emb2], 0) - 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) - - #Support multiple batches - embt = torch.tensor([timestep] * n).view(n, 1, 1, 1).float().to(get_torch_device()) - if "HalfTensor" in str(img0.type()): - embt = embt.half() - - mean_ = ( - torch.cat([img0, img1], 2) - .mean(1, keepdim=True) - .mean(2, keepdim=True) - .mean(3, keepdim=True) - ) - img0 = img0 - mean_ - img1 = img1 - mean_ - - img0_ = resize(img0, scale_factor=scale_factor) - img1_ = resize(img1, scale_factor=scale_factor) - - f0_1, f0_2, f0_3, f0_4 = self.encoder(img0_) - f1_1, f1_2, f1_3, f1_4 = self.encoder(img1_) - - out4 = self.decoder4(f0_4, f1_4, embt) - up_flow0_4 = out4[:, 0:2] - up_flow1_4 = out4[:, 2:4] - ft_3_ = out4[:, 4:] - - out3 = self.decoder3(ft_3_, f0_3, f1_3, up_flow0_4, up_flow1_4) - up_flow0_3 = out3[:, 0:2] + 2.0 * resize(up_flow0_4, scale_factor=2.0) - up_flow1_3 = out3[:, 2:4] + 2.0 * resize(up_flow1_4, scale_factor=2.0) - ft_2_ = out3[:, 4:] - - out2 = self.decoder2(ft_2_, f0_2, f1_2, up_flow0_3, up_flow1_3) - up_flow0_2 = out2[:, 0:2] + 2.0 * resize(up_flow0_3, scale_factor=2.0) - up_flow1_2 = out2[:, 2:4] + 2.0 * resize(up_flow1_3, scale_factor=2.0) - ft_1_ = out2[:, 4:] - - out1 = self.decoder1(ft_1_, f0_1, f1_1, up_flow0_2, up_flow1_2) - up_flow0_1 = out1[:, 0:2] + 2.0 * resize(up_flow0_2, scale_factor=2.0) - up_flow1_1 = out1[:, 2:4] + 2.0 * resize(up_flow1_2, scale_factor=2.0) - up_mask_1 = torch.sigmoid(out1[:, 4:5]) - up_res_1 = out1[:, 5:] - - up_flow0_1 = resize(up_flow0_1, scale_factor=(1.0 / scale_factor)) * ( - 1.0 / scale_factor - ) - up_flow1_1 = resize(up_flow1_1, scale_factor=(1.0 / scale_factor)) * ( - 1.0 / scale_factor - ) - up_mask_1 = resize(up_mask_1, scale_factor=(1.0 / scale_factor)) - up_res_1 = resize(up_res_1, scale_factor=(1.0 / scale_factor)) - - img0_warp = warp(img0, up_flow0_1) - img1_warp = warp(img1, up_flow1_1) - imgt_merge = up_mask_1 * img0_warp + (1 - up_mask_1) * img1_warp + mean_ - imgt_pred = imgt_merge + up_res_1 - imgt_pred = torch.clamp(imgt_pred, 0, 1) - return imgt_pred[:, :, :h, :w] diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/__init__.py deleted file mode 100644 index 2d7d3ce..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ifrnet/__init__.py +++ /dev/null @@ -1,57 +0,0 @@ -import torch -import pathlib -from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames -import typing -from comfy.model_management import get_torch_device -from vfi_utils import generic_frame_loop, InterpolationStateList - -MODEL_TYPE = pathlib.Path(__file__).parent.name -CKPT_NAMES = ["IFRNet_S_Vimeo90K.pth", "IFRNet_L_Vimeo90K.pth"] - -class IFRNet_VFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (CKPT_NAMES, ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}), - "multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}), - "scale_factor": ([0.25, 0.5, 1.0, 2.0, 4.0], {"default": 1.0}), - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames: typing.SupportsInt = 1, - multiplier: typing.SupportsInt = 2, - scale_factor: typing.SupportsFloat = 1.0, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - from .IFRNet_S_arch import IRFNet_S - from .IFRNet_L_arch import IRFNet_L - model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) - interpolation_model = IRFNet_S() if 'S' in ckpt_name else IRFNet_L() - interpolation_model.load_state_dict(torch.load(model_path)) - interpolation_model.eval().to(get_torch_device()) - frames = preprocess_frames(frames) - - def return_middle_frame(frame_0, frame_1, timestep, model, scale_factor): - return model(frame_0, frame_1, timestep, scale_factor) - - args = [interpolation_model, scale_factor] - out = postprocess_frames( - generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, - interpolation_states=optional_interpolation_states, dtype=torch.float32) - ) - return (out,) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ifunet/IFUNet_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ifunet/IFUNet_arch.py deleted file mode 100644 index 8206613..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ifunet/IFUNet_arch.py +++ /dev/null @@ -1,766 +0,0 @@ -""" -https://github.com/98mxr/IFUNet/blob/main/model/IFUNet.py -https://github.com/98mxr/IFUNet/blob/main/model/cbam.py -https://github.com/98mxr/IFUNet/blob/main/model/warplayer.py -https://github.com/98mxr/IFUNet/blob/5be535c8cff66d6fa1967252685719df4c0620e4/model/RIFE.py -https://github.com/98mxr/IFUNet/blob/main/model/rrdb.py -https://github.com/98mxr/IFUNet/blob/main/model/ResynNet.py -""" -import torch -import torch.nn as nn -import torch.nn.functional as F -from comfy.model_management import get_torch_device - -backwarp_tenGrid = {} -device = get_torch_device() - - -def conv(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, - ), - nn.PReLU(out_planes), - ) - - -def conv_bn(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=False, - ), - nn.BatchNorm2d(out_planes), - nn.PReLU(out_planes), - ) - - -class DegCNN(nn.Module): - def __init__(self): - super(DegCNN, self).__init__() - self.conv0 = conv(3, 32, 3, 2, 1) - self.conv1 = conv(32, 32, 3, 2, 1) - self.conv2 = conv(32, 32, 3, 2, 1) - self.conv3 = conv(32, 32, 3, 2, 1) - self.deconv = nn.Sequential( - nn.Dropout2d(0.95), - nn.ConvTranspose2d(4 * 32, 32, 4, 2, 1), - nn.PReLU(32), - nn.Conv2d(32, 3, 3, 1, 1), - nn.Sigmoid(), - ) - - def forward(self, x): - f0 = self.conv0(x) - f1 = self.conv1(f0) - f2 = self.conv2(f1) - f3 = self.conv3(f2) - f1 = F.interpolate(f1, scale_factor=2.0, mode="bilinear", align_corners=False) - f2 = F.interpolate(f2, scale_factor=4.0, mode="bilinear", align_corners=False) - f3 = F.interpolate(f3, scale_factor=8.0, mode="bilinear", align_corners=False) - return self.deconv(torch.cat((f0, f1, f2, f3), 1)) - - -class FlowBlock(nn.Module): - def __init__(self, in_planes, c=64): - super(FlowBlock, self).__init__() - self.conv0 = nn.Sequential( - conv_bn(in_planes, c // 2, 3, 2, 1), - conv_bn(c // 2, c, 3, 2, 1), - conv_bn(c, 2 * c, 3, 2, 1), - ) - self.convblock = nn.Sequential( - conv_bn(2 * c, 2 * c), - conv_bn(2 * c, 2 * c), - conv_bn(2 * c, 2 * c), - conv_bn(2 * c, 2 * c), - conv_bn(2 * c, 2 * c), - conv_bn(2 * c, 2 * c), - ) - self.lastconv = nn.ConvTranspose2d(2 * c, 4, 4, 2, 1) - - def forward(self, x, flow, 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) - feat = self.convblock(feat) + feat - tmp = self.lastconv(feat) - tmp = F.interpolate( - tmp, scale_factor=scale * 4, mode="bilinear", align_corners=False - ) - flow = tmp[:, :2] * scale * 4 - mask = tmp[:, 2:3] - return flow, mask - - -class ResynNet(nn.Module): - def __init__(self): - super(ResynNet, self).__init__() - self.block0 = FlowBlock(6, c=128) - self.block1 = FlowBlock(12, c=128) - self.block2 = FlowBlock(12, c=128) - self.degrad = DegCNN() - # Contextual Refinement context + decode - self.context0 = nn.Sequential( - conv(3, 16, 3, 2, 1), - conv(16, 32, 3, 2, 1), - ) - self.context1 = nn.Sequential( - conv(3, 16, 3, 2, 1), - conv(16, 32, 3, 2, 1), - ) - self.decode = nn.Sequential( - nn.ConvTranspose2d(64, 32, 4, 2, 1), - nn.ConvTranspose2d(32, 3, 4, 2, 1), - nn.Tanh(), - ) - - def calflow(self, img0, lowres, scale): - flow = None - stu = [self.block0, self.block1, self.block2] - for i in range(3): - if flow is not None: - flow_d, mask_d = stu[i]( - torch.cat((img0, lowres, warped_img0, mask), 1), - flow, - scale=scale[i], - ) - flow = flow + flow_d - mask = mask + mask_d - else: - flow, mask = stu[i](torch.cat((img0, lowres), 1), None, scale=scale[i]) - warped_img0 = warp(img0, flow) - flow_down = ( - F.interpolate(flow, scale_factor=0.25, mode="bilinear", align_corners=False) - * 0.25 - ) - c0 = warp(self.context0(img0), flow_down) - c1 = self.context1(warped_img0) - warped_img0 = warped_img0 + self.decode(torch.cat((c0, c1), 1)) - return flow, mask, torch.clamp(warped_img0, 0, 1) - - def forward( - self, x, deg=None, gt=None, scale=[4, 2, 1], training=False, blend=True - ): - if training: - deg = self.degrad(gt) - loss_cons = (gt - deg).abs().mean() - else: - loss_cons = torch.tensor([0]) - img_list = [] - N = x.shape[1] // 3 - for i in range(N): - img_list.append(x[:, i * 3 : i * 3 + 3]) - warped_list = [] - merged = [] - mask_list = [] - flow_list = [] - for i in range(N): - f, m, img = self.calflow(img_list[i], deg.detach(), scale) - mask_list.append(m) - warped_list.append(img) - flow_list.append(f) - if blend: - N += 1 - mask_list.append(m * 0) - warped_list.append(deg) - mask = F.softmax(torch.clamp(torch.cat(mask_list, 1), -4, 4), dim=1) - merged = 0 - for i in range(N): - merged += warped_list[i] * mask[:, i : i + 1] - return merged, loss_cons - - -def make_layer(basic_block, num_basic_block, **kwarg): - """Make layers by stacking the same blocks. - Args: - basic_block (nn.module): nn.module class for basic block. - num_basic_block (int): number of blocks. - Returns: - nn.Sequential: Stacked blocks in nn.Sequential. - """ - layers = [] - for _ in range(num_basic_block): - layers.append(basic_block(**kwarg)) - return nn.Sequential(*layers) - - -class ResidualDenseBlock(nn.Module): - """Residual Dense Block. - - Used in RRDB block in ESRGAN. - - Args: - num_feat (int): Channel number of intermediate features. - num_grow_ch (int): Channels for each growth. - """ - - def __init__(self, num_feat=64, num_grow_ch=32): - super(ResidualDenseBlock, self).__init__() - self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1) - self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1) - self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1) - self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1) - self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1) - - self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) - - # initialization - # default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1) - # 只能先取消，default_init_weights来自basicsr.arch_util - - def forward(self, x): - x1 = self.lrelu(self.conv1(x)) - x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1))) - x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1))) - x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1))) - x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1)) - # Emperically, we use 0.2 to scale the residual for better performance - # 原作者这么说我就这么听着吧 - return x5 * 0.2 + x - - -class RRDB(nn.Module): - """Residual in Residual Dense Block. - - Used in RRDB-Net in ESRGAN. - - Args: - num_feat (int): Channel number of intermediate features. - num_grow_ch (int): Channels for each growth. - """ - - def __init__(self, num_feat, num_grow_ch=32): - super(RRDB, self).__init__() - self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch) - self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch) - self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch) - - def forward(self, x): - out = self.rdb1(x) - out = self.rdb2(out) - out = self.rdb3(out) - # Emperically, we use 0.2 to scale the residual for better performance - # 原作者这么说我就这么听着吧 - return out * 0.2 + x - - -class RRDBNet(nn.Module): - """Networks consisting of Residual in Residual Dense Block, which is used - in ESRGAN. - - ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks. - - We extend ESRGAN for scale x2 and scale x1. - Note: This is one option for scale 1, scale 2 in RRDBNet. - We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size - and enlarge the channel size before feeding inputs into the main ESRGAN architecture. - - Args: - num_in_ch (int): Channel number of inputs. - num_out_ch (int): Channel number of outputs. - num_feat (int): Channel number of intermediate features. - Default: 64 - num_block (int): Block number in the trunk network. Defaults: 23 - num_grow_ch (int): Channels for each growth. Default: 32. - """ - - def __init__( - self, num_in_ch=16, num_out_ch=1, num_feat=64, num_block=6, num_grow_ch=32 - ): - super(RRDBNet, self).__init__() - self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1) - self.body = make_layer( - RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch - ) - self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - # upsample - self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) - self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) - - self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) - - def forward(self, img0, img1, warped_img0, warped_img1, flow): - x = torch.cat((img0, img1, warped_img0, warped_img1), 1) - x = F.interpolate(x, scale_factor=0.25, mode="bilinear", align_corners=False) - flow = ( - F.interpolate(flow, scale_factor=0.25, mode="bilinear", align_corners=False) - * 0.25 - ) - feat = torch.cat((x, flow), 1) - - feat = self.conv_first(feat) - body_feat = self.conv_body(self.body(feat)) - feat = feat + body_feat - # upsample，充分利用四倍放大 - feat = self.lrelu( - self.conv_up1(F.interpolate(feat, scale_factor=2.0, mode="nearest")) - ) - feat = self.lrelu( - self.conv_up2(F.interpolate(feat, scale_factor=2.0, mode="nearest")) - ) - out = self.conv_last(self.lrelu(self.conv_hr(feat))) - - out = torch.sigmoid(out) - return out - - -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) - return torch.nn.functional.grid_sample( - input=tenInput, - grid=g, - mode="bilinear", - padding_mode="border", - align_corners=True, - ) - - -class BasicConv(nn.Module): - def __init__( - self, - in_planes, - out_planes, - kernel_size, - stride=1, - padding=0, - dilation=1, - groups=1, - relu=True, - bn=True, - bias=False, - ): - super(BasicConv, self).__init__() - self.out_channels = out_planes - self.conv = nn.Conv2d( - in_planes, - out_planes, - kernel_size=kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - groups=groups, - bias=bias, - ) - self.bn = ( - nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True) - if bn - else None - ) - self.relu = nn.ReLU() if relu else None - - def forward(self, x): - x = self.conv(x) - if self.bn is not None: - x = self.bn(x) - if self.relu is not None: - x = self.relu(x) - return x - - -class Flatten(nn.Module): - def forward(self, x): - return x.view(x.size(0), -1) - - -class ChannelGate(nn.Module): - def __init__(self, gate_channels, reduction_ratio=16, pool_types=["avg", "max"]): - super(ChannelGate, self).__init__() - self.gate_channels = gate_channels - self.mlp = nn.Sequential( - Flatten(), - nn.Linear(gate_channels, gate_channels // reduction_ratio), - nn.ReLU(), - nn.Linear(gate_channels // reduction_ratio, gate_channels), - ) - self.pool_types = pool_types - - def forward(self, x): - channel_att_sum = None - for pool_type in self.pool_types: - if pool_type == "avg": - avg_pool = F.avg_pool2d( - x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)) - ) - channel_att_raw = self.mlp(avg_pool) - elif pool_type == "max": - max_pool = F.max_pool2d( - x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)) - ) - channel_att_raw = self.mlp(max_pool) - elif pool_type == "lp": - lp_pool = F.lp_pool2d( - x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)) - ) - channel_att_raw = self.mlp(lp_pool) - elif pool_type == "lse": - # LSE pool only - lse_pool = logsumexp_2d(x) - channel_att_raw = self.mlp(lse_pool) - - if channel_att_sum is None: - channel_att_sum = channel_att_raw - else: - channel_att_sum = channel_att_sum + channel_att_raw - - scale = F.sigmoid(channel_att_sum).unsqueeze(2).unsqueeze(3).expand_as(x) - return x * scale - - -def logsumexp_2d(tensor): - tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1) - s, _ = torch.max(tensor_flatten, dim=2, keepdim=True) - outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log() - return outputs - - -class ChannelPool(nn.Module): - def forward(self, x): - return torch.cat( - (torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1 - ) - - -class SpatialGate(nn.Module): - def __init__(self): - super(SpatialGate, self).__init__() - kernel_size = 7 - self.compress = ChannelPool() - self.spatial = BasicConv( - 2, 1, kernel_size, stride=1, padding=(kernel_size - 1) // 2, relu=False - ) - - def forward(self, x): - x_compress = self.compress(x) - x_out = self.spatial(x_compress) - scale = F.sigmoid(x_out) # broadcasting - return x * scale - - -class CBAM(nn.Module): - def __init__( - self, - gate_channels, - reduction_ratio=16, - pool_types=["avg", "max"], - no_spatial=False, - ): - super(CBAM, self).__init__() - self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types) - self.no_spatial = no_spatial - if not no_spatial: - self.SpatialGate = SpatialGate() - - def forward(self, x): - x_out = self.ChannelGate(x) - if not self.no_spatial: - x_out = self.SpatialGate(x_out) - return x_out - - -def conv(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, - ), - nn.PReLU(out_planes), - ) - - -class UNetConv(nn.Module): - def __init__(self, in_planes, out_planes, att=True): - super(UNetConv, self).__init__() - self.conv1 = conv(in_planes, out_planes, 3, 2, 1) - self.conv2 = conv(out_planes, out_planes, 3, 1, 1) - - if att: - self.cbam = CBAM(out_planes, 16) # 这一步导致了通道数最低为128 - else: - self.cbam = None - - def forward(self, x): - x = self.conv1(x) - x = self.conv2(x) - if self.cbam is not None: - x = self.cbam(x) - return x - - -class UpConv(nn.Module): - def __init__(self, in_planes, out_planes, att=True): - super(UpConv, self).__init__() - self.deconv = nn.Sequential( - nn.ConvTranspose2d(in_planes, in_planes // 2, 4, 2, 1), - nn.PReLU(in_planes // 2), - ) - - # 也许不需要这么卷积，我不确定 - self.conv1 = conv(in_planes, in_planes // 2, 3, 1, 1) - self.conv2 = conv(in_planes // 2, out_planes, 3, 1, 1) - - if att: - self.cbam = CBAM(out_planes, 16) - else: - self.cbam = None - - def forward(self, x1, x2): - x1 = self.deconv(x1) - y = self.conv1(torch.cat((x1, x2), 1)) - y = self.conv2(y) - if self.cbam is not None: - y = self.cbam(y) - return y - - -class FeatureNet(nn.Module): - def __init__(self, in_planes, out_planes): - super(FeatureNet, self).__init__() - # 处理IFBlock0时通道数问题 - self.conv0 = conv(7, in_planes, 1, 1, 0) - - self.conv1 = UNetConv(in_planes, out_planes // 8, att=False) - self.conv2 = UNetConv(out_planes // 8, out_planes // 4, att=True) - self.conv3 = UNetConv(out_planes // 4, out_planes // 2, att=True) - self.conv4 = UNetConv(out_planes // 2, out_planes, att=True) - self.conv5 = UNetConv(out_planes, 2 * out_planes, att=True) - - self.deconv5 = UpConv(2 * out_planes, out_planes, att=True) - self.deconv4 = UpConv(out_planes, out_planes // 2, att=False) - self.deconv3 = UpConv(out_planes // 2, out_planes // 4, att=False) - - def forward(self, x, level=0): - if x.shape[1] != 17: - x = self.conv0(x) - x2 = self.conv1(x) - x4 = self.conv2(x2) - x8 = self.conv3(x4) - x16 = self.conv4(x8) - x32 = self.conv5(x16) - y = self.deconv5(x32, x16) # 匹配IFBlock0通道和尺寸 - - # “早退机制”以期待用同一个UNet提取特征，不确定是否对训练产生影响 - if level != 0: - y = self.deconv4(y, x8) # 匹配IFBlock1通道和尺寸 - if level == 2: - y = self.deconv3(y, x4) # 匹配IFBlock2通道和尺寸 - return y - - -class IFBlock(nn.Module): - def __init__(self, c=64, level=0): - super(IFBlock, self).__init__() - self.convblock = nn.Sequential( - conv(c, c), - conv(c, c), - conv(c, c), - conv(c, c), - conv(c, c), - conv(c, c), - ) - self.flowconv = nn.Conv2d(c, 4, 3, 1, 1) - self.maskconvx16 = nn.Conv2d(c, 16 * 16 * 9, 1, 1, 0) - self.maskconvx8 = nn.Conv2d(c, 8 * 8 * 9, 1, 1, 0) - self.maskconvx4 = nn.Conv2d(c, 4 * 4 * 9, 1, 1, 0) - - self.level = level - assert self.level in [4, 8, 16], "Bitch" - - def mask_conv(self, x): - if self.level == 4: - return self.maskconvx4(x) - if self.level == 8: - return self.maskconvx8(x) - if self.level == 16: - return self.maskconvx16(x) - - def upsample_flow(self, flow, mask): - # 俺寻思俺懂了 - N, _, H, W = flow.shape - mask = mask.view(N, 1, 9, self.level, self.level, H, W) - mask = torch.softmax(mask, dim=2) - - up_flow = F.unfold(self.level * flow, [3, 3], padding=1) - up_flow = up_flow.view(N, 4, 9, 1, 1, H, W) - - up_flow = torch.sum(mask * up_flow, dim=2) - up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) - return up_flow.reshape(N, 4, self.level * H, self.level * W) - - def forward(self, x, scale): - x = self.convblock(x) + x # 类似ResNet的f(x) + x - tmp = self.flowconv(x) - up_mask = self.mask_conv(x) - flow_up = self.upsample_flow(tmp, up_mask) - flow = ( - F.interpolate( - flow_up, scale_factor=scale, mode="bilinear", align_corners=False - ) - * scale - ) - return flow - - -class IFUNet(nn.Module): - def __init__(self): - super(IFUNet, self).__init__() - # block0通道数必须为128的整倍数 - self.fmap = FeatureNet(in_planes=17, out_planes=256) - self.block0 = IFBlock(c=256, level=16) - self.block1 = IFBlock(c=128, level=8) - self.block2 = IFBlock(c=64, level=4) - - def forward(self, x, scale=1.0, timestep=0.5, ensemble=True): - 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]) - warped_img0 = img0 - warped_img1 = img1 - flow = None - block = [self.block0, self.block1, self.block2] - for i in range(3): - if flow != None: - x = torch.cat((img0, img1, timestep, warped_img0, warped_img1), 1) - flowtmp = flow - if scale != 1: - x = F.interpolate( - x, scale_factor=scale, mode="bilinear", align_corners=False - ) - flowtmp = ( - F.interpolate( - flow, - scale_factor=scale, - mode="bilinear", - align_corners=False, - ) - * scale - ) - x = torch.cat((x, flowtmp), 1) - # 期待UNet能提取到特征，不再需要ensemble - Fmap = self.fmap(x, level=i) - flow_d = block[i](Fmap, scale=1.0 / scale) - flow = flow + flow_d - - if ensemble: - x = torch.cat( - (img1, img0, 1 - timestep, warped_img0, warped_img1), 1 - ) - flowtmp = flow - if scale != 1: - x = F.interpolate( - x, scale_factor=scale, mode="bilinear", align_corners=False - ) - flowtmp = ( - F.interpolate( - flow, - scale_factor=scale, - mode="bilinear", - align_corners=False, - ) - * scale - ) - x = torch.cat((x, flowtmp), 1) - # 期待UNet能提取到特征，不再需要ensemble - Fmap = self.fmap(x, level=i) - flow_d = block[i](Fmap, scale=1.0 / scale) - flow2 = flow + flow_d - flow = (flow + flow2) / 2 - else: - x = torch.cat((img0, img1, timestep), 1) - if scale != 1: - x = F.interpolate( - x, scale_factor=scale, mode="bilinear", align_corners=False - ) - Fmap = self.fmap(x, level=i) - flow = block[i](Fmap, scale=1.0 / scale) - - if ensemble: - x = torch.cat((img1, img0, 1 - timestep), 1) - if scale != 1: - x = F.interpolate( - x, scale_factor=scale, mode="bilinear", align_corners=False - ) - Fmap = self.fmap(x, level=i) - flow2 = block[i](Fmap, scale=1.0 / scale) - flow = (flow + flow2) / 2 - - warped_img0 = warp(img0, flow[:, :2]) - warped_img1 = warp(img1, flow[:, 2:4]) - return flow, warped_img0, warped_img1 - - -class IFUNetModel(nn.Module): - def __init__(self, local_rank=-1): - super(IFUNetModel, self).__init__() - self.flownet = IFUNet() - self.fusionnet = RRDBNet() - self.refinenet = ResynNet() - - def forward(self, img0, img1, timestep=0.5, scale=1.0, ensemble=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) - - imgs = torch.cat((img0, img1), 1) - flow, warped_img0, warped_img1 = self.flownet(imgs, scale, timestep, ensemble) - mask = self.fusionnet(img0, img1, warped_img0, warped_img1, flow) - merged = warped_img0 * mask + warped_img1 * (1 - mask) - merged, _ = self.refinenet(imgs, deg=merged, scale=[4, 2, 1]) - return merged[:, :, :h, :w] diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ifunet/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ifunet/__init__.py deleted file mode 100644 index bec9867..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ifunet/__init__.py +++ /dev/null @@ -1,59 +0,0 @@ -import torch -from torch.utils.data import DataLoader -import pathlib -from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames, generic_frame_loop, InterpolationStateList -import typing -from comfy.model_management import get_torch_device - -MODEL_TYPE = pathlib.Path(__file__).parent.name -CKPT_NAMES = ["IFUNet.pth"] - -class IFUnet_VFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (CKPT_NAMES, ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}), - "multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}), - "scale_factor": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 100, "step": 0.1}), - "ensemble": ("BOOLEAN", {"default":True}) - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames: typing.SupportsInt = 1, - multiplier: typing.SupportsInt = 2, - scale_factor: typing.SupportsFloat = 1.0, - ensemble: bool = True, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - from .IFUNet_arch import IFUNetModel - model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) - interpolation_model = IFUNetModel() - interpolation_model.load_state_dict(torch.load(model_path)) - interpolation_model.eval().to(get_torch_device()) - frames = preprocess_frames(frames) - - def return_middle_frame(frame_0, frame_1, timestep, model, scale_factor, ensemble): - return model(frame_0, frame_1, timestep=timestep, scale=scale_factor, ensemble=ensemble) - - args = [interpolation_model, scale_factor, ensemble] - out = postprocess_frames( - generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, - interpolation_states=optional_interpolation_states, dtype=torch.float32) - ) - return (out,) - diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/m2m/M2M_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/m2m/M2M_arch.py deleted file mode 100644 index 536f915..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/m2m/M2M_arch.py +++ /dev/null @@ -1,1037 +0,0 @@ -""" -https://github.com/feinanshan/M2M_VFI/blob/main/Test/model/py -https://raw.githubusercontent.com/feinanshan/M2M_VFI/main/Test/model/py -https://github.com/feinanshan/M2M_VFI/blob/main/Test/model/py -https://github.com/feinanshan/M2M_VFI/blob/main/Test/model/py -https://github.com/feinanshan/M2M_VFI/blob/main/Test/model/m2m.py -""" - -import collections -import math -import os -import re -import torch -import typing -from vfi_models.ops import softsplat_func -from vfi_models.ops import costvol_func - -########################################################## - - -objBackwarpcache = {} - - -def backwarp(tenIn: torch.Tensor, tenFlow: torch.Tensor): - if ( - "grid" - + str(tenFlow.dtype) - + str(tenFlow.device) - + str(tenFlow.shape[2]) - + str(tenFlow.shape[3]) - not in objBackwarpcache - ): - tenHor = ( - torch.linspace( - start=-1.0, - end=1.0, - steps=tenFlow.shape[3], - dtype=tenFlow.dtype, - device=tenFlow.device, - ) - .view(1, 1, 1, -1) - .repeat(1, 1, tenFlow.shape[2], 1) - ) - tenVer = ( - torch.linspace( - start=-1.0, - end=1.0, - steps=tenFlow.shape[2], - dtype=tenFlow.dtype, - device=tenFlow.device, - ) - .view(1, 1, -1, 1) - .repeat(1, 1, 1, tenFlow.shape[3]) - ) - - objBackwarpcache[ - "grid" - + str(tenFlow.dtype) - + str(tenFlow.device) - + str(tenFlow.shape[2]) - + str(tenFlow.shape[3]) - ] = torch.cat([tenHor, tenVer], 1) - # end - - if tenFlow.shape[3] == tenFlow.shape[2]: - tenFlow = tenFlow * (2.0 / ((tenFlow.shape[3] and tenFlow.shape[2]) - 1.0)) - - elif tenFlow.shape[3] != tenFlow.shape[2]: - tenFlow = tenFlow * torch.tensor( - data=[2.0 / (tenFlow.shape[3] - 1.0), 2.0 / (tenFlow.shape[2] - 1.0)], - dtype=tenFlow.dtype, - device=tenFlow.device, - ).view(1, 2, 1, 1) - - # end - - return torch.nn.functional.grid_sample( - input=tenIn, - grid=( - objBackwarpcache[ - "grid" - + str(tenFlow.dtype) - + str(tenFlow.device) - + str(tenFlow.shape[2]) - + str(tenFlow.shape[3]) - ] - + tenFlow - ).permute(0, 2, 3, 1), - mode="bilinear", - padding_mode="zeros", - align_corners=True, - ) - - -# end - -########################################################## - - -class Basic(torch.nn.Module): - def __init__( - self, - strType: str, - intChans: typing.List[int], - objScratch: typing.Optional[typing.Dict] = None, - ): - super().__init__() - - self.strType = strType - self.netEvenize = None - self.netMain = None - self.netShortcut = None - - intIn = intChans[0] - intOut = intChans[-1] - netMain = [] - intChans = intChans.copy() - fltStride = 1.0 - - for intPart, strPart in enumerate(self.strType.split("+")[0].split("-")): - if strPart.startswith("evenize") == True and intPart == 0: - - class Evenize(torch.nn.Module): - def __init__(self, strPad): - super().__init__() - - self.strPad = strPad - - # end - - def forward(self, tenIn: torch.Tensor) -> torch.Tensor: - intPad = [0, 0, 0, 0] - - if tenIn.shape[3] % 2 != 0: - intPad[1] = 1 - if tenIn.shape[2] % 2 != 0: - intPad[3] = 1 - - if min(intPad) != 0 or max(intPad) != 0: - tenIn = torch.nn.functional.pad( - input=tenIn, - pad=intPad, - mode=self.strPad - if self.strPad != "zeros" - else "constant", - value=0.0, - ) - # end - - return tenIn - - # end - - # end - - strPad = "zeros" - - if "(" in strPart: - if "replpad" in strPart.split("(")[1].split(")")[0].split(","): - strPad = "replicate" - if "reflpad" in strPart.split("(")[1].split(")")[0].split(","): - strPad = "reflect" - # end - - self.netEvenize = Evenize(strPad) - - elif strPart.startswith("conv") == True: - intKsize = 3 - intPad = 1 - strPad = "zeros" - - if "(" in strPart: - intKsize = int(strPart.split("(")[1].split(")")[0].split(",")[0]) - intPad = int(math.floor(0.5 * (intKsize - 1))) - - if "replpad" in strPart.split("(")[1].split(")")[0].split(","): - strPad = "replicate" - if "reflpad" in strPart.split("(")[1].split(")")[0].split(","): - strPad = "reflect" - # end - - if "nopad" in self.strType.split("+"): - intPad = 0 - # end - - netMain += [ - torch.nn.Conv2d( - in_channels=intChans[0], - out_channels=intChans[1], - kernel_size=intKsize, - stride=1, - padding=intPad, - padding_mode=strPad, - bias="nobias" not in self.strType.split("+"), - ) - ] - intChans = intChans[1:] - fltStride *= 1.0 - - elif strPart.startswith("sconv") == True: - intKsize = 3 - intPad = 1 - strPad = "zeros" - - if "(" in strPart: - intKsize = int(strPart.split("(")[1].split(")")[0].split(",")[0]) - intPad = int(math.floor(0.5 * (intKsize - 1))) - - if "replpad" in strPart.split("(")[1].split(")")[0].split(","): - strPad = "replicate" - if "reflpad" in strPart.split("(")[1].split(")")[0].split(","): - strPad = "reflect" - # end - - if "nopad" in self.strType.split("+"): - intPad = 0 - # end - - netMain += [ - torch.nn.Conv2d( - in_channels=intChans[0], - out_channels=intChans[1], - kernel_size=intKsize, - stride=2, - padding=intPad, - padding_mode=strPad, - bias="nobias" not in self.strType.split("+"), - ) - ] - intChans = intChans[1:] - fltStride *= 2.0 - - elif strPart.startswith("up") == True: - - class Up(torch.nn.Module): - def __init__(self, strType): - super().__init__() - - self.strType = strType - - # end - - def forward(self, tenIn: torch.Tensor) -> torch.Tensor: - if self.strType == "nearest": - return torch.nn.functional.interpolate( - input=tenIn, - scale_factor=2.0, - mode="nearest-exact", - align_corners=False, - ) - - elif self.strType == "bilinear": - return torch.nn.functional.interpolate( - input=tenIn, - scale_factor=2.0, - mode="bilinear", - align_corners=False, - ) - - elif self.strType == "pyramid": - return pyramid(tenIn, None, "up") - - elif self.strType == "shuffle": - return torch.nn.functional.pixel_shuffle( - tenIn, upscale_factor=2 - ) # https://github.com/pytorch/pytorch/issues/62854 - - # end - - assert False # to make torchscript happy - - # end - - # end - - strType = "bilinear" - - if "(" in strPart: - if "nearest" in strPart.split("(")[1].split(")")[0].split(","): - strType = "nearest" - if "pyramid" in strPart.split("(")[1].split(")")[0].split(","): - strType = "pyramid" - if "shuffle" in strPart.split("(")[1].split(")")[0].split(","): - strType = "shuffle" - # end - - netMain += [Up(strType)] - fltStride *= 0.5 - - elif strPart.startswith("prelu") == True: - netMain += [ - torch.nn.PReLU( - num_parameters=1, - init=float(strPart.split("(")[1].split(")")[0].split(",")[0]), - ) - ] - fltStride *= 1.0 - - elif True: - assert False - - # end - # end - - self.netMain = torch.nn.Sequential(*netMain) - - for strPart in self.strType.split("+")[1:]: - if strPart.startswith("skip") == True: - if intIn == intOut and fltStride == 1.0: - self.netShortcut = torch.nn.Identity() - - elif intIn != intOut and fltStride == 1.0: - self.netShortcut = torch.nn.Conv2d( - in_channels=intIn, - out_channels=intOut, - kernel_size=1, - stride=1, - padding=0, - bias="nobias" not in self.strType.split("+"), - ) - - elif intIn == intOut and fltStride != 1.0: - - class Down(torch.nn.Module): - def __init__(self, fltScale): - super().__init__() - - self.fltScale = fltScale - - # end - - def forward(self, tenIn: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.interpolate( - input=tenIn, - scale_factor=self.fltScale, - mode="bilinear", - align_corners=False, - ) - - # end - - # end - - self.netShortcut = Down(1.0 / fltStride) - - elif intIn != intOut and fltStride != 1.0: - - class Down(torch.nn.Module): - def __init__(self, fltScale): - super().__init__() - - self.fltScale = fltScale - - # end - - def forward(self, tenIn: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.interpolate( - input=tenIn, - scale_factor=self.fltScale, - mode="bilinear", - align_corners=False, - ) - - # end - - # end - - self.netShortcut = torch.nn.Sequential( - Down(1.0 / fltStride), - torch.nn.Conv2d( - in_channels=intIn, - out_channels=intOut, - kernel_size=1, - stride=1, - padding=0, - bias="nobias" not in self.strType.split("+"), - ), - ) - - # end - - elif strPart.startswith("...") == True: - pass - - # end - # end - - assert len(intChans) == 1 - - # end - - def forward(self, tenIn: torch.Tensor) -> torch.Tensor: - if self.netEvenize is not None: - tenIn = self.netEvenize(tenIn) - # end - - tenOut = self.netMain(tenIn) - - if self.netShortcut is not None: - tenOut = tenOut + self.netShortcut(tenIn) - # end - - return tenOut - - # end - - -# end - - -########################################################## - - -class Network(torch.nn.Module): - def __init__(self): - super().__init__() - - class Extractor(torch.nn.Module): - def __init__(self): - super().__init__() - - self.netOne = Basic( - "evenize(replpad)-sconv(2)-prelu(0.25)-conv(3,replpad)-prelu(0.25)-conv(3,replpad)-prelu(0.25)", - [3, 32, 32, 32], - None, - ) - self.netTwo = Basic( - "evenize(replpad)-sconv(2)-prelu(0.25)-conv(3,replpad)-prelu(0.25)-conv(3,replpad)-prelu(0.25)", - [32, 32, 32, 32], - None, - ) - self.netThr = Basic( - "evenize(replpad)-sconv(2)-prelu(0.25)-conv(3,replpad)-prelu(0.25)-conv(3,replpad)-prelu(0.25)", - [32, 32, 32, 32], - None, - ) - - # end - - def forward(self, tenIn): - tenOne = self.netOne(tenIn) - tenTwo = self.netTwo(tenOne) - tenThr = self.netThr(tenTwo) - tenFou = torch.nn.functional.avg_pool2d( - input=tenThr, kernel_size=2, stride=2, count_include_pad=False - ) - tenFiv = torch.nn.functional.avg_pool2d( - input=tenFou, kernel_size=2, stride=2, count_include_pad=False - ) - - return [tenOne, tenTwo, tenThr, tenFou, tenFiv] - - # end - - # end - - class Decoder(torch.nn.Module): - def __init__(self, intChannels): - super().__init__() - - self.netCostacti = torch.nn.PReLU(num_parameters=1, init=0.25) - self.netMain = Basic( - "conv(3,replpad)-prelu(0.25)-conv(3,replpad)-prelu(0.25)-conv(3,replpad)-prelu(0.25)-conv(3,replpad)-prelu(0.25)-conv(3,replpad)-prelu(0.25)-conv(3,replpad)", - [intChannels, 128, 128, 96, 64, 32, 2], - None, - ) - - # end - - def forward(self, tenOne, tenTwo, tenFlow): - if tenFlow is not None: - tenFlow = 2.0 * torch.nn.functional.interpolate( - input=tenFlow, - scale_factor=2.0, - mode="bilinear", - align_corners=False, - ) - # end - - tenMain = [] - - if tenFlow is None: - tenMain.append(tenOne) - tenMain.append(self.netCostacti(costvol_func.apply(tenOne, tenTwo))) - - elif tenFlow is not None: - tenMain.append(tenOne) - tenMain.append( - self.netCostacti( - costvol_func.apply( - tenOne, backwarp(tenTwo, tenFlow.detach()) - ) - ) - ) - tenMain.append(tenFlow) - - # end - - return (tenFlow if tenFlow is not None else 0.0) + self.netMain( - torch.cat(tenMain, 1) - ) - - # end - - # end - - self.netExtractor = Extractor() - - self.netFiv = Decoder(32 + 81 + 0) - self.netFou = Decoder(32 + 81 + 2) - self.netThr = Decoder(32 + 81 + 2) - self.netTwo = Decoder(32 + 81 + 2) - self.netOne = Decoder(32 + 81 + 2) - - # end - - def bidir(self, tenOne, tenTwo): - tenOne, tenTwo = list( - zip( - *[ - torch.split(tenFeat, [tenOne.shape[0], tenTwo.shape[0]], 0) - for tenFeat in self.netExtractor(torch.cat([tenOne, tenTwo], 0)) - ] - ) - ) - - tenFwd = None - tenFwd = self.netFiv(tenOne[-1], tenTwo[-1], tenFwd) - tenFwd = self.netFou(tenOne[-2], tenTwo[-2], tenFwd) - tenFwd = self.netThr(tenOne[-3], tenTwo[-3], tenFwd) - tenFwd = self.netTwo(tenOne[-4], tenTwo[-4], tenFwd) - tenFwd = self.netOne(tenOne[-5], tenTwo[-5], tenFwd) - - tenBwd = None - tenBwd = self.netFiv(tenTwo[-1], tenOne[-1], tenBwd) - tenBwd = self.netFou(tenTwo[-2], tenOne[-2], tenBwd) - tenBwd = self.netThr(tenTwo[-3], tenOne[-3], tenBwd) - tenBwd = self.netTwo(tenTwo[-4], tenOne[-4], tenBwd) - tenBwd = self.netOne(tenTwo[-5], tenOne[-5], tenBwd) - - return tenFwd, tenBwd - - # end - - -# end - -########################################################## - - -def forwarp_mframe_mask( - tenIn1, tenFlow1, t1, tenIn2, tenFlow2, t2, tenMetric1=None, tenMetric2=None -): - def one_fdir(tenIn, tenFlow, td, tenMetric): - tenIn = torch.cat( - [ - tenIn * td * (tenMetric).clip(-20.0, 20.0).exp(), - td * (tenMetric).clip(-20.0, 20.0).exp(), - ], - 1, - ) - - tenOut = softsplat_func.apply(tenIn, tenFlow) - - return tenOut[:, :-1, :, :], tenOut[:, -1:, :, :] + 0.0000001 - - flow_num = tenFlow1.shape[0] - tenOut = 0 - tenNormalize = 0 - for idx in range(flow_num): - tenOutF, tenNormalizeF = one_fdir( - tenIn1[idx], tenFlow1[idx], t1[idx], tenMetric1[idx] - ) - tenOutB, tenNormalizeB = one_fdir( - tenIn2[idx], tenFlow2[idx], t2[idx], tenMetric2[idx] - ) - - tenOut += tenOutF + tenOutB - tenNormalize += tenNormalizeF + tenNormalizeB - - return tenOut / tenNormalize, tenNormalize < 0.00001 - - -################################################################### - -c = 16 - - -def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): - return torch.nn.Sequential( - torch.nn.Conv2d( - in_planes, - out_planes, - kernel_size=kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - bias=True, - ), - torch.nn.PReLU(out_planes), - ) - - -def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): - return torch.nn.Sequential( - torch.torch.nn.ConvTranspose2d( - in_channels=in_planes, - out_channels=out_planes, - kernel_size=4, - stride=2, - padding=1, - bias=True, - ), - torch.nn.PReLU(out_planes), - ) - - -class Conv2(torch.nn.Module): - def __init__(self, in_planes, out_planes, stride=2): - super(Conv2, self).__init__() - self.conv1 = conv(in_planes, out_planes, 3, stride, 1) - self.conv2 = conv(out_planes, out_planes, 3, 1, 1) - - def forward(self, x): - x = self.conv1(x) - x = self.conv2(x) - return x - - -class Conv2n(torch.nn.Module): - def __init__(self, in_planes, out_planes, stride=2): - super(Conv2n, self).__init__() - self.conv1 = conv(in_planes, in_planes, 3, stride, 1) - self.conv2 = conv(in_planes, in_planes, 3, 1, 1) - self.conv3 = conv(in_planes, in_planes, 1, 1, 0) - self.conv4 = conv(in_planes, out_planes, 1, 1, 0) - - def forward(self, x): - x = self.conv1(x) - x = self.conv2(x) - x = self.conv3(x) - x = self.conv4(x) - return x - - -##################################################### - - -class ImgPyramid(torch.nn.Module): - def __init__(self): - super(ImgPyramid, self).__init__() - self.conv1 = Conv2(3, c) - self.conv2 = Conv2(c, 2 * c) - self.conv3 = Conv2(2 * c, 4 * c) - self.conv4 = Conv2(4 * c, 8 * c) - - def forward(self, x): - x1 = self.conv1(x) - x2 = self.conv2(x1) - x3 = self.conv3(x2) - x4 = self.conv4(x3) - return [x1, x2, x3, x4] - - -class EncDec(torch.nn.Module): - def __init__(self, branch): - super(EncDec, self).__init__() - self.branch = branch - - self.down0 = Conv2(8, 2 * c) - self.down1 = Conv2(6 * c, 4 * c) - self.down2 = Conv2(12 * c, 8 * c) - self.down3 = Conv2(24 * c, 16 * c) - - self.up0 = deconv(48 * c, 8 * c) - self.up1 = deconv(16 * c, 4 * c) - self.up2 = deconv(8 * c, 2 * c) - self.up3 = deconv(4 * c, c) - self.conv = torch.nn.Conv2d(c, 2 * self.branch, 3, 1, 1) - - self.conv_m = torch.nn.Conv2d(c, 1, 3, 1, 1) - - # For Channel dimennsion - self.conv_C = torch.nn.Sequential( - torch.nn.AdaptiveAvgPool2d(1), - torch.nn.Conv2d( - 16 * c, - 16 * 16 * c, - kernel_size=(1, 1), - stride=(1, 1), - padding=(0, 0), - bias=True, - ), - torch.nn.Sigmoid(), - ) - - # For Height dimennsion - self.conv_H = torch.nn.Sequential( - torch.nn.AdaptiveAvgPool2d((None, 1)), - torch.nn.Conv2d( - 16 * c, 16, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=True - ), - torch.nn.Sigmoid(), - ) - - # For Width dimennsion - self.conv_W = torch.nn.Sequential( - torch.nn.AdaptiveAvgPool2d((1, None)), - torch.nn.Conv2d( - 16 * c, 16, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=True - ), - torch.nn.Sigmoid(), - ) - - self.sigmoid = torch.nn.Sigmoid() - - def forward(self, flow0, flow1, im0, im1, c0, c1): - N_, C_, H_, W_ = im0.shape - - wim1 = backwarp(im1, flow0) - wim0 = backwarp(im0, flow1) - s0_0 = self.down0(torch.cat((flow0, im0, wim1), 1)) - s1_0 = self.down0(torch.cat((flow1, im1, wim0), 1)) - - ######################################################################################### - flow0 = ( - torch.nn.functional.interpolate( - flow0, scale_factor=0.5, mode="bilinear", align_corners=False - ) - * 0.5 - ) - flow1 = ( - torch.nn.functional.interpolate( - flow1, scale_factor=0.5, mode="bilinear", align_corners=False - ) - * 0.5 - ) - - wf0 = backwarp(torch.cat((s0_0, c0[0]), 1), flow1) - wf1 = backwarp(torch.cat((s1_0, c1[0]), 1), flow0) - - s0_1 = self.down1(torch.cat((s0_0, c0[0], wf1), 1)) - s1_1 = self.down1(torch.cat((s1_0, c1[0], wf0), 1)) - - ######################################################################################### - flow0 = ( - torch.nn.functional.interpolate( - flow0, scale_factor=0.5, mode="bilinear", align_corners=False - ) - * 0.5 - ) - flow1 = ( - torch.nn.functional.interpolate( - flow1, scale_factor=0.5, mode="bilinear", align_corners=False - ) - * 0.5 - ) - - wf0 = backwarp(torch.cat((s0_1, c0[1]), 1), flow1) - wf1 = backwarp(torch.cat((s1_1, c1[1]), 1), flow0) - - s0_2 = self.down2(torch.cat((s0_1, c0[1], wf1), 1)) - s1_2 = self.down2(torch.cat((s1_1, c1[1], wf0), 1)) - - ######################################################################################### - flow0 = ( - torch.nn.functional.interpolate( - flow0, scale_factor=0.5, mode="bilinear", align_corners=False - ) - * 0.5 - ) - flow1 = ( - torch.nn.functional.interpolate( - flow1, scale_factor=0.5, mode="bilinear", align_corners=False - ) - * 0.5 - ) - - wf0 = backwarp(torch.cat((s0_2, c0[2]), 1), flow1) - wf1 = backwarp(torch.cat((s1_2, c1[2]), 1), flow0) - - s0_3 = self.down3(torch.cat((s0_2, c0[2], wf1), 1)) - s1_3 = self.down3(torch.cat((s1_2, c1[2], wf0), 1)) - - ######################################################################################### - - s0_3_c = self.conv_C(s0_3) - s0_3_c = s0_3_c.view(N_, 16, -1, 1, 1) - - s0_3_h = self.conv_H(s0_3) - s0_3_h = s0_3_h.view(N_, 16, 1, -1, 1) - - s0_3_w = self.conv_W(s0_3) - s0_3_w = s0_3_w.view(N_, 16, 1, 1, -1) - - cube0 = (s0_3_c * s0_3_h * s0_3_w).mean(1) - - s0_3 = s0_3 * cube0 - - s1_3_c = self.conv_C(s1_3) - s1_3_c = s1_3_c.view(N_, 16, -1, 1, 1) - - s1_3_h = self.conv_H(s1_3) - s1_3_h = s1_3_h.view(N_, 16, 1, -1, 1) - - s1_3_w = self.conv_W(s1_3) - s1_3_w = s1_3_w.view(N_, 16, 1, 1, -1) - - cube1 = (s1_3_c * s1_3_h * s1_3_w).mean(1) - - s1_3 = s1_3 * cube1 - - ######################################################################################### - flow0 = ( - torch.nn.functional.interpolate( - flow0, scale_factor=0.5, mode="bilinear", align_corners=False - ) - * 0.5 - ) - flow1 = ( - torch.nn.functional.interpolate( - flow1, scale_factor=0.5, mode="bilinear", align_corners=False - ) - * 0.5 - ) - - wf0 = backwarp(torch.cat((s0_3, c0[3]), 1), flow1) - wf1 = backwarp(torch.cat((s1_3, c1[3]), 1), flow0) - - x0 = self.up0(torch.cat((s0_3, c0[3], wf1), 1)) - x1 = self.up0(torch.cat((s1_3, c1[3], wf0), 1)) - - x0 = self.up1(torch.cat((s0_2, x0), 1)) - x1 = self.up1(torch.cat((s1_2, x1), 1)) - - x0 = self.up2(torch.cat((s0_1, x0), 1)) - x1 = self.up2(torch.cat((s1_1, x1), 1)) - - x0 = self.up3(torch.cat((s0_0, x0), 1)) - x1 = self.up3(torch.cat((s1_0, x1), 1)) - - m0 = self.sigmoid(self.conv_m(x0)) * 0.8 + 0.1 - m1 = self.sigmoid(self.conv_m(x1)) * 0.8 + 0.1 - - x0 = self.conv(x0) - x1 = self.conv(x1) - - return x0, x1, m0.repeat(1, self.branch, 1, 1), m1.repeat(1, self.branch, 1, 1) - - -class M2M_PWC(torch.nn.Module): - def __init__(self, ratio=4): - super(M2M_PWC, self).__init__() - self.branch = 4 - self.ratio = ratio - - self.netFlow = Network() - - self.paramAlpha = torch.nn.Parameter(10.0 * torch.ones(1, 1, 1, 1)) - - class MotionRefineNet(torch.nn.Module): - def __init__(self, branch): - super(MotionRefineNet, self).__init__() - self.branch = branch - self.img_pyramid = ImgPyramid() - self.motion_encdec = EncDec(branch) - - def forward(self, flow0, flow1, im0, im1, ratio): - flow0 = ratio * torch.nn.functional.interpolate( - input=flow0, - scale_factor=ratio, - mode="bilinear", - align_corners=False, - ) - flow1 = ratio * torch.nn.functional.interpolate( - input=flow1, - scale_factor=ratio, - mode="bilinear", - align_corners=False, - ) - - c0 = self.img_pyramid(im0) - c1 = self.img_pyramid(im1) - - flow_res = self.motion_encdec(flow0, flow1, im0, im1, c0, c1) - - flow0 = flow0.repeat(1, self.branch, 1, 1) + flow_res[0] - flow1 = flow1.repeat(1, self.branch, 1, 1) + flow_res[1] - - return flow0, flow1, flow_res[2], flow_res[3] - - self.MRN = MotionRefineNet(self.branch) - - def forward(self, im0, im1, fltTimes=[0.5], ratio=None): - if ratio is None: - ratio = self.ratio - - intWidth = im0.shape[3] and im1.shape[3] - intHeight = im0.shape[2] and im1.shape[2] - - intPadr = ((ratio * 16) - (intWidth % (ratio * 16))) % (ratio * 16) - intPadb = ((ratio * 16) - (intHeight % (ratio * 16))) % (ratio * 16) - - im0 = torch.nn.functional.pad( - input=im0, pad=[0, intPadr, 0, intPadb], mode="replicate" - ) - im1 = torch.nn.functional.pad( - input=im1, pad=[0, intPadr, 0, intPadb], mode="replicate" - ) - - N_, C_, H_, W_ = im0.shape - - outputs = [] - - with torch.set_grad_enabled(False): - tenStats = [im0, im1] - tenMean_ = sum([tenIn.mean([1, 2, 3], True) for tenIn in tenStats]) / len( - tenStats - ) - tenStd_ = ( - sum( - [ - tenIn.std([1, 2, 3], False, True).square() - + (tenMean_ - tenIn.mean([1, 2, 3], True)).square() - for tenIn in tenStats - ] - ) - / len(tenStats) - ).sqrt() - - im0_o = (im0 - tenMean_) / (tenStd_ + 0.0000001) - im1_o = (im1 - tenMean_) / (tenStd_ + 0.0000001) - - im0 = (im0 - tenMean_) / (tenStd_ + 0.0000001) - im1 = (im1 - tenMean_) / (tenStd_ + 0.0000001) - - im0_ = torch.nn.functional.interpolate( - input=im0, scale_factor=2.0 / ratio, mode="bilinear", align_corners=False - ) - im1_ = torch.nn.functional.interpolate( - input=im1, scale_factor=2.0 / ratio, mode="bilinear", align_corners=False - ) - - tenFwd, tenBwd = self.netFlow.bidir(im0_, im1_) - - tenFwd, tenBwd, WeiMF, WeiMB = self.MRN(tenFwd, tenBwd, im0, im1, ratio) - - for fltTime_ in fltTimes: - im0 = im0_o.repeat(1, self.branch, 1, 1) - im1 = im1_o.repeat(1, self.branch, 1, 1) - tenStd = tenStd_.repeat(1, self.branch, 1, 1) - tenMean = tenMean_.repeat(1, self.branch, 1, 1) - fltTime = fltTime_.repeat(1, self.branch, 1, 1) - - tenFwd = tenFwd.reshape(N_, self.branch, 2, H_, W_).view( - N_ * self.branch, 2, H_, W_ - ) - tenBwd = tenBwd.reshape(N_, self.branch, 2, H_, W_).view( - N_ * self.branch, 2, H_, W_ - ) - - WeiMF = WeiMF.reshape(N_, self.branch, 1, H_, W_).view( - N_ * self.branch, 1, H_, W_ - ) - WeiMB = WeiMB.reshape(N_, self.branch, 1, H_, W_).view( - N_ * self.branch, 1, H_, W_ - ) - - im0 = im0.reshape(N_, self.branch, 3, H_, W_).view( - N_ * self.branch, 3, H_, W_ - ) - im1 = im1.reshape(N_, self.branch, 3, H_, W_).view( - N_ * self.branch, 3, H_, W_ - ) - - tenStd = tenStd.reshape(N_, self.branch, 1, 1, 1).view( - N_ * self.branch, 1, 1, 1 - ) - tenMean = tenMean.reshape(N_, self.branch, 1, 1, 1).view( - N_ * self.branch, 1, 1, 1 - ) - fltTime = fltTime.reshape(N_, self.branch, 1, 1, 1).view( - N_ * self.branch, 1, 1, 1 - ) - - tenPhotoone = ( - ( - 1.0 - - ( - WeiMF - * (im0 - backwarp(im1, tenFwd).detach()).abs().mean([1], True) - ) - ) - .clip(0.001, None) - .square() - ) - tenPhototwo = ( - ( - 1.0 - - ( - WeiMB - * (im1 - backwarp(im0, tenBwd).detach()).abs().mean([1], True) - ) - ) - .clip(0.001, None) - .square() - ) - - t0 = fltTime - flow0 = tenFwd * t0 - metric0 = self.paramAlpha * tenPhotoone - - t1 = 1.0 - fltTime - flow1 = tenBwd * t1 - metric1 = self.paramAlpha * tenPhototwo - - flow0 = flow0.reshape(N_, self.branch, 2, H_, W_).permute(1, 0, 2, 3, 4) - flow1 = flow1.reshape(N_, self.branch, 2, H_, W_).permute(1, 0, 2, 3, 4) - - metric0 = metric0.reshape(N_, self.branch, 1, H_, W_).permute(1, 0, 2, 3, 4) - metric1 = metric1.reshape(N_, self.branch, 1, H_, W_).permute(1, 0, 2, 3, 4) - - im0 = im0.reshape(N_, self.branch, 3, H_, W_).permute(1, 0, 2, 3, 4) - im1 = im1.reshape(N_, self.branch, 3, H_, W_).permute(1, 0, 2, 3, 4) - - t0 = t0.reshape(N_, self.branch, 1, 1, 1).permute(1, 0, 2, 3, 4) - t1 = t1.reshape(N_, self.branch, 1, 1, 1).permute(1, 0, 2, 3, 4) - - tenOutput, mask = forwarp_mframe_mask( - im0, flow0, t1, im1, flow1, t0, metric0, metric1 - ) - - tenOutput = tenOutput + mask * (t1.mean(0) * im0_o + t0.mean(0) * im1_o) - - outputs.append((tenOutput * (tenStd_ + 0.0000001)) + tenMean_) - - return [output[:, :, :intHeight, :intWidth] for output in outputs] diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/m2m/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/m2m/__init__.py deleted file mode 100644 index 88e489e..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/m2m/__init__.py +++ /dev/null @@ -1,60 +0,0 @@ -import pathlib -import torch -from torch.utils.data import DataLoader -import pathlib -from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames -import typing -from comfy.model_management import get_torch_device -from vfi_utils import InterpolationStateList, generic_frame_loop - -MODEL_TYPE = pathlib.Path(__file__).parent.name -CKPT_NAMES = ["M2M.pth"] - - -class M2M_VFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (CKPT_NAMES, ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}), - "multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}), - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames: typing.SupportsInt = 1, - multiplier: typing.SupportsInt = 2, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - from .M2M_arch import M2M_PWC - model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) - interpolation_model = M2M_PWC() - interpolation_model.load_state_dict(torch.load(model_path)) - interpolation_model.eval().to(get_torch_device()) - frames = preprocess_frames(frames) - - def return_middle_frame(frame_0, frame_1, int_timestep, model): - tenSteps = [ - torch.FloatTensor([int_timestep] * len(frame_0)).view(len(frame_0), 1, 1, 1).to(get_torch_device()) - ] - return model(frame_0, frame_1, tenSteps)[0] - - args = [interpolation_model] - out = postprocess_frames( - generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, - interpolation_states=optional_interpolation_states, dtype=torch.float32) - ) - return (out,) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/__init__.py deleted file mode 100644 index 25b13b4..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -import torch.multiprocessing as mp - -if mp.current_process().name == "MainProcess": - import yaml - import os - from pathlib import Path - - config_path = Path(Path(__file__).parent.parent.parent.resolve(), "config.yaml") - - if os.path.exists(config_path): - config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader) - ops_backend = config["ops_backend"] - else: - ops_backend = "taichi" - - assert ops_backend in ["taichi", "cupy"] - - if ops_backend == "taichi": - from .taichi_ops import softsplat, ModuleSoftsplat, FunctionSoftsplat, softsplat_func, costvol_func, sepconv_func, init, batch_edt, FunctionAdaCoF, ModuleCorrelation, FunctionCorrelation, _FunctionCorrelation - else: - from .cupy_ops import softsplat, ModuleSoftsplat, FunctionSoftsplat, softsplat_func, costvol_func, sepconv_func, init, batch_edt, FunctionAdaCoF, ModuleCorrelation, FunctionCorrelation, _FunctionCorrelation - diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/__init__.py deleted file mode 100644 index 1009079..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -from .costvol import * -from .sepconv import * -from .softsplat import * -from .adacof import * -from .correlation import * -from comfy.model_management import is_nvidia, get_torch_device_name, get_torch_device - -def init(): - if not is_nvidia(): - raise NotImplementedError(f"CuPy ops backend only support CUDA device but found {get_torch_device_name(get_torch_device())} instead. Try Taichi ops backend by editing config.yaml") - return \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/adacof.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/adacof.py deleted file mode 100644 index 3784690..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/adacof.py +++ /dev/null @@ -1,491 +0,0 @@ -import torch -from .utils import cuda_kernel, cuda_launch, cuda_int32 -import math - -kernel_AdaCoF_updateOutput = """ - extern "C" __global__ void kernel_AdaCoF_updateOutput( - const int n, - const float* input, - const float* weight, - const float* offset_i, - const float* offset_j, - float* output - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - float dblOutput = 0.0; - - const int intSample = ( intIndex / SIZE_3(output) / SIZE_2(output) / SIZE_1(output) ) % SIZE_0(output); - const int c = ( intIndex / SIZE_3(output) / SIZE_2(output) ) % SIZE_1(output); - const int i = ( intIndex / SIZE_3(output) ) % SIZE_2(output); - const int j = ( intIndex ) % SIZE_3(output); - - for (int k = 0; k < F_SIZE; k += 1) { - for (int l = 0; l < F_SIZE; l += 1) { - float w = VALUE_4(weight, intSample, k*F_SIZE+l, i, j); - float alpha = VALUE_4(offset_i, intSample, k*F_SIZE+l, i, j); - float beta = VALUE_4(offset_j, intSample, k*F_SIZE+l, i, j); - int A = (int) alpha; - int B = (int) beta; - - int i_k_A = i+k*DILATION+A; - if(i_k_A < 0) - i_k_A = 0; - if(i_k_A > SIZE_2(input) - 1) - i_k_A = SIZE_2(input) - 1; - - int j_l_B = j+l*DILATION+B; - if(j_l_B < 0) - j_l_B = 0; - if(j_l_B > SIZE_3(input) - 1) - j_l_B = SIZE_3(input) - 1; - - int i_k_A_1 = i+k*DILATION+A+1; - if(i_k_A_1 < 0) - i_k_A_1 = 0; - if(i_k_A_1 > SIZE_2(input) - 1) - i_k_A_1 = SIZE_2(input) - 1; - - int j_l_B_1 = j+l*DILATION+B+1; - if(j_l_B_1 < 0) - j_l_B_1 = 0; - if(j_l_B_1 > SIZE_3(input) - 1) - j_l_B_1 = SIZE_3(input) - 1; - - dblOutput += w * ( - VALUE_4(input, intSample, c, i_k_A, j_l_B)*(1-(alpha-(float)A))*(1-(beta-(float)B)) + - VALUE_4(input, intSample, c, i_k_A_1, j_l_B)*(alpha-(float)A)*(1-(beta-(float)B)) + - VALUE_4(input, intSample, c, i_k_A, j_l_B_1)*(1-(alpha-(float)A))*(beta-(float)B) + - VALUE_4(input, intSample, c, i_k_A_1, j_l_B_1)*(alpha-(float)A)*(beta-(float)B) - ); - } - } - - output[intIndex] = dblOutput; - } } -""" - -kernel_AdaCoF_updateGradWeight = """ - extern "C" __global__ void kernel_AdaCoF_updateGradWeight( - const int n, - const float* gradLoss, - const float* input, - const float* offset_i, - const float* offset_j, - float* gradWeight - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - float floatOutput = 0.0; - - const int intSample = ( intIndex / SIZE_3(gradWeight) / SIZE_2(gradWeight) / SIZE_1(gradWeight) ) % SIZE_0(gradWeight); - const int intDepth = ( intIndex / SIZE_3(gradWeight) / SIZE_2(gradWeight) ) % SIZE_1(gradWeight); - const int i = ( intIndex / SIZE_3(gradWeight) ) % SIZE_2(gradWeight); - const int j = ( intIndex ) % SIZE_3(gradWeight); - - int k = intDepth / F_SIZE; - int l = intDepth % F_SIZE; - - for (int c = 0; c < 3; c++) - { - float delta = VALUE_4(gradLoss, intSample, c, i, j); - float alpha = VALUE_4(offset_i, intSample, k*F_SIZE+l, i, j); - float beta = VALUE_4(offset_j, intSample, k*F_SIZE+l, i, j); - int A = (int) alpha; - int B = (int) beta; - - int i_k_A = i+k*DILATION+A; - if(i_k_A < 0) - i_k_A = 0; - if(i_k_A > SIZE_2(input) - 1) - i_k_A = SIZE_2(input) - 1; - - int j_l_B = j+l*DILATION+B; - if(j_l_B < 0) - j_l_B = 0; - if(j_l_B > SIZE_3(input) - 1) - j_l_B = SIZE_3(input) - 1; - - int i_k_A_1 = i+k*DILATION+A+1; - if(i_k_A_1 < 0) - i_k_A_1 = 0; - if(i_k_A_1 > SIZE_2(input) - 1) - i_k_A_1 = SIZE_2(input) - 1; - - int j_l_B_1 = j+l*DILATION+B+1; - if(j_l_B_1 < 0) - j_l_B_1 = 0; - if(j_l_B_1 > SIZE_3(input) - 1) - j_l_B_1 = SIZE_3(input) - 1; - - floatOutput += delta * ( - VALUE_4(input, intSample, c, i_k_A, j_l_B)*(1-(alpha-(float)A))*(1-(beta-(float)B)) + - VALUE_4(input, intSample, c, i_k_A_1, j_l_B)*(alpha-(float)A)*(1-(beta-(float)B)) + - VALUE_4(input, intSample, c, i_k_A, j_l_B_1)*(1-(alpha-(float)A))*(beta-(float)B) + - VALUE_4(input, intSample, c, i_k_A_1, j_l_B_1)*(alpha-(float)A)*(beta-(float)B) - ); - } - - gradWeight[intIndex] = floatOutput; - } } -""" - -kernel_AdaCoF_updateGradAlpha = """ - extern "C" __global__ void kernel_AdaCoF_updateGradAlpha( - const int n, - const float* gradLoss, - const float* input, - const float* weight, - const float* offset_i, - const float* offset_j, - float* gradOffset_i - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - float floatOutput = 0.0; - - const int intSample = ( intIndex / SIZE_3(gradOffset_i) / SIZE_2(gradOffset_i) / SIZE_1(gradOffset_i) ) % SIZE_0(gradOffset_i); - const int intDepth = ( intIndex / SIZE_3(gradOffset_i) / SIZE_2(gradOffset_i) ) % SIZE_1(gradOffset_i); - const int i = ( intIndex / SIZE_3(gradOffset_i) ) % SIZE_2(gradOffset_i); - const int j = ( intIndex ) % SIZE_3(gradOffset_i); - - int k = intDepth / F_SIZE; - int l = intDepth % F_SIZE; - - for (int c = 0; c < 3; c++) - { - float delta = VALUE_4(gradLoss, intSample, c, i, j); - float w = VALUE_4(weight, intSample, k*F_SIZE+l, i, j); - float alpha = VALUE_4(offset_i, intSample, k*F_SIZE+l, i, j); - float beta = VALUE_4(offset_j, intSample, k*F_SIZE+l, i, j); - int A = (int) alpha; - int B = (int) beta; - - int i_k_A = i+k*DILATION+A; - if(i_k_A < 0) - i_k_A = 0; - if(i_k_A > SIZE_2(input) - 1) - i_k_A = SIZE_2(input) - 1; - - int j_l_B = j+l*DILATION+B; - if(j_l_B < 0) - j_l_B = 0; - if(j_l_B > SIZE_3(input) - 1) - j_l_B = SIZE_3(input) - 1; - - int i_k_A_1 = i+k*DILATION+A+1; - if(i_k_A_1 < 0) - i_k_A_1 = 0; - if(i_k_A_1 > SIZE_2(input) - 1) - i_k_A_1 = SIZE_2(input) - 1; - - int j_l_B_1 = j+l*DILATION+B+1; - if(j_l_B_1 < 0) - j_l_B_1 = 0; - if(j_l_B_1 > SIZE_3(input) - 1) - j_l_B_1 = SIZE_3(input) - 1; - - floatOutput += delta * w * ( - - VALUE_4(input, intSample, c, i_k_A, j_l_B)*(1-(beta-(float)B)) + - VALUE_4(input, intSample, c, i_k_A_1, j_l_B)*(1-(beta-(float)B)) - - VALUE_4(input, intSample, c, i_k_A, j_l_B_1)*(beta-(float)B) + - VALUE_4(input, intSample, c, i_k_A_1, j_l_B_1)*(beta-(float)B) - ); - } - - gradOffset_i[intIndex] = floatOutput; - } } -""" - -kernel_AdaCoF_updateGradBeta = """ - extern "C" __global__ void kernel_AdaCoF_updateGradBeta( - const int n, - const float* gradLoss, - const float* input, - const float* weight, - const float* offset_i, - const float* offset_j, - float* gradOffset_j - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - float floatOutput = 0.0; - - const int intSample = ( intIndex / SIZE_3(gradOffset_j) / SIZE_2(gradOffset_j) / SIZE_1(gradOffset_j) ) % SIZE_0(gradOffset_j); - const int intDepth = ( intIndex / SIZE_3(gradOffset_j) / SIZE_2(gradOffset_j) ) % SIZE_1(gradOffset_j); - const int i = ( intIndex / SIZE_3(gradOffset_j) ) % SIZE_2(gradOffset_j); - const int j = ( intIndex ) % SIZE_3(gradOffset_j); - - int k = intDepth / F_SIZE; - int l = intDepth % F_SIZE; - - for (int c = 0; c < 3; c++) - { - float delta = VALUE_4(gradLoss, intSample, c, i, j); - float w = VALUE_4(weight, intSample, k*F_SIZE+l, i, j); - float alpha = VALUE_4(offset_i, intSample, k*F_SIZE+l, i, j); - float beta = VALUE_4(offset_j, intSample, k*F_SIZE+l, i, j); - int A = (int) alpha; - int B = (int) beta; - - int i_k_A = i+k*DILATION+A; - if(i_k_A < 0) - i_k_A = 0; - if(i_k_A > SIZE_2(input) - 1) - i_k_A = SIZE_2(input) - 1; - - int j_l_B = j+l*DILATION+B; - if(j_l_B < 0) - j_l_B = 0; - if(j_l_B > SIZE_3(input) - 1) - j_l_B = SIZE_3(input) - 1; - - int i_k_A_1 = i+k*DILATION+A+1; - if(i_k_A_1 < 0) - i_k_A_1 = 0; - if(i_k_A_1 > SIZE_2(input) - 1) - i_k_A_1 = SIZE_2(input) - 1; - - int j_l_B_1 = j+l*DILATION+B+1; - if(j_l_B_1 < 0) - j_l_B_1 = 0; - if(j_l_B_1 > SIZE_3(input) - 1) - j_l_B_1 = SIZE_3(input) - 1; - - floatOutput += delta * w * ( - - VALUE_4(input, intSample, c, i_k_A, j_l_B)*(1-(alpha-(float)A)) - - VALUE_4(input, intSample, c, i_k_A_1, j_l_B)*(alpha-(float)A) + - VALUE_4(input, intSample, c, i_k_A, j_l_B_1)*(1-(alpha-(float)A)) + - VALUE_4(input, intSample, c, i_k_A_1, j_l_B_1)*(alpha-(float)A) - ); - } - - gradOffset_j[intIndex] = floatOutput; - } } -""" - -class FunctionAdaCoF(torch.autograd.Function): - # end - @staticmethod - def forward(ctx, input, weight, offset_i, offset_j, dilation): - ctx.save_for_backward(input, weight, offset_i, offset_j) - ctx.dilation = dilation - - intSample = input.size(0) - intInputDepth = input.size(1) - intInputHeight = input.size(2) - intInputWidth = input.size(3) - intFilterSize = int(math.sqrt(weight.size(1))) - intOutputHeight = weight.size(2) - intOutputWidth = weight.size(3) - - assert ( - intInputHeight - ((intFilterSize - 1) * dilation + 1) == intOutputHeight - 1 - ) - assert ( - intInputWidth - ((intFilterSize - 1) * dilation + 1) == intOutputWidth - 1 - ) - - assert input.is_contiguous() == True - assert weight.is_contiguous() == True - assert offset_i.is_contiguous() == True - assert offset_j.is_contiguous() == True - - output = input.new_zeros( - intSample, intInputDepth, intOutputHeight, intOutputWidth - ) - - if input.is_cuda == True: - - class Stream: - ptr = torch.cuda.current_stream().cuda_stream - - # end - - n = output.nelement() - cuda_launch( - cuda_kernel( - "kernel_AdaCoF_updateOutput", - kernel_AdaCoF_updateOutput, - { - "input": input, - "weight": weight, - "offset_i": offset_i, - "offset_j": offset_j, - "output": output, - }, - F_SIZE=str(intFilterSize), - DILATION=str(dilation) - ), - )( - grid=tuple([int((n + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - n, - input.data_ptr(), - weight.data_ptr(), - offset_i.data_ptr(), - offset_j.data_ptr(), - output.data_ptr(), - ], - stream=Stream, - ) - - elif input.is_cuda == False: - raise NotImplementedError() - - # end - - return output - - # end - @staticmethod - def backward(ctx, gradOutput): - input, weight, offset_i, offset_j = ctx.saved_tensors - dilation = ctx.dilation - - intSample = input.size(0) - intInputDepth = input.size(1) - intInputHeight = input.size(2) - intInputWidth = input.size(3) - intFilterSize = int(math.sqrt(weight.size(1))) - intOutputHeight = weight.size(2) - intOutputWidth = weight.size(3) - - assert ( - intInputHeight - ((intFilterSize - 1) * dilation + 1) == intOutputHeight - 1 - ) - assert ( - intInputWidth - ((intFilterSize - 1) * dilation + 1) == intOutputWidth - 1 - ) - - assert gradOutput.is_contiguous() == True - - gradInput = ( - input.new_zeros(intSample, intInputDepth, intInputHeight, intInputWidth) - if ctx.needs_input_grad[0] == True - else None - ) - gradWeight = ( - input.new_zeros( - intSample, intFilterSize**2, intOutputHeight, intOutputWidth - ) - if ctx.needs_input_grad[1] == True - else None - ) - gradOffset_i = ( - input.new_zeros( - intSample, intFilterSize**2, intOutputHeight, intOutputWidth - ) - if ctx.needs_input_grad[2] == True - else None - ) - gradOffset_j = ( - input.new_zeros( - intSample, intFilterSize**2, intOutputHeight, intOutputWidth - ) - if ctx.needs_input_grad[2] == True - else None - ) - - if input.is_cuda == True: - - class Stream: - ptr = torch.cuda.current_stream().cuda_stream - - # end - - # weight grad - n_w = gradWeight.nelement() - cuda_launch( - cuda_kernel( - "kernel_AdaCoF_updateGradWeight", - kernel_AdaCoF_updateGradWeight, - { - "gradLoss": gradOutput, - "input": input, - "offset_i": offset_i, - "offset_j": offset_j, - "gradWeight": gradWeight, - }, - F_SIZE=str(intFilterSize), - DILATION=str(dilation) - ), - )( - grid=tuple([int((n_w + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - n_w, - gradOutput.data_ptr(), - input.data_ptr(), - offset_i.data_ptr(), - offset_j.data_ptr(), - gradWeight.data_ptr(), - ], - stream=Stream, - ) - - # alpha grad - n_i = gradOffset_i.nelement() - cuda_launch( - cuda_kernel( - "kernel_AdaCoF_updateGradAlpha", - kernel_AdaCoF_updateGradAlpha, - { - "gradLoss": gradOutput, - "input": input, - "weight": weight, - "offset_i": offset_i, - "offset_j": offset_j, - "gradOffset_i": gradOffset_i, - }, - F_SIZE=str(intFilterSize), - DILATION=str(dilation) - ), - )( - grid=tuple([int((n_i + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - n_i, - gradOutput.data_ptr(), - input.data_ptr(), - weight.data_ptr(), - offset_i.data_ptr(), - offset_j.data_ptr(), - gradOffset_i.data_ptr(), - ], - stream=Stream, - ) - - # beta grad - n_j = gradOffset_j.nelement() - cuda_launch( - cuda_kernel( - "kernel_AdaCoF_updateGradBeta", - kernel_AdaCoF_updateGradBeta, - { - "gradLoss": gradOutput, - "input": input, - "weight": weight, - "offset_i": offset_i, - "offset_j": offset_j, - "gradOffset_j": gradOffset_j, - }, - F_SIZE=str(intFilterSize), - DILATION=str(dilation) - ), - )( - grid=tuple([int((n_j + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - n_j, - gradOutput.data_ptr(), - input.data_ptr(), - weight.data_ptr(), - offset_i.data_ptr(), - offset_j.data_ptr(), - gradOffset_j.data_ptr(), - ], - stream=Stream, - ) - - elif input.is_cuda == False: - raise NotImplementedError() - - # end - - return gradInput, gradWeight, gradOffset_i, gradOffset_j, None - -__all__ = ["FunctionAdaCoF"] diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/batch_edt.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/batch_edt.py deleted file mode 100644 index 3c3542c..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/batch_edt.py +++ /dev/null @@ -1,119 +0,0 @@ -############### DISTANCE TRANSFORM ############### -# img tensor: (bs,h,w) or (bs,1,h,w) -# returns same shape -# expects white lines, black whitespace -# defaults to diameter if empty image -from .utils import cuda_kernel, cuda_launch, cuda_int32, cuda_float32 -import torch - -_batch_edt_kernel = ( - "kernel_dt", - """ - extern "C" __global__ void kernel_dt( - const int bs, - const int h, - const int w, - const float diam2, - float* data, - float* output - ) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= bs*h*w) { - return; - } - int pb = idx / (h*w); - int pi = (idx - h*w*pb) / w; - int pj = (idx - h*w*pb - w*pi); - - float cost; - float mincost = diam2; - for (int j = 0; j < w; j++) { - cost = data[h*w*pb + w*pi + j] + (pj-j)*(pj-j); - if (cost < mincost) { - mincost = cost; - } - } - output[idx] = mincost; - return; - } -""", -) -_batch_edt = None - - -def batch_edt(img, block=1024): - # must initialize cuda/cupy after forking - global _batch_edt - if _batch_edt is None: - _batch_edt = cuda_launch(*_batch_edt_kernel) - - # bookkeeppingg - if len(img.shape) == 4: - assert img.shape[1] == 1 - img = img.squeeze(1) - expand = True - else: - expand = False - bs, h, w = img.shape - diam2 = h**2 + w**2 - odtype = img.dtype - grid = (img.nelement() + block - 1) // block - - # cupy implementation - if img.is_cuda: - # first pass, y-axis - data = ((1 - img.type(torch.float32)) * diam2).contiguous() - intermed = torch.zeros_like(data) - _batch_edt( - grid=(grid, 1, 1), - block=(block, 1, 1), # < 1024 - args=[ - cuda_int32(bs), - cuda_int32(h), - cuda_int32(w), - cuda_float32(diam2), - data.data_ptr(), - intermed.data_ptr(), - ], - ) - - # second pass, x-axis - intermed = intermed.permute(0, 2, 1).contiguous() - out = torch.zeros_like(intermed) - _batch_edt( - grid=(grid, 1, 1), - block=(block, 1, 1), - args=[ - cuda_int32(bs), - cuda_int32(w), - cuda_int32(h), - cuda_float32(diam2), - intermed.data_ptr(), - out.data_ptr(), - ], - ) - ans = out.permute(0, 2, 1).sqrt() - ans = ans.type(odtype) if odtype != ans.dtype else ans - - # default to scipy cpu implementation - else: - raise NotImplementedError() - """ sums = img.sum(dim=(1, 2)) - ans = torch.tensor( - np.stack( - [ - scipy.ndimage.morphology.distance_transform_edt(i) - if s != 0 - else np.ones_like(i) # change scipy behavior for empty image - * np.sqrt(diam2) - for i, s in zip(1 - img, sums) - ] - ), - dtype=odtype, - ) """ - - if expand: - ans = ans.unsqueeze(1) - return ans - -__all__ = ["batch_edt"] \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/correlation.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/correlation.py deleted file mode 100644 index d1e69e2..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/correlation.py +++ /dev/null @@ -1,413 +0,0 @@ -import torch -from .utils import cuda_kernel, cuda_launch, cuda_int32 - -kernel_Correlation_rearrange = """ - extern "C" __global__ void kernel_Correlation_rearrange( - const int n, - const float* input, - float* output - ) { - int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; - - if (intIndex >= n) { - return; - } - - int intSample = blockIdx.z; - int intChannel = blockIdx.y; - - float fltValue = input[(((intSample * SIZE_1(input)) + intChannel) * SIZE_2(input) * SIZE_3(input)) + intIndex]; - - __syncthreads(); - - int intPaddedY = (intIndex / SIZE_3(input)) + 4; - int intPaddedX = (intIndex % SIZE_3(input)) + 4; - int intRearrange = ((SIZE_3(input) + 8) * intPaddedY) + intPaddedX; - - output[(((intSample * SIZE_1(output) * SIZE_2(output)) + intRearrange) * SIZE_1(input)) + intChannel] = fltValue; - } -""" - -kernel_Correlation_updateOutput = """ - extern "C" __global__ void kernel_Correlation_updateOutput( - const int n, - const float* rbot0, - const float* rbot1, - float* top - ) { - extern __shared__ char patch_data_char[]; - - float *patch_data = (float *)patch_data_char; - - // First (upper left) position of kernel upper-left corner in current center position of neighborhood in image 1 - int x1 = blockIdx.x + 4; - int y1 = blockIdx.y + 4; - int item = blockIdx.z; - int ch_off = threadIdx.x; - - // Load 3D patch into shared shared memory - for (int j = 0; j < 1; j++) { // HEIGHT - for (int i = 0; i < 1; i++) { // WIDTH - int ji_off = (j + i) * SIZE_3(rbot0); - for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS - int idx1 = ((item * SIZE_1(rbot0) + y1+j) * SIZE_2(rbot0) + x1+i) * SIZE_3(rbot0) + ch; - int idxPatchData = ji_off + ch; - patch_data[idxPatchData] = rbot0[idx1]; - } - } - } - - __syncthreads(); - - __shared__ float sum[32]; - - // Compute correlation - for (int top_channel = 0; top_channel < SIZE_1(top); top_channel++) { - sum[ch_off] = 0; - - int s2o = top_channel % 9 - 4; - int s2p = top_channel / 9 - 4; - - for (int j = 0; j < 1; j++) { // HEIGHT - for (int i = 0; i < 1; i++) { // WIDTH - int ji_off = (j + i) * SIZE_3(rbot0); - for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS - int x2 = x1 + s2o; - int y2 = y1 + s2p; - - int idxPatchData = ji_off + ch; - int idx2 = ((item * SIZE_1(rbot0) + y2+j) * SIZE_2(rbot0) + x2+i) * SIZE_3(rbot0) + ch; - - sum[ch_off] += patch_data[idxPatchData] * rbot1[idx2]; - } - } - } - - __syncthreads(); - - if (ch_off == 0) { - float total_sum = 0; - for (int idx = 0; idx < 32; idx++) { - total_sum += sum[idx]; - } - const int sumelems = SIZE_3(rbot0); - const int index = ((top_channel*SIZE_2(top) + blockIdx.y)*SIZE_3(top))+blockIdx.x; - top[index + item*SIZE_1(top)*SIZE_2(top)*SIZE_3(top)] = total_sum / (float)sumelems; - } - } - } -""" - -kernel_Correlation_updateGradFirst = """ - #define ROUND_OFF 50000 - - extern "C" __global__ void kernel_Correlation_updateGradFirst( - const int n, - const int intSample, - const float* rbot0, - const float* rbot1, - const float* gradOutput, - float* gradFirst, - float* gradSecond - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - int n = intIndex % SIZE_1(gradFirst); // channels - int l = (intIndex / SIZE_1(gradFirst)) % SIZE_3(gradFirst) + 4; // w-pos - int m = (intIndex / SIZE_1(gradFirst) / SIZE_3(gradFirst)) % SIZE_2(gradFirst) + 4; // h-pos - - // round_off is a trick to enable integer division with ceil, even for negative numbers - // We use a large offset, for the inner part not to become negative. - const int round_off = ROUND_OFF; - const int round_off_s1 = round_off; - - // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior: - int xmin = (l - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4) - int ymin = (m - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4) - - // Same here: - int xmax = (l - 4 + round_off_s1) - round_off; // floor (l - 4) - int ymax = (m - 4 + round_off_s1) - round_off; // floor (m - 4) - - float sum = 0; - if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) { - xmin = max(0,xmin); - xmax = min(SIZE_3(gradOutput)-1,xmax); - - ymin = max(0,ymin); - ymax = min(SIZE_2(gradOutput)-1,ymax); - - for (int p = -4; p <= 4; p++) { - for (int o = -4; o <= 4; o++) { - // Get rbot1 data: - int s2o = o; - int s2p = p; - int idxbot1 = ((intSample * SIZE_1(rbot0) + (m+s2p)) * SIZE_2(rbot0) + (l+s2o)) * SIZE_3(rbot0) + n; - float bot1tmp = rbot1[idxbot1]; // rbot1[l+s2o,m+s2p,n] - - // Index offset for gradOutput in following loops: - int op = (p+4) * 9 + (o+4); // index[o,p] - int idxopoffset = (intSample * SIZE_1(gradOutput) + op); - - for (int y = ymin; y <= ymax; y++) { - for (int x = xmin; x <= xmax; x++) { - int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p] - sum += gradOutput[idxgradOutput] * bot1tmp; - } - } - } - } - } - const int sumelems = SIZE_1(gradFirst); - const int bot0index = ((n * SIZE_2(gradFirst)) + (m-4)) * SIZE_3(gradFirst) + (l-4); - gradFirst[bot0index + intSample*SIZE_1(gradFirst)*SIZE_2(gradFirst)*SIZE_3(gradFirst)] = sum / (float)sumelems; - } } -""" - -kernel_Correlation_updateGradSecond = """ - #define ROUND_OFF 50000 - - extern "C" __global__ void kernel_Correlation_updateGradSecond( - const int n, - const int intSample, - const float* rbot0, - const float* rbot1, - const float* gradOutput, - float* gradFirst, - float* gradSecond - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - int n = intIndex % SIZE_1(gradSecond); // channels - int l = (intIndex / SIZE_1(gradSecond)) % SIZE_3(gradSecond) + 4; // w-pos - int m = (intIndex / SIZE_1(gradSecond) / SIZE_3(gradSecond)) % SIZE_2(gradSecond) + 4; // h-pos - - // round_off is a trick to enable integer division with ceil, even for negative numbers - // We use a large offset, for the inner part not to become negative. - const int round_off = ROUND_OFF; - const int round_off_s1 = round_off; - - float sum = 0; - for (int p = -4; p <= 4; p++) { - for (int o = -4; o <= 4; o++) { - int s2o = o; - int s2p = p; - - //Get X,Y ranges and clamp - // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior: - int xmin = (l - 4 - s2o + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o) - int ymin = (m - 4 - s2p + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o) - - // Same here: - int xmax = (l - 4 - s2o + round_off_s1) - round_off; // floor (l - 4 - s2o) - int ymax = (m - 4 - s2p + round_off_s1) - round_off; // floor (m - 4 - s2p) - - if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) { - xmin = max(0,xmin); - xmax = min(SIZE_3(gradOutput)-1,xmax); - - ymin = max(0,ymin); - ymax = min(SIZE_2(gradOutput)-1,ymax); - - // Get rbot0 data: - int idxbot0 = ((intSample * SIZE_1(rbot0) + (m-s2p)) * SIZE_2(rbot0) + (l-s2o)) * SIZE_3(rbot0) + n; - float bot0tmp = rbot0[idxbot0]; // rbot1[l+s2o,m+s2p,n] - - // Index offset for gradOutput in following loops: - int op = (p+4) * 9 + (o+4); // index[o,p] - int idxopoffset = (intSample * SIZE_1(gradOutput) + op); - - for (int y = ymin; y <= ymax; y++) { - for (int x = xmin; x <= xmax; x++) { - int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p] - sum += gradOutput[idxgradOutput] * bot0tmp; - } - } - } - } - } - const int sumelems = SIZE_1(gradSecond); - const int bot1index = ((n * SIZE_2(gradSecond)) + (m-4)) * SIZE_3(gradSecond) + (l-4); - gradSecond[bot1index + intSample*SIZE_1(gradSecond)*SIZE_2(gradSecond)*SIZE_3(gradSecond)] = sum / (float)sumelems; - } } -""" - - -class _FunctionCorrelation(torch.autograd.Function): - @staticmethod - def forward(self, first, second): - rbot0 = first.new_zeros( - [first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1]] - ) - rbot1 = first.new_zeros( - [first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1]] - ) - - self.save_for_backward(first, second, rbot0, rbot1) - - first = first.contiguous() - assert first.is_cuda == True - second = second.contiguous() - assert second.is_cuda == True - - output = first.new_zeros([first.shape[0], 81, first.shape[2], first.shape[3]]) - - if first.is_cuda == True: - n = first.shape[2] * first.shape[3] - cuda_launch( - cuda_kernel( - "kernel_Correlation_rearrange", kernel_Correlation_rearrange, {"input": first, "output": rbot0} - ), - )( - grid=tuple([int((n + 16 - 1) / 16), first.shape[1], first.shape[0]]), - block=tuple([16, 1, 1]), - args=[n, first.data_ptr(), rbot0.data_ptr()], - ) - - n = second.shape[2] * second.shape[3] - cuda_launch( - cuda_kernel( - "kernel_Correlation_rearrange", kernel_Correlation_rearrange, {"input": second, "output": rbot1} - ), - )( - grid=tuple([int((n + 16 - 1) / 16), second.shape[1], second.shape[0]]), - block=tuple([16, 1, 1]), - args=[n, second.data_ptr(), rbot1.data_ptr()], - ) - - n = output.shape[1] * output.shape[2] * output.shape[3] - cuda_launch( - cuda_kernel( - "kernel_Correlation_updateOutput", - kernel_Correlation_updateOutput, - {"rbot0": rbot0, "rbot1": rbot1, "top": output}, - ), - )( - grid=tuple([output.shape[3], output.shape[2], output.shape[0]]), - block=tuple([32, 1, 1]), - shared_mem=first.shape[1] * 4, - args=[n, rbot0.data_ptr(), rbot1.data_ptr(), output.data_ptr()], - ) - - elif first.is_cuda == False: - raise NotImplementedError() - - # end - - return output - - # end - - @staticmethod - def backward(self, gradOutput): - first, second, rbot0, rbot1 = self.saved_tensors - - gradOutput = gradOutput.contiguous() - assert gradOutput.is_cuda == True - - gradFirst = ( - first.new_zeros( - [first.shape[0], first.shape[1], first.shape[2], first.shape[3]] - ) - if self.needs_input_grad[0] == True - else None - ) - gradSecond = ( - first.new_zeros( - [first.shape[0], first.shape[1], first.shape[2], first.shape[3]] - ) - if self.needs_input_grad[1] == True - else None - ) - - if first.is_cuda == True: - if gradFirst is not None: - for intSample in range(first.shape[0]): - n = first.shape[1] * first.shape[2] * first.shape[3] - cuda_launch( - cuda_kernel( - "kernel_Correlation_updateGradFirst", - kernel_Correlation_updateGradFirst, - { - "rbot0": rbot0, - "rbot1": rbot1, - "gradOutput": gradOutput, - "gradFirst": gradFirst, - "gradSecond": None, - }, - ), - )( - grid=tuple([int((n + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - n, - intSample, - rbot0.data_ptr(), - rbot1.data_ptr(), - gradOutput.data_ptr(), - gradFirst.data_ptr(), - None, - ], - ) - # end - # end - - if gradSecond is not None: - for intSample in range(first.shape[0]): - n = first.shape[1] * first.shape[2] * first.shape[3] - cuda_launch( - cuda_kernel( - "kernel_Correlation_updateGradSecond", - kernel_Correlation_updateGradSecond, - { - "rbot0": rbot0, - "rbot1": rbot1, - "gradOutput": gradOutput, - "gradFirst": None, - "gradSecond": gradSecond, - }, - ), - )( - grid=tuple([int((n + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - n, - intSample, - rbot0.data_ptr(), - rbot1.data_ptr(), - gradOutput.data_ptr(), - None, - gradSecond.data_ptr(), - ], - ) - # end - # end - - elif first.is_cuda == False: - raise NotImplementedError() - - # end - - return gradFirst, gradSecond - - # end - - -# end - - -def FunctionCorrelation(tenFirst, tenSecond): - return _FunctionCorrelation.apply(tenFirst, tenSecond) - - -# end - - -class ModuleCorrelation(torch.nn.Module): - def __init__(self): - super(ModuleCorrelation, self).__init__() - - # end - - def forward(self, tenFirst, tenSecond): - return _FunctionCorrelation.apply(tenFirst, tenSecond) - - # end - -__all__ = ["_FunctionCorrelation", "FunctionCorrelation", "ModuleCorrelation"] diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/costvol.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/costvol.py deleted file mode 100644 index 070dd48..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/costvol.py +++ /dev/null @@ -1,317 +0,0 @@ -from .utils import cuda_kernel, cuda_launch, cuda_int32 -import torch, collections - -costvol_out = """ - extern "C" __global__ void __launch_bounds__(512) costvol_out( - const int n, - const {{type}}* __restrict__ tenOne, - const {{type}}* __restrict__ tenTwo, - {{type}}* __restrict__ tenOut - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - const int intN = ( intIndex / SIZE_3(tenOut) / SIZE_2(tenOut) ) % SIZE_0(tenOut); - const int intC = -1; - const int intY = ( intIndex / SIZE_3(tenOut) ) % SIZE_2(tenOut); - const int intX = ( intIndex ) % SIZE_3(tenOut); - - {{type}} fltOne[{{intChans}}]; - - for (int intValue = 0; intValue < SIZE_1(tenOne); intValue += 1) { - fltOne[intValue] = VALUE_4(tenOne, intN, intValue, intY, intX); - } - - int intOffset = OFFSET_4(tenOut, intN, 0, intY, intX); - - for (int intOy = intY - 4; intOy <= intY + 4; intOy += 1) { - for (int intOx = intX - 4; intOx <= intX + 4; intOx += 1) { - {{type}} fltValue = 0.0f; - - if ((intOy >= 0) && (intOy < SIZE_2(tenOut)) && (intOx >= 0) && (intOx < SIZE_3(tenOut))) { - for (int intValue = 0; intValue < SIZE_1(tenOne); intValue += 1) { - fltValue += abs(fltOne[intValue] - VALUE_4(tenTwo, intN, intValue, intOy, intOx)); - } - } else { - for (int intValue = 0; intValue < SIZE_1(tenOne); intValue += 1) { - fltValue += abs(fltOne[intValue]); - } - } - - tenOut[intOffset] = fltValue / SIZE_1(tenOne); - intOffset += SIZE_2(tenOut) * SIZE_3(tenOut); - } - } - } } -""" - -costvol_onegrad = """ - extern "C" __global__ void __launch_bounds__(512) costvol_onegrad( - const int n, - const {{type}}* __restrict__ tenOne, - const {{type}}* __restrict__ tenTwo, - const {{type}}* __restrict__ tenOutgrad, - {{type}}* __restrict__ tenOnegrad, - {{type}}* __restrict__ tenTwograd - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - const int intN = ( intIndex / SIZE_3(tenOnegrad) / SIZE_2(tenOnegrad) ) % SIZE_0(tenOnegrad); - const int intC = -1; - const int intY = ( intIndex / SIZE_3(tenOnegrad) ) % SIZE_2(tenOnegrad); - const int intX = ( intIndex ) % SIZE_3(tenOnegrad); - - {{type}} fltOne[{{intChans}}]; - - for (int intValue = 0; intValue < SIZE_1(tenOne); intValue += 1) { - fltOne[intValue] = VALUE_4(tenOne, intN, intValue, intY, intX); - } - - int intOffset = OFFSET_4(tenOutgrad, intN, 0, intY, intX); - - for (int intOy = intY - 4; intOy <= intY + 4; intOy += 1) { - for (int intOx = intX - 4; intOx <= intX + 4; intOx += 1) { - if ((intOy >= 0) && (intOy < SIZE_2(tenOutgrad)) && (intOx >= 0) && (intOx < SIZE_3(tenOutgrad))) { - for (int intValue = 0; intValue < SIZE_1(tenOne); intValue += 1) { - if (fltOne[intValue] - VALUE_4(tenTwo, intN, intValue, intOy, intOx) >= 0.0f) { - tenOnegrad[OFFSET_4(tenOnegrad, intN, intValue, intY, intX)] += +tenOutgrad[intOffset] / SIZE_1(tenOne); - } else { - tenOnegrad[OFFSET_4(tenOnegrad, intN, intValue, intY, intX)] += -tenOutgrad[intOffset] / SIZE_1(tenOne); - } - } - } else { - for (int intValue = 0; intValue < SIZE_1(tenOne); intValue += 1) { - if (fltOne[intValue] >= 0.0f) { - tenOnegrad[OFFSET_4(tenOnegrad, intN, intValue, intY, intX)] += +tenOutgrad[intOffset] / SIZE_1(tenOne); - } else { - tenOnegrad[OFFSET_4(tenOnegrad, intN, intValue, intY, intX)] += -tenOutgrad[intOffset] / SIZE_1(tenOne); - } - } - } - - intOffset += SIZE_2(tenOutgrad) * SIZE_3(tenOutgrad); - } - } - } } -""" - -costvol_twograd = """ - extern "C" __global__ void __launch_bounds__(512) costvol_twograd( - const int n, - const {{type}}* __restrict__ tenOne, - const {{type}}* __restrict__ tenTwo, - const {{type}}* __restrict__ tenOutgrad, - {{type}}* __restrict__ tenOnegrad, - {{type}}* __restrict__ tenTwograd - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - const int intN = ( intIndex / SIZE_3(tenTwograd) / SIZE_2(tenTwograd) ) % SIZE_0(tenTwograd); - const int intC = -1; - const int intY = ( intIndex / SIZE_3(tenTwograd) ) % SIZE_2(tenTwograd); - const int intX = ( intIndex ) % SIZE_3(tenTwograd); - - {{type}} fltOne[{{intChans}}]; - - for (int intValue = 0; intValue < SIZE_1(tenOne); intValue += 1) { - fltOne[intValue] = VALUE_4(tenOne, intN, intValue, intY, intX); - } - - int intOffset = OFFSET_4(tenOutgrad, intN, 0, intY, intX); - - for (int intOy = intY - 4; intOy <= intY + 4; intOy += 1) { - for (int intOx = intX - 4; intOx <= intX + 4; intOx += 1) { - if ((intOy >= 0) && (intOy < SIZE_2(tenOutgrad)) && (intOx >= 0) && (intOx < SIZE_3(tenOutgrad))) { - for (int intValue = 0; intValue < SIZE_1(tenOne); intValue += 1) { - if (fltOne[intValue] - VALUE_4(tenTwo, intN, intValue, intOy, intOx) >= 0.0f) { - atomicAdd(&tenTwograd[OFFSET_4(tenTwograd, intN, intValue, intOy, intOx)], -tenOutgrad[intOffset] / SIZE_1(tenOne)); - } else { - atomicAdd(&tenTwograd[OFFSET_4(tenTwograd, intN, intValue, intOy, intOx)], +tenOutgrad[intOffset] / SIZE_1(tenOne)); - } - } - } else { - // ... - } - - intOffset += SIZE_2(tenOutgrad) * SIZE_3(tenOutgrad); - } - } - } } -""" - -class costvol_func(torch.autograd.Function): - @staticmethod - @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) - def forward(self, tenOne, tenTwo): - tenOut = tenOne.new_empty( - [tenOne.shape[0], 81, tenOne.shape[2], tenOne.shape[3]] - ) - - cuda_launch( - cuda_kernel( - "costvol_out", - costvol_out, - { - "intChans": tenOne.shape[1], - "tenOne": tenOne, - "tenTwo": tenTwo, - "tenOut": tenOut, - }, - ) - )( - grid=tuple( - [ - int( - ( - (tenOut.shape[0] * tenOut.shape[2] * tenOut.shape[3]) - + 512 - - 1 - ) - / 512 - ), - 1, - 1, - ] - ), - block=tuple([512, 1, 1]), - args=[ - cuda_int32(tenOut.shape[0] * tenOut.shape[2] * tenOut.shape[3]), - tenOne.data_ptr(), - tenTwo.data_ptr(), - tenOut.data_ptr(), - ], - stream=collections.namedtuple("Stream", "ptr")( - torch.cuda.current_stream().cuda_stream - ), - ) - - self.save_for_backward(tenOne, tenTwo) - - return tenOut - - # end - - @staticmethod - @torch.cuda.amp.custom_bwd - def backward(self, tenOutgrad): - tenOne, tenTwo = self.saved_tensors - - tenOutgrad = tenOutgrad.contiguous() - assert tenOutgrad.is_cuda == True - - tenOnegrad = ( - tenOne.new_zeros( - [tenOne.shape[0], tenOne.shape[1], tenOne.shape[2], tenOne.shape[3]] - ) - if self.needs_input_grad[0] == True - else None - ) - tenTwograd = ( - tenTwo.new_zeros( - [tenTwo.shape[0], tenTwo.shape[1], tenTwo.shape[2], tenTwo.shape[3]] - ) - if self.needs_input_grad[1] == True - else None - ) - - if tenOnegrad is not None: - cuda_launch( - cuda_kernel( - "costvol_onegrad", - costvol_onegrad, - { - "intChans": tenOne.shape[1], - "tenOne": tenOne, - "tenTwo": tenTwo, - "tenOutgrad": tenOutgrad, - "tenOnegrad": tenOnegrad, - "tenTwograd": tenTwograd, - }, - ) - )( - grid=tuple( - [ - int( - ( - ( - tenOnegrad.shape[0] - * tenOnegrad.shape[2] - * tenOnegrad.shape[3] - ) - + 512 - - 1 - ) - / 512 - ), - 1, - 1, - ] - ), - block=tuple([512, 1, 1]), - args=[ - cuda_int32( - tenOnegrad.shape[0] * tenOnegrad.shape[2] * tenOnegrad.shape[3] - ), - tenOne.data_ptr(), - tenTwo.data_ptr(), - tenOutgrad.data_ptr(), - tenOnegrad.data_ptr(), - tenTwograd.data_ptr(), - ], - stream=collections.namedtuple("Stream", "ptr")( - torch.cuda.current_stream().cuda_stream - ), - ) - # end - - if tenTwograd is not None: - cuda_launch( - cuda_kernel( - "costvol_twograd", - costvol_twograd, - { - "intChans": tenOne.shape[1], - "tenOne": tenOne, - "tenTwo": tenTwo, - "tenOutgrad": tenOutgrad, - "tenOnegrad": tenOnegrad, - "tenTwograd": tenTwograd, - }, - ) - )( - grid=tuple( - [ - int( - ( - ( - tenTwograd.shape[0] - * tenTwograd.shape[2] - * tenTwograd.shape[3] - ) - + 512 - - 1 - ) - / 512 - ), - 1, - 1, - ] - ), - block=tuple([512, 1, 1]), - args=[ - cuda_int32( - tenTwograd.shape[0] * tenTwograd.shape[2] * tenTwograd.shape[3] - ), - tenOne.data_ptr(), - tenTwo.data_ptr(), - tenOutgrad.data_ptr(), - tenOnegrad.data_ptr(), - tenTwograd.data_ptr(), - ], - stream=collections.namedtuple("Stream", "ptr")( - torch.cuda.current_stream().cuda_stream - ), - ) - # end - - return tenOnegrad, tenTwograd, None, None - - # end - - -# end - -__all__ = ["costvol_func"] \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/sepconv.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/sepconv.py deleted file mode 100644 index c334cdc..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/sepconv.py +++ /dev/null @@ -1,332 +0,0 @@ -import torch -from .utils import cuda_launch, cuda_kernel, cuda_int32 - -sepconv_vergrad = """ - extern "C" __global__ void __launch_bounds__(512) sepconv_vergrad( - const int n, - const {{type}}* __restrict__ tenIn, - const {{type}}* __restrict__ tenVer, - const {{type}}* __restrict__ tenHor, - const {{type}}* __restrict__ tenOutgrad, - {{type}}* __restrict__ tenIngrad, - {{type}}* __restrict__ tenVergrad, - {{type}}* __restrict__ tenHorgrad - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - const int intN = ( intIndex / SIZE_3(tenVergrad) / SIZE_2(tenVergrad) / SIZE_1(tenVergrad) ) % SIZE_0(tenVergrad); - const int intC = ( intIndex / SIZE_3(tenVergrad) / SIZE_2(tenVergrad) ) % SIZE_1(tenVergrad); - const int intY = ( intIndex / SIZE_3(tenVergrad) ) % SIZE_2(tenVergrad); - const int intX = ( intIndex ) % SIZE_3(tenVergrad); - - {{type}} fltVergrad = 0.0; - - {{type}} fltKahanc = 0.0; - {{type}} fltKahany = 0.0; - {{type}} fltKahant = 0.0; - - for (int intI = 0; intI < SIZE_1(tenIn); intI += 1) { - for (int intFx = 0; intFx < SIZE_1(tenHor); intFx += 1) { - fltKahany = VALUE_4(tenHor, intN, intFx, intY, intX) * VALUE_4(tenIn, intN, intI, intY + intC, intX + intFx) * VALUE_4(tenOutgrad, intN, intI, intY, intX); - fltKahany = fltKahany - fltKahanc; - fltKahant = fltVergrad + fltKahany; - fltKahanc = (fltKahant - fltVergrad) - fltKahany; - fltVergrad = fltKahant; - } - } - - tenVergrad[intIndex] = fltVergrad; - } } -""" - -sepconv_ingrad = """ - extern "C" __global__ void __launch_bounds__(512) sepconv_ingrad( - const int n, - const {{type}}* __restrict__ tenIn, - const {{type}}* __restrict__ tenVer, - const {{type}}* __restrict__ tenHor, - const {{type}}* __restrict__ tenOutgrad, - {{type}}* __restrict__ tenIngrad, - {{type}}* __restrict__ tenVergrad, - {{type}}* __restrict__ tenHorgrad - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - const int intN = ( intIndex / SIZE_3(tenIngrad) / SIZE_2(tenIngrad) / SIZE_1(tenIngrad) ) % SIZE_0(tenIngrad); - const int intC = ( intIndex / SIZE_3(tenIngrad) / SIZE_2(tenIngrad) ) % SIZE_1(tenIngrad); - const int intY = ( intIndex / SIZE_3(tenIngrad) ) % SIZE_2(tenIngrad); - const int intX = ( intIndex ) % SIZE_3(tenIngrad); - - {{type}} fltIngrad = 0.0; - - {{type}} fltKahanc = 0.0; - {{type}} fltKahany = 0.0; - {{type}} fltKahant = 0.0; - - for (int intFy = 0; intFy < SIZE_1(tenVer); intFy += 1) { - int intKy = intY + intFy - (SIZE_1(tenVer) - 1); - - if (intKy < 0) { continue; } - if (intKy >= SIZE_2(tenVer)) { continue; } - - for (int intFx = 0; intFx < SIZE_1(tenHor); intFx += 1) { - int intKx = intX + intFx - (SIZE_1(tenHor) - 1); - - if (intKx < 0) { continue; } - if (intKx >= SIZE_3(tenHor)) { continue; } - - fltKahany = VALUE_4(tenVer, intN, (SIZE_1(tenVer) - 1) - intFy, intKy, intKx) * VALUE_4(tenHor, intN, (SIZE_1(tenHor) - 1) - intFx, intKy, intKx) * VALUE_4(tenOutgrad, intN, intC, intKy, intKx); - fltKahany = fltKahany - fltKahanc; - fltKahant = fltIngrad + fltKahany; - fltKahanc = (fltKahant - fltIngrad) - fltKahany; - fltIngrad = fltKahant; - } - } - - tenIngrad[intIndex] = fltIngrad; - } } -""" - -sepconv_out = """ - extern "C" __global__ void __launch_bounds__(512) sepconv_out( - const int n, - const {{type}}* __restrict__ tenIn, - const {{type}}* __restrict__ tenVer, - const {{type}}* __restrict__ tenHor, - {{type}}* __restrict__ tenOut - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - const int intN = ( intIndex / SIZE_3(tenOut) / SIZE_2(tenOut) / SIZE_1(tenOut) ) % SIZE_0(tenOut); - const int intC = ( intIndex / SIZE_3(tenOut) / SIZE_2(tenOut) ) % SIZE_1(tenOut); - const int intY = ( intIndex / SIZE_3(tenOut) ) % SIZE_2(tenOut); - const int intX = ( intIndex ) % SIZE_3(tenOut); - - {{type}} fltOut = 0.0; - - {{type}} fltKahanc = 0.0; - {{type}} fltKahany = 0.0; - {{type}} fltKahant = 0.0; - - for (int intFy = 0; intFy < SIZE_1(tenVer); intFy += 1) { - for (int intFx = 0; intFx < SIZE_1(tenHor); intFx += 1) { - fltKahany = VALUE_4(tenIn, intN, intC, intY + intFy, intX + intFx) * VALUE_4(tenVer, intN, intFy, intY, intX) * VALUE_4(tenHor, intN, intFx, intY, intX); - fltKahany = fltKahany - fltKahanc; - fltKahant = fltOut + fltKahany; - fltKahanc = (fltKahant - fltOut) - fltKahany; - fltOut = fltKahant; - } - } - - tenOut[intIndex] = fltOut; - } } -""" - -sepconv_horgrad = """ - extern "C" __global__ void __launch_bounds__(512) sepconv_horgrad( - const int n, - const {{type}}* __restrict__ tenIn, - const {{type}}* __restrict__ tenVer, - const {{type}}* __restrict__ tenHor, - const {{type}}* __restrict__ tenOutgrad, - {{type}}* __restrict__ tenIngrad, - {{type}}* __restrict__ tenVergrad, - {{type}}* __restrict__ tenHorgrad - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - const int intN = ( intIndex / SIZE_3(tenHorgrad) / SIZE_2(tenHorgrad) / SIZE_1(tenHorgrad) ) % SIZE_0(tenHorgrad); - const int intC = ( intIndex / SIZE_3(tenHorgrad) / SIZE_2(tenHorgrad) ) % SIZE_1(tenHorgrad); - const int intY = ( intIndex / SIZE_3(tenHorgrad) ) % SIZE_2(tenHorgrad); - const int intX = ( intIndex ) % SIZE_3(tenHorgrad); - - {{type}} fltHorgrad = 0.0; - - {{type}} fltKahanc = 0.0; - {{type}} fltKahany = 0.0; - {{type}} fltKahant = 0.0; - - for (int intI = 0; intI < SIZE_1(tenIn); intI += 1) { - for (int intFy = 0; intFy < SIZE_1(tenVer); intFy += 1) { - fltKahany = VALUE_4(tenVer, intN, intFy, intY, intX) * VALUE_4(tenIn, intN, intI, intY + intFy, intX + intC) * VALUE_4(tenOutgrad, intN, intI, intY, intX); - fltKahany = fltKahany - fltKahanc; - fltKahant = fltHorgrad + fltKahany; - fltKahanc = (fltKahant - fltHorgrad) - fltKahany; - fltHorgrad = fltKahant; - } - } - - tenHorgrad[intIndex] = fltHorgrad; - } } -""" - -class sepconv_func(torch.autograd.Function): - @staticmethod - @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) - def forward(self, tenIn, tenVer, tenHor): - tenOut = tenIn.new_empty( - [ - tenIn.shape[0], - tenIn.shape[1], - tenVer.shape[2] and tenHor.shape[2], - tenVer.shape[3] and tenHor.shape[3], - ] - ) - - if tenIn.is_cuda == True: - cuda_launch( - cuda_kernel( - "sepconv_out", - sepconv_out, - { - "tenIn": tenIn, - "tenVer": tenVer, - "tenHor": tenHor, - "tenOut": tenOut, - }, - ) - )( - grid=tuple([int((tenOut.nelement() + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - cuda_int32(tenOut.nelement()), - tenIn.data_ptr(), - tenVer.data_ptr(), - tenHor.data_ptr(), - tenOut.data_ptr(), - ], - ) - - elif tenIn.is_cuda != True: - assert False - - # end - - self.save_for_backward(tenIn, tenVer, tenHor) - - return tenOut - - # end - - @staticmethod - @torch.cuda.amp.custom_bwd - def backward(self, tenOutgrad): - tenIn, tenVer, tenHor = self.saved_tensors - - tenOutgrad = tenOutgrad.contiguous() - assert tenOutgrad.is_cuda == True - - tenIngrad = ( - tenIn.new_empty( - [tenIn.shape[0], tenIn.shape[1], tenIn.shape[2], tenIn.shape[3]] - ) - if self.needs_input_grad[0] == True - else None - ) - tenVergrad = ( - tenVer.new_empty( - [tenVer.shape[0], tenVer.shape[1], tenVer.shape[2], tenVer.shape[3]] - ) - if self.needs_input_grad[1] == True - else None - ) - tenHorgrad = ( - tenHor.new_empty( - [tenHor.shape[0], tenHor.shape[1], tenHor.shape[2], tenHor.shape[3]] - ) - if self.needs_input_grad[2] == True - else None - ) - - if tenIngrad is not None: - cuda_launch( - cuda_kernel( - "sepconv_ingrad", - sepconv_ingrad, - { - "tenIn": tenIn, - "tenVer": tenVer, - "tenHor": tenHor, - "tenOutgrad": tenOutgrad, - "tenIngrad": tenIngrad, - "tenVergrad": tenVergrad, - "tenHorgrad": tenHorgrad, - }, - ) - )( - grid=tuple([int((tenIngrad.nelement() + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - cuda_int32(tenIngrad.nelement()), - tenIn.data_ptr(), - tenVer.data_ptr(), - tenHor.data_ptr(), - tenOutgrad.data_ptr(), - tenIngrad.data_ptr(), - None, - None, - ], - ) - # end - - if tenVergrad is not None: - cuda_launch( - cuda_kernel( - "sepconv_vergrad", - sepconv_vergrad, - { - "tenIn": tenIn, - "tenVer": tenVer, - "tenHor": tenHor, - "tenOutgrad": tenOutgrad, - "tenIngrad": tenIngrad, - "tenVergrad": tenVergrad, - "tenHorgrad": tenHorgrad, - }, - ) - )( - grid=tuple([int((tenVergrad.nelement() + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - cuda_int32(tenVergrad.nelement()), - tenIn.data_ptr(), - tenVer.data_ptr(), - tenHor.data_ptr(), - tenOutgrad.data_ptr(), - None, - tenVergrad.data_ptr(), - None, - ], - ) - # end - - if tenHorgrad is not None: - cuda_launch( - cuda_kernel( - "sepconv_horgrad", - sepconv_horgrad, - { - "tenIn": tenIn, - "tenVer": tenVer, - "tenHor": tenHor, - "tenOutgrad": tenOutgrad, - "tenIngrad": tenIngrad, - "tenVergrad": tenVergrad, - "tenHorgrad": tenHorgrad, - }, - ) - )( - grid=tuple([int((tenHorgrad.nelement() + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - cuda_int32(tenHorgrad.nelement()), - tenIn.data_ptr(), - tenVer.data_ptr(), - tenHor.data_ptr(), - tenOutgrad.data_ptr(), - None, - None, - tenHorgrad.data_ptr(), - ], - ) - # end - - return tenIngrad, tenVergrad, tenHorgrad - - # end - - -# end -__all__ = ["sepconv_func"] \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/softsplat.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/softsplat.py deleted file mode 100644 index 4a2ae47..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/softsplat.py +++ /dev/null @@ -1,440 +0,0 @@ -import torch -from .utils import cuda_launch, cuda_kernel, cuda_int32 -import cupy -import collections - -softsplat_flowgrad = """ - extern "C" __global__ void __launch_bounds__(512) softsplat_flowgrad( - const int n, - const {{type}}* __restrict__ tenIn, - const {{type}}* __restrict__ tenFlow, - const {{type}}* __restrict__ tenOutgrad, - {{type}}* __restrict__ tenIngrad, - {{type}}* __restrict__ tenFlowgrad - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - const int intN = ( intIndex / SIZE_3(tenFlowgrad) / SIZE_2(tenFlowgrad) / SIZE_1(tenFlowgrad) ) % SIZE_0(tenFlowgrad); - const int intC = ( intIndex / SIZE_3(tenFlowgrad) / SIZE_2(tenFlowgrad) ) % SIZE_1(tenFlowgrad); - const int intY = ( intIndex / SIZE_3(tenFlowgrad) ) % SIZE_2(tenFlowgrad); - const int intX = ( intIndex ) % SIZE_3(tenFlowgrad); - - assert(SIZE_1(tenFlow) == 2); - - {{type}} fltFlowgrad = 0.0f; - - {{type}} fltX = ({{type}}) (intX) + VALUE_4(tenFlow, intN, 0, intY, intX); - {{type}} fltY = ({{type}}) (intY) + VALUE_4(tenFlow, intN, 1, intY, intX); - - if (isfinite(fltX) == false) { return; } - if (isfinite(fltY) == false) { return; } - - int intNorthwestX = (int) (floor(fltX)); - int intNorthwestY = (int) (floor(fltY)); - int intNortheastX = intNorthwestX + 1; - int intNortheastY = intNorthwestY; - int intSouthwestX = intNorthwestX; - int intSouthwestY = intNorthwestY + 1; - int intSoutheastX = intNorthwestX + 1; - int intSoutheastY = intNorthwestY + 1; - - {{type}} fltNorthwest = 0.0f; - {{type}} fltNortheast = 0.0f; - {{type}} fltSouthwest = 0.0f; - {{type}} fltSoutheast = 0.0f; - - if (intC == 0) { - fltNorthwest = (({{type}}) (-1.0f)) * (({{type}}) (intSoutheastY) - fltY); - fltNortheast = (({{type}}) (+1.0f)) * (({{type}}) (intSouthwestY) - fltY); - fltSouthwest = (({{type}}) (-1.0f)) * (fltY - ({{type}}) (intNortheastY)); - fltSoutheast = (({{type}}) (+1.0f)) * (fltY - ({{type}}) (intNorthwestY)); - - } else if (intC == 1) { - fltNorthwest = (({{type}}) (intSoutheastX) - fltX) * (({{type}}) (-1.0f)); - fltNortheast = (fltX - ({{type}}) (intSouthwestX)) * (({{type}}) (-1.0f)); - fltSouthwest = (({{type}}) (intNortheastX) - fltX) * (({{type}}) (+1.0f)); - fltSoutheast = (fltX - ({{type}}) (intNorthwestX)) * (({{type}}) (+1.0f)); - - } - - for (int intChannel = 0; intChannel < SIZE_1(tenOutgrad); intChannel += 1) { - {{type}} fltIn = VALUE_4(tenIn, intN, intChannel, intY, intX); - - if ((intNorthwestX >= 0) && (intNorthwestX < SIZE_3(tenOutgrad)) && (intNorthwestY >= 0) && (intNorthwestY < SIZE_2(tenOutgrad))) { - fltFlowgrad += VALUE_4(tenOutgrad, intN, intChannel, intNorthwestY, intNorthwestX) * fltIn * fltNorthwest; - } - - if ((intNortheastX >= 0) && (intNortheastX < SIZE_3(tenOutgrad)) && (intNortheastY >= 0) && (intNortheastY < SIZE_2(tenOutgrad))) { - fltFlowgrad += VALUE_4(tenOutgrad, intN, intChannel, intNortheastY, intNortheastX) * fltIn * fltNortheast; - } - - if ((intSouthwestX >= 0) && (intSouthwestX < SIZE_3(tenOutgrad)) && (intSouthwestY >= 0) && (intSouthwestY < SIZE_2(tenOutgrad))) { - fltFlowgrad += VALUE_4(tenOutgrad, intN, intChannel, intSouthwestY, intSouthwestX) * fltIn * fltSouthwest; - } - - if ((intSoutheastX >= 0) && (intSoutheastX < SIZE_3(tenOutgrad)) && (intSoutheastY >= 0) && (intSoutheastY < SIZE_2(tenOutgrad))) { - fltFlowgrad += VALUE_4(tenOutgrad, intN, intChannel, intSoutheastY, intSoutheastX) * fltIn * fltSoutheast; - } - } - - tenFlowgrad[intIndex] = fltFlowgrad; - } } -""" - -softsplat_ingrad = """ - extern "C" __global__ void __launch_bounds__(512) softsplat_ingrad( - const int n, - const {{type}}* __restrict__ tenIn, - const {{type}}* __restrict__ tenFlow, - const {{type}}* __restrict__ tenOutgrad, - {{type}}* __restrict__ tenIngrad, - {{type}}* __restrict__ tenFlowgrad - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - const int intN = ( intIndex / SIZE_3(tenIngrad) / SIZE_2(tenIngrad) / SIZE_1(tenIngrad) ) % SIZE_0(tenIngrad); - const int intC = ( intIndex / SIZE_3(tenIngrad) / SIZE_2(tenIngrad) ) % SIZE_1(tenIngrad); - const int intY = ( intIndex / SIZE_3(tenIngrad) ) % SIZE_2(tenIngrad); - const int intX = ( intIndex ) % SIZE_3(tenIngrad); - - assert(SIZE_1(tenFlow) == 2); - - {{type}} fltIngrad = 0.0f; - - {{type}} fltX = ({{type}}) (intX) + VALUE_4(tenFlow, intN, 0, intY, intX); - {{type}} fltY = ({{type}}) (intY) + VALUE_4(tenFlow, intN, 1, intY, intX); - - if (isfinite(fltX) == false) { return; } - if (isfinite(fltY) == false) { return; } - - int intNorthwestX = (int) (floor(fltX)); - int intNorthwestY = (int) (floor(fltY)); - int intNortheastX = intNorthwestX + 1; - int intNortheastY = intNorthwestY; - int intSouthwestX = intNorthwestX; - int intSouthwestY = intNorthwestY + 1; - int intSoutheastX = intNorthwestX + 1; - int intSoutheastY = intNorthwestY + 1; - - {{type}} fltNorthwest = (({{type}}) (intSoutheastX) - fltX) * (({{type}}) (intSoutheastY) - fltY); - {{type}} fltNortheast = (fltX - ({{type}}) (intSouthwestX)) * (({{type}}) (intSouthwestY) - fltY); - {{type}} fltSouthwest = (({{type}}) (intNortheastX) - fltX) * (fltY - ({{type}}) (intNortheastY)); - {{type}} fltSoutheast = (fltX - ({{type}}) (intNorthwestX)) * (fltY - ({{type}}) (intNorthwestY)); - - if ((intNorthwestX >= 0) && (intNorthwestX < SIZE_3(tenOutgrad)) && (intNorthwestY >= 0) && (intNorthwestY < SIZE_2(tenOutgrad))) { - fltIngrad += VALUE_4(tenOutgrad, intN, intC, intNorthwestY, intNorthwestX) * fltNorthwest; - } - - if ((intNortheastX >= 0) && (intNortheastX < SIZE_3(tenOutgrad)) && (intNortheastY >= 0) && (intNortheastY < SIZE_2(tenOutgrad))) { - fltIngrad += VALUE_4(tenOutgrad, intN, intC, intNortheastY, intNortheastX) * fltNortheast; - } - - if ((intSouthwestX >= 0) && (intSouthwestX < SIZE_3(tenOutgrad)) && (intSouthwestY >= 0) && (intSouthwestY < SIZE_2(tenOutgrad))) { - fltIngrad += VALUE_4(tenOutgrad, intN, intC, intSouthwestY, intSouthwestX) * fltSouthwest; - } - - if ((intSoutheastX >= 0) && (intSoutheastX < SIZE_3(tenOutgrad)) && (intSoutheastY >= 0) && (intSoutheastY < SIZE_2(tenOutgrad))) { - fltIngrad += VALUE_4(tenOutgrad, intN, intC, intSoutheastY, intSoutheastX) * fltSoutheast; - } - - tenIngrad[intIndex] = fltIngrad; - } } -""" - -softsplat_out = """ - extern "C" __global__ void __launch_bounds__(512) softsplat_out( - const int n, - const {{type}}* __restrict__ tenIn, - const {{type}}* __restrict__ tenFlow, - {{type}}* __restrict__ tenOut - ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { - const int intN = ( intIndex / SIZE_3(tenOut) / SIZE_2(tenOut) / SIZE_1(tenOut) ) % SIZE_0(tenOut); - const int intC = ( intIndex / SIZE_3(tenOut) / SIZE_2(tenOut) ) % SIZE_1(tenOut); - const int intY = ( intIndex / SIZE_3(tenOut) ) % SIZE_2(tenOut); - const int intX = ( intIndex ) % SIZE_3(tenOut); - - assert(SIZE_1(tenFlow) == 2); - - {{type}} fltX = ({{type}}) (intX) + VALUE_4(tenFlow, intN, 0, intY, intX); - {{type}} fltY = ({{type}}) (intY) + VALUE_4(tenFlow, intN, 1, intY, intX); - - if (isfinite(fltX) == false) { return; } - if (isfinite(fltY) == false) { return; } - - {{type}} fltIn = VALUE_4(tenIn, intN, intC, intY, intX); - - int intNorthwestX = (int) (floor(fltX)); - int intNorthwestY = (int) (floor(fltY)); - int intNortheastX = intNorthwestX + 1; - int intNortheastY = intNorthwestY; - int intSouthwestX = intNorthwestX; - int intSouthwestY = intNorthwestY + 1; - int intSoutheastX = intNorthwestX + 1; - int intSoutheastY = intNorthwestY + 1; - - {{type}} fltNorthwest = (({{type}}) (intSoutheastX) - fltX) * (({{type}}) (intSoutheastY) - fltY); - {{type}} fltNortheast = (fltX - ({{type}}) (intSouthwestX)) * (({{type}}) (intSouthwestY) - fltY); - {{type}} fltSouthwest = (({{type}}) (intNortheastX) - fltX) * (fltY - ({{type}}) (intNortheastY)); - {{type}} fltSoutheast = (fltX - ({{type}}) (intNorthwestX)) * (fltY - ({{type}}) (intNorthwestY)); - - if ((intNorthwestX >= 0) && (intNorthwestX < SIZE_3(tenOut)) && (intNorthwestY >= 0) && (intNorthwestY < SIZE_2(tenOut))) { - atomicAdd(&tenOut[OFFSET_4(tenOut, intN, intC, intNorthwestY, intNorthwestX)], fltIn * fltNorthwest); - } - - if ((intNortheastX >= 0) && (intNortheastX < SIZE_3(tenOut)) && (intNortheastY >= 0) && (intNortheastY < SIZE_2(tenOut))) { - atomicAdd(&tenOut[OFFSET_4(tenOut, intN, intC, intNortheastY, intNortheastX)], fltIn * fltNortheast); - } - - if ((intSouthwestX >= 0) && (intSouthwestX < SIZE_3(tenOut)) && (intSouthwestY >= 0) && (intSouthwestY < SIZE_2(tenOut))) { - atomicAdd(&tenOut[OFFSET_4(tenOut, intN, intC, intSouthwestY, intSouthwestX)], fltIn * fltSouthwest); - } - - if ((intSoutheastX >= 0) && (intSoutheastX < SIZE_3(tenOut)) && (intSoutheastY >= 0) && (intSoutheastY < SIZE_2(tenOut))) { - atomicAdd(&tenOut[OFFSET_4(tenOut, intN, intC, intSoutheastY, intSoutheastX)], fltIn * fltSoutheast); - } - } } -""" - - -# end - -class softsplat_func(torch.autograd.Function): - @staticmethod - @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) - def forward(self, tenIn, tenFlow): - tenOut = tenIn.new_zeros( - [tenIn.shape[0], tenIn.shape[1], tenIn.shape[2], tenIn.shape[3]] - ) - - if tenIn.is_cuda == True: - cuda_launch( - cuda_kernel( - "softsplat_out", - softsplat_out, - {"tenIn": tenIn, "tenFlow": tenFlow, "tenOut": tenOut}, - ) - )( - grid=tuple([int((tenOut.nelement() + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - cuda_int32(tenOut.nelement()), - tenIn.data_ptr(), - tenFlow.data_ptr(), - tenOut.data_ptr(), - ], - stream=collections.namedtuple("Stream", "ptr")( - torch.cuda.current_stream().cuda_stream - ), - ) - - elif tenIn.is_cuda != True: - assert False - - # end - - self.save_for_backward(tenIn, tenFlow) - - return tenOut - - # end - - @staticmethod - @torch.cuda.amp.custom_bwd - def backward(self, tenOutgrad): - tenIn, tenFlow = self.saved_tensors - - tenOutgrad = tenOutgrad.contiguous() - assert tenOutgrad.is_cuda == True - - tenIngrad = ( - tenIn.new_zeros( - [tenIn.shape[0], tenIn.shape[1], tenIn.shape[2], tenIn.shape[3]] - ) - if self.needs_input_grad[0] == True - else None - ) - tenFlowgrad = ( - tenFlow.new_zeros( - [tenFlow.shape[0], tenFlow.shape[1], tenFlow.shape[2], tenFlow.shape[3]] - ) - if self.needs_input_grad[1] == True - else None - ) - - if tenIngrad is not None: - cuda_launch( - cuda_kernel( - "softsplat_ingrad", - softsplat_ingrad, - { - "tenIn": tenIn, - "tenFlow": tenFlow, - "tenOutgrad": tenOutgrad, - "tenIngrad": tenIngrad, - "tenFlowgrad": tenFlowgrad, - }, - ) - )( - grid=tuple([int((tenIngrad.nelement() + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - cuda_int32(tenIngrad.nelement()), - tenIn.data_ptr(), - tenFlow.data_ptr(), - tenOutgrad.data_ptr(), - tenIngrad.data_ptr(), - None, - ], - stream=collections.namedtuple("Stream", "ptr")( - torch.cuda.current_stream().cuda_stream - ), - ) - # end - - if tenFlowgrad is not None: - cuda_launch( - cuda_kernel( - "softsplat_flowgrad", - softsplat_flowgrad, - { - "tenIn": tenIn, - "tenFlow": tenFlow, - "tenOutgrad": tenOutgrad, - "tenIngrad": tenIngrad, - "tenFlowgrad": tenFlowgrad, - }, - ) - )( - grid=tuple([int((tenFlowgrad.nelement() + 512 - 1) / 512), 1, 1]), - block=tuple([512, 1, 1]), - args=[ - cuda_int32(tenFlowgrad.nelement()), - tenIn.data_ptr(), - tenFlow.data_ptr(), - tenOutgrad.data_ptr(), - None, - tenFlowgrad.data_ptr(), - ], - stream=collections.namedtuple("Stream", "ptr")( - torch.cuda.current_stream().cuda_stream - ), - ) - # end - - return tenIngrad, tenFlowgrad - - # end - - -def FunctionSoftsplat(tenInput, tenFlow, tenMetric, strType): - assert tenMetric is None or tenMetric.shape[1] == 1 - assert strType in ["summation", "average", "linear", "softmax"] - - if strType == "average": - tenInput = torch.cat( - [ - tenInput, - tenInput.new_ones( - tenInput.shape[0], 1, tenInput.shape[2], tenInput.shape[3] - ), - ], - 1, - ) - - elif strType == "linear": - tenInput = torch.cat([tenInput * tenMetric, tenMetric], 1) - - elif strType == "softmax": - tenInput = torch.cat([tenInput * tenMetric.exp(), tenMetric.exp()], 1) - - # end - - tenOutput = softsplat_func.apply(tenInput, tenFlow) - - if strType != "summation": - tenNormalize = tenOutput[:, -1:, :, :] - - tenNormalize[tenNormalize == 0.0] = 1.0 - - tenOutput = tenOutput[:, :-1, :, :] / tenNormalize - # end - - return tenOutput - - -# end - - -class ModuleSoftsplat(torch.nn.Module): - def __init__(self, strType): - super().__init__() - - self.strType = strType - - # end - - def forward(self, tenInput, tenFlow, tenMetric): - return FunctionSoftsplat(tenInput, tenFlow, tenMetric, self.strType) - - # end - - -# end - - - -def softsplat( - tenIn: torch.Tensor, tenFlow: torch.Tensor, tenMetric: torch.Tensor, strMode: str -): - assert strMode.split("-")[0] in ["sum", "avg", "linear", "soft"] - - if strMode == "sum": - assert tenMetric is None - if strMode == "avg": - assert tenMetric is None - if strMode.split("-")[0] == "linear": - assert tenMetric is not None - if strMode.split("-")[0] == "soft": - assert tenMetric is not None - - if strMode == "avg": - tenIn = torch.cat( - [ - tenIn, - tenIn.new_ones([tenIn.shape[0], 1, tenIn.shape[2], tenIn.shape[3]]), - ], - 1, - ) - - elif strMode.split("-")[0] == "linear": - tenIn = torch.cat([tenIn * tenMetric, tenMetric], 1) - - elif strMode.split("-")[0] == "soft": - tenIn = torch.cat([tenIn * tenMetric.exp(), tenMetric.exp()], 1) - - # end - - tenOut = softsplat_func.apply(tenIn, tenFlow) - - if strMode.split("-")[0] in ["avg", "linear", "soft"]: - tenNormalize = tenOut[:, -1:, :, :] - - if len(strMode.split("-")) == 1: - tenNormalize = tenNormalize + 0.0000001 - - elif strMode.split("-")[1] == "addeps": - tenNormalize = tenNormalize + 0.0000001 - - elif strMode.split("-")[1] == "zeroeps": - tenNormalize[tenNormalize == 0.0] = 1.0 - - elif strMode.split("-")[1] == "clipeps": - tenNormalize = tenNormalize.clip(0.0000001, None) - - # end - - tenOut = tenOut[:, :-1, :, :] / tenNormalize - # end - - return tenOut - - -# end - -__all__ = ["FunctionSoftsplat", "ModuleSoftsplat", "softsplat", "softsplat_func"] diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/utils.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/utils.py deleted file mode 100644 index ec29a48..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/cupy_ops/utils.py +++ /dev/null @@ -1,242 +0,0 @@ -import cupy -import os -import re -import torch -import typing -from pathlib import Path -import platform - -########################################################## - - -objCudacache = {} - - -def cuda_int32(intIn: int): - return cupy.int32(intIn) - - -# end - - -def cuda_float32(fltIn: float): - return cupy.float32(fltIn) - - -# end - - -def cuda_kernel(strFunction: str, strKernel: str, objVariables: typing.Dict, **replace_kwargs): - if "device" not in objCudacache: - objCudacache["device"] = torch.cuda.get_device_name() - # end - - strKey = strFunction - - for strVariable in objVariables: - objValue = objVariables[strVariable] - - strKey += strVariable - - if objValue is None: - continue - - elif type(objValue) == int: - strKey += str(objValue) - - elif type(objValue) == float: - strKey += str(objValue) - - elif type(objValue) == bool: - strKey += str(objValue) - - elif type(objValue) == str: - strKey += objValue - - elif type(objValue) == torch.Tensor: - strKey += str(objValue.dtype) - strKey += str(objValue.shape) - strKey += str(objValue.stride()) - - elif True: - print(strVariable, type(objValue)) - assert False - - # end - # end - - strKey += objCudacache["device"] - - if strKey not in objCudacache: - for strVariable in objVariables: - objValue = objVariables[strVariable] - - if objValue is None: - continue - - elif type(objValue) == int: - strKernel = strKernel.replace("{{" + strVariable + "}}", str(objValue)) - - elif type(objValue) == float: - strKernel = strKernel.replace("{{" + strVariable + "}}", str(objValue)) - - elif type(objValue) == bool: - strKernel = strKernel.replace("{{" + strVariable + "}}", str(objValue)) - - elif type(objValue) == str: - strKernel = strKernel.replace("{{" + strVariable + "}}", objValue) - - elif type(objValue) == torch.Tensor and objValue.dtype == torch.uint8: - strKernel = strKernel.replace("{{type}}", "unsigned char") - - elif type(objValue) == torch.Tensor and objValue.dtype == torch.float16: - strKernel = strKernel.replace("{{type}}", "half") - - elif type(objValue) == torch.Tensor and objValue.dtype == torch.float32: - strKernel = strKernel.replace("{{type}}", "float") - - elif type(objValue) == torch.Tensor and objValue.dtype == torch.float64: - strKernel = strKernel.replace("{{type}}", "double") - - elif type(objValue) == torch.Tensor and objValue.dtype == torch.int32: - strKernel = strKernel.replace("{{type}}", "int") - - elif type(objValue) == torch.Tensor and objValue.dtype == torch.int64: - strKernel = strKernel.replace("{{type}}", "long") - - elif type(objValue) == torch.Tensor: - print(strVariable, objValue.dtype) - assert False - - elif True: - print(strVariable, type(objValue)) - assert False - - # end - # end - - while True: - objMatch = re.search("(SIZE_)([0-4])(\()([^\)]*)(\))", strKernel) - - if objMatch is None: - break - # end - - intArg = int(objMatch.group(2)) - - strTensor = objMatch.group(4) - intSizes = objVariables[strTensor].size() - - strKernel = strKernel.replace(objMatch.group(), str(intSizes[intArg])) - # end - - while True: - objMatch = re.search("(OFFSET_)([0-4])(\()([^\)]+)(\))", strKernel) - - if objMatch is None: - break - # end - - intArgs = int(objMatch.group(2)) - strArgs = objMatch.group(4).split(",") - - strTensor = strArgs[0] - intStrides = objVariables[strTensor].stride() - strIndex = [ - "((" - + strArgs[intArg + 1].replace("{", "(").replace("}", ")").strip() - + ")*" - + str(intStrides[intArg]) - + ")" - for intArg in range(intArgs) - ] - - strKernel = strKernel.replace( - objMatch.group(0), "(" + str.join("+", strIndex) + ")" - ) - # end - - while True: - objMatch = re.search("(VALUE_)([0-4])(\()", strKernel) - - if objMatch is None: - break - # end - - intStart = objMatch.span()[1] - intStop = objMatch.span()[1] - intParentheses = 1 - - while True: - intParentheses += 1 if strKernel[intStop] == "(" else 0 - intParentheses -= 1 if strKernel[intStop] == ")" else 0 - - if intParentheses == 0: - break - # end - - intStop += 1 - # end - - intArgs = int(objMatch.group(2)) - strArgs = strKernel[intStart:intStop].split(",") - - assert intArgs == len(strArgs) - 1 - - strTensor = strArgs[0] - intStrides = objVariables[strTensor].stride() - - strIndex = [] - - for intArg in range(intArgs): - strIndex.append( - "((" - + strArgs[intArg + 1].replace("{", "(").replace("}", ")").strip() - + ")*" - + str(intStrides[intArg]) - + ")" - ) - # end - - strKernel = strKernel.replace( - "VALUE_" + str(intArgs) + "(" + strKernel[intStart:intStop] + ")", - strTensor + "[" + str.join("+", strIndex) + "]", - ) - # end - - for replace_key, value in replace_kwargs.items(): - strKernel = strKernel.replace(replace_key, value) - - objCudacache[strKey] = {"strFunction": strFunction, "strKernel": strKernel} - # end - - return strKey - - -# end -def get_cuda_home_path(): - if "CUDA_HOME" in os.environ: - return os.environ["CUDA_HOME"] - import torch - torch_lib_path = Path(torch.__file__).parent / "lib" - torch_lib_path = str(torch_lib_path.resolve()) - if os.path.exists(torch_lib_path): - nvrtc = filter(lambda lib_file: "nvrtc-builtins" in lib_file, os.listdir(torch_lib_path)) - nvrtc = list(nvrtc) - return torch_lib_path if len(nvrtc) > 0 else None - -@cupy.memoize(for_each_device=True) -def cuda_launch(strKey: str): - if True:#"CUDA_HOME" not in os.environ: - cuda_home = get_cuda_home_path() - if cuda_home is not None: - os.environ["CUDA_HOME"] = cuda_home - os.environ["CUDA_PATH"] = cuda_home - else: - os.environ["CUDA_HOME"] = "/usr/local/cuda/" - os.environ["CUDA_PATH"] = "/usr/local/cuda/" - # print(objCudacache[strKey]['strKernel']) - # return cupy.cuda.compile_with_cache(objCudacache[strKey]['strKernel'], tuple(['-I ' + os.environ['CUDA_HOME'], '-I ' + os.environ['CUDA_HOME'] + '/include'])).get_function(objCudacache[strKey]['strFunction']) - return cupy.RawModule(code=objCudacache[strKey]["strKernel"]).get_function( - objCudacache[strKey]["strFunction"] - ) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/__init__.py deleted file mode 100644 index e36bb6c..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/__init__.py +++ /dev/null @@ -1,150 +0,0 @@ -import comfy.model_management as model_management -import torch -import torch.multiprocessing as mp -from .worker_process import f -from .utils import to_shared_memory - -parent_conn, child_conn, process = None, None, None -device = model_management.get_torch_device() - -def req_to_taichi_process(op_name, *tensors): - global parent_conn, child_conn, process - if parent_conn is None: - mp.set_start_method('spawn', force=True) - parent_conn, child_conn = mp.Pipe() - process = mp.Process(target=f, args=(child_conn, device)) - process.start() - - tensors = to_shared_memory(tensors) - parent_conn.send((op_name, tensors)) - result = parent_conn.recv() - del tensors - - if type(result) not in [tuple, list]: - raise Exception(result) - - return [tensor.to(device) for tensor in result] - -def softsplat( - tenIn: torch.Tensor, tenFlow: torch.Tensor, tenMetric: torch.Tensor, strMode: str -): - assert strMode.split("-")[0] in ["sum", "avg", "linear", "soft"] - - if strMode == "sum": - assert tenMetric is None - if strMode == "avg": - assert tenMetric is None - if strMode.split("-")[0] == "linear": - assert tenMetric is not None - if strMode.split("-")[0] == "soft": - assert tenMetric is not None - - if strMode == "avg": - tenIn = torch.cat( - [ - tenIn, - tenIn.new_ones([tenIn.shape[0], 1, tenIn.shape[2], tenIn.shape[3]]), - ], - 1, - ) - - elif strMode.split("-")[0] == "linear": - tenIn = torch.cat([tenIn * tenMetric, tenMetric], 1) - - elif strMode.split("-")[0] == "soft": - tenIn = torch.cat([tenIn * tenMetric.exp(), tenMetric.exp()], 1) - - # end - - tenOut = req_to_taichi_process("softsplat_out", tenIn, tenFlow)[0] - - if strMode.split("-")[0] in ["avg", "linear", "soft"]: - tenNormalize = tenOut[:, -1:, :, :] - - if len(strMode.split("-")) == 1: - tenNormalize = tenNormalize + 0.0000001 - - elif strMode.split("-")[1] == "addeps": - tenNormalize = tenNormalize + 0.0000001 - - elif strMode.split("-")[1] == "zeroeps": - tenNormalize[tenNormalize == 0.0] = 1.0 - - elif strMode.split("-")[1] == "clipeps": - tenNormalize = tenNormalize.clip(0.0000001, None) - - # end - - tenOut = tenOut[:, :-1, :, :] / tenNormalize - # end - - return tenOut - -def FunctionSoftsplat(tenInput, tenFlow, tenMetric, strType): - assert tenMetric is None or tenMetric.shape[1] == 1 - assert strType in ["summation", "average", "linear", "softmax"] - - if strType == "average": - tenInput = torch.cat( - [ - tenInput, - tenInput.new_ones( - tenInput.shape[0], 1, tenInput.shape[2], tenInput.shape[3] - ), - ], - 1, - ) - - elif strType == "linear": - tenInput = torch.cat([tenInput * tenMetric, tenMetric], 1) - - elif strType == "softmax": - tenInput = torch.cat([tenInput * tenMetric.exp(), tenMetric.exp()], 1) - - # end - - tenOutput = req_to_taichi_process("softsplat_out", tenInput, tenFlow)[0] - - if strType != "summation": - tenNormalize = tenOutput[:, -1:, :, :] - - tenNormalize[tenNormalize == 0.0] = 1.0 - - tenOutput = tenOutput[:, :-1, :, :] / tenNormalize - # end - - return tenOutput - - -# end - - -class ModuleSoftsplat(torch.nn.Module): - def __init__(self, strType): - super(self).__init__() - - self.strType = strType - - # end - - def forward(self, tenInput, tenFlow, tenMetric): - return FunctionSoftsplat(tenInput, tenFlow, tenMetric, self.strType) - -def softsplat_func(tenIn, tenFlow): - return req_to_taichi_process("softsplat_out", tenIn, tenFlow)[0] - -class costvol_func: - @staticmethod - def apply(tenOne, tenTwo): - return req_to_taichi_process("costvol_out", tenOne, tenTwo)[0] - -class sepconv_func: - @staticmethod - def apply(tenIn, tenVer, tenHor): - return req_to_taichi_process("sepconv_out", tenIn, tenVer, tenHor)[0] - -def init(): - one_sample = torch.ones(1, 3, 16, 16, dtype=torch.float32, device=device) - softsplat_func(one_sample, one_sample) - costvol_func.apply(one_sample, one_sample) - sepconv_func.apply(one_sample, one_sample, one_sample) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/adacof.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/adacof.py deleted file mode 100644 index acf2672..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/adacof.py +++ /dev/null @@ -1,6 +0,0 @@ -import torch -class FunctionAdaCoF(torch.autograd.Function): - # end - @staticmethod - def forward(ctx, input, weight, offset_i, offset_j, dilation): - raise NotImplementedError() diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/batch_edt.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/batch_edt.py deleted file mode 100644 index c1fe1fd..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/batch_edt.py +++ /dev/null @@ -1,2 +0,0 @@ -def batch_edt(img, block=1024): - raise NotImplementedError() \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/correlation.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/correlation.py deleted file mode 100644 index 3782d1a..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/correlation.py +++ /dev/null @@ -1,15 +0,0 @@ -import torch - -class _FunctionCorrelation(torch.autograd.Function): - @staticmethod - def forward(self, first, second): - raise NotImplementedError() - -def FunctionCorrelation(tenFirst, tenSecond): - raise NotImplementedError() - return _FunctionCorrelation.apply(tenFirst, tenSecond) - -class ModuleCorrelation(torch.nn.Module): - def __init__(self): - raise NotImplementedError() - super(ModuleCorrelation, self).__init__() \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/costvol.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/costvol.py deleted file mode 100644 index a82b727..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/costvol.py +++ /dev/null @@ -1,26 +0,0 @@ -import taichi as ti -import taichi.math as tm - -""" @ti.kernel -def costvol_out(tenOne: ti.types.ndarray(), tltOne: ti.types.ndarray(), tenTwo: ti.types.ndarray(), tenOut: ti.types.ndarray()): - N, C, H, W = tenOut.shape - for i, ch, y, x in ti.ndrange(N, C, H, W): - for intValue in range(tenOne.shape[1]): - tltOne[intValue] = tenOne[i, intValue, y, x] - - tenOut_ch = 0 - for intOy in range(y - 4, y + 4 + 1): - for intOx in range(x - 4, x + 4 + 1): - point = tm.ivec2(intOx, intOy) - fltValue = 0.0 - for intValue in range(ch): - if (point.y >= 0) and (point.y < H) and (point.x >= 0) and (point.x < W): - fltValue += ti.abs(tltOne[intValue] - tenTwo[i, intValue, point.y, point.x]) - else: - fltValue += ti.abs(tltOne[intValue]) - - tenOut[i, tenOut_ch, y, x] = fltValue / tenOne.shape[1] - tenOut_ch += 1 """ - -def worker_interface(op_name, tensors): - raise NotImplementedError(op_name) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/raw_softsplat.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/raw_softsplat.py deleted file mode 100644 index 0cfd3fd..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/raw_softsplat.py +++ /dev/null @@ -1,126 +0,0 @@ -#Seperate taichi kernels to another file so that comfy.model_management won't be called in the new process - -import taichi as ti -import taichi.math as tm - -@ti.func -def put_to_tenOut(tenOut: ti.types.ndarray(), fltIn: ti.i32, flt: ti.i32, pos:tm.uvec2, i:ti.i32, ch:ti.i32): - N, C, H, W = tenOut.shape - if (pos.x >= 0) and (pos.x < W) and (pos.y >= 0) and (pos.y < H): - tenOut[i, ch, pos.y, pos.x] += fltIn * flt -@ti.kernel -def softsplat_out(tenIn: ti.types.ndarray(), tenFlow: ti.types.ndarray(), tenOut: ti.types.ndarray()): - N, C, H, W = tenIn.shape - for i, ch, y, x in ti.ndrange(N, C, H, W): - fltX = x + tenFlow[i, 0, y, x] - fltY = y + tenFlow[i, 1, y, x] - fltIn = tenIn[i, ch, y, x] - - northWest = tm.ivec2(ti.floor(fltX), ti.floor(fltY)) - northEast = northWest + [1, 0] - southWest = northWest + [0, 1] - southEast = northWest + [1, 1] - - fltNorthwest = (southEast.x - fltX) * (southEast.y - fltY) - fltNortheast = (fltX - southWest.x) * (southWest.y - fltY) - fltSouthwest = (northEast.x - fltX) * (fltY - northEast.y) - fltSoutheast = (fltX - northWest.x) * (fltY - northWest.y) - - put_to_tenOut(tenOut, fltIn, fltNorthwest, northWest, i, ch) - put_to_tenOut(tenOut, fltIn, fltNortheast, northEast, i, ch) - put_to_tenOut(tenOut, fltIn, fltSouthwest, southWest, i, ch) - put_to_tenOut(tenOut, fltIn, fltSoutheast, southEast, i, ch) - -@ti.func -def add_to_fltFlowgrad(fltFlowgrad, tenOutgrad, fltIn, flt, pos, i, ch): - N, C, H, W = tenOutgrad.shape - if (pos.x >= 0) and (pos.x < W) and (pos.y >= 0) and (pos.y < H): - fltFlowgrad += tenOutgrad[i, ch, pos.y, pos.x] * fltIn * flt - -@ti.kernel -def softsplat_flowgrad( - tenIn: ti.types.ndarray(), - tenFlow: ti.types.ndarray(), - tenOutgrad: ti.types.ndarray(), - tenIngrad: ti.types.ndarray(), - tenFlowgrad: ti.types.ndarray() -): - N, C, H, W = tenFlowgrad.shape - for i, ch, y, x in ti.ndrange(N, C, H, W): - fltFlowgrad = 0.0 - fltX = x + tenFlow[i, 0, y, x] - fltY = y + tenFlow[i, 1, y, x] - - northWest = tm.vec2(ti.floor(fltX, dtype=ti.i32), ti.floor(fltY, dtype=ti.i32)) - northEast = tm.vec2(northWest.x + 1, northWest.y) - southWest = tm.vec2(northWest.x, northWest.y + 1) - southEast = tm.vec2(northWest.x + 1, northWest.y + 1) - - if ch == 0: - fltNorthwest = -1.0 * (southEast.y - fltY) - fltNortheast = +1.0 * (southWest.y - fltY) - fltSouthwest = -1.0 * (fltY - northEast.y) - fltSoutheast = +1.0 * (fltY - northWest.y) - - elif ch == 1: - fltNorthwest = -1.0 * (southEast.x - fltX) - fltNortheast = -1.0 * (fltX - southWest.x) - fltSouthwest = +1.0 * (northEast.x - fltX) - fltSoutheast = +1.0 * (fltX - northWest.x) - - for outgrad_ch in ti.ndrange(tenOutgrad.shape[1]): - fltIn = tenIn[i, outgrad_ch, y, x] - add_to_fltFlowgrad(fltFlowgrad, tenOutgrad, fltIn, fltNorthwest, northWest, i, outgrad_ch) - add_to_fltFlowgrad(fltFlowgrad, tenOutgrad, fltIn, fltNortheast, northEast, i, outgrad_ch) - add_to_fltFlowgrad(fltFlowgrad, tenOutgrad, fltIn, fltSouthwest, southWest, i, outgrad_ch) - add_to_fltFlowgrad(fltFlowgrad, tenOutgrad, fltIn, fltSoutheast, southEast, i, outgrad_ch) - - tenFlowgrad[i] = fltFlowgrad #Is 'i' the same as intIndex? - -@ti.func -def add_to_fltIngrad(fltIngrad, tenOutgrad, flt, pos, i, ch): - N, C, H, W = tenOutgrad.shape - if (pos.x >= 0) and (pos.x < W) and (pos.y >= 0) and (pos.y < H): - fltIngrad += tenOutgrad[i, ch, pos.y, pos.x] * flt -@ti.kernel -def softsplat_ingrad( - tenIn: ti.types.ndarray(), - tenFlow: ti.types.ndarray(), - tenOutgrad: ti.types.ndarray(), - tenIngrad: ti.types.ndarray(), - tenFlowgrad: ti.types.ndarray() -): - N, C, H, W = tenIngrad.shape - for i, ch, y, x in ti.ndrange(N, C, H, W): - fltIngrad = 0.0 - fltX = x + tenFlow[i, 0, y, x] - fltY = y + tenFlow[i, 1, y, x] - - northWest = tm.vec2(ti.floor(fltX, dtype=ti.i32), ti.floor(fltY, dtype=ti.i32)) - northEast = tm.vec2(northWest.x + 1, northWest.y) - southWest = tm.vec2(northWest.x, northWest.y + 1) - southEast = tm.vec2(northWest.x + 1, northWest.y + 1) - - fltNorthwest = (southEast.x - fltX) * (southEast.y - fltY) - fltNortheast = (fltX - southWest.x) * (southWest.y - fltY) - fltSouthwest = (northEast.x - fltX) * (fltY - northEast.y) - fltSoutheast = (fltX - northWest.x) * (fltY - northWest.y) - - add_to_fltIngrad(fltIngrad, tenOutgrad, fltNorthwest, northWest, i, ch) - add_to_fltIngrad(fltIngrad, tenOutgrad, fltNortheast, northEast, i, ch) - add_to_fltIngrad(fltIngrad, tenOutgrad, fltSouthwest, southWest, i, ch) - add_to_fltIngrad(fltIngrad, tenOutgrad, fltSoutheast, southEast, i, ch) - tenIngrad[i] = fltIngrad - -# end - -def worker_interface(op_name, tensors): - if op_name == "softsplat_out": - tenIn, tenFlow = tensors - tenOut = tenIn.new_zeros(tenIn.shape) - softsplat_out(tenIn, tenFlow, tenOut) - return (tenOut, ) - - raise NotImplementedError(op_name) - -__all__ = ["worker_interface"] \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/sepconv.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/sepconv.py deleted file mode 100644 index f9d8909..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/sepconv.py +++ /dev/null @@ -1,39 +0,0 @@ -import taichi as ti -import taichi.math as tm -from functools import reduce - -@ti.kernel -def sepconv_out(tenIn: ti.types.ndarray(), tenVer: ti.types.ndarray(), tenHor: ti.types.ndarray(), tenOut: ti.types.ndarray()): - N, C, H, W = tenIn.shape - intIndex = 0 - for i, ch, y, x in ti.ndrange(N, C, H, W): - fltOut, fltKahanc, fltKahany, fltKahant = 0.0, 0.0, 0.0, 0.0 - for intFy, intFx in ti.ndrange(tenVer.shape[1], tenHor.shape[1]): - fltKahany = tenIn[i, ch, y + intFy, x + intFx] * tenVer[i, intFy, y, x] * tenHor[i, intFx, y, x] - fltKahany = fltKahany - fltKahanc - fltKahant = fltOut + fltKahany - fltKahanc = (fltKahant - fltOut) - fltKahany - fltOut = fltKahant - tenOut[intIndex] = fltOut - intIndex += 1 - - -def worker_interface(op_name, tensors): - if op_name == "sepconv_out": - tenIn, tenVer, tenHor = tensors - real_tenOut_shape = [ - tenIn.shape[0], - tenIn.shape[1], - tenVer.shape[2] and tenHor.shape[2], - tenVer.shape[3] and tenHor.shape[3], - ] - tenOut = tenIn.new_zeros([ - int(reduce(lambda a, b: a * b, real_tenOut_shape)) - ]) - sepconv_out(tenIn, tenVer, tenHor, tenOut) - tenOut = tenOut.view(*real_tenOut_shape) - return (tenOut, ) - - raise NotImplementedError(op_name) - -__all__ = ["worker_interface"] \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/utils.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/utils.py deleted file mode 100644 index ee42b2f..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/utils.py +++ /dev/null @@ -1,11 +0,0 @@ -import platform -import torch -def to_shared_memory(tensors: tuple[torch.Tensor]): - return [tensor.cpu() for tensor in tensors if tensor is not None] - """ if platform.system() == "Windows": - return [tensor.cpu() for tensor in tensors if tensor is not None] - - return [tensor.share_memory_() for tensor in tensors if tensor is not None] """ - -def to_device(tensors: tuple[torch.Tensor], device: torch.device): - return [tensor.to(device) for tensor in tensors if tensor is not None] \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/worker_process.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/worker_process.py deleted file mode 100644 index 586d06a..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/ops/taichi_ops/worker_process.py +++ /dev/null @@ -1,26 +0,0 @@ -import torch.multiprocessing as mp -import torch -from .raw_softsplat import worker_interface as raw_softsplat -from .costvol import worker_interface as costvol -from .sepconv import worker_interface as sepconv -from .utils import to_shared_memory, to_device -import taichi as ti -import traceback - -def f(child_conn, device: torch.DeviceObjType): - ti.init(arch=ti.gpu) - while True: - op_name, tensors = child_conn.recv() - tensors = to_device(tensors, device) - try: - if "softsplat" in op_name: - result = raw_softsplat(op_name, tensors) - elif "costvol" in op_name: - result = costvol(op_name, tensors) - elif "sepconv" in op_name: - result = sepconv(op_name, tensors) - else: - raise NotImplementedError(op_name) - child_conn.send(to_shared_memory(result)) - except: - child_conn.send(traceback.format_exc()) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/rife/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/rife/__init__.py deleted file mode 100644 index 9c9e82d..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/rife/__init__.py +++ /dev/null @@ -1,107 +0,0 @@ -import torch -from torch.utils.data import DataLoader -import pathlib -from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames, generic_frame_loop, InterpolationStateList -import typing -from comfy.model_management import get_torch_device -import re -from functools import cmp_to_key -from packaging import version - -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" - #Arch 4.10 doesn't work due to state dict mismatch - #TODO: Investigating and fix it - #"rife410.pth": "4.10", - #"rife411.pth": "4.10", - #"rife412.pth": "4.10" -} - -class RIFE_VFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": ( - sorted(list(CKPT_NAME_VER_DICT.keys()), key=lambda ckpt_name: version.parse(CKPT_NAME_VER_DICT[ckpt_name])), - {"default": "rife47.pth"} - ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}), - "multiplier": ("INT", {"default": 2, "min": 1}), - "fast_mode": ("BOOLEAN", {"default":True}), - "ensemble": ("BOOLEAN", {"default":True}), - "scale_factor": ([0.25, 0.5, 1.0, 2.0, 4.0], {"default": 1.0}) - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames = 10, - multiplier: typing.SupportsInt = 2, - fast_mode = False, - ensemble = False, - scale_factor = 1.0, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - """ - Perform video frame interpolation using a given checkpoint model. - - Args: - ckpt_name (str): The name of the checkpoint model to use. - frames (torch.Tensor): A tensor containing input video frames. - clear_cache_after_n_frames (int, optional): The number of frames to process before clearing CUDA cache - to prevent memory overflow. Defaults to 10. Lower numbers are safer but mean more processing time. - How high you should set it depends on how many input frames there are, input resolution (after upscaling), - how many times you want to multiply them, and how long you're willing to wait for the process to complete. - multiplier (int, optional): The multiplier for each input frame. 60 input frames * 2 = 120 output frames. Defaults to 2. - - Returns: - tuple: A tuple containing the output interpolated frames. - - Note: - This method interpolates frames in a video sequence using a specified checkpoint model. - It processes each frame sequentially, generating interpolated frames between them. - - To prevent memory overflow, it clears the CUDA cache after processing a specified number of frames. - """ - from .rife_arch import IFNet - 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(get_torch_device()) - frames = preprocess_frames(frames) - - def return_middle_frame(frame_0, frame_1, timestep, model, scale_list, in_fast_mode, in_ensemble): - return model(frame_0, frame_1, timestep, scale_list, in_fast_mode, in_ensemble) - - scale_list = [8 / scale_factor, 4 / scale_factor, 2 / scale_factor, 1 / scale_factor] - - args = [interpolation_model, scale_list, fast_mode, ensemble] - out = postprocess_frames( - generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, - interpolation_states=optional_interpolation_states, dtype=torch.float32) - ) - return (out,) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/rife/rife_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/rife/rife_arch.py deleted file mode 100644 index 01eeed8..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/rife/rife_arch.py +++ /dev/null @@ -1,581 +0,0 @@ -""" -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 -from comfy.model_management import get_torch_device - -device = get_torch_device() -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=0.5, - 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) - - 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]) - - 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.0-4.6 - if self.arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]: - flow, mask = block[i]( - torch.cat((img0[:, :3], img1[:, :3], timestep), 1), - None, - scale=scale_list[i], - ) - if ensemble: - f1, m1 = block[i]( - torch.cat((img1[:, :3], img0[:, :3], 1 - timestep), 1), - None, - scale=scale_list[i], - ) - flow = (flow + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2 - mask = (mask + (-m1)) / 2 - - # 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.0-4.6 - if self.arch_ver in ["4.0", "4.2", "4.3", "4.5", "4.6"]: - f0, m0 = block[i]( - torch.cat( - (warped_img0[:, :3], warped_img1[:, :3], timestep, mask), 1 - ), - flow, - scale=scale_list[i], - ) - - if self.arch_ver in ["4.0"]: - if ( - i == 1 - and f0[:, :2].abs().max() > 32 - and f0[:, 2:4].abs().max() > 32 - and not training - ): - for k in range(4): - scale_list[k] *= 2 - flow, mask = block[0]( - torch.cat((img0[:, :3], img1[:, :3], timestep), 1), - None, - scale=scale_list[0], - ) - warped_img0 = warp(img0, flow[:, :2]) - warped_img1 = warp(img1, flow[:, 2:4]) - f0, m0 = block[i]( - torch.cat( - ( - warped_img0[:, :3], - warped_img1[:, :3], - timestep, - mask, - ), - 1, - ), - flow, - scale=scale_list[i], - ) - - # 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) - return merged[3][:, :, :h, :w] \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/sepconv/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/sepconv/__init__.py deleted file mode 100644 index eb62e76..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/sepconv/__init__.py +++ /dev/null @@ -1,56 +0,0 @@ -import torch -from torch.utils.data import DataLoader -import pathlib -from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames -import typing -from comfy.model_management import soft_empty_cache, get_torch_device -from vfi_utils import InterpolationStateList, generic_frame_loop - -MODEL_TYPE = pathlib.Path(__file__).parent.name -CKPT_NAMES = ["sepconv.pth"] - - -class SepconvVFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (CKPT_NAMES, ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}), - "multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}) - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames = 10, - multiplier: typing.SupportsInt = 2, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - from .sepconv_enhanced import Network - model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) - interpolation_model = Network() - interpolation_model.load_state_dict(torch.load(model_path)) - interpolation_model.eval().to(get_torch_device()) - frames = preprocess_frames(frames) - - def return_middle_frame(frame_0, frame_1, timestep, model): - return model(frame_0, frame_1) - - args = [interpolation_model] - out = postprocess_frames( - generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, - interpolation_states=optional_interpolation_states, use_timestep=False, dtype=torch.float32) - ) - return (out,) diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/sepconv/sepconv_enhanced.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/sepconv/sepconv_enhanced.py deleted file mode 100644 index e5a747e..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/sepconv/sepconv_enhanced.py +++ /dev/null @@ -1,748 +0,0 @@ -""" -23-nov-21 -https://github.com/sniklaus/revisiting-sepconv/blob/fea509d98157170df1fb35bf615bd41d98858e1a/run.py -https://github.com/sniklaus/revisiting-sepconv/blob/fea509d98157170df1fb35bf615bd41d98858e1a/sepconv/sepconv.py -Deleted stuffs about arguments_strModel and getopt -""" -#!/usr/bin/env python -import torch -import typing -from comfy.model_management import get_torch_device - -########################################################## -from vfi_models.ops import sepconv_func -########################################################## - - - - - -import torch - -import math -import numpy -import os -import PIL -import PIL.Image -import sys -import typing - -########################################################## - -assert ( - int(str("").join(torch.__version__.split(".")[0:2])) >= 13 -) # requires at least pytorch version 1.3.0 - -torch.set_grad_enabled( - False -) # make sure to not compute gradients for computational performance - -torch.backends.cudnn.enabled = ( - True # make sure to use cudnn for computational performance -) - -########################################################## - -########################################################## - - -class Basic(torch.nn.Module): - def __init__( - self, - strType: str, - intChans: typing.List[int], - objScratch: typing.Optional[typing.Dict] = None, - ): - super().__init__() - - self.strType = strType - self.netEvenize = None - self.netMain = None - self.netShortcut = None - - intIn = intChans[0] - intOut = intChans[-1] - netMain = [] - intChans = intChans.copy() - fltStride = 1.0 - - for intPart, strPart in enumerate(self.strType.split("+")[0].split("-")): - if strPart.startswith("conv") == True: - intKsize = 3 - intPad = 1 - strPad = "zeros" - - if "(" in strPart: - intKsize = int(strPart.split("(")[1].split(")")[0].split(",")[0]) - intPad = int(math.floor(0.5 * (intKsize - 1))) - - if "replpad" in strPart.split("(")[1].split(")")[0].split(","): - strPad = "replicate" - if "reflpad" in strPart.split("(")[1].split(")")[0].split(","): - strPad = "reflect" - # end - - if "nopad" in self.strType.split("+"): - intPad = 0 - # end - - netMain += [ - torch.nn.Conv2d( - in_channels=intChans[0], - out_channels=intChans[1], - kernel_size=intKsize, - stride=1, - padding=intPad, - padding_mode=strPad, - bias="nobias" not in self.strType.split("+"), - ) - ] - intChans = intChans[1:] - fltStride *= 1.0 - - elif strPart.startswith("sconv") == True: - intKsize = 3 - intPad = 1 - strPad = "zeros" - - if "(" in strPart: - intKsize = int(strPart.split("(")[1].split(")")[0].split(",")[0]) - intPad = int(math.floor(0.5 * (intKsize - 1))) - - if "replpad" in strPart.split("(")[1].split(")")[0].split(","): - strPad = "replicate" - if "reflpad" in strPart.split("(")[1].split(")")[0].split(","): - strPad = "reflect" - # end - - if "nopad" in self.strType.split("+"): - intPad = 0 - # end - - netMain += [ - torch.nn.Conv2d( - in_channels=intChans[0], - out_channels=intChans[1], - kernel_size=intKsize, - stride=2, - padding=intPad, - padding_mode=strPad, - bias="nobias" not in self.strType.split("+"), - ) - ] - intChans = intChans[1:] - fltStride *= 2.0 - - elif strPart.startswith("up") == True: - - class Up(torch.nn.Module): - def __init__(self, strType): - super().__init__() - - self.strType = strType - - # end - - def forward(self, tenIn: torch.Tensor) -> torch.Tensor: - if self.strType == "nearest": - return torch.nn.functional.interpolate( - input=tenIn, - scale_factor=2.0, - mode="nearest", - align_corners=False, - ) - - elif self.strType == "bilinear": - return torch.nn.functional.interpolate( - input=tenIn, - scale_factor=2.0, - mode="bilinear", - align_corners=False, - ) - - elif self.strType == "pyramid": - return pyramid(tenIn, None, "up") - - elif self.strType == "shuffle": - return torch.nn.functional.pixel_shuffle( - tenIn, upscale_factor=2 - ) # https://github.com/pytorch/pytorch/issues/62854 - - # end - - assert False # to make torchscript happy - - # end - - # end - - strType = "bilinear" - - if "(" in strPart: - if "nearest" in strPart.split("(")[1].split(")")[0].split(","): - strType = "nearest" - if "pyramid" in strPart.split("(")[1].split(")")[0].split(","): - strType = "pyramid" - if "shuffle" in strPart.split("(")[1].split(")")[0].split(","): - strType = "shuffle" - # end - - netMain += [Up(strType)] - fltStride *= 0.5 - - elif strPart.startswith("prelu") == True: - netMain += [ - torch.nn.PReLU( - num_parameters=1, - init=float(strPart.split("(")[1].split(")")[0].split(",")[0]), - ) - ] - fltStride *= 1.0 - - elif True: - assert False - - # end - # end - - self.netMain = torch.nn.Sequential(*netMain) - - for strPart in self.strType.split("+")[1:]: - if strPart.startswith("skip") == True: - if intIn == intOut and fltStride == 1.0: - self.netShortcut = torch.nn.Identity() - - elif intIn != intOut and fltStride == 1.0: - self.netShortcut = torch.nn.Conv2d( - in_channels=intIn, - out_channels=intOut, - kernel_size=1, - stride=1, - padding=0, - bias="nobias" not in self.strType.split("+"), - ) - - elif intIn == intOut and fltStride != 1.0: - - class Down(torch.nn.Module): - def __init__(self, fltScale): - super().__init__() - - self.fltScale = fltScale - - # end - - def forward(self, tenIn: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.interpolate( - input=tenIn, - scale_factor=self.fltScale, - mode="bilinear", - align_corners=False, - ) - - # end - - # end - - self.netShortcut = Down(1.0 / fltStride) - - elif intIn != intOut and fltStride != 1.0: - - class Down(torch.nn.Module): - def __init__(self, fltScale): - super().__init__() - - self.fltScale = fltScale - - # end - - def forward(self, tenIn: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.interpolate( - input=tenIn, - scale_factor=self.fltScale, - mode="bilinear", - align_corners=False, - ) - - # end - - # end - - self.netShortcut = torch.nn.Sequential( - Down(1.0 / fltStride), - torch.nn.Conv2d( - in_channels=intIn, - out_channels=intOut, - kernel_size=1, - stride=1, - padding=0, - bias="nobias" not in self.strType.split("+"), - ), - ) - - # end - - elif strPart.startswith("...") == True: - pass - - # end - # end - - assert len(intChans) == 1 - - # end - - def forward(self, tenIn: torch.Tensor) -> torch.Tensor: - if self.netEvenize is not None: - tenIn = self.netEvenize(tenIn) - # end - - tenOut = self.netMain(tenIn) - - if self.netShortcut is not None: - tenOut = tenOut + self.netShortcut(tenIn) - # end - - return tenOut - - # end - - -# end - - -class Encode(torch.nn.Module): - objScratch: typing.Dict[str, typing.List[int]] = None - - def __init__( - self, - intIns: typing.List[int], - intOuts: typing.List[int], - strHor: str, - strVer: str, - objScratch: typing.Dict[str, typing.List[int]], - ): - super().__init__() - - assert len(intIns) == len(intOuts) - assert len(intOuts) == len(intIns) - - self.intRows = len(intIns) and len(intOuts) - self.intIns = intIns.copy() - self.intOuts = intOuts.copy() - self.strHor = strHor - self.strVer = strVer - self.objScratch = objScratch - - self.netHor = torch.nn.ModuleList() - self.netVer = torch.nn.ModuleList() - - for intRow in range(self.intRows): - netHor = torch.nn.Identity() - netVer = torch.nn.Identity() - - if self.intOuts[intRow] != 0: - if self.intIns[intRow] != 0: - netHor = Basic( - self.strHor, - [ - self.intIns[intRow], - self.intOuts[intRow], - self.intOuts[intRow], - ], - objScratch, - ) - # end - - if intRow != 0: - netVer = Basic( - self.strVer, - [ - self.intOuts[intRow - 1], - self.intOuts[intRow], - self.intOuts[intRow], - ], - objScratch, - ) - # end - # end - - self.netHor.append(netHor) - self.netVer.append(netVer) - # end - - # end - - def forward(self, tenIns: typing.List[torch.Tensor]) -> typing.List[torch.Tensor]: - intRow = 0 - for netHor in self.netHor: - if self.intOuts[intRow] != 0: - if self.intIns[intRow] != 0: - tenIns[intRow] = netHor(tenIns[intRow]) - # end - # end - intRow += 1 - # end - - intRow = 0 - for netVer in self.netVer: - if self.intOuts[intRow] != 0: - if intRow != 0: - tenIns[intRow] = tenIns[intRow] + netVer(tenIns[intRow - 1]) - # end - # end - intRow += 1 - # end - - for intRow, tenIn in enumerate(tenIns): - self.objScratch["levelshape" + str(intRow)] = tenIn.shape - # end - - return tenIns - - # end - - -# end - - -class Decode(torch.nn.Module): - objScratch: typing.Dict[str, typing.List[int]] = None - - def __init__( - self, - intIns: typing.List[int], - intOuts: typing.List[int], - strHor: str, - strVer: str, - objScratch: typing.Dict[str, typing.List[int]], - ): - super().__init__() - - assert len(intIns) == len(intOuts) - assert len(intOuts) == len(intIns) - - self.intRows = len(intIns) and len(intOuts) - self.intIns = intIns.copy() - self.intOuts = intOuts.copy() - self.strHor = strHor - self.strVer = strVer - self.objScratch = objScratch - - self.netHor = torch.nn.ModuleList() - self.netVer = torch.nn.ModuleList() - - for intRow in range(self.intRows - 1, -1, -1): - netHor = torch.nn.Identity() - netVer = torch.nn.Identity() - - if self.intOuts[intRow] != 0: - if self.intIns[intRow] != 0: - netHor = Basic( - self.strHor, - [ - self.intIns[intRow], - self.intOuts[intRow], - self.intOuts[intRow], - ], - objScratch, - ) - # end - - if intRow != self.intRows - 1: - netVer = Basic( - self.strVer, - [ - self.intOuts[intRow + 1], - self.intOuts[intRow], - self.intOuts[intRow], - ], - objScratch, - ) - # end - # end - - self.netHor.append(netHor) - self.netVer.append(netVer) - # end - - # end - - def forward(self, tenIns: typing.List[torch.Tensor]) -> typing.List[torch.Tensor]: - intRow = self.intRows - 1 - for netHor in self.netHor: - if self.intOuts[intRow] != 0: - if self.intIns[intRow] != 0: - tenIns[intRow] = netHor(tenIns[intRow]) - # end - # end - intRow -= 1 - # end - - intRow = self.intRows - 1 - for netVer in self.netVer: - if self.intOuts[intRow] != 0: - if intRow != self.intRows - 1: - tenVer = netVer(tenIns[intRow + 1]) - - if "levelshape" + str(intRow) in self.objScratch: - if ( - tenVer.shape[2] - == self.objScratch["levelshape" + str(intRow)][2] + 1 - ): - tenVer = torch.nn.functional.pad( - input=tenVer, - pad=[0, 0, 0, -1], - mode="constant", - value=0.0, - ) - if ( - tenVer.shape[3] - == self.objScratch["levelshape" + str(intRow)][3] + 1 - ): - tenVer = torch.nn.functional.pad( - input=tenVer, - pad=[0, -1, 0, 0], - mode="constant", - value=0.0, - ) - # end - - tenIns[intRow] = tenIns[intRow] + tenVer - # end - # end - intRow -= 1 - # end - - return tenIns - - # end - - -# end - -########################################################## - - -class Network(torch.nn.Module): - def __init__(self): - super().__init__() - - self.intEncdec = [1, 1] - self.intChannels = [32, 64, 128, 256, 512] - - self.objScratch = {} - - self.netInput = torch.nn.Conv2d( - in_channels=3, - out_channels=int(round(0.5 * self.intChannels[0])), - kernel_size=3, - stride=1, - padding=1, - padding_mode="zeros", - ) - - self.netEncode = torch.nn.Sequential( - *( - [ - Encode( - [0] * len(self.intChannels), - self.intChannels, - "prelu(0.25)-conv(3)-prelu(0.25)-conv(3)+skip", - "prelu(0.25)-sconv(3)-prelu(0.25)-conv(3)", - self.objScratch, - ) - ] - + [ - Encode( - self.intChannels, - self.intChannels, - "prelu(0.25)-conv(3)-prelu(0.25)-conv(3)+skip", - "prelu(0.25)-sconv(3)-prelu(0.25)-conv(3)", - self.objScratch, - ) - for intEncdec in range(1, self.intEncdec[0]) - ] - ) - ) - - self.netDecode = torch.nn.Sequential( - *( - [ - Decode( - [0] + self.intChannels[1:], - [0] + self.intChannels[1:], - "prelu(0.25)-conv(3)-prelu(0.25)-conv(3)+skip", - "prelu(0.25)-up(bilinear)-conv(3)-prelu(0.25)-conv(3)", - self.objScratch, - ) - for intEncdec in range(0, self.intEncdec[1]) - ] - ) - ) - - self.netVerone = Basic( - "up(bilinear)-conv(3)-prelu(0.25)-conv(3)", - [self.intChannels[1], self.intChannels[1], 51], - ) - self.netVertwo = Basic( - "up(bilinear)-conv(3)-prelu(0.25)-conv(3)", - [self.intChannels[1], self.intChannels[1], 51], - ) - self.netHorone = Basic( - "up(bilinear)-conv(3)-prelu(0.25)-conv(3)", - [self.intChannels[1], self.intChannels[1], 51], - ) - self.netHortwo = Basic( - "up(bilinear)-conv(3)-prelu(0.25)-conv(3)", - [self.intChannels[1], self.intChannels[1], 51], - ) - - # self.load_state_dict(torch.hub.load_state_dict_from_url(url='http://content.sniklaus.com/resepconv/network-' + arguments_strModel + '.pytorch', file_name='resepconv-' + arguments_strModel)) - - # end - - def forward(self, x1, x2): - # padding if needed - intWidth = x1.shape[3] - intHeight = x1.shape[2] - - intPadr = (2 - (intWidth % 2)) % 2 - intPadb = (2 - (intHeight % 2)) % 2 - - tenOne = torch.nn.functional.pad( - input=x1, pad=[0, intPadr, 0, intPadb], mode="replicate" - ) - tenTwo = torch.nn.functional.pad( - input=x2, pad=[0, intPadr, 0, intPadb], mode="replicate" - ) - #### - - tenSeq = [tenOne, tenTwo] - - with torch.set_grad_enabled(False): - tenStack = torch.stack(tenSeq, 1) - tenMean = ( - tenStack.view(tenStack.shape[0], -1) - .mean(1, True) - .view(tenStack.shape[0], 1, 1, 1) - ) - tenStd = ( - tenStack.view(tenStack.shape[0], -1) - .std(1, True) - .view(tenStack.shape[0], 1, 1, 1) - ) - tenSeq = [ - (tenFrame - tenMean) / (tenStd + 0.0000001) for tenFrame in tenSeq - ] - tenSeq = [tenFrame.detach() for tenFrame in tenSeq] - # end - - tenOut = self.netDecode( - self.netEncode( - [torch.cat([self.netInput(tenSeq[0]), self.netInput(tenSeq[1])], 1)] - + ([0.0] * (len(self.intChannels) - 1)) - ) - )[1] - - tenOne = torch.nn.functional.pad( - input=tenOne, - pad=[ - int(math.floor(0.5 * 51)), - int(math.floor(0.5 * 51)), - int(math.floor(0.5 * 51)), - int(math.floor(0.5 * 51)), - ], - mode="replicate", - ) - tenTwo = torch.nn.functional.pad( - input=tenTwo, - pad=[ - int(math.floor(0.5 * 51)), - int(math.floor(0.5 * 51)), - int(math.floor(0.5 * 51)), - int(math.floor(0.5 * 51)), - ], - mode="replicate", - ) - - tenOne = torch.cat( - [ - tenOne, - tenOne.new_ones([tenOne.shape[0], 1, tenOne.shape[2], tenOne.shape[3]]), - ], - 1, - ).detach() - tenTwo = torch.cat( - [ - tenTwo, - tenTwo.new_ones([tenTwo.shape[0], 1, tenTwo.shape[2], tenTwo.shape[3]]), - ], - 1, - ).detach() - - tenVerone = self.netVerone(tenOut) - tenVertwo = self.netVertwo(tenOut) - tenHorone = self.netHorone(tenOut) - tenHortwo = self.netHortwo(tenOut) - - tenOut = sepconv_func.apply(tenOne, tenVerone, tenHorone) + sepconv_func.apply( - tenTwo, tenVertwo, tenHortwo - ) - - tenNormalize = tenOut[:, -1:, :, :] - tenNormalize[tenNormalize.abs() < 0.01] = 1.0 - tenOut = tenOut[:, :-1, :, :] / tenNormalize - - # crop if needed - return tenOut[:, :, :intHeight, :intWidth] - - # end - - -# end - -netNetwork = None - -########################################################## - - -def estimate(tenOne, tenTwo): - global netNetwork - - if netNetwork is None: - netNetwork = Network().to(get_torch_device()).eval() - # end - - assert tenOne.shape[1] == tenTwo.shape[1] - assert tenOne.shape[2] == tenTwo.shape[2] - - intWidth = tenOne.shape[2] - intHeight = tenOne.shape[1] - - assert ( - intWidth <= 1280 - ) # while our approach works with larger images, we do not recommend it unless you are aware of the implications - assert ( - intHeight <= 720 - ) # while our approach works with larger images, we do not recommend it unless you are aware of the implications - - tenPreprocessedOne = tenOne.to(get_torch_device()).view(1, 3, intHeight, intWidth) - tenPreprocessedTwo = tenTwo.to(get_torch_device()).view(1, 3, intHeight, intWidth) - - intPadr = (2 - (intWidth % 2)) % 2 - intPadb = (2 - (intHeight % 2)) % 2 - - tenPreprocessedOne = torch.nn.functional.pad( - input=tenPreprocessedOne, pad=[0, intPadr, 0, intPadb], mode="replicate" - ) - tenPreprocessedTwo = torch.nn.functional.pad( - input=tenPreprocessedTwo, pad=[0, intPadr, 0, intPadb], mode="replicate" - ) - - return netNetwork([tenPreprocessedOne, tenPreprocessedTwo])[ - 0, :, :intHeight, :intWidth - ].cpu() - - -# end diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/stmfnet/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/stmfnet/__init__.py deleted file mode 100644 index 1613ab9..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/stmfnet/__init__.py +++ /dev/null @@ -1,100 +0,0 @@ -import torch -from comfy.model_management import get_torch_device, soft_empty_cache -import numpy as np -import typing -from vfi_utils import InterpolationStateList, load_file_from_github_release, preprocess_frames, postprocess_frames, assert_batch_size -import pathlib -import warnings -import gc - -MODEL_TYPE = pathlib.Path(__file__).parent.name -device = get_torch_device() - -class STMFNet_VFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (["stmfnet.pth"], ), - "frames": ("IMAGE", ), - "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}), - "multiplier": ("INT", {"default": 2, "min": 2, "max": 2}), #TODO: Implement recursively invoking interpolator for multi-frame interpolation - "duplicate_first_last_frames": ("BOOLEAN", {"default": False}) - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ) - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - #Reference: https://github.com/danier97/ST-MFNet/blob/main/interpolate_yuv.py#L93 - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - clear_cache_after_n_frames = 10, - multiplier: typing.SupportsInt = 2, - duplicate_first_last_frames: bool = False, - optional_interpolation_states: InterpolationStateList = None, - **kwargs - ): - from .stmfnet_arch import STMFNet_Model - if multiplier != 2: - warnings.warn("Currently, ST-MFNet only supports 2x interpolation. The process will continue but please set multiplier=2 afterward") - - assert_batch_size(frames, batch_size=4, vfi_name="ST-MFNet") - interpolation_states = optional_interpolation_states - model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) - model = STMFNet_Model() - model.load_state_dict(torch.load(model_path)) - model = model.eval().to(device) - - frames = preprocess_frames(frames) - number_of_frames_processed_since_last_cleared_cuda_cache = 0 - output_frames = [] - for frame_itr in range(len(frames) - 3): - #Does skipping frame i+1 make sanse in this case? - if interpolation_states is not None and interpolation_states.is_frame_skipped(frame_itr) and interpolation_states.is_frame_skipped(frame_itr + 1): - continue - - #Ensure that input frames are in fp32 - the same dtype as model - frame0, frame1, frame2, frame3 = ( - frames[frame_itr:frame_itr+1].float(), - frames[frame_itr+1:frame_itr+2].float(), - frames[frame_itr+2:frame_itr+3].float(), - frames[frame_itr+3:frame_itr+4].float() - ) - new_frame = model(frame0.to(device), frame1.to(device), frame2.to(device), frame3.to(device)).detach().cpu() - number_of_frames_processed_since_last_cleared_cuda_cache += 2 - - if frame_itr == 0: - output_frames.append(frame0) - if duplicate_first_last_frames: - output_frames.append(frame0) # repeat the first frame - output_frames.append(frame1) - output_frames.append(new_frame) - output_frames.append(frame2) - if frame_itr == len(frames) - 4: - output_frames.append(frame3) - if duplicate_first_last_frames: - output_frames.append(frame3) # repeat the last frame - - # Try to avoid a memory overflow by clearing cuda cache regularly - if number_of_frames_processed_since_last_cleared_cuda_cache >= clear_cache_after_n_frames: - print("Comfy-VFI: Clearing cache...", end = ' ') - soft_empty_cache() - number_of_frames_processed_since_last_cleared_cuda_cache = 0 - print("Done cache clearing") - gc.collect() - - dtype = torch.float32 - output_frames = [frame.cpu().to(dtype=dtype) for frame in output_frames] #Ensure all frames are in cpu - out = torch.cat(output_frames, dim=0) - # clear cache for courtesy - print("Comfy-VFI: Final clearing cache...", end = ' ') - soft_empty_cache() - print("Done cache clearing") - return (postprocess_frames(out), ) \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/stmfnet/stmfnet_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/stmfnet/stmfnet_arch.py deleted file mode 100644 index 1e495e0..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/stmfnet/stmfnet_arch.py +++ /dev/null @@ -1,2858 +0,0 @@ -# https://github.com/danielism97/ST-MFNet/blob/main/models/stmfnet.py -# https://github.com/danielism97/ST-MFNet/blob/main/models/misc/pwcnet.py -# https://github.com/danielism97/ST-MFNet/blob/main/models/misc/correlation/correlation.py -# https://github.com/danielism97/ST-MFNet/blob/main/models/misc/gridnet.py -# https://github.com/danielism97/ST-MFNet/blob/main/models/feature.py -# https://github.com/danielism97/ST-MFNet/blob/main/utility.py -# https://github.com/danielism97/ST-MFNet/blob/main/models/misc/resnet_3D.py -# https://github.com/danielism97/ST-MFNet/blob/main/cupy_module/adacof.py -# https://github.com/danielism97/ST-MFNet/blob/main/cupy_module/softsplat.py -# https://github.com/danielism97/ST-MFNet/blob/main/models/misc/__init__.py -# https://github.com/danielism97/ST-MFNet/blob/main/models/misc/pwcnet.py -# https://github.com/danielism97/ST-MFNet/blob/main/models/misc/correlation/correlation.py -from torch.nn import functional as F -from torch.utils.model_zoo import load_url as load_state_dict_from_url -import cv2 -import math -import numpy -import numpy as np -import PIL -import PIL.Image -import re -import sys -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -import torch.optim.lr_scheduler as lrs -from vfi_models.ops import FunctionCorrelation, FunctionAdaCoF, ModuleSoftsplat -from vfi_utils import get_ckpt_container_path -import pathlib -MODEL_TYPE = pathlib.Path(__file__).parent.name - -#Simple way to reduce oranges on VSCode bar -def identity(x): - return x - - -def backwarp(tenInput, tenFlow): - backwarp_tenGrid = {} - backwarp_tenPartial = {} - if str(tenFlow.shape) not in backwarp_tenGrid: - tenHor = ( - torch.linspace( - -1.0 + (1.0 / tenFlow.shape[3]), - 1.0 - (1.0 / tenFlow.shape[3]), - tenFlow.shape[3], - ) - .view(1, 1, 1, -1) - .expand(-1, -1, tenFlow.shape[2], -1) - ) - tenVer = ( - torch.linspace( - -1.0 + (1.0 / tenFlow.shape[2]), - 1.0 - (1.0 / tenFlow.shape[2]), - tenFlow.shape[2], - ) - .view(1, 1, -1, 1) - .expand(-1, -1, -1, tenFlow.shape[3]) - ) - - backwarp_tenGrid[str(tenFlow.shape)] = torch.cat([tenHor, tenVer], 1).cuda() - # end - - if str(tenFlow.shape) not in backwarp_tenPartial: - backwarp_tenPartial[str(tenFlow.shape)] = tenFlow.new_ones( - [tenFlow.shape[0], 1, tenFlow.shape[2], tenFlow.shape[3]] - ) - # end - - tenFlow = torch.cat( - [ - tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), - tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0), - ], - 1, - ) - tenInput = torch.cat([tenInput, backwarp_tenPartial[str(tenFlow.shape)]], 1) - - tenOutput = torch.nn.functional.grid_sample( - input=tenInput, - grid=(backwarp_tenGrid[str(tenFlow.shape)] + tenFlow).permute(0, 2, 3, 1), - mode="bilinear", - padding_mode="zeros", - align_corners=False, - ) - - tenMask = tenOutput[:, -1:, :, :] - tenMask[tenMask > 0.999] = 1.0 - tenMask[tenMask < 1.0] = 0.0 - - return tenOutput[:, :-1, :, :] * tenMask - - -# end - -########################################################## - - -class PWCNet(torch.nn.Module): - def __init__(self): - super(PWCNet, self).__init__() - - class Extractor(torch.nn.Module): - def __init__(self): - super(Extractor, self).__init__() - - self.netOne = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=3, - out_channels=16, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=16, - out_channels=16, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=16, - out_channels=16, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netTwo = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=16, - out_channels=32, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=32, - out_channels=32, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=32, - out_channels=32, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netThr = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=32, - out_channels=64, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=64, - out_channels=64, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=64, - out_channels=64, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netFou = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=64, - out_channels=96, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=96, - out_channels=96, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=96, - out_channels=96, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netFiv = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=96, - out_channels=128, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=128, - out_channels=128, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=128, - out_channels=128, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netSix = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=128, - out_channels=196, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=196, - out_channels=196, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=196, - out_channels=196, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - # end - - def forward(self, tenInput): - tenOne = self.netOne(tenInput) - tenTwo = self.netTwo(tenOne) - tenThr = self.netThr(tenTwo) - tenFou = self.netFou(tenThr) - tenFiv = self.netFiv(tenFou) - tenSix = self.netSix(tenFiv) - - return [tenOne, tenTwo, tenThr, tenFou, tenFiv, tenSix] - - # end - - # end - - class Decoder(torch.nn.Module): - def __init__(self, intLevel): - super(Decoder, self).__init__() - - intPrevious = [ - None, - None, - 81 + 32 + 2 + 2, - 81 + 64 + 2 + 2, - 81 + 96 + 2 + 2, - 81 + 128 + 2 + 2, - 81, - None, - ][intLevel + 1] - intCurrent = [ - None, - None, - 81 + 32 + 2 + 2, - 81 + 64 + 2 + 2, - 81 + 96 + 2 + 2, - 81 + 128 + 2 + 2, - 81, - None, - ][intLevel + 0] - - if intLevel < 6: - self.netUpflow = torch.nn.ConvTranspose2d( - in_channels=2, - out_channels=2, - kernel_size=4, - stride=2, - padding=1, - ) - if intLevel < 6: - self.netUpfeat = torch.nn.ConvTranspose2d( - in_channels=intPrevious + 128 + 128 + 96 + 64 + 32, - out_channels=2, - kernel_size=4, - stride=2, - padding=1, - ) - if intLevel < 6: - self.fltBackwarp = [None, None, None, 5.0, 2.5, 1.25, 0.625, None][ - intLevel + 1 - ] - - self.netOne = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent, - out_channels=128, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netTwo = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent + 128, - out_channels=128, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netThr = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent + 128 + 128, - out_channels=96, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netFou = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent + 128 + 128 + 96, - out_channels=64, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netFiv = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent + 128 + 128 + 96 + 64, - out_channels=32, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netSix = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent + 128 + 128 + 96 + 64 + 32, - out_channels=2, - kernel_size=3, - stride=1, - padding=1, - ) - ) - - # end - - def forward(self, tenFirst, tenSecond, objPrevious): - tenFlow = None - tenFeat = None - - if objPrevious is None: - tenFlow = None - tenFeat = None - - tenVolume = torch.nn.functional.leaky_relu( - input=FunctionCorrelation( - tenFirst=tenFirst, tenSecond=tenSecond - ), - negative_slope=0.1, - inplace=False, - ) - - tenFeat = torch.cat([tenVolume], 1) - - elif objPrevious is not None: - tenFlow = self.netUpflow(objPrevious["tenFlow"]) - tenFeat = self.netUpfeat(objPrevious["tenFeat"]) - - tenVolume = torch.nn.functional.leaky_relu( - input=FunctionCorrelation( - tenFirst=tenFirst, - tenSecond=backwarp( - tenInput=tenSecond, tenFlow=tenFlow * self.fltBackwarp - ), - ), - negative_slope=0.1, - inplace=False, - ) - - tenFeat = torch.cat([tenVolume, tenFirst, tenFlow, tenFeat], 1) - - # end - - tenFeat = torch.cat([self.netOne(tenFeat), tenFeat], 1) - tenFeat = torch.cat([self.netTwo(tenFeat), tenFeat], 1) - tenFeat = torch.cat([self.netThr(tenFeat), tenFeat], 1) - tenFeat = torch.cat([self.netFou(tenFeat), tenFeat], 1) - tenFeat = torch.cat([self.netFiv(tenFeat), tenFeat], 1) - - tenFlow = self.netSix(tenFeat) - - return {"tenFlow": tenFlow, "tenFeat": tenFeat} - - # end - - # end - - class Refiner(torch.nn.Module): - def __init__(self): - super(Refiner, self).__init__() - - self.netMain = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=81 + 32 + 2 + 2 + 128 + 128 + 96 + 64 + 32, - out_channels=128, - kernel_size=3, - stride=1, - padding=1, - dilation=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=128, - out_channels=128, - kernel_size=3, - stride=1, - padding=2, - dilation=2, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=128, - out_channels=128, - kernel_size=3, - stride=1, - padding=4, - dilation=4, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=128, - out_channels=96, - kernel_size=3, - stride=1, - padding=8, - dilation=8, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=96, - out_channels=64, - kernel_size=3, - stride=1, - padding=16, - dilation=16, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=64, - out_channels=32, - kernel_size=3, - stride=1, - padding=1, - dilation=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=32, - out_channels=2, - kernel_size=3, - stride=1, - padding=1, - dilation=1, - ), - ) - - # end - - def forward(self, tenInput): - return self.netMain(tenInput) - - # end - - # end - - self.netExtractor = Extractor() - - self.netTwo = Decoder(2) - self.netThr = Decoder(3) - self.netFou = Decoder(4) - self.netFiv = Decoder(5) - self.netSix = Decoder(6) - - self.netRefiner = Refiner() - - self.load_state_dict( - { - strKey.replace("module", "net"): tenWeight - for strKey, tenWeight in torch.hub.load_state_dict_from_url( - url="http://content.sniklaus.com/github/pytorch-pwc/network-" - + "default" - + ".pytorch", - model_dir=get_ckpt_container_path(MODEL_TYPE) - ).items() - } - ) - - # end - - def forward(self, tenFirst, tenSecond, *args): - # optionally pass pre-extracted feature pyramid in as args - if len(args) == 0: - tenFirst = self.netExtractor(tenFirst) - tenSecond = self.netExtractor(tenSecond) - else: - tenFirst, tenSecond = args - - objEstimate = self.netSix(tenFirst[-1], tenSecond[-1], None) - objEstimate = self.netFiv(tenFirst[-2], tenSecond[-2], objEstimate) - objEstimate = self.netFou(tenFirst[-3], tenSecond[-3], objEstimate) - objEstimate = self.netThr(tenFirst[-4], tenSecond[-4], objEstimate) - objEstimate = self.netTwo(tenFirst[-5], tenSecond[-5], objEstimate) - - return objEstimate["tenFlow"] + self.netRefiner(objEstimate["tenFeat"]) - - # end - - def extract_pyramid(self, tenFirst, tenSecond): - return self.netExtractor(tenFirst), self.netExtractor(tenSecond) - - def extract_pyramid_single(self, tenFirst): - return self.netExtractor(tenFirst) - - -# end - -netNetwork = None - -########################################################## - - -def estimate(tenFirst, tenSecond): - global netNetwork - - if netNetwork is None: - netNetwork = Network().cuda().eval() - # end - - assert tenFirst.shape[1] == tenSecond.shape[1] - assert tenFirst.shape[2] == tenSecond.shape[2] - - intWidth = tenFirst.shape[2] - intHeight = tenFirst.shape[1] - - assert ( - intWidth == 1024 - ) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue - assert ( - intHeight == 436 - ) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue - - tenPreprocessedFirst = tenFirst.cuda().view(1, 3, intHeight, intWidth) - tenPreprocessedSecond = tenSecond.cuda().view(1, 3, intHeight, intWidth) - - intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0)) - intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0)) - - tenPreprocessedFirst = torch.nn.functional.interpolate( - input=tenPreprocessedFirst, - size=(intPreprocessedHeight, intPreprocessedWidth), - mode="bilinear", - align_corners=False, - ) - tenPreprocessedSecond = torch.nn.functional.interpolate( - input=tenPreprocessedSecond, - size=(intPreprocessedHeight, intPreprocessedWidth), - mode="bilinear", - align_corners=False, - ) - - tenFlow = 20.0 * torch.nn.functional.interpolate( - input=netNetwork(tenPreprocessedFirst, tenPreprocessedSecond), - size=(intHeight, intWidth), - mode="bilinear", - align_corners=False, - ) - - tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth) - tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight) - - return tenFlow[0, :, :, :].cpu() - - -# end - - -class Upsampler_8tap(nn.Module): - def __init__(self): - super(Upsampler_8tap, self).__init__() - filt_8tap = torch.tensor([[-1, 4, -11, 40, 40, -11, 4, -1]]).div(64) - self.filter = nn.Parameter(filt_8tap.repeat(3, 1, 1, 1), requires_grad=False) - - def forward(self, im): - b, c, h, w = im.shape - im_up = torch.zeros(b, c, h * 2, w * 2).to(im.device) - im_up[:, :, ::2, ::2] = im - - p = (8 - 1) // 2 - im_up_row = F.conv2d( - F.pad(im, pad=(p, p + 1, 0, 0), mode="reflect"), self.filter, groups=3 - ) - im_up[:, :, 0::2, 1::2] = im_up_row - im_up_col = torch.transpose( - F.conv2d( - F.pad(torch.transpose(im, 2, 3), pad=(p, p + 1, 0, 0), mode="reflect"), - self.filter, - groups=3, - ), - 2, - 3, - ) - im_up[:, :, 1::2, 0::2] = im_up_col - im_up_cross = F.conv2d( - F.pad(im_up[:, :, 1::2, ::2], pad=(p, p + 1, 0, 0), mode="reflect"), - self.filter, - groups=3, - ) - im_up[:, :, 1::2, 1::2] = im_up_cross - return im_up - -# end - - -model_urls = { - "r3d_18": "https://download.pytorch.org/models/r3d_18-b3b3357e.pth", - "mc3_18": "https://download.pytorch.org/models/mc3_18-a90a0ba3.pth", - "r2plus1d_18": "https://download.pytorch.org/models/r2plus1d_18-91a641e6.pth", -} - - -class Conv3DSimple(nn.Conv3d): - def __init__(self, in_planes, out_planes, midplanes=None, stride=1, padding=1): - super(Conv3DSimple, self).__init__( - in_channels=in_planes, - out_channels=out_planes, - kernel_size=(3, 3, 3), - stride=stride, - padding=padding, - bias=False, - ) - - @staticmethod - def get_downsample_stride(stride, temporal_stride): - if temporal_stride: - return (temporal_stride, stride, stride) - else: - return (stride, stride, stride) - - -class Conv2Plus1D(nn.Sequential): - def __init__(self, in_planes, out_planes, midplanes, stride=1, padding=1): - super(Conv2Plus1D, self).__init__( - nn.Conv3d( - in_planes, - midplanes, - kernel_size=(1, 3, 3), - stride=(1, stride, stride), - padding=(0, padding, padding), - bias=False, - ), - batchnorm(midplanes), - nn.ReLU(inplace=True), - nn.Conv3d( - midplanes, - out_planes, - kernel_size=(3, 1, 1), - stride=(stride, 1, 1), - padding=(padding, 0, 0), - bias=False, - ), - ) - - @staticmethod - def get_downsample_stride(stride): - return stride, stride, stride - - -class Conv3DNoTemporal(nn.Conv3d): - def __init__(self, in_planes, out_planes, midplanes=None, stride=1, padding=1): - super(Conv3DNoTemporal, self).__init__( - in_channels=in_planes, - out_channels=out_planes, - kernel_size=(1, 3, 3), - stride=(1, stride, stride), - padding=(0, padding, padding), - bias=False, - ) - - @staticmethod - def get_downsample_stride(stride): - return 1, stride, stride - - -class SEGating(nn.Module): - def __init__(self, inplanes, reduction=16): - super().__init__() - - self.pool = nn.AdaptiveAvgPool3d(1) - self.attn_layer = nn.Sequential( - nn.Conv3d(inplanes, inplanes, kernel_size=1, stride=1, bias=True), - nn.Sigmoid(), - ) - - def forward(self, x): - out = self.pool(x) - y = self.attn_layer(out) - return x * y - - -class BasicBlock(nn.Module): - expansion = 1 - - def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None): - midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes) - - super(BasicBlock, self).__init__() - self.conv1 = nn.Sequential( - conv_builder(inplanes, planes, midplanes, stride), - batchnorm(planes), - nn.ReLU(inplace=True), - ) - self.conv2 = nn.Sequential( - conv_builder(planes, planes, midplanes), batchnorm(planes) - ) - self.fg = SEGating(planes) ## Feature Gating, from FLAVR - self.relu = nn.ReLU(inplace=True) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - out = self.conv1(x) - out = self.conv2(out) - out = self.fg(out) - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class Bottleneck(nn.Module): - expansion = 4 - - def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None): - super(Bottleneck, self).__init__() - midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes) - - # 1x1x1 - self.conv1 = nn.Sequential( - nn.Conv3d(inplanes, planes, kernel_size=1, bias=False), - batchnorm(planes), - nn.ReLU(inplace=True), - ) - # Second kernel - self.conv2 = nn.Sequential( - conv_builder(planes, planes, midplanes, stride), - batchnorm(planes), - nn.ReLU(inplace=True), - ) - - # 1x1x1 - self.conv3 = nn.Sequential( - nn.Conv3d(planes, planes * self.expansion, kernel_size=1, bias=False), - batchnorm(planes * self.expansion), - ) - self.relu = nn.ReLU(inplace=True) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.conv2(out) - out = self.conv3(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class BasicStem(nn.Sequential): - """The default conv-batchnorm-relu stem""" - - def __init__(self, outplanes=32): - super(BasicStem, self).__init__( - nn.Conv3d( - 3, - outplanes, - kernel_size=(3, 7, 7), - stride=(1, 2, 2), - padding=(1, 3, 3), - bias=False, - ), - batchnorm(outplanes), - nn.ReLU(inplace=True), - ) - - -class R2Plus1dStem(nn.Sequential): - """R(2+1)D stem is different than the default one as it uses separated 3D convolution""" - - def __init__(self): - super(R2Plus1dStem, self).__init__( - nn.Conv3d( - 3, - 45, - kernel_size=(1, 7, 7), - stride=(1, 2, 2), - padding=(0, 3, 3), - bias=False, - ), - batchnorm(45), - nn.ReLU(inplace=True), - nn.Conv3d( - 45, - 64, - kernel_size=(3, 1, 1), - stride=(1, 1, 1), - padding=(1, 0, 0), - bias=False, - ), - batchnorm(64), - nn.ReLU(inplace=True), - ) - - -class VideoResNet(nn.Module): - def __init__( - self, - block, - conv_makers, - layers, - stem, - zero_init_residual=False, - channels=[32, 64, 96, 128], - ): - """Generic resnet video generator. - - Args: - block (nn.Module): resnet building block - conv_makers (list(functions)): generator function for each layer - layers (List[int]): number of blocks per layer - stem (nn.Module, optional): Resnet stem, if None, defaults to conv-bn-relu. Defaults to None. - zero_init_residual (bool, optional): Zero init bottleneck residual BN. Defaults to False. - """ - super(VideoResNet, self).__init__() - self.inplanes = channels[0] # output channel of first stem - - self.stem = stem() - - self.layer1 = self._make_layer( - block, conv_makers[0], channels[0], layers[0], stride=1 - ) - self.layer2 = self._make_layer( - block, conv_makers[1], channels[1], layers[1], stride=2, temporal_stride=1 - ) - self.layer3 = self._make_layer( - block, conv_makers[2], channels[2], layers[2], stride=2, temporal_stride=1 - ) - self.layer4 = self._make_layer( - block, conv_makers[3], channels[3], layers[3], stride=1, temporal_stride=1 - ) - - # init weights - self._initialize_weights() - - if zero_init_residual: - for m in self.modules(): - if isinstance(m, Bottleneck): - nn.init.constant_(m.bn3.weight, 0) - - def forward(self, x): - tensorConv0 = self.stem(x) - tensorConv1 = self.layer1(tensorConv0) - tensorConv2 = self.layer2(tensorConv1) - tensorConv3 = self.layer3(tensorConv2) - tensorConv4 = self.layer4(tensorConv3) - return tensorConv0, tensorConv1, tensorConv2, tensorConv3, tensorConv4 - - def _make_layer( - self, block, conv_builder, planes, blocks, stride=1, temporal_stride=None - ): - downsample = None - - if stride != 1 or self.inplanes != planes * block.expansion: - ds_stride = conv_builder.get_downsample_stride(stride, temporal_stride) - downsample = nn.Sequential( - nn.Conv3d( - self.inplanes, - planes * block.expansion, - kernel_size=1, - stride=ds_stride, - bias=False, - ), - batchnorm(planes * block.expansion), - ) - stride = ds_stride - - layers = [] - layers.append(block(self.inplanes, planes, conv_builder, stride, downsample)) - - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes, conv_builder)) - - return nn.Sequential(*layers) - - def _initialize_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv3d): - nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.BatchNorm3d): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Linear): - nn.init.normal_(m.weight, 0, 0.01) - nn.init.constant_(m.bias, 0) - - -def _video_resnet(arch, pretrained=False, progress=True, **kwargs): - model = VideoResNet(**kwargs) - - if pretrained: - state_dict = load_state_dict_from_url(model_urls[arch], progress=progress, model_dir=get_ckpt_container_path(MODEL_TYPE)) - model.load_state_dict(state_dict) - return model - - -def r3d_18(bn=False, pretrained=False, progress=True, **kwargs): - """Construct 18 layer Resnet3D model as in - https://arxiv.org/abs/1711.11248 - - Args: - pretrained (bool): If True, returns a model pre-trained on Kinetics-400 - progress (bool): If True, displays a progress bar of the download to stderr - - Returns: - nn.Module: R3D-18 network - """ - - global batchnorm - if bn: - batchnorm = nn.BatchNorm3d - else: - batchnorm = identity - - return _video_resnet( - "r3d_18", - pretrained, - progress, - block=BasicBlock, - conv_makers=[Conv3DSimple] * 4, - layers=[2, 2, 2, 2], - stem=BasicStem, - **kwargs, - ) - - -def mc3_18(bn=False, pretrained=False, progress=True, **kwargs): - """Constructor for 18 layer Mixed Convolution network as in - https://arxiv.org/abs/1711.11248 - - Args: - pretrained (bool): If True, returns a model pre-trained on Kinetics-400 - progress (bool): If True, displays a progress bar of the download to stderr - - Returns: - nn.Module: MC3 Network definition - """ - global batchnorm - if bn: - batchnorm = nn.BatchNorm3d - else: - batchnorm = identity - - return _video_resnet( - "mc3_18", - pretrained, - progress, - block=BasicBlock, - conv_makers=[Conv3DSimple] + [Conv3DNoTemporal] * 3, - layers=[2, 2, 2, 2], - stem=BasicStem, - **kwargs, - ) - - -def r2plus1d_18(bn=False, pretrained=False, progress=True, **kwargs): - """Constructor for the 18 layer deep R(2+1)D network as in - https://arxiv.org/abs/1711.11248 - - Args: - pretrained (bool): If True, returns a model pre-trained on Kinetics-400 - progress (bool): If True, displays a progress bar of the download to stderr - - Returns: - nn.Module: R(2+1)D-18 network - """ - - global batchnorm - if bn: - batchnorm = nn.BatchNorm3d - else: - batchnorm = identity - - return _video_resnet( - "r2plus1d_18", - pretrained, - progress, - block=BasicBlock, - conv_makers=[Conv2Plus1D] * 4, - layers=[2, 2, 2, 2], - stem=R2Plus1dStem, - **kwargs, - ) - - -class upConv3D(nn.Module): - def __init__(self, in_ch, out_ch, kernel_size, stride, padding, upmode="transpose"): - super().__init__() - self.upmode = upmode - if self.upmode == "transpose": - self.upconv = nn.ModuleList( - [ - nn.ConvTranspose3d( - in_ch, - out_ch, - kernel_size=kernel_size, - stride=stride, - padding=padding, - ), - SEGating(out_ch), - batchnorm(out_ch), - ] - ) - else: - self.upconv = nn.ModuleList( - [ - nn.Upsample( - mode="trilinear", scale_factor=(1, 2, 2), align_corners=False - ), - nn.Conv3d(in_ch, out_ch, kernel_size=1, stride=1), - SEGating(out_ch), - batchnorm(out_ch), - ] - ) - self.upconv = nn.Sequential(*self.upconv) - - def forward(self, x): - return self.upconv(x) - - -class Conv_3d(nn.Module): - def __init__(self, in_ch, out_ch, kernel_size, stride=1, padding=0, bias=True): - super().__init__() - self.conv = nn.Sequential( - nn.Conv3d( - in_ch, - out_ch, - kernel_size=kernel_size, - stride=stride, - padding=padding, - bias=bias, - ), - SEGating(out_ch), - batchnorm(out_ch), - ) - - def forward(self, x): - return self.conv(x) - - -def make_optimizer(args, my_model): - trainable = filter(lambda x: x.requires_grad, my_model.parameters()) - - if args.optimizer == "SGD": - optimizer_function = optim.SGD - kwargs = {"momentum": 0.9} - elif args.optimizer == "ADAM": - optimizer_function = optim.Adam - kwargs = {"betas": (0.9, 0.999), "eps": 1e-08} - elif args.optimizer == "ADAMax": - optimizer_function = optim.Adamax - kwargs = {"betas": (0.9, 0.999), "eps": 1e-08} - elif args.optimizer == "RMSprop": - optimizer_function = optim.RMSprop - kwargs = {"eps": 1e-08} - - kwargs["lr"] = args.lr - kwargs["weight_decay"] = args.weight_decay - - return optimizer_function(trainable, **kwargs) - - -def make_scheduler(args, my_optimizer): - if args.decay_type == "step": - scheduler = lrs.StepLR(my_optimizer, step_size=args.lr_decay, gamma=args.gamma) - elif args.decay_type.find("step") >= 0: - milestones = args.decay_type.split("_") - milestones.pop(0) - milestones = list(map(lambda x: int(x), milestones)) - scheduler = lrs.MultiStepLR( - my_optimizer, milestones=milestones, gamma=args.gamma - ) - elif args.decay_type == "plateau": - scheduler = lrs.ReduceLROnPlateau( - my_optimizer, - mode="max", - factor=args.gamma, - patience=args.patience, - threshold=0.01, # metric to be used is psnr - threshold_mode="abs", - verbose=True, - ) - - return scheduler - - -def gaussian_kernel(sz, sigma): - k = torch.arange(-(sz - 1) / 2, (sz + 1) / 2) - k = torch.exp(-1.0 / (2 * sigma**2) * k**2) - k = k.reshape(-1, 1) * k.reshape(1, -1) - k = k / torch.sum(k) - return k - - -def moduleNormalize(frame): - return torch.cat( - [ - (frame[:, 0:1, :, :] - 0.4631), - (frame[:, 1:2, :, :] - 0.4352), - (frame[:, 2:3, :, :] - 0.3990), - ], - 1, - ) - - -class FoldUnfold: - """ - Class to handle folding tensor frame into batch of patches and back to frame again - Thanks to Charlie Tan (charlie.tan.2019@bristol.ac.uk) for the earier version. - """ - - def __init__(self, height, width, patch_size, overlap): - if height % 2 or width % 2 or patch_size % 2 or overlap % 2: - print( - "only defined for even values of height, width, patch_size size and overlap, odd values will reconstruct incorrectly" - ) - return - - self.height = height - self.width = width - - self.patch_size = patch_size - self.overlap = overlap - self.stride = patch_size - overlap - - def fold_to_patches(self, *frames): - """ - args: frames -- list of (1,3,H,W) tensors - returns: list of (B,3,h,w) image patches - """ - - # number of blocks in each direction - n_blocks_h = (self.height // (self.stride)) + 1 - n_blocks_w = (self.width // (self.stride)) + 1 - - # how much to pad each edge by - self.pad_h = (self.stride * n_blocks_h + self.overlap - self.height) // 2 - self.pad_w = (self.stride * n_blocks_w + self.overlap - self.width) // 2 - self.height_pad = self.height + 2 * self.pad_h - self.width_pad = self.width + 2 * self.pad_w - - # pad the frames and unfold into patches - patches_list = [] - for i in range(len(frames)): - padded = F.pad( - frames[i], - (self.pad_w, self.pad_w, self.pad_h, self.pad_h), - mode="reflect", - ) - unfolded = F.unfold(padded, self.patch_size, stride=self.stride) - patches = unfolded.permute(2, 1, 0).reshape( - -1, 3, self.patch_size, self.patch_size - ) - patches_list.append(patches) - - return patches_list - - def unfold_to_frame(self, patches): - """ - args: patches -- tensor of shape (B,3,h,w) - returns: frame -- tensor of shape (1,3,H,W) - """ - - # reshape and permute back into [frames, chans * patch_size ** 2, num_patches] as expected by fold - frame_unfold = patches.reshape(-1, 3 * self.patch_size**2, 1).permute(2, 1, 0) - - # fold into tensor of shape pad_shape - frame_fold = F.fold( - frame_unfold, - (self.height_pad, self.width_pad), - self.patch_size, - stride=self.stride, - ) - - # unfold sums overlaps instead of averaging so tensor of ones unfolded and - # folded to track overlaps and take mean of overlapping pixels - ones = torch.ones_like(frame_fold) - ones_unfold = F.unfold(ones, self.patch_size, stride=self.stride) - - # divisor is tensor of shape pad_shape where each element is the number of values that have overlapped - # 1 = no overlaps - divisor = F.fold( - ones_unfold, - (self.height_pad, self.width_pad), - self.patch_size, - stride=self.stride, - ) - - # divide reconstructed frame by divisor - frame_div = frame_fold / divisor - - # crop frame to remove the padded areas - frame_crop = frame_div[ - :, :, self.pad_h : -self.pad_h, self.pad_w : -self.pad_w - ].clone() - - return frame_crop - - -def read_frame_yuv2rgb(stream, width, height, iFrame, bit_depth, pix_fmt="420"): - if pix_fmt == "420": - multiplier = 1 - uv_factor = 2 - elif pix_fmt == "444": - multiplier = 2 - uv_factor = 1 - else: - print("Pixel format {} is not supported".format(pix_fmt)) - return - - if bit_depth == 8: - datatype = np.uint8 - stream.seek(iFrame * 1.5 * width * height * multiplier) - Y = np.fromfile(stream, dtype=datatype, count=width * height).reshape( - (height, width) - ) - - # read chroma samples and upsample since original is 4:2:0 sampling - U = np.fromfile( - stream, dtype=datatype, count=(width // uv_factor) * (height // uv_factor) - ).reshape((height // uv_factor, width // uv_factor)) - V = np.fromfile( - stream, dtype=datatype, count=(width // uv_factor) * (height // uv_factor) - ).reshape((height // uv_factor, width // uv_factor)) - - else: - datatype = np.uint16 - stream.seek(iFrame * 3 * width * height * multiplier) - Y = np.fromfile(stream, dtype=datatype, count=width * height).reshape( - (height, width) - ) - - U = np.fromfile( - stream, dtype=datatype, count=(width // uv_factor) * (height // uv_factor) - ).reshape((height // uv_factor, width // uv_factor)) - V = np.fromfile( - stream, dtype=datatype, count=(width // uv_factor) * (height // uv_factor) - ).reshape((height // uv_factor, width // uv_factor)) - - if pix_fmt == "420": - yuv = np.empty((height * 3 // 2, width), dtype=datatype) - yuv[0:height, :] = Y - - yuv[height : height + height // 4, :] = U.reshape(-1, width) - yuv[height + height // 4 :, :] = V.reshape(-1, width) - - if bit_depth != 8: - yuv = (yuv / (2**bit_depth - 1) * 255).astype(np.uint8) - - # convert to rgb - rgb = cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB_I420) - - else: - yvu = np.stack([Y, V, U], axis=2) - if bit_depth != 8: - yvu = (yvu / (2**bit_depth - 1) * 255).astype(np.uint8) - rgb = cv2.cvtColor(yvu, cv2.COLOR_YCrCb2RGB) - - return rgb - - -def quantize(imTensor): - return imTensor.clamp(0.0, 1.0).mul(255).round() - - -def tensor2rgb(tensor): - """ - Convert GPU Tensor to RGB image (numpy array) - """ - out = [] - for b in range(tensor.shape[0]): - out.append( - np.moveaxis(quantize(tensor[b]).cpu().detach().numpy(), 0, 2).astype( - np.uint8 - ) - ) - return np.array(out) # (B,H,W,C) - - -class Identity(nn.Module): - def __init__(self, *args): - super(Identity, self).__init__() - - def forward(self, x): - return x - - -class SEBlock(nn.Module): - def __init__(self, input_dim, reduction=16): - super(SEBlock, self).__init__() - mid = int(input_dim / reduction) - self.avg_pool = nn.AdaptiveAvgPool2d(1) - self.fc = nn.Sequential( - nn.Linear(input_dim, mid), - nn.ReLU(inplace=True), - nn.Linear(mid, input_dim), - nn.Sigmoid(), - ) - - def forward(self, x): - b, c, _, _ = x.size() - y = self.avg_pool(x).view(b, c) - y = self.fc(y).view(b, c, 1, 1) - return x * y - - -class ResNextBlock(nn.Module): - def __init__( - self, down, cin, cout, ks, stride=1, groups=32, base_width=4, norm_layer=None - ): - super(ResNextBlock, self).__init__() - if norm_layer is None or norm_layer == "batch": - norm_layer = nn.BatchNorm2d - elif norm_layer == "identity": - norm_layer = Identity - width = int(cout * (base_width / 64.0)) * groups - # Both self.conv2 and self.downsample layers downsample the input when stride != 1 - self.conv1 = nn.Conv2d(cin, width, kernel_size=1, stride=1, bias=False) - self.bn1 = norm_layer(width) - if down: - self.conv2 = nn.Conv2d( - width, - width, - kernel_size=ks, - stride=stride, - padding=(ks - 1) // 2, - groups=groups, - bias=False, - ) - else: - self.conv2 = nn.ConvTranspose2d( - width, - width, - kernel_size=ks, - stride=stride, - padding=(ks - stride) // 2, - groups=groups, - bias=False, - ) - self.bn2 = norm_layer(width) - self.conv3 = nn.Conv2d(width, cout, kernel_size=1, stride=1, bias=False) - self.bn3 = norm_layer(cout) - self.relu = nn.ReLU(inplace=True) - self.downsample = None - if stride != 1 or cin != cout: - if down: - self.downsample = nn.Sequential( - nn.Conv2d(cin, cout, kernel_size=1, stride=stride, bias=False), - norm_layer(cout), - ) - else: - self.downsample = nn.Sequential( - # ks = stride here s.t. resolution can be kept - nn.ConvTranspose2d( - cin, cout, kernel_size=2, stride=stride, bias=False - ), - norm_layer(cout), - ) - self.stride = stride - - def forward(self, x): - identity = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - out = self.relu(out) - - out = self.conv3(out) - out = self.bn3(out) - - if self.downsample is not None: - identity = self.downsample(x) - - out += identity - out = self.relu(out) - - return out - - -class MultiScaleResNextBlock(nn.Module): - def __init__(self, down, cin, cout, ks_s, ks_l, stride, norm_layer): - super(MultiScaleResNextBlock, self).__init__() - self.resnext_small = ResNextBlock( - down, cin, cout // 2, ks_s, stride, norm_layer=norm_layer - ) - self.resnext_large = ResNextBlock( - down, cin, cout // 2, ks_l, stride, norm_layer=norm_layer - ) - self.attention = SEBlock(cout) - - def forward(self, tensorCombine): - out_small = self.resnext_small(tensorCombine) - out_large = self.resnext_large(tensorCombine) - out = torch.cat([out_small, out_large], 1) - out = self.attention(out) - return out - - -class UMultiScaleResNext(nn.Module): - def __init__( - self, channels=[64, 128, 256, 512], norm_layer="batch", inplanes=6, **kwargs - ): - super(UMultiScaleResNext, self).__init__() - self.conv1 = MultiScaleResNextBlock( - True, inplanes, channels[0], ks_s=3, ks_l=7, stride=2, norm_layer=norm_layer - ) - self.conv2 = MultiScaleResNextBlock( - True, - channels[0], - channels[1], - ks_s=3, - ks_l=7, - stride=2, - norm_layer=norm_layer, - ) - self.conv3 = MultiScaleResNextBlock( - True, - channels[1], - channels[2], - ks_s=3, - ks_l=5, - stride=2, - norm_layer=norm_layer, - ) - self.conv4 = MultiScaleResNextBlock( - True, - channels[2], - channels[3], - ks_s=3, - ks_l=5, - stride=2, - norm_layer=norm_layer, - ) - - self.deconv4 = MultiScaleResNextBlock( - True, - channels[3], - channels[3], - ks_s=3, - ks_l=5, - stride=1, - norm_layer=norm_layer, - ) - self.deconv3 = MultiScaleResNextBlock( - False, - channels[3], - channels[2], - ks_s=4, - ks_l=6, - stride=2, - norm_layer=norm_layer, - ) - self.deconv2 = MultiScaleResNextBlock( - False, - channels[2], - channels[1], - ks_s=4, - ks_l=8, - stride=2, - norm_layer=norm_layer, - ) - self.deconv1 = MultiScaleResNextBlock( - False, - channels[1], - channels[0], - ks_s=4, - ks_l=8, - stride=2, - norm_layer=norm_layer, - ) - - def forward(self, im0, im2): - tensorJoin = torch.cat([im0, im2], 1) # (B,6,H,W) - - tensorConv1 = self.conv1(tensorJoin) - tensorConv2 = self.conv2(tensorConv1) - tensorConv3 = self.conv3(tensorConv2) - tensorConv4 = self.conv4(tensorConv3) - - tensorDeconv4 = self.deconv4(tensorConv4) - tensorDeconv3 = self.deconv3(tensorDeconv4 + tensorConv4) - tensorDeconv2 = self.deconv2(tensorDeconv3 + tensorConv3) - tensorDeconv1 = self.deconv1(tensorDeconv2 + tensorConv2) - - return tensorDeconv1 - - -class MultiInputGridNet(nn.Module): - def __init__(self, in_chs, out_chs, grid_chs=(32, 64, 96), n_row=3, n_col=6): - super(MultiInputGridNet, self).__init__() - - self.n_row = n_row - self.n_col = n_col - self.n_chs = grid_chs - assert ( - len(grid_chs) == self.n_row - ), "should give num channels for each row (scale stream)" - assert ( - len(in_chs) == self.n_row - ), "should give input channels for each row (scale stream)" - - for r, n_ch in enumerate(self.n_chs): - setattr(self, f"lateral_{r}_0", LateralBlock(in_chs[r], n_ch)) - for c in range(1, self.n_col): - setattr(self, f"lateral_{r}_{c}", LateralBlock(n_ch, n_ch)) - - for r, (in_ch, out_ch) in enumerate(zip(self.n_chs[:-1], self.n_chs[1:])): - for c in range(int(self.n_col / 2)): - setattr(self, f"down_{r}_{c}", DownSamplingBlock(in_ch, out_ch)) - - for r, (in_ch, out_ch) in enumerate(zip(self.n_chs[1:], self.n_chs[:-1])): - for c in range(int(self.n_col / 2)): - setattr(self, f"up_{r}_{c}", UpSamplingBlock(in_ch, out_ch)) - - self.lateral_final = LateralBlock(self.n_chs[0], out_chs) - - def forward(self, *args): - assert len(args) == self.n_row - - # extensible, memory-efficient - cur_col = list(args) - for c in range(int(self.n_col / 2)): - for r in range(self.n_row): - cur_col[r] = getattr(self, f"lateral_{r}_{c}")(cur_col[r]) - if r != 0: - cur_col[r] += getattr(self, f"down_{r-1}_{c}")(cur_col[r - 1]) - - for c in range(int(self.n_col / 2), self.n_col): - for r in range(self.n_row - 1, -1, -1): - cur_col[r] = getattr(self, f"lateral_{r}_{c}")(cur_col[r]) - if r != self.n_row - 1: - cur_col[r] += getattr(self, f"up_{r}_{c-int(self.n_col/2)}")( - cur_col[r + 1] - ) - - return self.lateral_final(cur_col[0]) - - -class MIMOGridNet(nn.Module): - def __init__( - self, in_chs, out_chs, grid_chs=(32, 64, 96), n_row=3, n_col=6, outrow=(0, 1, 2) - ): - super(MIMOGridNet, self).__init__() - - self.n_row = n_row - self.n_col = n_col - self.n_chs = grid_chs - self.outrow = outrow - assert ( - len(grid_chs) == self.n_row - ), "should give num channels for each row (scale stream)" - assert ( - len(in_chs) == self.n_row - ), "should give input channels for each row (scale stream)" - assert len(out_chs) == len( - self.outrow - ), "should give out channels for each output row (scale stream)" - - for r, n_ch in enumerate(self.n_chs): - setattr(self, f"lateral_{r}_0", LateralBlock(in_chs[r], n_ch)) - for c in range(1, self.n_col): - setattr(self, f"lateral_{r}_{c}", LateralBlock(n_ch, n_ch)) - - for r, (in_ch, out_ch) in enumerate(zip(self.n_chs[:-1], self.n_chs[1:])): - for c in range(int(self.n_col / 2)): - setattr(self, f"down_{r}_{c}", DownSamplingBlock(in_ch, out_ch)) - - for r, (in_ch, out_ch) in enumerate(zip(self.n_chs[1:], self.n_chs[:-1])): - for c in range(int(self.n_col / 2)): - setattr(self, f"up_{r}_{c}", UpSamplingBlock(in_ch, out_ch)) - - for i, r in enumerate(outrow): - setattr(self, f"lateral_final_{r}", LateralBlock(self.n_chs[r], out_chs[i])) - - def forward(self, *args): - assert len(args) == self.n_row - - # extensible, memory-efficient - cur_col = list(args) - for c in range(int(self.n_col / 2)): - for r in range(self.n_row): - cur_col[r] = getattr(self, f"lateral_{r}_{c}")(cur_col[r]) - if r != 0: - cur_col[r] += getattr(self, f"down_{r-1}_{c}")(cur_col[r - 1]) - - for c in range(int(self.n_col / 2), self.n_col): - for r in range(self.n_row - 1, -1, -1): - cur_col[r] = getattr(self, f"lateral_{r}_{c}")(cur_col[r]) - if r != self.n_row - 1: - cur_col[r] += getattr(self, f"up_{r}_{c-int(self.n_col/2)}")( - cur_col[r + 1] - ) - - out = [] - for r in self.outrow: - out.append(getattr(self, f"lateral_final_{r}")(cur_col[r])) - - return out - - -class GeneralGridNet(nn.Module): - def __init__(self, in_chs, out_chs, grid_chs=(32, 64, 96), n_row=3, n_col=6): - super(GeneralGridNet, self).__init__() - - self.n_row = n_row - self.n_col = n_col - self.n_chs = grid_chs - assert ( - len(grid_chs) == self.n_row - ), "should give num channels for each row (scale stream)" - - for r, n_ch in enumerate(self.n_chs): - if r == 0: - setattr(self, f"lateral_{r}_0", LateralBlock(in_chs, n_ch)) - for c in range(1, self.n_col): - setattr(self, f"lateral_{r}_{c}", LateralBlock(n_ch, n_ch)) - - for r, (in_ch, out_ch) in enumerate(zip(self.n_chs[:-1], self.n_chs[1:])): - for c in range(int(self.n_col / 2)): - setattr(self, f"down_{r}_{c}", DownSamplingBlock(in_ch, out_ch)) - - for r, (in_ch, out_ch) in enumerate(zip(self.n_chs[1:], self.n_chs[:-1])): - for c in range(int(self.n_col / 2)): - setattr(self, f"up_{r}_{c}", UpSamplingBlock(in_ch, out_ch)) - - self.lateral_final = LateralBlock(self.n_chs[0], out_chs) - - def forward(self, x): - cur_col = [x] + [None] * (self.n_row - 1) - for c in range(int(self.n_col / 2)): - for r in range(self.n_row): - if cur_col[r] != None: - cur_col[r] = getattr(self, f"lateral_{r}_{c}")(cur_col[r]) - else: - cur_col[r] = 0.0 - if r != 0: - cur_col[r] += getattr(self, f"down_{r-1}_{c}")(cur_col[r - 1]) - - for c in range(int(self.n_col / 2), self.n_col): - for r in range(self.n_row - 1, -1, -1): - cur_col[r] = getattr(self, f"lateral_{r}_{c}")(cur_col[r]) - if r != self.n_row - 1: - cur_col[r] += getattr(self, f"up_{r}_{c-int(self.n_col/2)}")( - cur_col[r + 1] - ) - - return self.lateral_final(cur_col[0]) - - -class GridNet(nn.Module): - def __init__(self, in_chs, out_chs, grid_chs=(32, 64, 96)): - super(GridNet, self).__init__() - - self.n_row = 3 - self.n_col = 6 - self.n_chs = grid_chs - assert ( - len(grid_chs) == self.n_row - ), "should give num channels for each row (scale stream)" - - self.lateral_init = LateralBlock(in_chs, self.n_chs[0]) - - for r, n_ch in enumerate(self.n_chs): - for c in range(self.n_col - 1): - setattr(self, f"lateral_{r}_{c}", LateralBlock(n_ch, n_ch)) - - for r, (in_ch, out_ch) in enumerate(zip(self.n_chs[:-1], self.n_chs[1:])): - for c in range(int(self.n_col / 2)): - setattr(self, f"down_{r}_{c}", DownSamplingBlock(in_ch, out_ch)) - - for r, (in_ch, out_ch) in enumerate(zip(self.n_chs[1:], self.n_chs[:-1])): - for c in range(int(self.n_col / 2)): - setattr(self, f"up_{r}_{c}", UpSamplingBlock(in_ch, out_ch)) - - self.lateral_final = LateralBlock(self.n_chs[0], out_chs) - - def forward(self, x): - state_00 = self.lateral_init(x) - state_10 = self.down_0_0(state_00) - state_20 = self.down_1_0(state_10) - - state_01 = self.lateral_0_0(state_00) - state_11 = self.down_0_1(state_01) + self.lateral_1_0(state_10) - state_21 = self.down_1_1(state_11) + self.lateral_2_0(state_20) - - state_02 = self.lateral_0_1(state_01) - state_12 = self.down_0_2(state_02) + self.lateral_1_1(state_11) - state_22 = self.down_1_2(state_12) + self.lateral_2_1(state_21) - - state_23 = self.lateral_2_2(state_22) - state_13 = self.up_1_0(state_23) + self.lateral_1_2(state_12) - state_03 = self.up_0_0(state_13) + self.lateral_0_2(state_02) - - state_24 = self.lateral_2_3(state_23) - state_14 = self.up_1_1(state_24) + self.lateral_1_3(state_13) - state_04 = self.up_0_1(state_14) + self.lateral_0_3(state_03) - - state_25 = self.lateral_2_4(state_24) - state_15 = self.up_1_2(state_25) + self.lateral_1_4(state_14) - state_05 = self.up_0_2(state_15) + self.lateral_0_4(state_04) - - return self.lateral_final(state_05) - - -class LateralBlock(nn.Module): - def __init__(self, ch_in, ch_out): - super(LateralBlock, self).__init__() - self.f = nn.Sequential( - nn.PReLU(), - nn.Conv2d(ch_in, ch_out, kernel_size=3, padding=1), - nn.PReLU(), - nn.Conv2d(ch_out, ch_out, kernel_size=3, padding=1), - ) - if ch_in != ch_out: - self.conv = nn.Conv2d(ch_in, ch_out, kernel_size=3, padding=1) - - def forward(self, x): - fx = self.f(x) - if fx.shape[1] != x.shape[1]: - x = self.conv(x) - return fx + x - - -class DownSamplingBlock(nn.Module): - def __init__(self, ch_in, ch_out): - super(DownSamplingBlock, self).__init__() - self.f = nn.Sequential( - nn.PReLU(), - nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=2, padding=1), - nn.PReLU(), - nn.Conv2d(ch_out, ch_out, kernel_size=3, padding=1), - ) - - def forward(self, x): - return self.f(x) - - -class UpSamplingBlock(nn.Module): - def __init__(self, ch_in, ch_out): - super(UpSamplingBlock, self).__init__() - self.f = nn.Sequential( - nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False), - nn.PReLU(), - nn.Conv2d(ch_in, ch_out, kernel_size=3, padding=1), - nn.PReLU(), - nn.Conv2d(ch_out, ch_out, kernel_size=3, padding=1), - ) - - def forward(self, x): - return self.f(x) - -# end - - -class Network(torch.nn.Module): - def __init__(self): - super(Network, self).__init__() - - class Extractor(torch.nn.Module): - def __init__(self): - super(Extractor, self).__init__() - - self.netOne = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=3, - out_channels=16, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=16, - out_channels=16, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=16, - out_channels=16, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netTwo = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=16, - out_channels=32, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=32, - out_channels=32, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=32, - out_channels=32, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netThr = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=32, - out_channels=64, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=64, - out_channels=64, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=64, - out_channels=64, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netFou = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=64, - out_channels=96, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=96, - out_channels=96, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=96, - out_channels=96, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netFiv = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=96, - out_channels=128, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=128, - out_channels=128, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=128, - out_channels=128, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netSix = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=128, - out_channels=196, - kernel_size=3, - stride=2, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=196, - out_channels=196, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=196, - out_channels=196, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - # end - - def forward(self, tenInput): - tenOne = self.netOne(tenInput) - tenTwo = self.netTwo(tenOne) - tenThr = self.netThr(tenTwo) - tenFou = self.netFou(tenThr) - tenFiv = self.netFiv(tenFou) - tenSix = self.netSix(tenFiv) - - return [tenOne, tenTwo, tenThr, tenFou, tenFiv, tenSix] - - # end - - # end - - class Decoder(torch.nn.Module): - def __init__(self, intLevel): - super(Decoder, self).__init__() - - intPrevious = [ - None, - None, - 81 + 32 + 2 + 2, - 81 + 64 + 2 + 2, - 81 + 96 + 2 + 2, - 81 + 128 + 2 + 2, - 81, - None, - ][intLevel + 1] - intCurrent = [ - None, - None, - 81 + 32 + 2 + 2, - 81 + 64 + 2 + 2, - 81 + 96 + 2 + 2, - 81 + 128 + 2 + 2, - 81, - None, - ][intLevel + 0] - - if intLevel < 6: - self.netUpflow = torch.nn.ConvTranspose2d( - in_channels=2, - out_channels=2, - kernel_size=4, - stride=2, - padding=1, - ) - if intLevel < 6: - self.netUpfeat = torch.nn.ConvTranspose2d( - in_channels=intPrevious + 128 + 128 + 96 + 64 + 32, - out_channels=2, - kernel_size=4, - stride=2, - padding=1, - ) - if intLevel < 6: - self.fltBackwarp = [None, None, None, 5.0, 2.5, 1.25, 0.625, None][ - intLevel + 1 - ] - - self.netOne = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent, - out_channels=128, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netTwo = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent + 128, - out_channels=128, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netThr = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent + 128 + 128, - out_channels=96, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netFou = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent + 128 + 128 + 96, - out_channels=64, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netFiv = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent + 128 + 128 + 96 + 64, - out_channels=32, - kernel_size=3, - stride=1, - padding=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - ) - - self.netSix = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=intCurrent + 128 + 128 + 96 + 64 + 32, - out_channels=2, - kernel_size=3, - stride=1, - padding=1, - ) - ) - - # end - - def forward(self, tenFirst, tenSecond, objPrevious): - tenFlow = None - tenFeat = None - - if objPrevious is None: - tenFlow = None - tenFeat = None - - tenVolume = torch.nn.functional.leaky_relu( - input=FunctionCorrelation( - tenFirst=tenFirst, tenSecond=tenSecond - ), - negative_slope=0.1, - inplace=False, - ) - - tenFeat = torch.cat([tenVolume], 1) - - elif objPrevious is not None: - tenFlow = self.netUpflow(objPrevious["tenFlow"]) - tenFeat = self.netUpfeat(objPrevious["tenFeat"]) - - tenVolume = torch.nn.functional.leaky_relu( - input=FunctionCorrelation( - tenFirst=tenFirst, - tenSecond=backwarp( - tenInput=tenSecond, tenFlow=tenFlow * self.fltBackwarp - ), - ), - negative_slope=0.1, - inplace=False, - ) - - tenFeat = torch.cat([tenVolume, tenFirst, tenFlow, tenFeat], 1) - - # end - - tenFeat = torch.cat([self.netOne(tenFeat), tenFeat], 1) - tenFeat = torch.cat([self.netTwo(tenFeat), tenFeat], 1) - tenFeat = torch.cat([self.netThr(tenFeat), tenFeat], 1) - tenFeat = torch.cat([self.netFou(tenFeat), tenFeat], 1) - tenFeat = torch.cat([self.netFiv(tenFeat), tenFeat], 1) - - tenFlow = self.netSix(tenFeat) - - return {"tenFlow": tenFlow, "tenFeat": tenFeat} - - # end - - # end - - class Refiner(torch.nn.Module): - def __init__(self): - super(Refiner, self).__init__() - - self.netMain = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=81 + 32 + 2 + 2 + 128 + 128 + 96 + 64 + 32, - out_channels=128, - kernel_size=3, - stride=1, - padding=1, - dilation=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=128, - out_channels=128, - kernel_size=3, - stride=1, - padding=2, - dilation=2, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=128, - out_channels=128, - kernel_size=3, - stride=1, - padding=4, - dilation=4, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=128, - out_channels=96, - kernel_size=3, - stride=1, - padding=8, - dilation=8, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=96, - out_channels=64, - kernel_size=3, - stride=1, - padding=16, - dilation=16, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=64, - out_channels=32, - kernel_size=3, - stride=1, - padding=1, - dilation=1, - ), - torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), - torch.nn.Conv2d( - in_channels=32, - out_channels=2, - kernel_size=3, - stride=1, - padding=1, - dilation=1, - ), - ) - - # end - - def forward(self, tenInput): - return self.netMain(tenInput) - - # end - - # end - - self.netExtractor = Extractor() - - self.netTwo = Decoder(2) - self.netThr = Decoder(3) - self.netFou = Decoder(4) - self.netFiv = Decoder(5) - self.netSix = Decoder(6) - - self.netRefiner = Refiner() - - self.load_state_dict( - { - strKey.replace("module", "net"): tenWeight - for strKey, tenWeight in torch.hub.load_state_dict_from_url( - url="http://content.sniklaus.com/github/pytorch-pwc/network-" - + "default" - + ".pytorch", - model_dir=get_ckpt_container_path(MODEL_TYPE) - ).items() - } - ) - - # end - - def forward(self, tenFirst, tenSecond, *args): - # optionally pass pre-extracted feature pyramid in as args - if len(args) == 0: - tenFirst = self.netExtractor(tenFirst) - tenSecond = self.netExtractor(tenSecond) - else: - tenFirst, tenSecond = args - - objEstimate = self.netSix(tenFirst[-1], tenSecond[-1], None) - objEstimate = self.netFiv(tenFirst[-2], tenSecond[-2], objEstimate) - objEstimate = self.netFou(tenFirst[-3], tenSecond[-3], objEstimate) - objEstimate = self.netThr(tenFirst[-4], tenSecond[-4], objEstimate) - objEstimate = self.netTwo(tenFirst[-5], tenSecond[-5], objEstimate) - - return objEstimate["tenFlow"] + self.netRefiner(objEstimate["tenFeat"]) - - # end - - def extract_pyramid(self, tenFirst, tenSecond): - return self.netExtractor(tenFirst), self.netExtractor(tenSecond) - - def extract_pyramid_single(self, tenFirst): - return self.netExtractor(tenFirst) - - -# end - -netNetwork = None - -########################################################## - - -def estimate(tenFirst, tenSecond): - global netNetwork - - if netNetwork is None: - netNetwork = Network().cuda().eval() - # end - - assert tenFirst.shape[1] == tenSecond.shape[1] - assert tenFirst.shape[2] == tenSecond.shape[2] - - intWidth = tenFirst.shape[2] - intHeight = tenFirst.shape[1] - - assert ( - intWidth == 1024 - ) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue - assert ( - intHeight == 436 - ) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue - - tenPreprocessedFirst = tenFirst.cuda().view(1, 3, intHeight, intWidth) - tenPreprocessedSecond = tenSecond.cuda().view(1, 3, intHeight, intWidth) - - intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0)) - intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0)) - - tenPreprocessedFirst = torch.nn.functional.interpolate( - input=tenPreprocessedFirst, - size=(intPreprocessedHeight, intPreprocessedWidth), - mode="bilinear", - align_corners=False, - ) - tenPreprocessedSecond = torch.nn.functional.interpolate( - input=tenPreprocessedSecond, - size=(intPreprocessedHeight, intPreprocessedWidth), - mode="bilinear", - align_corners=False, - ) - - tenFlow = 20.0 * torch.nn.functional.interpolate( - input=netNetwork(tenPreprocessedFirst, tenPreprocessedSecond), - size=(intHeight, intWidth), - mode="bilinear", - align_corners=False, - ) - - tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth) - tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight) - - return tenFlow[0, :, :, :].cpu() - - -# end - - -class UNet3d_18(nn.Module): - def __init__(self, channels=[32, 64, 96, 128], bn=True): - super(UNet3d_18, self).__init__() - growth = 2 # since concatenating previous outputs - upmode = "transpose" # use transposeConv to upsample - - self.channels = channels - - self.lrelu = nn.LeakyReLU(0.2, True) - - self.encoder = r3d_18(bn=bn, channels=channels) - - self.decoder = nn.Sequential( - Conv_3d( - channels[::-1][0], - channels[::-1][1], - kernel_size=3, - padding=1, - bias=True, - ), - upConv3D( - channels[::-1][1] * growth, - channels[::-1][2], - kernel_size=(3, 4, 4), - stride=(1, 2, 2), - padding=(1, 1, 1), - upmode=upmode, - ), - upConv3D( - channels[::-1][2] * growth, - channels[::-1][3], - kernel_size=(3, 4, 4), - stride=(1, 2, 2), - padding=(1, 1, 1), - upmode=upmode, - ), - Conv_3d( - channels[::-1][3] * growth, - channels[::-1][3], - kernel_size=3, - padding=1, - bias=True, - ), - upConv3D( - channels[::-1][3] * growth, - channels[::-1][3], - kernel_size=(3, 4, 4), - stride=(1, 2, 2), - padding=(1, 1, 1), - upmode=upmode, - ), - ) - - self.feature_fuse = nn.Sequential( - *( - [ - nn.Conv2d( - channels[::-1][3] * 5, - channels[::-1][3], - kernel_size=1, - stride=1, - bias=False, - ) - ] - + [nn.BatchNorm2d(channels[::-1][3]) if bn else Identity] - ) - ) - - self.outconv = nn.Sequential( - nn.ReflectionPad2d(3), - nn.Conv2d(channels[::-1][3], 3, kernel_size=7, stride=1, padding=0), - ) - - def forward(self, im1, im3, im5, im7, im4_tilde): - images = torch.stack((im1, im3, im4_tilde, im5, im7), dim=2) - - x_0, x_1, x_2, x_3, x_4 = self.encoder(images) - - dx_3 = self.lrelu(self.decoder[0](x_4)) - dx_3 = torch.cat([dx_3, x_3], dim=1) - - dx_2 = self.lrelu(self.decoder[1](dx_3)) - dx_2 = torch.cat([dx_2, x_2], dim=1) - - dx_1 = self.lrelu(self.decoder[2](dx_2)) - dx_1 = torch.cat([dx_1, x_1], dim=1) - - dx_0 = self.lrelu(self.decoder[3](dx_1)) - dx_0 = torch.cat([dx_0, x_0], dim=1) - - dx_out = self.lrelu(self.decoder[4](dx_0)) - dx_out = torch.cat(torch.unbind(dx_out, 2), 1) - - out = self.lrelu(self.feature_fuse(dx_out)) - out = self.outconv(out) - - return out - - -class KernelEstimation(torch.nn.Module): - def __init__(self, kernel_size): - super(KernelEstimation, self).__init__() - self.kernel_size = kernel_size - - def Subnet_offset(ks): - return torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=ks, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), - torch.nn.Conv2d( - in_channels=ks, out_channels=ks, kernel_size=3, stride=1, padding=1 - ), - ) - - def Subnet_weight(ks): - return torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=ks, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True), - torch.nn.Conv2d( - in_channels=ks, out_channels=ks, kernel_size=3, stride=1, padding=1 - ), - torch.nn.Softmax(dim=1), - ) - - def Subnet_offset_ds(ks): - return torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=ks, kernel_size=3, stride=1, padding=1 - ), - ) - - def Subnet_weight_ds(ks): - return torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=ks, kernel_size=3, stride=1, padding=1 - ), - torch.nn.Softmax(dim=1), - ) - - def Subnet_offset_us(ks): - return torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=ks, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Upsample(scale_factor=4, mode="bilinear", align_corners=True), - torch.nn.Conv2d( - in_channels=ks, out_channels=ks, kernel_size=3, stride=1, padding=1 - ), - ) - - def Subnet_weight_us(ks): - return torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Conv2d( - in_channels=64, out_channels=ks, kernel_size=3, stride=1, padding=1 - ), - torch.nn.ReLU(inplace=False), - torch.nn.Upsample(scale_factor=4, mode="bilinear", align_corners=True), - torch.nn.Conv2d( - in_channels=ks, out_channels=ks, kernel_size=3, stride=1, padding=1 - ), - torch.nn.Softmax(dim=1), - ) - - self.moduleWeight1_ds = Subnet_weight_ds(self.kernel_size**2) - self.moduleAlpha1_ds = Subnet_offset_ds(self.kernel_size**2) - self.moduleBeta1_ds = Subnet_offset_ds(self.kernel_size**2) - self.moduleWeight2_ds = Subnet_weight_ds(self.kernel_size**2) - self.moduleAlpha2_ds = Subnet_offset_ds(self.kernel_size**2) - self.moduleBeta2_ds = Subnet_offset_ds(self.kernel_size**2) - - self.moduleWeight1 = Subnet_weight(self.kernel_size**2) - self.moduleAlpha1 = Subnet_offset(self.kernel_size**2) - self.moduleBeta1 = Subnet_offset(self.kernel_size**2) - self.moduleWeight2 = Subnet_weight(self.kernel_size**2) - self.moduleAlpha2 = Subnet_offset(self.kernel_size**2) - self.moduleBeta2 = Subnet_offset(self.kernel_size**2) - - self.moduleWeight1_us = Subnet_weight_us(self.kernel_size**2) - self.moduleAlpha1_us = Subnet_offset_us(self.kernel_size**2) - self.moduleBeta1_us = Subnet_offset_us(self.kernel_size**2) - self.moduleWeight2_us = Subnet_weight_us(self.kernel_size**2) - self.moduleAlpha2_us = Subnet_offset_us(self.kernel_size**2) - self.moduleBeta2_us = Subnet_offset_us(self.kernel_size**2) - - def forward(self, tensorCombine): - # Frame 0 - Weight1_ds = self.moduleWeight1_ds(tensorCombine) - Weight1 = self.moduleWeight1(tensorCombine) - Weight1_us = self.moduleWeight1_us(tensorCombine) - Alpha1_ds = self.moduleAlpha1_ds(tensorCombine) - Alpha1 = self.moduleAlpha1(tensorCombine) - Alpha1_us = self.moduleAlpha1_us(tensorCombine) - Beta1_ds = self.moduleBeta1_ds(tensorCombine) - Beta1 = self.moduleBeta1(tensorCombine) - Beta1_us = self.moduleBeta1_us(tensorCombine) - - # Frame 2 - Weight2_ds = self.moduleWeight2_ds(tensorCombine) - Weight2 = self.moduleWeight2(tensorCombine) - Weight2_us = self.moduleWeight2_us(tensorCombine) - Alpha2_ds = self.moduleAlpha2_ds(tensorCombine) - Alpha2 = self.moduleAlpha2(tensorCombine) - Alpha2_us = self.moduleAlpha2_us(tensorCombine) - Beta2_ds = self.moduleBeta2_ds(tensorCombine) - Beta2 = self.moduleBeta2(tensorCombine) - Beta2_us = self.moduleBeta2_us(tensorCombine) - - return ( - Weight1_ds, - Alpha1_ds, - Beta1_ds, - Weight2_ds, - Alpha2_ds, - Beta2_ds, - Weight1, - Alpha1, - Beta1, - Weight2, - Alpha2, - Beta2, - Weight1_us, - Alpha1_us, - Beta1_us, - Weight2_us, - Alpha2_us, - Beta2_us, - ) - - -class STMFNet_Model(torch.nn.Module): - def __init__(self): - super(STMFNet_Model, self).__init__() - - class Metric(torch.nn.Module): - def __init__(self): - super(Metric, self).__init__() - self.paramScale = torch.nn.Parameter(-torch.ones(1, 1, 1, 1)) - - def forward(self, tenFirst, tenSecond, tenFlow): - return self.paramScale * F.l1_loss( - input=tenFirst, - target=backwarp(tenSecond, tenFlow), - reduction="none", - ).mean(1, True) - - self.kernel_size = 5 - self.dilation = 1 - self.featc = [64, 128, 256, 512] - self.featnorm = "batch" - self.finetune_pwc = False - - self.kernel_pad = int(((self.kernel_size - 1) * self.dilation) / 2.0) - - self.feature_extractor = UMultiScaleResNext( - self.featc, norm_layer=self.featnorm - ) - - self.get_kernel = KernelEstimation(self.kernel_size) - - self.modulePad = torch.nn.ReplicationPad2d( - [self.kernel_pad, self.kernel_pad, self.kernel_pad, self.kernel_pad] - ) - - self.moduleAdaCoF = FunctionAdaCoF.apply - - self.gauss_kernel = torch.nn.Parameter( - gaussian_kernel(5, 0.5).repeat(3, 1, 1, 1), requires_grad=False - ) - - self.upsampler = Upsampler_8tap() - - self.scale_synthesis = MIMOGridNet( - (6, 6 + 6, 6), (3,), grid_chs=(32, 64, 96), n_row=3, n_col=4, outrow=(1,) - ) - - self.flow_estimator = PWCNet() - - self.softsplat = ModuleSoftsplat(strType="softmax") - - self.metric = Metric() - - self.dyntex_generator = UNet3d_18(bn=self.featnorm) - - # freeze weights of PWCNet if not finetuning it - if not self.finetune_pwc: - for param in self.flow_estimator.parameters(): - param.requires_grad = False - - def forward(self, I0, I1, I2, I3): - h0 = int(list(I1.size())[2]) - w0 = int(list(I1.size())[3]) - h2 = int(list(I2.size())[2]) - w2 = int(list(I2.size())[3]) - if h0 != h2 or w0 != w2: - sys.exit("Frame sizes do not match") - - h_padded = False - w_padded = False - if h0 % 128 != 0: - pad_h = 128 - (h0 % 128) - I0 = F.pad(I0, (0, 0, 0, pad_h), mode="reflect") - I1 = F.pad(I1, (0, 0, 0, pad_h), mode="reflect") - I2 = F.pad(I2, (0, 0, 0, pad_h), mode="reflect") - I3 = F.pad(I3, (0, 0, 0, pad_h), mode="reflect") - h_padded = True - - if w0 % 128 != 0: - pad_w = 128 - (w0 % 128) - I0 = F.pad(I0, (0, pad_w, 0, 0), mode="reflect") - I1 = F.pad(I1, (0, pad_w, 0, 0), mode="reflect") - I2 = F.pad(I2, (0, pad_w, 0, 0), mode="reflect") - I3 = F.pad(I3, (0, pad_w, 0, 0), mode="reflect") - w_padded = True - - feats = self.feature_extractor(moduleNormalize(I1), moduleNormalize(I2)) - kernelest = self.get_kernel(feats) - Weight1_ds, Alpha1_ds, Beta1_ds, Weight2_ds, Alpha2_ds, Beta2_ds = kernelest[:6] - Weight1, Alpha1, Beta1, Weight2, Alpha2, Beta2 = kernelest[6:12] - Weight1_us, Alpha1_us, Beta1_us, Weight2_us, Alpha2_us, Beta2_us = kernelest[ - 12: - ] - - # Original scale - tensorAdaCoF1 = ( - self.moduleAdaCoF(self.modulePad(I1), Weight1, Alpha1, Beta1, self.dilation) - * 1.0 - ) - tensorAdaCoF2 = ( - self.moduleAdaCoF(self.modulePad(I2), Weight2, Alpha2, Beta2, self.dilation) - * 1.0 - ) - - # 1/2 downsampled version - c, h, w = I1.shape[1:] - p = (self.gauss_kernel.shape[-1] - 1) // 2 - I1_blur = F.conv2d( - F.pad(I1, pad=(p, p, p, p), mode="reflect"), self.gauss_kernel, groups=c - ) - I2_blur = F.conv2d( - F.pad(I2, pad=(p, p, p, p), mode="reflect"), self.gauss_kernel, groups=c - ) - I1_ds = F.interpolate( - I1_blur, size=(h // 2, w // 2), mode="bilinear", align_corners=False - ) - I2_ds = F.interpolate( - I2_blur, size=(h // 2, w // 2), mode="bilinear", align_corners=False - ) - tensorAdaCoF1_ds = ( - self.moduleAdaCoF( - self.modulePad(I1_ds), Weight1_ds, Alpha1_ds, Beta1_ds, self.dilation - ) - * 1.0 - ) - tensorAdaCoF2_ds = ( - self.moduleAdaCoF( - self.modulePad(I2_ds), Weight2_ds, Alpha2_ds, Beta2_ds, self.dilation - ) - * 1.0 - ) - - # x2 upsampled version - I1_us = self.upsampler(I1) - I2_us = self.upsampler(I2) - tensorAdaCoF1_us = ( - self.moduleAdaCoF( - self.modulePad(I1_us), Weight1_us, Alpha1_us, Beta1_us, self.dilation - ) - * 1.0 - ) - tensorAdaCoF2_us = ( - self.moduleAdaCoF( - self.modulePad(I2_us), Weight2_us, Alpha2_us, Beta2_us, self.dilation - ) - * 1.0 - ) - - # use softsplat for refinement - pyramid0, pyramid2 = self.flow_estimator.extract_pyramid(I1, I2) - flow_0_2 = 20 * self.flow_estimator(I1, I2, pyramid0, pyramid2) - flow_0_2 = F.interpolate( - flow_0_2, size=(h, w), mode="bilinear", align_corners=False - ) - flow_2_0 = 20 * self.flow_estimator(I2, I1, pyramid2, pyramid0) - flow_2_0 = F.interpolate( - flow_2_0, size=(h, w), mode="bilinear", align_corners=False - ) - metric_0_2 = self.metric(I1, I2, flow_0_2) - metric_2_0 = self.metric(I2, I1, flow_2_0) - tensorSoftsplat0 = self.softsplat(I1, 0.5 * flow_0_2, metric_0_2) - tensorSoftsplat2 = self.softsplat(I2, 0.5 * flow_2_0, metric_2_0) - - # synthesize multiple scales - tensorCombine_us = torch.cat([tensorAdaCoF1_us, tensorAdaCoF2_us], dim=1) - tensorCombine = torch.cat( - [tensorAdaCoF1, tensorAdaCoF2, tensorSoftsplat0, tensorSoftsplat2], dim=1 - ) - tensorCombine_ds = torch.cat([tensorAdaCoF1_ds, tensorAdaCoF2_ds], dim=1) - output_tilde = self.scale_synthesis( - tensorCombine_us, tensorCombine, tensorCombine_ds - )[0] - - # generate dynamic texture - dyntex = self.dyntex_generator(I0, I1, I2, I3, output_tilde) - output = output_tilde + dyntex - - if h_padded: - output = output[:, :, 0:h0, :] - if w_padded: - output = output[:, :, :, 0:w0] - - if self.training: - return {"frame1": output} - else: - return output \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/xvfi/__init__.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/xvfi/__init__.py deleted file mode 100644 index e52fc1a..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/xvfi/__init__.py +++ /dev/null @@ -1,115 +0,0 @@ -import argparse -import torch -import torch.nn as nn -import torch.nn.functional as F -import einops -from torch.utils.data import DataLoader -import pathlib -from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames, InterpolationStateList -import typing -from comfy.model_management import get_torch_device - -CKPT_CONFIGS = { - "XVFInet_X4K1000FPS_exp1_latest.pt": { - "module_scale_factor": 4, - "S_trn": 3, - "S_tst": 5 - }, - "XVFInet_Vimeo_exp1_latest.pt": { - "module_scale_factor": 2, - "S_trn": 1, - "S_tst": 1 - } -} - -class XVFI_Inference(nn.Module): - def __init__(self, model_path, model_config) -> None: - super(XVFI_Inference, self).__init__() - from .xvfi_arch import XVFInet, weights_init - model_config = model_config - args = argparse.Namespace( - gpu=get_torch_device(), - nf=64, - **model_config, - img_ch=3, - ) - self.model = XVFInet(args).apply(weights_init).to(get_torch_device()) - self.model.load_state_dict(torch.load(model_path, map_location=get_torch_device())["state_dict_Model"]) - - def forward(self, I0, I1, timestep): - #"Real" inference is called "test_custom" in the original repo - #https://github.com/JihyongOh/XVFI/blob/main/utils.py#L434 - #https://github.com/JihyongOh/XVFI/blob/main/main.py#L336 - - x = torch.stack([I0, I1], dim=0) - x = einops.rearrange(x, "t b c h w -> b c t h w") - return self.model(x, timestep, is_training=False) - -MODEL_TYPE = pathlib.Path(__file__).parent.name - -class XVFI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (list(CKPT_CONFIGS.keys()), ), - "frames": ("IMAGE", ), - "batch_size": ("INT", {"default": 1, "min": 1, "max": 100}), - "multipler": ("INT", {"default": 2, "min": 2, "max": 1000}), - }, - "optional": { - "optional_interpolation_states": ("INTERPOLATION_STATES", ), - } - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "ComfyUI-Frame-Interpolation/VFI" - - def vfi( - self, - ckpt_name: typing.AnyStr, - frames: torch.Tensor, - batch_size: typing.SupportsInt = 1, - multipler: typing.SupportsInt = 2, - optional_interpolation_states: InterpolationStateList = None - ): - model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name) - ckpt_config = CKPT_CONFIGS[ckpt_name] - global model - model = XVFI_Inference(model_path, ckpt_config) - - frames = preprocess_frames(frames) - #https://github.com/JihyongOh/XVFI/blob/main/main.py#L314 - divide = 2 ** (ckpt_config["S_tst"]) * ckpt_config["module_scale_factor"] * 4 - B, C, H, W = frames.size() - H_padding = (divide - H % divide) % divide - W_padding = (divide - W % divide) % divide - if H_padding != 0 or W_padding != 0: - frames = F.pad(frames, (0, W_padding, 0, H_padding), "constant") - - frame_dict = { - str(i): frames[i].unsqueeze(0) for i in range(frames.shape[0]) - } - - if optional_interpolation_states is None: - interpolation_states = [True] * (frames.shape[0] - 1) - else: - interpolation_states = optional_interpolation_states - - enabled_former_idxs = [i for i, state in enumerate(interpolation_states) if state] - former_idxs_loader = DataLoader(enabled_former_idxs, batch_size=batch_size) - - for former_idxs_batch in former_idxs_loader: - for middle_i in range(1, multipler): - _middle_frames = model( - frames[former_idxs_batch], - frames[former_idxs_batch + 1], - timestep=torch.tensor([middle_i/multipler]).repeat(len(former_idxs_batch)).unsqueeze(1).to(get_torch_device()) - ) - for i, former_idx in enumerate(former_idxs_batch): - frame_dict[f'{former_idx}.{middle_i}'] = _middle_frames[i].unsqueeze(0) - - out_frames = torch.cat([frame_dict[key] for key in sorted(frame_dict.keys())], dim=0)[:, :, :H, :W] - return (postprocess_frames(out_frames), ) - diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_models/xvfi/xvfi_arch.py b/imports/ComfyUI_Frame_Interpolation/vfi_models/xvfi/xvfi_arch.py deleted file mode 100644 index dae013e..0000000 --- a/imports/ComfyUI_Frame_Interpolation/vfi_models/xvfi/xvfi_arch.py +++ /dev/null @@ -1,506 +0,0 @@ -import functools, random -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.autograd import Variable -import numpy as np -from torch.nn import init -from comfy.model_management import get_torch_device - - - -class XVFInet(nn.Module): - - def __init__(self, args): - super(XVFInet, self).__init__() - self.args = args - self.device = get_torch_device() - self.nf = args.nf - self.scale = args.module_scale_factor - self.vfinet = VFInet(args) - self.lrelu = nn.ReLU() - self.in_channels = 3 - self.channel_converter = nn.Sequential( - nn.Conv3d(self.in_channels, self.nf, [1, 3, 3], [1, 1, 1], [0, 1, 1]), - nn.ReLU()) - - self.rec_ext_ds_module = [self.channel_converter] - self.rec_ext_ds = nn.Conv3d(self.nf, self.nf, [1, 3, 3], [1, 2, 2], [0, 1, 1]) - for _ in range(int(np.log2(self.scale))): - self.rec_ext_ds_module.append(self.rec_ext_ds) - self.rec_ext_ds_module.append(nn.ReLU()) - self.rec_ext_ds_module.append(nn.Conv3d(self.nf, self.nf, [1, 3, 3], 1, [0, 1, 1])) - self.rec_ext_ds_module.append(RResBlock2D_3D(args, T_reduce_flag=False)) - self.rec_ext_ds_module = nn.Sequential(*self.rec_ext_ds_module) - - self.rec_ctx_ds = nn.Conv3d(self.nf, self.nf, [1, 3, 3], [1, 2, 2], [0, 1, 1]) - - print("The lowest scale depth for training (S_trn): ", self.args.S_trn) - print("The lowest scale depth for test (S_tst): ", self.args.S_tst) - - def forward(self, x, t_value, is_training=True): - ''' - x shape : [B,C,T,H,W] - t_value shape : [B,1] ############### - ''' - B, C, T, H, W = x.size() - B2, C2 = t_value.size() - assert C2 == 1, "t_value shape is [B,]" - assert T % 2 == 0, "T must be an even number" - t_value = t_value.view(B, 1, 1, 1) - - flow_l = None - feat_x = self.rec_ext_ds_module(x) - feat_x_list = [feat_x] - self.lowest_depth_level = self.args.S_trn if is_training else self.args.S_tst - for level in range(1, self.lowest_depth_level+1): - feat_x = self.rec_ctx_ds(feat_x) - feat_x_list.append(feat_x) - - if is_training: - out_l_list = [] - flow_refine_l_list = [] - out_l, flow_l, flow_refine_l = self.vfinet(x, feat_x_list[self.args.S_trn], flow_l, t_value, level=self.args.S_trn, is_training=True) - out_l_list.append(out_l) - flow_refine_l_list.append(flow_refine_l) - for level in range(self.args.S_trn-1, 0, -1): ## self.args.S_trn, self.args.S_trn-1, ..., 1. level 0 is not included - out_l, flow_l = self.vfinet(x, feat_x_list[level], flow_l, t_value, level=level, is_training=True) - out_l_list.append(out_l) - out_l, flow_l, flow_refine_l, occ_0_l0 = self.vfinet(x, feat_x_list[0], flow_l, t_value, level=0, is_training=True) - out_l_list.append(out_l) - flow_refine_l_list.append(flow_refine_l) - return out_l_list[::-1], flow_refine_l_list[::-1], occ_0_l0, torch.mean(x, dim=2) # out_l_list should be reversed. [out_l0, out_l1, ...] - - else: # Testing - for level in range(self.args.S_tst, 0, -1): ## self.args.S_tst, self.args.S_tst-1, ..., 1. level 0 is not included - flow_l = self.vfinet(x, feat_x_list[level], flow_l, t_value, level=level, is_training=False) - out_l = self.vfinet(x, feat_x_list[0], flow_l, t_value, level=0, is_training=False) - return out_l - - -class VFInet(nn.Module): - - def __init__(self, args): - super(VFInet, self).__init__() - self.args = args - self.device = get_torch_device() - self.nf = args.nf - self.scale = args.module_scale_factor - self.in_channels = 3 - - self.conv_flow_bottom = nn.Sequential( - nn.Conv2d(2*self.nf, 2*self.nf, [4,4], 2, [1,1]), - nn.ReLU(), - nn.Conv2d(2*self.nf, 4*self.nf, [4,4], 2, [1,1]), - nn.ReLU(), - nn.UpsamplingNearest2d(scale_factor=2), - nn.Conv2d(4 * self.nf, 2 * self.nf, [3, 3], 1, [1, 1]), - nn.ReLU(), - nn.UpsamplingNearest2d(scale_factor=2), - nn.Conv2d(2 * self.nf, self.nf, [3, 3], 1, [1, 1]), - nn.ReLU(), - nn.Conv2d(self.nf, 6, [3,3], 1, [1,1]), - ) - - self.conv_flow1 = nn.Conv2d(2*self.nf, self.nf, [3, 3], 1, [1, 1]) - - self.conv_flow2 = nn.Sequential( - nn.Conv2d(2*self.nf + 4, 2 * self.nf, [4, 4], 2, [1, 1]), - nn.ReLU(), - nn.Conv2d(2 * self.nf, 4 * self.nf, [4, 4], 2, [1, 1]), - nn.ReLU(), - nn.UpsamplingNearest2d(scale_factor=2), - nn.Conv2d(4 * self.nf, 2 * self.nf, [3, 3], 1, [1, 1]), - nn.ReLU(), - nn.UpsamplingNearest2d(scale_factor=2), - nn.Conv2d(2 * self.nf, self.nf, [3, 3], 1, [1, 1]), - nn.ReLU(), - nn.Conv2d(self.nf, 6, [3, 3], 1, [1, 1]), - ) - - self.conv_flow3 = nn.Sequential( - nn.Conv2d(4 + self.nf * 4, self.nf, [1, 1], 1, [0, 0]), - nn.ReLU(), - nn.Conv2d(self.nf, 2 * self.nf, [4, 4], 2, [1, 1]), - nn.ReLU(), - nn.Conv2d(2 * self.nf, 4 * self.nf, [4, 4], 2, [1, 1]), - nn.ReLU(), - nn.UpsamplingNearest2d(scale_factor=2), - nn.Conv2d(4 * self.nf, 2 * self.nf, [3, 3], 1, [1, 1]), - nn.ReLU(), - nn.UpsamplingNearest2d(scale_factor=2), - nn.Conv2d(2 * self.nf, self.nf, [3, 3], 1, [1, 1]), - nn.ReLU(), - nn.Conv2d(self.nf, 4, [3, 3], 1, [1, 1]), - ) - - self.refine_unet = RefineUNet(args) - self.lrelu = nn.ReLU() - - def forward(self, x, feat_x, flow_l_prev, t_value, level, is_training): - ''' - x shape : [B,C,T,H,W] - t_value shape : [B,1] ############### - ''' - B, C, T, H, W = x.size() - assert T % 2 == 0, "T must be an even number" - - ####################### For a single level - l = 2 ** level - x_l = x.permute(0,2,1,3,4) - x_l = x_l.contiguous().view(B * T, C, H, W) - - if level == 0: - pass - else: - x_l = F.interpolate(x_l, scale_factor=(1.0 / l, 1.0 / l), mode='bicubic', align_corners=False) - ''' - Down pixel-shuffle - ''' - x_l = x_l.view(B, T, C, H//l, W//l) - x_l = x_l.permute(0,2,1,3,4) - - B, C, T, H, W = x_l.size() - - ## Feature extraction - feat0_l = feat_x[:,:,0,:,:] - feat1_l = feat_x[:,:,1,:,:] - - ## Flow estimation - if flow_l_prev is None: - flow_l_tmp = self.conv_flow_bottom(torch.cat((feat0_l, feat1_l), dim=1)) - flow_l = flow_l_tmp[:,:4,:,:] - else: - up_flow_l_prev = 2.0*F.interpolate(flow_l_prev.detach(), scale_factor=(2,2), mode='bilinear', align_corners=False) - warped_feat1_l = self.bwarp(feat1_l, up_flow_l_prev[:,:2,:,:]) - warped_feat0_l = self.bwarp(feat0_l, up_flow_l_prev[:,2:,:,:]) - flow_l_tmp = self.conv_flow2(torch.cat([self.conv_flow1(torch.cat([feat0_l, warped_feat1_l],dim=1)), self.conv_flow1(torch.cat([feat1_l, warped_feat0_l],dim=1)), up_flow_l_prev],dim=1)) - flow_l = flow_l_tmp[:,:4,:,:] + up_flow_l_prev - - if not is_training and level!=0: - return flow_l - - flow_01_l = flow_l[:,:2,:,:] - flow_10_l = flow_l[:,2:,:,:] - z_01_l = torch.sigmoid(flow_l_tmp[:,4:5,:,:]) - z_10_l = torch.sigmoid(flow_l_tmp[:,5:6,:,:]) - - ## Complementary Flow Reversal (CFR) - flow_forward, norm0_l = self.z_fwarp(flow_01_l, t_value * flow_01_l, z_01_l) ## Actually, F (t) -> (t+1). Translation only. Not normalized yet - flow_backward, norm1_l = self.z_fwarp(flow_10_l, (1-t_value) * flow_10_l, z_10_l) ## Actually, F (1-t) -> (-t). Translation only. Not normalized yet - - flow_t0_l = -(1-t_value) * ((t_value)*flow_forward) + (t_value) * ((t_value)*flow_backward) # The numerator of Eq.(1) in the paper. - flow_t1_l = (1-t_value) * ((1-t_value)*flow_forward) - (t_value) * ((1-t_value)*flow_backward) # The numerator of Eq.(2) in the paper. - - norm_l = (1-t_value)*norm0_l + t_value*norm1_l - mask_ = (norm_l.detach() > 0).type(norm_l.type()) - flow_t0_l = (1-mask_) * flow_t0_l + mask_ * (flow_t0_l.clone() / (norm_l.clone() + (1-mask_))) # Divide the numerator with denominator in Eq.(1) - flow_t1_l = (1-mask_) * flow_t1_l + mask_ * (flow_t1_l.clone() / (norm_l.clone() + (1-mask_))) # Divide the numerator with denominator in Eq.(2) - - ## Feature warping - warped0_l = self.bwarp(feat0_l, flow_t0_l) - warped1_l = self.bwarp(feat1_l, flow_t1_l) - - ## Flow refinement - flow_refine_l = torch.cat([feat0_l, warped0_l, warped1_l, feat1_l, flow_t0_l, flow_t1_l], dim=1) - flow_refine_l = self.conv_flow3(flow_refine_l) + torch.cat([flow_t0_l, flow_t1_l], dim=1) - flow_t0_l = flow_refine_l[:, :2, :, :] - flow_t1_l = flow_refine_l[:, 2:4, :, :] - - warped0_l = self.bwarp(feat0_l, flow_t0_l) - warped1_l = self.bwarp(feat1_l, flow_t1_l) - - ## Flow upscale - flow_t0_l = self.scale * F.interpolate(flow_t0_l, scale_factor=(self.scale, self.scale), mode='bilinear',align_corners=False) - flow_t1_l = self.scale * F.interpolate(flow_t1_l, scale_factor=(self.scale, self.scale), mode='bilinear',align_corners=False) - - ## Image warping and blending - warped_img0_l = self.bwarp(x_l[:,:,0,:,:], flow_t0_l) - warped_img1_l = self.bwarp(x_l[:,:,1,:,:], flow_t1_l) - - refine_out = self.refine_unet(torch.cat([F.pixel_shuffle(torch.cat([feat0_l, feat1_l, warped0_l, warped1_l],dim=1), self.scale), x_l[:,:,0,:,:], x_l[:,:,1,:,:], warped_img0_l, warped_img1_l, flow_t0_l, flow_t1_l],dim=1)) - occ_0_l = torch.sigmoid(refine_out[:, 0:1, :, :]) - occ_1_l = 1-occ_0_l - - out_l = (1-t_value)*occ_0_l*warped_img0_l + t_value*occ_1_l*warped_img1_l - out_l = out_l / ( (1-t_value)*occ_0_l + t_value*occ_1_l ) + refine_out[:, 1:4, :, :] - - if not is_training and level==0: - return out_l - - if is_training: - if flow_l_prev is None: - # if level == self.args.S_trn: - return out_l, flow_l, flow_refine_l[:, 0:4, :, :] - elif level != 0: - return out_l, flow_l - else: # level==0 - return out_l, flow_l, flow_refine_l[:, 0:4, :, :], occ_0_l - - def bwarp(self, x, flo): - ''' - x: [B, C, H, W] (im2) - flo: [B, 2, H, W] flow - ''' - B, C, H, W = x.size() - # mesh grid - xx = torch.arange(0, W).view(1, 1, 1, W).expand(B, 1, H, W) - yy = torch.arange(0, H).view(1, 1, H, 1).expand(B, 1, H, W) - grid = torch.cat((xx, yy), 1).float() - - grid = grid.to(self.device) - vgrid = torch.autograd.Variable(grid) + flo - - # scale grid to [-1,1] - vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :].clone() / max(W - 1, 1) - 1.0 - vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :].clone() / max(H - 1, 1) - 1.0 - - vgrid = vgrid.permute(0, 2, 3, 1) # [B,H,W,2] - output = nn.functional.grid_sample(x, vgrid, align_corners=True) - mask = torch.autograd.Variable(torch.ones(x.size())).to(self.device) - mask = nn.functional.grid_sample(mask, vgrid, align_corners=True) - - # mask[mask<0.9999] = 0 - # mask[mask>0] = 1 - mask = mask.masked_fill_(mask < 0.999, 0) - mask = mask.masked_fill_(mask > 0, 1) - - return output * mask - - def fwarp(self, img, flo): - - """ - -img: image (N, C, H, W) - -flo: optical flow (N, 2, H, W) - elements of flo is in [0, H] and [0, W] for dx, dy - https://github.com/lyh-18/EQVI/blob/EQVI-master/models/forward_warp_gaussian.py - """ - - # (x1, y1) (x1, y2) - # +---------------+ - # | | - # | o(x, y) | - # | | - # | | - # | | - # | | - # +---------------+ - # (x2, y1) (x2, y2) - - N, C, _, _ = img.size() - - # translate start-point optical flow to end-point optical flow - y = flo[:, 0:1:, :] - x = flo[:, 1:2, :, :] - - x = x.repeat(1, C, 1, 1) - y = y.repeat(1, C, 1, 1) - - # Four point of square (x1, y1), (x1, y2), (x2, y1), (y2, y2) - x1 = torch.floor(x) - x2 = x1 + 1 - y1 = torch.floor(y) - y2 = y1 + 1 - - # firstly, get gaussian weights - w11, w12, w21, w22 = self.get_gaussian_weights(x, y, x1, x2, y1, y2) - - # secondly, sample each weighted corner - img11, o11 = self.sample_one(img, x1, y1, w11) - img12, o12 = self.sample_one(img, x1, y2, w12) - img21, o21 = self.sample_one(img, x2, y1, w21) - img22, o22 = self.sample_one(img, x2, y2, w22) - - imgw = img11 + img12 + img21 + img22 - o = o11 + o12 + o21 + o22 - - return imgw, o - - - def z_fwarp(self, img, flo, z): - """ - -img: image (N, C, H, W) - -flo: optical flow (N, 2, H, W) - elements of flo is in [0, H] and [0, W] for dx, dy - modified from https://github.com/lyh-18/EQVI/blob/EQVI-master/models/forward_warp_gaussian.py - """ - - # (x1, y1) (x1, y2) - # +---------------+ - # | | - # | o(x, y) | - # | | - # | | - # | | - # | | - # +---------------+ - # (x2, y1) (x2, y2) - - N, C, _, _ = img.size() - - # translate start-point optical flow to end-point optical flow - y = flo[:, 0:1:, :] - x = flo[:, 1:2, :, :] - - x = x.repeat(1, C, 1, 1) - y = y.repeat(1, C, 1, 1) - - # Four point of square (x1, y1), (x1, y2), (x2, y1), (y2, y2) - x1 = torch.floor(x) - x2 = x1 + 1 - y1 = torch.floor(y) - y2 = y1 + 1 - - # firstly, get gaussian weights - w11, w12, w21, w22 = self.get_gaussian_weights(x, y, x1, x2, y1, y2, z+1e-5) - - # secondly, sample each weighted corner - img11, o11 = self.sample_one(img, x1, y1, w11) - img12, o12 = self.sample_one(img, x1, y2, w12) - img21, o21 = self.sample_one(img, x2, y1, w21) - img22, o22 = self.sample_one(img, x2, y2, w22) - - imgw = img11 + img12 + img21 + img22 - o = o11 + o12 + o21 + o22 - - return imgw, o - - - def get_gaussian_weights(self, x, y, x1, x2, y1, y2, z=1.0): - # z 0.0 ~ 1.0 - w11 = z * torch.exp(-((x - x1) ** 2 + (y - y1) ** 2)) - w12 = z * torch.exp(-((x - x1) ** 2 + (y - y2) ** 2)) - w21 = z * torch.exp(-((x - x2) ** 2 + (y - y1) ** 2)) - w22 = z * torch.exp(-((x - x2) ** 2 + (y - y2) ** 2)) - - return w11, w12, w21, w22 - - def sample_one(self, img, shiftx, shifty, weight): - """ - Input: - -img (N, C, H, W) - -shiftx, shifty (N, c, H, W) - """ - - N, C, H, W = img.size() - - # flatten all (all restored as Tensors) - flat_shiftx = shiftx.view(-1) - flat_shifty = shifty.view(-1) - flat_basex = torch.arange(0, H, requires_grad=False).view(-1, 1)[None, None].to(self.device).long().repeat(N, C,1,W).view(-1) - flat_basey = torch.arange(0, W, requires_grad=False).view(1, -1)[None, None].to(self.device).long().repeat(N, C,H,1).view(-1) - flat_weight = weight.view(-1) - flat_img = img.contiguous().view(-1) - - # The corresponding positions in I1 - idxn = torch.arange(0, N, requires_grad=False).view(N, 1, 1, 1).to(self.device).long().repeat(1, C, H, W).view(-1) - idxc = torch.arange(0, C, requires_grad=False).view(1, C, 1, 1).to(self.device).long().repeat(N, 1, H, W).view(-1) - idxx = flat_shiftx.long() + flat_basex - idxy = flat_shifty.long() + flat_basey - - # recording the inside part the shifted - mask = idxx.ge(0) & idxx.lt(H) & idxy.ge(0) & idxy.lt(W) - - # Mask off points out of boundaries - ids = (idxn * C * H * W + idxc * H * W + idxx * W + idxy) - ids_mask = torch.masked_select(ids, mask).clone().to(self.device) - - # Note here! accmulate fla must be true for proper bp - img_warp = torch.zeros([N * C * H * W, ]).to(self.device) - img_warp.put_(ids_mask, torch.masked_select(flat_img * flat_weight, mask), accumulate=True) - - one_warp = torch.zeros([N * C * H * W, ]).to(self.device) - one_warp.put_(ids_mask, torch.masked_select(flat_weight, mask), accumulate=True) - - return img_warp.view(N, C, H, W), one_warp.view(N, C, H, W) - -class RefineUNet(nn.Module): - def __init__(self, args): - super(RefineUNet, self).__init__() - self.args = args - self.scale = args.module_scale_factor - self.nf = args.nf - self.conv1 = nn.Conv2d(self.nf, self.nf, [3,3], 1, [1,1]) - self.conv2 = nn.Conv2d(self.nf, self.nf, [3,3], 1, [1,1]) - self.lrelu = nn.ReLU() - self.NN = nn.UpsamplingNearest2d(scale_factor=2) - - self.enc1 = nn.Conv2d((4*self.nf)//self.scale//self.scale + 4*args.img_ch + 4, self.nf, [4, 4], 2, [1, 1]) - self.enc2 = nn.Conv2d(self.nf, 2*self.nf, [4, 4], 2, [1, 1]) - self.enc3 = nn.Conv2d(2*self.nf, 4*self.nf, [4, 4], 2, [1, 1]) - self.dec0 = nn.Conv2d(4*self.nf, 4*self.nf, [3, 3], 1, [1, 1]) - self.dec1 = nn.Conv2d(4*self.nf + 2*self.nf, 2*self.nf, [3, 3], 1, [1, 1]) ## input concatenated with enc2 - self.dec2 = nn.Conv2d(2*self.nf + self.nf, self.nf, [3, 3], 1, [1, 1]) ## input concatenated with enc1 - self.dec3 = nn.Conv2d(self.nf, 1+args.img_ch, [3, 3], 1, [1, 1]) ## input added with warped image - - def forward(self, concat): - enc1 = self.lrelu(self.enc1(concat)) - enc2 = self.lrelu(self.enc2(enc1)) - out = self.lrelu(self.enc3(enc2)) - - out = self.lrelu(self.dec0(out)) - out = self.NN(out) - - out = torch.cat((out,enc2),dim=1) - out = self.lrelu(self.dec1(out)) - - out = self.NN(out) - out = torch.cat((out,enc1),dim=1) - out = self.lrelu(self.dec2(out)) - - out = self.NN(out) - out = self.dec3(out) - return out - -class ResBlock2D_3D(nn.Module): - ## Shape of input [B,C,T,H,W] - ## Shape of output [B,C,T,H,W] - def __init__(self, args): - super(ResBlock2D_3D, self).__init__() - self.args = args - self.nf = args.nf - - self.conv3x3_1 = nn.Conv3d(self.nf, self.nf, [1,3,3], 1, [0,1,1]) - self.conv3x3_2 = nn.Conv3d(self.nf, self.nf, [1,3,3], 1, [0,1,1]) - self.lrelu = nn.ReLU() - - def forward(self, x): - ''' - x shape : [B,C,T,H,W] - ''' - B, C, T, H, W = x.size() - - out = self.conv3x3_2(self.lrelu(self.conv3x3_1(x))) - - return x + out - -class RResBlock2D_3D(nn.Module): - - def __init__(self, args, T_reduce_flag=False): - super(RResBlock2D_3D, self).__init__() - self.args = args - self.nf = args.nf - self.T_reduce_flag = T_reduce_flag - self.resblock1 = ResBlock2D_3D(self.args) - self.resblock2 = ResBlock2D_3D(self.args) - if T_reduce_flag: - self.reduceT_conv = nn.Conv3d(self.nf, self.nf, [3,1,1], 1, [0,0,0]) - - def forward(self, x): - ''' - x shape : [B,C,T,H,W] - ''' - out = self.resblock1(x) - out = self.resblock2(out) - if self.T_reduce_flag: - return self.reduceT_conv(out + x) - else: - return out + x - -def weights_init(m): - classname = m.__class__.__name__ - if (classname.find('Conv2d') != -1) or (classname.find('Conv3d') != -1): - init.xavier_normal_(m.weight) - # init.kaiming_normal_(m.weight, nonlinearity='relu') - if hasattr(m, 'bias') and m.bias is not None: - init.zeros_(m.bias) \ No newline at end of file diff --git a/imports/ComfyUI_Frame_Interpolation/vfi_utils.py b/imports/ComfyUI_Frame_Interpolation/vfi_utils.py index 7504637..9c762b0 100644 --- a/imports/ComfyUI_Frame_Interpolation/vfi_utils.py +++ b/imports/ComfyUI_Frame_Interpolation/vfi_utils.py @@ -24,7 +24,7 @@ else: raise Exception("config.yaml file is neccessary, plz recreate the config file by downloading it from https://github.com/Fannovel16/ComfyUI-Frame-Interpolation") DEVICE = get_torch_device() -class InterpolationStateList(): +class InterpolationStateListImport(): def __init__(self, frame_indices: typing.List[int], is_skip_list: bool): self.frame_indices = frame_indices @@ -35,7 +35,7 @@ class InterpolationStateList(): return self.is_skip_list and is_frame_in_list or not self.is_skip_list and not is_frame_in_list -class MakeInterpolationStateList: +class MakeInterpolationStateListImport: @classmethod def INPUT_TYPES(s): return { @@ -52,7 +52,7 @@ class MakeInterpolationStateList: def create_options(self, frame_indices: str, is_skip_list: bool): frame_indices_list = [int(item) for item in frame_indices.split(',')] - interpolation_state_list = InterpolationStateList( + interpolation_state_list = InterpolationStateListImport( frame_indices=frame_indices_list, is_skip_list=is_skip_list, ) @@ -127,7 +127,7 @@ def _generic_frame_loop( multiplier: typing.Union[typing.SupportsInt, typing.List], return_middle_frame_function, *return_middle_frame_function_args, - interpolation_states: InterpolationStateList = None, + interpolation_states: InterpolationStateListImport = None, use_timestep=True, dtype=torch.float16, final_logging=True): @@ -213,7 +213,7 @@ def generic_frame_loop( multiplier: typing.Union[typing.SupportsInt, typing.List], return_middle_frame_function, *return_middle_frame_function_args, - interpolation_states: InterpolationStateList = None, + interpolation_states: InterpolationStateListImport = None, use_timestep=True, dtype=torch.float32): @@ -255,7 +255,7 @@ def generic_frame_loop( return output_frames raise NotImplementedError(f"multipiler of {type(multiplier)}") -class FloatToInt: +class FloatToIntImport: @classmethod def INPUT_TYPES(s): return {