File size: 4,494 Bytes
82ea528 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
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), ) |