115 lines
5.4 KiB
Python
115 lines
5.4 KiB
Python
"""PromptForge Frame Interpolate: RIFE / FILM frame interpolation in one node that backs off instead of crashing.
|
|
|
|
It uses ComfyUI's own interpolation models and node (no other custom node pack is needed). The checkpoint list holds
|
|
the models in `models/frame_interpolation` (RIFE 4.25+ and FILM; older RIFE .pth files such as rife49 are not supported by ComfyUI's loader). When the whole clip does not fit
|
|
in memory, the clip is interpolated in shorter overlapping runs (down to a few frames), and what is still too much is
|
|
cross-faded instead of failing.
|
|
"""
|
|
import gc
|
|
|
|
import psutil
|
|
import torch
|
|
|
|
import comfy.model_management as model_management
|
|
import comfy.model_patcher
|
|
import comfy.storage
|
|
import comfy.utils
|
|
import folder_paths
|
|
from comfy_extras.nodes_frame_interpolation import FrameInterpolate, FrameInterpolationModelLoader
|
|
|
|
BLEND = "blend (no model, instant)"
|
|
NO_CKPT = "(no frame interpolation model found)"
|
|
MIN_RUN = 4
|
|
RUN = 24
|
|
|
|
|
|
def _checkpoints():
|
|
return {name: folder_paths.get_full_path("frame_interpolation", name) for name in folder_paths.get_filename_list("frame_interpolation")}
|
|
|
|
|
|
def _load(path):
|
|
state = comfy.utils.load_torch_file(path, safe_load=True)
|
|
model = FrameInterpolationModelLoader._detect_and_load(state)
|
|
dtype = torch.float16 if model_management.should_use_fp16(model_management.get_torch_device()) else torch.float32
|
|
model.eval().to(dtype)
|
|
return comfy.model_patcher.CoreModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device(),
|
|
fast_disk=comfy.storage.state_dict_fast_disk(state))
|
|
|
|
|
|
def _blend(frames, multiplier):
|
|
out = [frames[:1]]
|
|
for index in range(frames.shape[0] - 1):
|
|
steps = torch.arange(1, multiplier + 1, dtype=frames.dtype).view(-1, 1, 1, 1) / multiplier
|
|
out.append(torch.lerp(frames[index:index + 1], frames[index + 1:index + 2], steps))
|
|
return torch.cat(out)
|
|
|
|
|
|
class PROMPTFORGE_FrameInterpolate:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
names = list(_checkpoints()) or [NO_CKPT]
|
|
preferred = next((name for name in names if "rife" in name.lower()), names[0])
|
|
return {
|
|
"required": {
|
|
"frames": ("IMAGE",),
|
|
"multiplier": ("INT", {"default": 2, "min": 2, "max": 8, "tooltip": "2 doubles the frame count, 3 triples it, and so on."}),
|
|
"ckpt_name": (names + [BLEND], {"default": preferred, "tooltip": "RIFE or FILM model. blend is a plain cross-fade: instant, no model, ghosting on motion."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "interpolate"
|
|
CATEGORY = "🔮 PromptForge/Video"
|
|
DESCRIPTION = "RIFE / FILM frame interpolation with ComfyUI's own models that interpolates in shorter runs, or cross-fades, instead of running out of memory."
|
|
|
|
def interpolate(self, frames, multiplier, ckpt_name):
|
|
count = int(frames.shape[0])
|
|
if count < 2:
|
|
return (frames,)
|
|
frames = frames[..., :3]
|
|
total = (count - 1) * multiplier + 1
|
|
available = psutil.virtual_memory().available
|
|
values = total * frames.shape[1] * frames.shape[2] * 3
|
|
dtype = torch.float32 if values * 4 < 0.1 * available else torch.float16
|
|
if values * 2 > 0.9 * available:
|
|
raise ValueError(f"[PromptForge] the interpolated clip needs {values * 2 / 1024 ** 3:.1f} GB of RAM even in half precision and only {available / 1024 ** 3:.1f} GB is free. Use a lower multiplier or fewer frames.")
|
|
out = torch.empty((total, *frames.shape[1:]), dtype=dtype)
|
|
if ckpt_name == BLEND:
|
|
out.copy_(_blend(frames, multiplier))
|
|
return (out,)
|
|
paths = _checkpoints()
|
|
if ckpt_name not in paths:
|
|
raise ValueError("[PromptForge] no frame interpolation model was found. Put a RIFE 4.25+ or FILM model in models/frame_interpolation (for example rife_v4.26.safetensors), or choose blend.")
|
|
patcher = _load(paths[ckpt_name])
|
|
|
|
start, span, filled = 0, min(count, RUN), 0
|
|
pbar = comfy.utils.ProgressBar(count - 1)
|
|
while start < count - 1:
|
|
model_management.throw_exception_if_processing_interrupted()
|
|
end = min(count, start + span)
|
|
run = frames[start:end]
|
|
try:
|
|
result = FrameInterpolate.execute(patcher, run, multiplier).args[0]
|
|
except model_management.OOM_EXCEPTION:
|
|
gc.collect()
|
|
model_management.soft_empty_cache()
|
|
if span > MIN_RUN:
|
|
span = max(MIN_RUN, span // 2)
|
|
print(f"[PromptForge] frame interpolation ran out of memory; continuing in runs of {span} frames.")
|
|
continue
|
|
print("[PromptForge] frame interpolation ran out of memory even in short runs; cross-fading the rest.")
|
|
result = _blend(run, multiplier)
|
|
result = result if filled == 0 else result[1:]
|
|
out[filled:filled + len(result)] = result
|
|
filled += len(result)
|
|
pbar.update(end - 1 - start)
|
|
start = end - 1
|
|
del patcher
|
|
gc.collect()
|
|
model_management.soft_empty_cache()
|
|
return (out,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"PROMPTFORGE_FrameInterpolate": PROMPTFORGE_FrameInterpolate}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"PROMPTFORGE_FrameInterpolate": "PromptForge Frame Interpolate (low memory)"}
|