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 <noreply@anthropic.com>
This commit is contained in:
18
inference.py
18
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"]
|
||||
|
||||
Reference in New Issue
Block a user