Upload batch_loop_clean_rife_fill.py

batch_loop_clean_rife_fill.py

@@ -0,0 +1,568 @@
+from __future__ import annotations
+
+import os
+import sys
+import importlib
+import threading
+from typing import List, Tuple, Optional
+
+import torch
+import torch.nn.functional as F
+
+
+# ============================================================
+# Hardcoded HYBRID compare settings (exactly as requested)
+# ============================================================
+_DOWNSCALE_MAX = 256
+_BLUR_SIGMA = 1.2
+_HIST_BINS = 32
+_SCALE = 1000.0
+
+_W_PIXEL = 1.00
+_W_SSIM  = 1.00
+_W_EDGE  = 0.50
+_W_HIST  = 0.20
+
+
+# ============================================================
+# Lazy RIFE import (mirrors your wrapper behavior)
+# ============================================================
+_IMPORT_LOCK = threading.Lock()
+_RIFE_CLASS = None
+
+_HARDCODED_CKPT_NAME = "rife47.pth"
+_HARDCODED_CLEAR_CACHE_AFTER_N_FRAMES = 10
+_HARDCODED_FAST_MODE = True
+_HARDCODED_ENSEMBLE = True
+_HARDCODED_SCALE_FACTOR = 1.0
+
+
+def _lazy_get_rife_class():
+    """
+    Lazily import ComfyUI-Frame-Interpolation's RIFE_VFI class.
+    Expected folder:
+      ComfyUI/custom_nodes/ComfyUI-Frame-Interpolation
+    """
+    global _RIFE_CLASS
+    if _RIFE_CLASS is not None:
+        return _RIFE_CLASS
+
+    with _IMPORT_LOCK:
+        if _RIFE_CLASS is not None:
+            return _RIFE_CLASS
+
+        this_dir = os.path.dirname(os.path.abspath(__file__))
+        custom_nodes_dir = os.path.abspath(os.path.join(this_dir, ".."))
+        cfi_dir = os.path.join(custom_nodes_dir, "ComfyUI-Frame-Interpolation")
+
+        if not os.path.isdir(cfi_dir):
+            raise FileNotFoundError(
+                f"Could not find ComfyUI-Frame-Interpolation folder at:\n  {cfi_dir}\n"
+                f"Expected it at:\n  {os.path.join(custom_nodes_dir, 'ComfyUI-Frame-Interpolation')}"
+            )
+
+        if cfi_dir not in sys.path:
+            sys.path.insert(0, cfi_dir)
+
+        rife_mod = importlib.import_module("vfi_models.rife")
+        rife_cls = getattr(rife_mod, "RIFE_VFI", None)
+        if rife_cls is None:
+            raise ImportError("vfi_models.rife imported, but RIFE_VFI class was not found.")
+
+        _RIFE_CLASS = rife_cls
+        return _RIFE_CLASS
+
+
+def _run_rife(frames_bhwc: torch.Tensor, multiplier: int) -> torch.Tensor:
+    """
+    frames_bhwc: [2,H,W,C]
+    returns: [multiplier+1, H, W, C] (usually includes originals at ends)
+    """
+    RIFE_VFI = _lazy_get_rife_class()
+    rife_node = RIFE_VFI()
+
+    out = rife_node.vfi(
+        ckpt_name=_HARDCODED_CKPT_NAME,
+        frames=frames_bhwc,
+        clear_cache_after_n_frames=_HARDCODED_CLEAR_CACHE_AFTER_N_FRAMES,
+        multiplier=int(multiplier),
+        fast_mode=_HARDCODED_FAST_MODE,
+        ensemble=_HARDCODED_ENSEMBLE,
+        scale_factor=_HARDCODED_SCALE_FACTOR,
+        optional_interpolation_states=None,
+    )
+
+    # Some versions may return (IMAGE,) or (IMAGE, states). We only want the IMAGE.
+    if isinstance(out, (tuple, list)):
+        return out[0]
+    return out
+
+
+# ============================================================
+# Image helpers
+# ============================================================
+
+def _bhwc_to_nchw(img: torch.Tensor) -> torch.Tensor:
+    if img.dim() != 4:
+        raise ValueError(f"Expected IMAGE tensor [B,H,W,C], got {tuple(img.shape)}")
+    return img.permute(0, 3, 1, 2).contiguous()
+
+def _nchw_to_bhwc(img: torch.Tensor) -> torch.Tensor:
+    if img.dim() != 4:
+        raise ValueError(f"Expected NCHW tensor [B,C,H,W], got {tuple(img.shape)}")
+    return img.permute(0, 2, 3, 1).contiguous()
+
+def _drop_alpha_if_any(x: torch.Tensor) -> torch.Tensor:
+    if x.shape[1] > 3:
+        return x[:, :3, :, :].contiguous()
+    return x
+
+def _ensure_3ch(x: torch.Tensor) -> torch.Tensor:
+    if x.shape[1] == 1:
+        return x.repeat(1, 3, 1, 1)
+    return x
+
+def _to_luma(x: torch.Tensor) -> torch.Tensor:
+    if x.shape[1] == 1:
+        return x
+    r = x[:, 0:1, :, :]
+    g = x[:, 1:2, :, :]
+    b = x[:, 2:3, :, :]
+    return (0.2989 * r + 0.5870 * g + 0.1140 * b)
+
+def _resize_max(x: torch.Tensor, max_size: int) -> torch.Tensor:
+    if max_size <= 0:
+        return x
+    b, c, h, w = x.shape
+    m = max(h, w)
+    if m <= max_size:
+        return x
+    scale = max_size / float(m)
+    nh = max(1, int(round(h * scale)))
+    nw = max(1, int(round(w * scale)))
+    return F.interpolate(x, size=(nh, nw), mode="bilinear", align_corners=False)
+
+def _gaussian_blur(x: torch.Tensor, sigma: float) -> torch.Tensor:
+    if sigma <= 0:
+        return x
+
+    radius = int(max(1, round(3.0 * sigma)))
+    ksize = 2 * radius + 1
+    device = x.device
+    dtype = x.dtype
+
+    coords = torch.arange(-radius, radius + 1, device=device, dtype=dtype)
+    kernel1d = torch.exp(-(coords * coords) / (2.0 * sigma * sigma))
+    kernel1d = kernel1d / (kernel1d.sum() + 1e-12)
+
+    c = x.shape[1]
+    kh = kernel1d.view(1, 1, 1, ksize).repeat(c, 1, 1, 1)
+    kv = kernel1d.view(1, 1, ksize, 1).repeat(c, 1, 1, 1)
+
+    out = F.conv2d(x, kh, padding=(0, radius), groups=c)
+    out = F.conv2d(out, kv, padding=(radius, 0), groups=c)
+    return out
+
+def _sobel_edges(y: torch.Tensor) -> torch.Tensor:
+    device = y.device
+    dtype = y.dtype
+    c = y.shape[1]
+
+    kx = torch.tensor(
+        [[-1, 0, 1],
+         [-2, 0, 2],
+         [-1, 0, 1]],
+        device=device, dtype=dtype
+    ) / 8.0
+
+    ky = torch.tensor(
+        [[-1, -2, -1],
+         [ 0,  0,  0],
+         [ 1,  2,  1]],
+        device=device, dtype=dtype
+    ) / 8.0
+
+    kx = kx.view(1, 1, 3, 3).repeat(c, 1, 1, 1)
+    ky = ky.view(1, 1, 3, 3).repeat(c, 1, 1, 1)
+
+    gx = F.conv2d(y, kx, padding=1, groups=c)
+    gy = F.conv2d(y, ky, padding=1, groups=c)
+    return torch.sqrt(gx * gx + gy * gy + 1e-12)
+
+
+# ============================================================
+# SSIM (vectorized for batch of pairs)
+# ============================================================
+
+def _make_ssim_kernel(device, dtype, window_size: int = 11, sigma: float = 1.5):
+    radius = window_size // 2
+    coords = torch.arange(window_size, device=device, dtype=dtype) - radius
+    g = torch.exp(-(coords * coords) / (2.0 * sigma * sigma))
+    g = g / (g.sum() + 1e-12)
+    w2d = (g[:, None] * g[None, :]).view(1, 1, window_size, window_size)
+    return w2d, radius
+
+def _ssim_batch_luma(x: torch.Tensor, y: torch.Tensor, w2d: torch.Tensor, radius: int) -> torch.Tensor:
+    """
+    x,y: [N,1,H,W]
+    returns: [N] ssim values
+    """
+    C1 = (0.01) ** 2
+    C2 = (0.03) ** 2
+
+    mu_x = F.conv2d(x, w2d, padding=radius, groups=1)
+    mu_y = F.conv2d(y, w2d, padding=radius, groups=1)
+
+    mu_x2 = mu_x * mu_x
+    mu_y2 = mu_y * mu_y
+    mu_xy = mu_x * mu_y
+
+    sigma_x2 = F.conv2d(x * x, w2d, padding=radius, groups=1) - mu_x2
+    sigma_y2 = F.conv2d(y * y, w2d, padding=radius, groups=1) - mu_y2
+    sigma_xy = F.conv2d(x * y, w2d, padding=radius, groups=1) - mu_xy
+
+    num = (2.0 * mu_xy + C1) * (2.0 * sigma_xy + C2)
+    den = (mu_x2 + mu_y2 + C1) * (sigma_x2 + sigma_y2 + C2)
+    ssim_map = num / (den + 1e-12)
+    return ssim_map.mean(dim=[1, 2, 3])
+
+
+# ============================================================
+# Histogram (per-frame) + chi2 between frames
+# ============================================================
+
+def _compute_histograms(rgb_resized: torch.Tensor, bins: int) -> torch.Tensor:
+    """
+    rgb_resized: [B,3,H,W] in [0,1]
+    returns hist: [B,3,bins] normalized
+    Uses torch.histc. If device histc fails, falls back to CPU.
+    """
+    eps = 1e-12
+    B = rgb_resized.shape[0]
+    device = rgb_resized.device
+
+    try:
+        h = torch.zeros((B, 3, bins), device=device, dtype=torch.float32)
+        for i in range(B):
+            for c in range(3):
+                hc = torch.histc(rgb_resized[i, c], bins=bins, min=0.0, max=1.0)
+                hc = hc / (hc.sum() + eps)
+                h[i, c] = hc
+        return h
+    except Exception:
+        rgb_cpu = rgb_resized.detach().float().cpu()
+        h_cpu = torch.zeros((B, 3, bins), device="cpu", dtype=torch.float32)
+        for i in range(B):
+            for c in range(3):
+                hc = torch.histc(rgb_cpu[i, c], bins=bins, min=0.0, max=1.0)
+                hc = hc / (hc.sum() + eps)
+                h_cpu[i, c] = hc
+        return h_cpu.to(device)
+
+
+def _chi2_from_hist(h1: torch.Tensor, h2: torch.Tensor) -> torch.Tensor:
+    """
+    h1,h2: [...,3,bins]
+    returns: [...] chi2 distance averaged across channels
+    """
+    eps = 1e-12
+    diff2 = (h1 - h2) ** 2
+    denom = (h1 + h2 + eps)
+    chi = 0.5 * torch.sum(diff2 / denom, dim=-1)  # sum over bins -> [...,3]
+    return torch.mean(chi, dim=-1)                # avg over channels -> [...]
+
+
+# ============================================================
+# Preprocess + HYBRID scores
+# ============================================================
+
+class _Pre:
+    def __init__(self, rgb_resized, rgb_blur, luma_blur, edges, hist, w2d, radius):
+        self.rgb_resized = rgb_resized  # [B,3,h,w]
+        self.rgb_blur = rgb_blur        # [B,3,h,w]
+        self.luma_blur = luma_blur      # [B,1,h,w]
+        self.edges = edges              # [B,1,h,w]
+        self.hist = hist                # [B,3,bins]
+        self.w2d = w2d
+        self.radius = radius
+
+def _preprocess(images_bhwc: torch.Tensor) -> _Pre:
+    x = _bhwc_to_nchw(images_bhwc)
+    x = _drop_alpha_if_any(x).clamp(0.0, 1.0)
+    x = _ensure_3ch(x)
+
+    rgb_resized = _resize_max(x, _DOWNSCALE_MAX)
+    rgb_blur = _gaussian_blur(rgb_resized, _BLUR_SIGMA)
+    luma_blur = _to_luma(rgb_blur)
+    edges = _sobel_edges(luma_blur)
+
+    hist = _compute_histograms(rgb_resized, _HIST_BINS)
+    w2d, radius = _make_ssim_kernel(device=luma_blur.device, dtype=luma_blur.dtype)
+    return _Pre(rgb_resized, rgb_blur, luma_blur, edges, hist, w2d, radius)
+
+def _hybrid_scores_adj(pre: _Pre) -> torch.Tensor:
+    """
+    returns scores for adjacent pairs: [B-1] (scaled by _SCALE)
+    """
+    B = pre.rgb_blur.shape[0]
+    if B <= 1:
+        return torch.zeros((0,), device=pre.rgb_blur.device, dtype=torch.float32)
+
+    # Pixel MAE on blurred RGB
+    pix = torch.mean(torch.abs(pre.rgb_blur[:-1] - pre.rgb_blur[1:]), dim=[1, 2, 3])  # [B-1]
+
+    # SSIM diff on blurred luma
+    ssim_vals = _ssim_batch_luma(pre.luma_blur[:-1], pre.luma_blur[1:], pre.w2d, pre.radius)  # [B-1]
+    ssim_diff = torch.clamp(1.0 - ssim_vals, min=0.0)
+
+    # Edge MAE
+    ed = torch.mean(torch.abs(pre.edges[:-1] - pre.edges[1:]), dim=[1, 2, 3])
+
+    # Hist chi2
+    hist = _chi2_from_hist(pre.hist[:-1], pre.hist[1:])  # [B-1]
+
+    score = (_W_PIXEL * pix) + (_W_SSIM * ssim_diff) + (_W_EDGE * ed) + (_W_HIST * hist)
+    return score * _SCALE
+
+def _hybrid_score_pair(pre: _Pre, i: int, j: int) -> float:
+    pix = torch.mean(torch.abs(pre.rgb_blur[i] - pre.rgb_blur[j]))
+    ssim_val = _ssim_batch_luma(pre.luma_blur[i:i+1], pre.luma_blur[j:j+1], pre.w2d, pre.radius)[0]
+    ssim_diff = torch.clamp(1.0 - ssim_val, min=0.0)
+    ed = torch.mean(torch.abs(pre.edges[i] - pre.edges[j]))
+    hist = _chi2_from_hist(pre.hist[i:i+1], pre.hist[j:j+1])[0]
+    score = (_W_PIXEL * pix) + (_W_SSIM * ssim_diff) + (_W_EDGE * ed) + (_W_HIST * hist)
+    return float(score.item() * _SCALE)
+
+def _hybrid_scores_to_anchor(pre: _Pre, anchor_idx: int, cand_indices: List[int]) -> torch.Tensor:
+    """
+    returns [N] scores (scaled) between anchor and each candidate
+    """
+    device = pre.rgb_blur.device
+    if len(cand_indices) == 0:
+        return torch.zeros((0,), device=device, dtype=torch.float32)
+
+    idx = torch.tensor(cand_indices, device=device, dtype=torch.long)
+
+    # gather candidates
+    rgb_c = pre.rgb_blur.index_select(0, idx)     # [N,3,h,w]
+    luma_c = pre.luma_blur.index_select(0, idx)   # [N,1,h,w]
+    edge_c = pre.edges.index_select(0, idx)       # [N,1,h,w]
+    hist_c = pre.hist.index_select(0, idx)        # [N,3,bins]
+
+    rgb_a = pre.rgb_blur[anchor_idx].unsqueeze(0).expand_as(rgb_c)
+    luma_a = pre.luma_blur[anchor_idx].unsqueeze(0).expand_as(luma_c)
+    edge_a = pre.edges[anchor_idx].unsqueeze(0).expand_as(edge_c)
+    hist_a = pre.hist[anchor_idx].unsqueeze(0).expand_as(hist_c)
+
+    pix = torch.mean(torch.abs(rgb_c - rgb_a), dim=[1, 2, 3])  # [N]
+    ssim_vals = _ssim_batch_luma(luma_a, luma_c, pre.w2d, pre.radius)
+    ssim_diff = torch.clamp(1.0 - ssim_vals, min=0.0)
+    ed = torch.mean(torch.abs(edge_c - edge_a), dim=[1, 2, 3])
+    hist = _chi2_from_hist(hist_c, hist_a)
+
+    score = (_W_PIXEL * pix) + (_W_SSIM * ssim_diff) + (_W_EDGE * ed) + (_W_HIST * hist)
+    return score * _SCALE
+
+
+# ============================================================
+# The Node
+# ============================================================
+
+class LoopCleanRifeFill51:
+    """
+    1) Remove frozen tail
+    2) Remove frozen frames across whole batch (dedup pass)
+    3) Crop to looping segment [anchor .. best_end]
+    4) Repeatedly insert RIFE interpolated frames into highest-diff adjacent gap
+    5) Stop at target_frames
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+
+                # you explicitly wanted this configurable
+                "loop_anchor": ("INT", {"default": 9, "min": 0, "max": 4096, "step": 1}),
+
+                # tail search window for loop end matching
+                "loop_tail_search": ("INT", {"default": 15, "min": 1, "max": 512, "step": 1}),
+
+                # keep as input (default 3.0), since you might tune this per dataset
+                "freeze_threshold": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 1000.0, "step": 0.1}),
+
+                # final length
+                "target_frames": ("INT", {"default": 51, "min": 1, "max": 1000, "step": 1}),
+
+                # multiplier behavior (you mainly asked 2 or 3)
+                "max_multiplier": ("INT", {"default": 3, "min": 2, "max": 8, "step": 1}),
+                "big_gap_threshold": ("FLOAT", {"default": 20.0, "min": 0.0, "max": 10000.0, "step": 0.5}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("images",)
+    FUNCTION = "process"
+    CATEGORY = "image/analysis"
+
+    def process(
+        self,
+        images: torch.Tensor,
+        loop_anchor: int,
+        loop_tail_search: int,
+        freeze_threshold: float,
+        target_frames: int,
+        max_multiplier: int,
+        big_gap_threshold: float,
+    ):
+        # ---------------------------
+        # Basic sanity
+        # ---------------------------
+        if images.dim() != 4:
+            raise ValueError(f"Expected IMAGE [B,H,W,C], got {tuple(images.shape)}")
+        B = images.shape[0]
+        if B <= 1:
+            # If it's a single frame, just repeat to target (rare / your "should never happen")
+            if target_frames > 1:
+                images = images.repeat(target_frames, 1, 1, 1)
+            return (images[:target_frames],)
+
+        # =====================================================
+        # 1) Remove frozen tail
+        # =====================================================
+        pre = _preprocess(images)
+        scores_adj = _hybrid_scores_adj(pre)  # [B-1]
+        keep_last = images.shape[0] - 1
+        while keep_last > 0 and float(scores_adj[keep_last - 1].item()) < freeze_threshold:
+            keep_last -= 1
+        images = images[: keep_last + 1]
+
+        if images.shape[0] <= 1:
+            if target_frames > 1:
+                images = images.repeat(target_frames, 1, 1, 1)
+            return (images[:target_frames],)
+
+        # =====================================================
+        # 2) Remove frozen frames across entire batch (dedup)
+        #    Remove only ONE of a frozen pair => drop the later one.
+        # =====================================================
+        pre = _preprocess(images)
+        keep: List[int] = [0]
+        last_kept = 0
+        for i in range(1, images.shape[0]):
+            sc = _hybrid_score_pair(pre, last_kept, i)
+            if sc >= freeze_threshold:
+                keep.append(i)
+                last_kept = i
+
+        keep_t = torch.tensor(keep, device=images.device, dtype=torch.long)
+        images = images.index_select(0, keep_t)
+
+        if images.shape[0] <= 1:
+            if target_frames > 1:
+                images = images.repeat(target_frames, 1, 1, 1)
+            return (images[:target_frames],)
+
+        # =====================================================
+        # 3) Crop to looping segment using anchor + closest end
+        # =====================================================
+        L = images.shape[0]
+        anchor = int(max(0, min(loop_anchor, L - 1)))
+
+        # Candidates are from the last N frames, but must be > anchor.
+        tail_start = max(anchor + 1, L - int(loop_tail_search))
+        if tail_start <= L - 1:
+            cand = list(range(L - 1, tail_start - 1, -1))  # reverse from end
+            pre = _preprocess(images)
+            scores = _hybrid_scores_to_anchor(pre, anchor, cand)  # [N]
+            best_k = int(torch.argmin(scores).item())
+            end_idx = cand[best_k]
+            if end_idx <= anchor:
+                # fallback: keep from anchor to end
+                images = images[anchor:]
+            else:
+                images = images[anchor : end_idx + 1]
+        else:
+            # No candidates after anchor; fallback: keep from anchor to end
+            images = images[anchor:]
+
+        if images.shape[0] <= 1:
+            if target_frames > 1:
+                images = images.repeat(target_frames, 1, 1, 1)
+            return (images[:target_frames],)
+
+        # =====================================================
+        # 4+5) Insert RIFE interpolations into highest gap until target_frames
+        # =====================================================
+        # Clamp max_multiplier (at least 2)
+        max_multiplier = int(max(2, max_multiplier))
+
+        safety = 0
+        while images.shape[0] < target_frames:
+            safety += 1
+            if safety > 500:
+                # Prevent infinite loops in pathological cases
+                break
+
+            n = images.shape[0]
+            if n < 2:
+                break
+
+            pre = _preprocess(images)
+            scores_adj = _hybrid_scores_adj(pre)  # [n-1]
+            if scores_adj.numel() == 0:
+                break
+
+            # Highest-diff adjacent pair
+            idx = int(torch.argmax(scores_adj).item())
+            max_score = float(scores_adj[idx].item())
+
+            remaining = target_frames - n
+
+            # Choose multiplier (mostly 2, sometimes 3+ if gap is large and we have room)
+            m = 2
+            if remaining >= 2 and max_multiplier >= 3 and max_score >= big_gap_threshold:
+                m = 3
+
+            # If we still have lots of room, allow higher multipliers up to max_multiplier
+            # (optional, but useful if the batch got really short)
+            # Inserts (m-1) frames.
+            if remaining >= 3 and max_multiplier > 3 and max_score >= big_gap_threshold:
+                # try to use as much as we can without overshooting
+                m = min(max_multiplier, remaining + 1)
+
+            # Never overshoot target
+            if (m - 1) > remaining:
+                m = remaining + 1
+            m = int(max(2, m))
+
+            # Run RIFE on the pair (batch of 2)
+            pair = images[idx : idx + 2]  # [2,H,W,C]
+            rife_out = _run_rife(pair, multiplier=m)  # [m+1,H,W,C] typically
+
+            # Take only the inserted frames (exclude first and last originals)
+            inserted = rife_out[1:-1]  # [m-1,H,W,C]
+            if inserted.shape[0] == 0:
+                # fallback: if something weird happens, just stop
+                break
+
+            # If we would overshoot due to some mismatch, clamp inserted
+            if inserted.shape[0] > remaining:
+                inserted = inserted[:remaining]
+
+            # Insert into the batch between idx and idx+1
+            images = torch.cat([images[:idx+1], inserted, images[idx+1:]], dim=0)
+
+        # If we somehow overshot (shouldn't), clamp
+        images = images[:target_frames]
+        return (images,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "LoopCleanRifeFill51": LoopCleanRifeFill51,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "LoopCleanRifeFill51": "Loop Clean + RIFE Fill to 51 (Hybrid hardcoded)",
+}