fix: add BIM-VFI artifact-safe mode

README.md

@@ -117,8 +117,11 @@ Loads the BiM-VFI checkpoint. Auto-downloads from Google Drive on first use to `
 | Input | Description |
 |-------|-------------|
 | **model_path** | Checkpoint from `models/bim-vfi/` |
-| **auto_pyr_level** | Auto pyramid level by resolution (<540p=3, 540p=5, 1080p=6, 4K=7) |
+| **auto_pyr_level** | Official automatic pyramid policy (below 1080p=5, 1080p=6, 4K=7) |
 | **pyr_level** | Manual pyramid level (3–7), used when auto is off |
+| **artifact_safe_mode** | Disables the RGB refinement residual to suppress wrong-edge/halo artifacts caused by flow misalignment in blurry or large-motion shots. Off preserves official behavior and can retain more detail on easy shots |
+
+`artifact_safe_mode` implements the workaround recommended by the official BIM-VFI maintainer for [wrong-edge artifacts caused by severely misaligned warped inputs](https://github.com/KAIST-VICLab/BiM-VFI/issues/1). Enable it selectively for affected footage.
 
 #### BIM-VFI Interpolate
 

bim_vfi_arch/bim_vfi.py

@@ -12,13 +12,17 @@ from ..utils.padder import InputPadder
 
 
 class BiMVFI(nn.Module):
-    def __init__(self, pyr_level=3, feat_channels=32, **kwargs):
+    def __init__(self, pyr_level=3, feat_channels=32,
+                 artifact_safe_mode=False, **kwargs):
         super(BiMVFI, self).__init__()
         self.pyr_level = pyr_level
         self.mfe = ResNetPyramid(feat_channels)
         self.cfe = ResNetPyramid(feat_channels)
         self.bimfn = BiMFN(feat_channels)
-        self.sn = SynthesisNetwork(feat_channels)
+        self.sn = SynthesisNetwork(
+            feat_channels,
+            use_rgb_refine_residual=not artifact_safe_mode,
+        )
         self.feat_channels = feat_channels
         self.caun = CAUN(feat_channels)
 

bim_vfi_arch/sn.py

@@ -5,8 +5,9 @@ from .backwarp import backwarp
 
 
 class SynthesisNetwork(nn.Module):
-    def __init__(self, feat_channels):
+    def __init__(self, feat_channels, use_rgb_refine_residual=True):
         super(SynthesisNetwork, self).__init__()
+        self.use_rgb_refine_residual = use_rgb_refine_residual
         input_channels = 6 + 1
         self.conv_down1 = nn.Sequential(
             nn.Conv2d(input_channels, feat_channels, 7, padding=3),
@@ -59,6 +60,13 @@ class SynthesisNetwork(nn.Module):
             warped_img1 = backwarp(i1, flow_t1)
             return warped_img0, warped_img1, warped_c0, warped_c1
 
+    def merge_warped_images(self, warped_img0, warped_img1,
+                            blending_mask, refine_res):
+        merged_img = warped_img0 * blending_mask + warped_img1 * (1 - blending_mask)
+        if self.use_rgb_refine_residual:
+            merged_img = merged_img + refine_res
+        return merged_img
+
     def forward(self, i0, i1, c0_pyr, c1_pyr, bi_flow_pyr, occ):
         warped_img0, warped_img1, warped_c0, warped_c1 = \
             self.get_warped_representations(
@@ -82,7 +90,9 @@ class SynthesisNetwork(nn.Module):
         occ_res = refine[:, 3:]
         occ_out = occ + occ_res
         blending_mask = torch.sigmoid(occ_out)
-        merged_img = (warped_img0 * blending_mask + warped_img1 * (1 - blending_mask)) + refine_res
+        merged_img = self.merge_warped_images(
+            warped_img0, warped_img1, blending_mask, refine_res
+        )
         interp_img = merged_img
 
         extra_dict = {}

example_workflows/tween_speed_bim_model_lab.json

@@ -161,7 +161,7 @@
       "id": 6,
       "type": "LoadBIMVFIModel",
       "pos": [390, 790],
-      "size": [300, 125],
+      "size": [300, 150],
       "flags": {},
       "order": 5,
       "mode": 0,
@@ -173,7 +173,7 @@
         "aux_id": "ComfyUI-Tween.git",
         "Node name for S&R": "LoadBIMVFIModel"
       },
-      "widgets_values": ["bim_vfi.pth", true, 3],
+      "widgets_values": ["bim_vfi.pth", true, 3, false],
       "color": "#28384a",
       "bgcolor": "#36506b"
     },

inference.py

@@ -17,12 +17,18 @@ logger = logging.getLogger("Tween")
 class BiMVFIModel:
     """Clean inference wrapper around BiMVFI for ComfyUI integration."""
 
-    def __init__(self, checkpoint_path, pyr_level=3, auto_pyr_level=True, device="cpu"):
+    def __init__(self, checkpoint_path, pyr_level=3, auto_pyr_level=True,
+                 artifact_safe_mode=False, device="cpu"):
         self.pyr_level = pyr_level
         self.auto_pyr_level = auto_pyr_level
+        self.artifact_safe_mode = artifact_safe_mode
         self.device = device
 
-        self.model = BiMVFI(pyr_level=pyr_level, feat_channels=32)
+        self.model = BiMVFI(
+            pyr_level=pyr_level,
+            feat_channels=32,
+            artifact_safe_mode=artifact_safe_mode,
+        )
         self._load_checkpoint(checkpoint_path)
         self.model.eval()
         self.model.to(device)
@@ -61,10 +67,11 @@ class BiMVFIModel:
                 return 7
             elif h >= 1080:
                 return 6
-            elif h >= 540:
-                return 5
             else:
-                return 3
+                # Match the official video inference path. Level 3 can miss
+                # large motion even at low resolutions; it remains available
+                # as a manual speed/quality tradeoff.
+                return 5
         return self.pyr_level
 
     @torch.no_grad()

nodes.py

@@ -289,13 +289,19 @@ class LoadBIMVFIModel:
                 }),
                 "auto_pyr_level": ("BOOLEAN", {
                     "default": True,
-                    "tooltip": "Automatically select pyramid level based on input resolution: <540p=3, 540p=5, 1080p=6, 4K=7. Disable to use manual pyr_level.",
+                    "tooltip": "Use the official pyramid policy: below 1080p=5, 1080p=6, 4K=7. Disable to use manual pyr_level.",
                 }),
                 "pyr_level": ("INT", {
                     "default": 3, "min": 3, "max": 7, "step": 1,
                     "tooltip": "Manual pyramid levels for coarse-to-fine processing. Only used when auto_pyr_level is disabled. More levels = captures larger motion but slower.",
                 }),
-            }
+            },
+            "optional": {
+                "artifact_safe_mode": ("BOOLEAN", {
+                    "default": False,
+                    "tooltip": "Disable BIM-VFI's RGB refinement residual. This can remove wrong-edge/halo artifacts when optical flow is misaligned by blur or large motion, but may reduce detail on easy shots.",
+                }),
+            },
         }
 
     RETURN_TYPES = ("BIM_VFI_MODEL",)
@@ -303,7 +309,8 @@ class LoadBIMVFIModel:
     FUNCTION = "load_model"
     CATEGORY = "video/BIM-VFI"
 
-    def load_model(self, model_path, auto_pyr_level, pyr_level):
+    def load_model(self, model_path, auto_pyr_level, pyr_level,
+                   artifact_safe_mode=False):
         full_path = os.path.join(MODEL_DIR, model_path)
 
         if not os.path.exists(full_path):
@@ -314,11 +321,15 @@ class LoadBIMVFIModel:
             checkpoint_path=full_path,
             pyr_level=pyr_level,
             auto_pyr_level=auto_pyr_level,
+            artifact_safe_mode=artifact_safe_mode,
             device="cpu",
         )
 
         mode = "auto" if auto_pyr_level else f"manual ({pyr_level})"
-        logger.info(f"BIM-VFI model loaded (pyr_level={mode})")
+        synthesis = "artifact-safe" if artifact_safe_mode else "official"
+        logger.info(
+            f"BIM-VFI model loaded (pyr_level={mode}, synthesis={synthesis})"
+        )
         return (wrapper,)
 
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-tween"
 description = "Video frame interpolation nodes for ComfyUI using BIM-VFI, EMA-VFI, SGM-VFI, GIMM-VFI, SPEED, and LDF-VFI."
-version = "1.2.0"
+version = "1.2.1"
 license = "Apache-2.0"
 requires-python = ">=3.10"
 dependencies = [

tests/test_bim_vfi.py

@@ -0,0 +1,82 @@
+import importlib
+from pathlib import Path
+import sys
+import types
+
+import torch
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "_tween_bim_tests"
+
+
+def _load_inference_module():
+    package = types.ModuleType(PACKAGE_NAME)
+    package.__path__ = [str(REPO_ROOT)]
+    package.__package__ = PACKAGE_NAME
+    sys.modules.setdefault(PACKAGE_NAME, package)
+    return importlib.import_module(f"{PACKAGE_NAME}.inference")
+
+
+def test_auto_pyramid_levels_match_official_video_inference():
+    model_class = _load_inference_module().BiMVFIModel
+    model = model_class.__new__(model_class)
+    model.auto_pyr_level = True
+
+    assert model._get_pyr_level(240) == 5
+    assert model._get_pyr_level(539) == 5
+    assert model._get_pyr_level(720) == 5
+    assert model._get_pyr_level(1079) == 5
+    assert model._get_pyr_level(1080) == 6
+    assert model._get_pyr_level(2159) == 6
+    assert model._get_pyr_level(2160) == 7
+
+
+def test_manual_pyramid_level_remains_available():
+    model_class = _load_inference_module().BiMVFIModel
+    model = model_class.__new__(model_class)
+    model.auto_pyr_level = False
+    model.pyr_level = 3
+
+    assert model._get_pyr_level(720) == 3
+
+
+def test_artifact_safe_mode_removes_only_rgb_refinement_residual():
+    module = _load_inference_module()
+    official = module.BiMVFI(pyr_level=3, feat_channels=1)
+    artifact_safe = module.BiMVFI(
+        pyr_level=3,
+        feat_channels=1,
+        artifact_safe_mode=True,
+    )
+
+    assert official.sn.use_rgb_refine_residual is True
+    assert artifact_safe.sn.use_rgb_refine_residual is False
+
+    official_keys = official.state_dict().keys()
+    artifact_safe_keys = artifact_safe.state_dict().keys()
+    assert official_keys == artifact_safe_keys
+
+    warped0 = torch.full((1, 3, 2, 2), 0.2)
+    warped1 = torch.full((1, 3, 2, 2), 0.8)
+    mask = torch.full((1, 1, 2, 2), 0.25)
+    residual = torch.full((1, 3, 2, 2), 0.1)
+    blend = warped0 * mask + warped1 * (1 - mask)
+
+    assert torch.equal(
+        official.sn.merge_warped_images(warped0, warped1, mask, residual),
+        blend + residual,
+    )
+    assert torch.equal(
+        artifact_safe.sn.merge_warped_images(warped0, warped1, mask, residual),
+        blend,
+    )
+
+
+def test_demo_exposes_artifact_safe_mode_without_enabling_it():
+    import json
+
+    workflow_path = REPO_ROOT / "example_workflows" / "tween_speed_bim_model_lab.json"
+    workflow = json.loads(workflow_path.read_text(encoding="utf-8"))
+    loader = next(node for node in workflow["nodes"] if node["type"] == "LoadBIMVFIModel")
+
+    assert loader["widgets_values"] == ["bim_vfi.pth", True, 3, False]