Upload salia_special_rife_batch.py

salia_special_rife_batch.py

@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import os
+import sys
+import importlib
+import threading
+from typing import Dict, Optional, Tuple
+
+import torch
+
+# -----------------------------------------------------------------------------
+# Lazy import wrapper around rife_lazy.py (which itself lazy-loads Frame-Interpolation)
+# -----------------------------------------------------------------------------
+_RIFE_IMPORT_LOCK = threading.Lock()
+_RIFE_LAZY_MOD = None
+_RIFE_VFI_CLASS = None
+
+
+def _lazy_import_rife_lazy_module():
+    """
+    Import rife_lazy.py only when this node is executed.
+    rife_lazy.py is expected to be in the same folder as this file.
+    """
+    global _RIFE_LAZY_MOD
+    if _RIFE_LAZY_MOD is not None:
+        return _RIFE_LAZY_MOD
+
+    with _RIFE_IMPORT_LOCK:
+        if _RIFE_LAZY_MOD is not None:
+            return _RIFE_LAZY_MOD
+
+        this_dir = os.path.dirname(os.path.abspath(__file__))
+        if this_dir not in sys.path:
+            sys.path.insert(0, this_dir)
+
+        _RIFE_LAZY_MOD = importlib.import_module("rife_lazy")
+        return _RIFE_LAZY_MOD
+
+
+def _lazy_get_rife_vfi_class():
+    """
+    Uses rife_lazy._lazy_get_rife_class() to obtain the real RIFE_VFI class.
+    """
+    global _RIFE_VFI_CLASS
+    if _RIFE_VFI_CLASS is not None:
+        return _RIFE_VFI_CLASS
+
+    mod = _lazy_import_rife_lazy_module()
+
+    with _RIFE_IMPORT_LOCK:
+        if _RIFE_VFI_CLASS is not None:
+            return _RIFE_VFI_CLASS
+        _RIFE_VFI_CLASS = mod._lazy_get_rife_class()
+        return _RIFE_VFI_CLASS
+
+
+def _unwrap_image_output(result):
+    """
+    Many ComfyUI nodes return tuples. We only want the IMAGE output.
+    """
+    if isinstance(result, (tuple, list)):
+        return result[0]
+    return result
+
+
+def _normalize_to_batch(images: torch.Tensor) -> torch.Tensor:
+    """
+    Accept [H,W,C] and convert to [1,H,W,C].
+    """
+    if isinstance(images, torch.Tensor) and images.dim() == 3:
+        return images.unsqueeze(0)
+    return images
+
+
+def _torch_inference_context():
+    """
+    Use torch.inference_mode if available, otherwise fall back to torch.no_grad.
+    """
+    return torch.inference_mode() if hasattr(torch, "inference_mode") else torch.no_grad()
+
+
+# -----------------------------------------------------------------------------
+# Node
+# -----------------------------------------------------------------------------
+class SALIA_SPECIAL_BATCH_RIFE47:
+    """
+    Input:  IMAGE batch
+
+    Pipeline:
+      1) Custom_Batch_Output split:
+         - Batch_UP = [7] + [9..25] + [27..31] + [33..36]
+         - Rife_x3  = [37, 4]
+      2) Rife_x3 -> RIFE(rife47, mult=3) => RifeOutput
+      3) almost_final_batch = concat(Batch_UP, RifeOutput)
+
+      4) Extra inserts (indices refer to almost_final_batch *before* inserts):
+         - DUO_14_15:      interpolate (14,15) with mult=3 => keep middle 2 frames => insert after 14
+         - SINGLE_25_26:   interpolate (25,26) with mult=2 => keep middle 1 frame => insert after 25
+         - SINGLE_30_1:    interpolate (30, 1) with mult=2 => keep middle 1 frame => insert after 30
+
+    Output: IMAGE batch
+    """
+
+    CATEGORY = "salia_online/VFI"
+    FUNCTION = "process"
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("IMAGE",)
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"images": ("IMAGE",)}}
+
+    # ---- Custom_Batch_Output indices (0-based) ----
+    _BATCH_UP_INDICES = (
+        [7]
+        + list(range(9, 26))    # 9..25
+        + list(range(27, 32))   # 27..31
+        + list(range(33, 37))   # 33..36
+    )
+    _RIFE_X3_INDICES = [37, 4]
+
+    # ---- Insert specs (based on almost_final_batch BEFORE inserts) ----
+    _DUO_PAIR: Tuple[int, int] = (14, 15)
+    _DUO_MULT: int = 3
+
+    _SINGLE_25_26_PAIR: Tuple[int, int] = (25, 26)
+    _SINGLE_25_26_MULT: int = 2
+
+    _SINGLE_30_1_PAIR: Tuple[int, int] = (30, 1)
+    _SINGLE_30_1_MULT: int = 2
+
+    def _call_rife(self, rife_node, rife_lazy_mod, frames: torch.Tensor, multiplier: int) -> torch.Tensor:
+        """
+        Call the underlying RIFE_VFI node with the hardcoded params from rife_lazy.py.
+        """
+        result = rife_node.vfi(
+            ckpt_name=getattr(rife_lazy_mod, "_HARDCODED_CKPT_NAME", "rife47.pth"),
+            frames=frames,
+            clear_cache_after_n_frames=getattr(rife_lazy_mod, "_HARDCODED_CLEAR_CACHE_AFTER_N_FRAMES", 10),
+            multiplier=int(multiplier),
+            fast_mode=getattr(rife_lazy_mod, "_HARDCODED_FAST_MODE", True),
+            ensemble=getattr(rife_lazy_mod, "_HARDCODED_ENSEMBLE", True),
+            scale_factor=getattr(rife_lazy_mod, "_HARDCODED_SCALE_FACTOR", 1.0),
+            optional_interpolation_states=None,
+        )
+        out = _unwrap_image_output(result)
+        if not isinstance(out, torch.Tensor):
+            raise TypeError("RIFE output was not a torch.Tensor")
+        out = _normalize_to_batch(out)
+        if out.dim() != 4:
+            raise ValueError(f"RIFE output must be [B,H,W,C], got shape: {tuple(out.shape)}")
+        return out
+
+    def _rife_middle_frames(
+        self,
+        rife_node,
+        rife_lazy_mod,
+        base: torch.Tensor,
+        i: int,
+        j: int,
+        multiplier: int,
+    ) -> Optional[torch.Tensor]:
+        """
+        Interpolate between base[i] and base[j], then discard endpoints and return only middle frames:
+          - mult=3 => output len 4 => middle 2 frames
+          - mult=2 => output len 3 => middle 1 frame
+        """
+        b = int(base.shape[0])
+        if i < 0 or j < 0 or i >= b or j >= b:
+            return None
+
+        device = base.device
+        idx = torch.tensor([i, j], dtype=torch.long, device=device)
+        pair = torch.index_select(base, 0, idx)
+
+        out = self._call_rife(rife_node, rife_lazy_mod, pair, multiplier)
+
+        # keep only middle frames
+        if out.shape[0] < 3:
+            return None
+        mid = out[1:-1]
+        if mid.shape[0] == 0:
+            return None
+        return mid
+
+    def _insert_after_indices(self, base: torch.Tensor, inserts: Dict[int, torch.Tensor]) -> torch.Tensor:
+        """
+        inserts maps "base_index -> batch_of_frames_to_insert_after_that_index"
+        This avoids index-shift confusion because inserts are applied relative to the original base.
+        """
+        parts = []
+        for i in range(int(base.shape[0])):
+            parts.append(base[i : i + 1])
+            extra = inserts.get(i, None)
+            if extra is not None:
+                extra = _normalize_to_batch(extra)
+                parts.append(extra)
+        return torch.cat(parts, dim=0)
+
+    def process(self, images):
+        # Safety / validation
+        if not isinstance(images, torch.Tensor):
+            return (images,)
+
+        images = _normalize_to_batch(images)
+        if images.dim() != 4:
+            return (images,)
+
+        b = int(images.shape[0])
+        # Must be able to address index 37 (needs B >= 38)
+        if b < 38:
+            return (images,)
+
+        device = images.device
+
+        # Custom_Batch_Output split
+        idx_up = torch.tensor(self._BATCH_UP_INDICES, dtype=torch.long, device=device)
+        idx_rife = torch.tensor(self._RIFE_X3_INDICES, dtype=torch.long, device=device)
+
+        batch_up = torch.index_select(images, 0, idx_up)
+        rife_x3 = torch.index_select(images, 0, idx_rife)
+
+        # Lazy-load RIFE
+        rife_lazy_mod = _lazy_import_rife_lazy_module()
+        RIFE_VFI = _lazy_get_rife_vfi_class()
+        rife_node = RIFE_VFI()
+
+        with _torch_inference_context():
+            # Rife_x3 -> RIFE mult=3
+            rife_output = self._call_rife(rife_node, rife_lazy_mod, rife_x3, multiplier=3)
+
+            # almost_final_batch = Batch_UP + RifeOutput
+            almost_final = torch.cat([batch_up, rife_output], dim=0)
+
+            # Need indices up to 30 and also index 1 for the last interpolation
+            if int(almost_final.shape[0]) <= 30:
+                return (almost_final,)
+
+            # Build inserts from the original almost_final (pre-insert)
+            inserts: Dict[int, torch.Tensor] = {}
+
+            duo_mid = self._rife_middle_frames(
+                rife_node, rife_lazy_mod, almost_final,
+                i=self._DUO_PAIR[0], j=self._DUO_PAIR[1], multiplier=self._DUO_MULT
+            )
+            if duo_mid is not None:
+                inserts[self._DUO_PAIR[0]] = duo_mid  # insert after first index of the pair
+
+            single_25_26_mid = self._rife_middle_frames(
+                rife_node, rife_lazy_mod, almost_final,
+                i=self._SINGLE_25_26_PAIR[0], j=self._SINGLE_25_26_PAIR[1], multiplier=self._SINGLE_25_26_MULT
+            )
+            if single_25_26_mid is not None:
+                inserts[self._SINGLE_25_26_PAIR[0]] = single_25_26_mid
+
+            single_30_1_mid = self._rife_middle_frames(
+                rife_node, rife_lazy_mod, almost_final,
+                i=self._SINGLE_30_1_PAIR[0], j=self._SINGLE_30_1_PAIR[1], multiplier=self._SINGLE_30_1_MULT
+            )
+            if single_30_1_mid is not None:
+                inserts[self._SINGLE_30_1_PAIR[0]] = single_30_1_mid
+
+            # Apply inserts without index-shift mistakes
+            final_batch = self._insert_after_indices(almost_final, inserts)
+            return (final_batch,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "SALIA_SPECIAL_BATCH_RIFE47": SALIA_SPECIAL_BATCH_RIFE47,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SALIA_SPECIAL_BATCH_RIFE47": "Special Batch + RIFE Inserts (rife47)",
+}