From 69a4aebfe72b2271d2b4aa219e3ccb56949725bf Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Thu, 12 Feb 2026 18:51:29 +0100 Subject: [PATCH] Add auto_pyr_level toggle to select pyramid level by resolution When enabled (default), automatically picks the optimal pyr_level based on input height: <540p=3, 540p=5, 1080p=6, 4K=7. When disabled, uses the manual pyr_level value. Co-Authored-By: Claude Opus 4.6 --- inference.py | 18 ++++++++++++++++-- nodes.py | 12 +++++++++--- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/inference.py b/inference.py index a550c20..d0cdcbb 100644 --- a/inference.py +++ b/inference.py @@ -5,8 +5,9 @@ from .bim_vfi_arch import BiMVFI class BiMVFIModel: """Clean inference wrapper around BiMVFI for ComfyUI integration.""" - def __init__(self, checkpoint_path, pyr_level=3, device="cpu"): + def __init__(self, checkpoint_path, pyr_level=3, auto_pyr_level=True, device="cpu"): self.pyr_level = pyr_level + self.auto_pyr_level = auto_pyr_level self.device = device self.model = BiMVFI(pyr_level=pyr_level, feat_channels=32) @@ -58,12 +59,25 @@ class BiMVFIModel: img0 = frame0.to(device) img1 = frame1.to(device) + if self.auto_pyr_level: + _, _, h, _ = img0.shape + if h >= 2160: + pyr_level = 7 + elif h >= 1080: + pyr_level = 6 + elif h >= 540: + pyr_level = 5 + else: + pyr_level = 3 + else: + pyr_level = self.pyr_level + time_step_tensor = torch.tensor([time_step], device=device).view(1, 1, 1, 1) result_dict = self.model( img0=img0, img1=img1, time_step=time_step_tensor, - pyr_level=self.pyr_level, + pyr_level=pyr_level, ) interp = result_dict["imgt_pred"] diff --git a/nodes.py b/nodes.py index 4f325a9..5df0eb3 100644 --- a/nodes.py +++ b/nodes.py @@ -57,9 +57,13 @@ class LoadBIMVFIModel: "default": MODEL_FILENAME, "tooltip": "Checkpoint file from models/bim-vfi/. Auto-downloads on first use if missing.", }), + "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.", + }), "pyr_level": ("INT", { "default": 3, "min": 3, "max": 7, "step": 1, - "tooltip": "Pyramid levels for coarse-to-fine processing. More levels = captures larger motion but slower. Recommended: 3-5 for <540p, 5-6 for 1080p, 6-7 for 4K.", + "tooltip": "Manual pyramid levels for coarse-to-fine processing. Only used when auto_pyr_level is disabled. More levels = captures larger motion but slower.", }), } } @@ -69,7 +73,7 @@ class LoadBIMVFIModel: FUNCTION = "load_model" CATEGORY = "video/BIM-VFI" - def load_model(self, model_path, pyr_level): + def load_model(self, model_path, auto_pyr_level, pyr_level): full_path = os.path.join(MODEL_DIR, model_path) if not os.path.exists(full_path): @@ -79,10 +83,12 @@ class LoadBIMVFIModel: wrapper = BiMVFIModel( checkpoint_path=full_path, pyr_level=pyr_level, + auto_pyr_level=auto_pyr_level, device="cpu", ) - logger.info(f"BIM-VFI model loaded (pyr_level={pyr_level})") + mode = "auto" if auto_pyr_level else f"manual ({pyr_level})" + logger.info(f"BIM-VFI model loaded (pyr_level={mode})") return (wrapper,)