From e644cb101532d19f18db19808f623f33927c52c7 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 16 Aug 2026 00:19:56 +0200 Subject: [PATCH] fix: add BIM-VFI artifact-safe mode --- README.md | 5 +- bim_vfi_arch/bim_vfi.py | 8 +- bim_vfi_arch/sn.py | 14 +++- .../tween_speed_bim_model_lab.json | 4 +- inference.py | 17 ++-- nodes.py | 19 ++++- pyproject.toml | 2 +- tests/test_bim_vfi.py | 82 +++++++++++++++++++ 8 files changed, 134 insertions(+), 17 deletions(-) create mode 100644 tests/test_bim_vfi.py diff --git a/README.md b/README.md index ec257be..078034f 100644 --- a/README.md +++ b/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 diff --git a/bim_vfi_arch/bim_vfi.py b/bim_vfi_arch/bim_vfi.py index fe3ca7b..2acb6ca 100644 --- a/bim_vfi_arch/bim_vfi.py +++ b/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) diff --git a/bim_vfi_arch/sn.py b/bim_vfi_arch/sn.py index 9aa268a..af99f6a 100644 --- a/bim_vfi_arch/sn.py +++ b/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 = {} diff --git a/example_workflows/tween_speed_bim_model_lab.json b/example_workflows/tween_speed_bim_model_lab.json index 37c75fe..f77a664 100644 --- a/example_workflows/tween_speed_bim_model_lab.json +++ b/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" }, diff --git a/inference.py b/inference.py index c9d71c9..f16e126 100644 --- a/inference.py +++ b/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() diff --git a/nodes.py b/nodes.py index e28cd76..2890736 100644 --- a/nodes.py +++ b/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,) diff --git a/pyproject.toml b/pyproject.toml index 06b6c64..8fe4de1 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/tests/test_bim_vfi.py b/tests/test_bim_vfi.py new file mode 100644 index 0000000..7efe1e0 --- /dev/null +++ b/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]