mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-22 03:51:21 +08:00
85 lines
3.1 KiB
Python
85 lines
3.1 KiB
Python
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,)
|