|
|
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): |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
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), ) |
|
|
|
|
|
|
