165 lines
6.2 KiB
Python
165 lines
6.2 KiB
Python
# nodes.py -- ComfyUI-Lee-RIFE
|
|
#
|
|
# GPU-batched RIFE frame interpolation as a self-contained ComfyUI node.
|
|
#
|
|
# Why batched: the common RIFE node pushes frame-pairs through ONE AT A TIME,
|
|
# copies each pair CPU<->GPU and flushes the CUDA cache every few frames. On a
|
|
# fast card the RIFE forward is tiny, so the GPU idles between those serial
|
|
# steps (low GPU utilisation). This node batches the pairs and skips the
|
|
# per-frame flush, keeping the GPU busy. It also enables RIFE's real
|
|
# bidirectional "ensemble" (averaging a forward and a time-reversed pass).
|
|
#
|
|
# The architecture (rife_arch.py) and the weights (models/rife49.pth) come from
|
|
# RIFE / Practical-RIFE and ComfyUI-Frame-Interpolation, both MIT-licensed.
|
|
# See THIRD-PARTY-NOTICES.md.
|
|
|
|
import os
|
|
import torch
|
|
|
|
from .rife_arch import IFNet
|
|
|
|
try:
|
|
from comfy.model_management import get_torch_device, soft_empty_cache
|
|
except Exception: # allows importing/testing outside ComfyUI
|
|
def get_torch_device():
|
|
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
def soft_empty_cache(*a, **k):
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
MODELS_DIR = os.path.join(HERE, "models")
|
|
|
|
# RIFE checkpoint -> architecture version
|
|
CKPT_ARCH = {
|
|
"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",
|
|
}
|
|
|
|
|
|
def _available_ckpts():
|
|
found = [n for n in CKPT_ARCH if os.path.isfile(os.path.join(MODELS_DIR, n))]
|
|
# Fallback so the node still registers even before the weights are in place.
|
|
return found or ["rife49.pth"]
|
|
|
|
|
|
class LeeRIFEGPU:
|
|
_cache = {} # (ckpt_name, dtype_str) -> loaded model
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
ckpts = _available_ckpts()
|
|
default = "rife49.pth" if "rife49.pth" in ckpts else ckpts[0]
|
|
return {
|
|
"required": {
|
|
"frames": ("IMAGE",),
|
|
"ckpt_name": (ckpts, {"default": default}),
|
|
"multiplier": ("INT", {"default": 2, "min": 1, "max": 16}),
|
|
"ensemble": ("BOOLEAN", {"default": True}),
|
|
"scale_factor": ([0.25, 0.5, 1.0, 2.0, 4.0], {"default": 1.0}),
|
|
"precision": (["fp16", "fp32"], {"default": "fp16"}),
|
|
"batch_size": ("INT", {"default": 16, "min": 1, "max": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "interpolate"
|
|
CATEGORY = "Lee/RIFE"
|
|
|
|
def _get_model(self, ckpt_name, dtype):
|
|
key = (ckpt_name, str(dtype))
|
|
if key in self._cache:
|
|
return self._cache[key]
|
|
path = os.path.join(MODELS_DIR, ckpt_name)
|
|
if not os.path.isfile(path):
|
|
raise FileNotFoundError(
|
|
f"{ckpt_name} not found in {MODELS_DIR}. "
|
|
f"Place the RIFE weights file there (see models/HOW-TO-ADD-WEIGHTS.txt)."
|
|
)
|
|
arch = CKPT_ARCH.get(ckpt_name, "4.7")
|
|
model = IFNet(arch_ver=arch)
|
|
try:
|
|
sd = torch.load(path, map_location="cpu", weights_only=True)
|
|
except Exception:
|
|
sd = torch.load(path, map_location="cpu", weights_only=False)
|
|
if isinstance(sd, dict) and "state_dict" in sd:
|
|
sd = sd["state_dict"]
|
|
model.load_state_dict(sd)
|
|
model.eval().to(get_torch_device())
|
|
if dtype == torch.float16:
|
|
model.half()
|
|
print(f"[ComfyUI-Lee-RIFE] loaded {ckpt_name} (arch {arch}) on "
|
|
f"{get_torch_device()} / {dtype}")
|
|
self._cache[key] = model
|
|
return model
|
|
|
|
@torch.no_grad()
|
|
def interpolate(self, frames, ckpt_name, multiplier, ensemble,
|
|
scale_factor, precision, batch_size):
|
|
device = get_torch_device()
|
|
dtype = torch.float16 if precision == "fp16" else torch.float32
|
|
model = self._get_model(ckpt_name, dtype)
|
|
|
|
N = frames.shape[0]
|
|
if N < 2 or multiplier < 2:
|
|
return (frames,)
|
|
|
|
# ComfyUI IMAGE is (N,H,W,C) in [0,1]; RIFE wants (N,C,H,W).
|
|
# Stays on CPU here; only the per-batch chunks are moved to the GPU.
|
|
seq = frames.permute(0, 3, 1, 2).contiguous()
|
|
img0_all = seq[:-1]
|
|
img1_all = seq[1:]
|
|
P = img0_all.shape[0]
|
|
scale_list = [8 / scale_factor, 4 / scale_factor,
|
|
2 / scale_factor, 1 / scale_factor]
|
|
|
|
# One batched pass per intermediate timestep k/multiplier.
|
|
mids = [] # list over k, each tensor (P,C,H,W) on CPU / float32
|
|
for k in range(1, multiplier):
|
|
t = k / multiplier
|
|
parts = []
|
|
s = 0
|
|
chunk = int(batch_size)
|
|
while s < P:
|
|
a = img0_all[s:s + chunk].to(device, dtype)
|
|
b = img1_all[s:s + chunk].to(device, dtype)
|
|
try:
|
|
r = model(a, b, timestep=t, scale_list=list(scale_list),
|
|
training=False, fastmode=True, ensemble=ensemble)
|
|
except RuntimeError as e:
|
|
del a, b
|
|
if "out of memory" not in str(e).lower():
|
|
raise
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
if chunk == 1:
|
|
raise
|
|
chunk = max(1, chunk // 2) # auto-shrink and retry
|
|
continue
|
|
parts.append(r.float().clamp(0, 1).cpu())
|
|
del a, b, r
|
|
s += chunk
|
|
mids.append(torch.cat(parts, 0))
|
|
|
|
# Interleave originals + midpoints, pad to multiplier*N output frames.
|
|
seq_cpu = seq.float().cpu()
|
|
out = []
|
|
for i in range(N - 1):
|
|
out.append(seq_cpu[i])
|
|
for k in range(multiplier - 1):
|
|
out.append(mids[k][i])
|
|
out.append(seq_cpu[N - 1])
|
|
while len(out) < multiplier * N:
|
|
out.append(seq_cpu[N - 1])
|
|
|
|
result = torch.stack(out, 0).permute(0, 2, 3, 1).contiguous() # (M*N,H,W,C)
|
|
soft_empty_cache()
|
|
return (result[..., :3],)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"LeeRIFEGPU": LeeRIFEGPU}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"LeeRIFEGPU": "RIFE GPU Fast (Lee)"}
|