Add tiling, batch inference, interpolation states and multiplier schedule

Ported and adapted from Fannovel16/ComfyUI-Frame-Interpolation:

- gc.collect() after soft_empty_cache() to reduce RAM fragmentation
- Explicit float32 cast in postprocess_frames for fp16 engine safety
- InterpolationStateList class + MakeInterpolationStateList node to
  skip/keep specific frame pairs during interpolation
- Per-pair multiplier schedule (comma-separated list) overriding uniform
  multiplier, with automatic padding to match the number of pairs
- Automatic tiling with cosine feathering when frame dimensions exceed
  the TensorRT engine optimization profile (fixes 1440x2560 with max=1920)
- Padding with replicate mode when frame dimensions are below profile min
- get_input_profile_bounds() on Engine to retrieve min/opt/max shapes
- Batch inference: process multiple (pair, timestep) tasks per GPU call
  with dynamic batch axis in ONNX export and engine profile
- Backward compatible: all new features are opt-in via optional inputs

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

__init__.py

@@ -5,7 +5,7 @@ import subprocess
 import re
 from pathlib import Path
 from comfy.model_management import get_torch_device
-from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife, logger
+from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife, logger, InterpolationStateList, infer_tiled
 from .trt_utilities import Engine
 from .utilities import download_file, ColoredLogger
 import folder_paths
@@ -383,6 +383,11 @@ class AutoLoadRifeTensorrtModel:
                 "model": (model_options, {"default": model_default, "tooltip": model_tooltip}),
                 "precision": (precision_options, {"default": precision_default, "tooltip": precision_tooltip}),
                 "resolution_profile": (profile_options, {"default": profile_default, "tooltip": profile_tooltip}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 16, "step": 1,
+                                       "tooltip": "Max batch size for the engine profile. "
+                                                  "Higher values allow batched inference (multiple frame pairs per GPU call) "
+                                                  "for better throughput, but use more VRAM and require a rebuild. "
+                                                  "Set to 1 for the original behaviour."}),
             },
             "optional": {
                 "custom_config": ("RIFE_RESOLUTION_CONFIG", {"tooltip": "Custom resolution config (used when profile='custom')"}),
@@ -395,7 +400,7 @@ class AutoLoadRifeTensorrtModel:
     DESCRIPTION = "Load RIFE tensorrt models, they will be built automatically if not found."
     FUNCTION = "load_rife_tensorrt_model"
 
-    def load_rife_tensorrt_model(self, model, precision, resolution_profile, custom_config=None):
+    def load_rife_tensorrt_model(self, model, precision, resolution_profile, batch_size=1, custom_config=None):
         tensorrt_models_dir = os.path.join(folder_paths.models_dir, "tensorrt", "rife")
         onnx_models_dir = os.path.join(folder_paths.models_dir, "onnx")
 
@@ -426,7 +431,9 @@ class AutoLoadRifeTensorrtModel:
 
         # Build tensorrt model path with detailed naming (includes profile)
         engine_channel = 3
-        engine_min_batch, engine_opt_batch, engine_max_batch = 1, 1, 1
+        engine_min_batch = 1
+        engine_opt_batch = max(1, batch_size)
+        engine_max_batch = max(1, batch_size)
         engine_min_h, engine_opt_h, engine_max_h = dim_min, dim_opt, dim_max
         engine_min_w, engine_opt_w, engine_max_w = dim_min, dim_opt, dim_max
         tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{model}_{precision}_{profile_name}_{engine_min_batch}x{engine_channel}x{engine_min_h}x{engine_min_w}_{engine_opt_batch}x{engine_channel}x{engine_opt_h}x{engine_opt_w}_{engine_max_batch}x{engine_channel}x{engine_max_h}x{engine_max_w}_{tensorrt.__version__}.trt")
@@ -463,6 +470,38 @@ class AutoLoadRifeTensorrtModel:
 
         return (engine,)
 
+
+class MakeInterpolationStateList:
+    """Build an InterpolationStateList to skip or keep specific frame pairs.
+    Ported from Fannovel16/ComfyUI-Frame-Interpolation.
+    """
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "frame_indices": ("STRING", {"multiline": True, "default": "1,2,3",
+                                             "tooltip": "Comma-separated list of frame-pair indices (0-based). "
+                                                        "Pair i is formed by frame i and frame i+1."}),
+                "is_skip_list": ("BOOLEAN", {"default": True,
+                                             "tooltip": "True = skip interpolation for the listed pairs. "
+                                                        "False = interpolate ONLY the listed pairs."}),
+            },
+        }
+
+    RETURN_TYPES = ("INTERPOLATION_STATES",)
+    FUNCTION = "create_options"
+    CATEGORY = "⚡️ TensorRT/RIFE"
+    DESCRIPTION = "Build a skip/keep list controlling which frame pairs get interpolated by RIFE."
+
+    def create_options(self, frame_indices: str, is_skip_list: bool):
+        frame_indices_list = [int(item.strip()) for item in frame_indices.split(',') if item.strip() != '']
+        interpolation_state_list = InterpolationStateList(
+            frame_indices=frame_indices_list,
+            is_skip_list=is_skip_list,
+        )
+        return (interpolation_state_list,)
+
+
 class AutoRifeTensorrt:
     @classmethod
     def INPUT_TYPES(s):
@@ -471,8 +510,19 @@ class AutoRifeTensorrt:
                 "frames": ("IMAGE", {"tooltip": "Input frames for video frame interpolation"}),
                 "rife_trt_model": ("RIFE_TRT_MODEL", {"tooltip": "Tensorrt model built and loaded"}),
                 "clear_cache_after_n_frames": ("INT", {"default": 100, "min": 1, "max": 1000, "tooltip": "Clear CUDA cache after processing this many frames"}),
-                "multiplier": ("INT", {"default": 2, "min": 1, "tooltip": "Frame interpolation multiplier"}),
+                "multiplier": ("INT", {"default": 2, "min": 1, "tooltip": "Frame interpolation multiplier (uniform). Ignored if multiplier_list is connected."}),
                 "keep_model_loaded": ("BOOLEAN", {"default": False, "tooltip": "Keep model loaded in memory after processing"}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 16,
+                                       "tooltip": "Number of interpolation tasks per GPU call. "
+                                                  "Higher values improve throughput but use more VRAM. "
+                                                  "Must be <= the batch_size used when building the engine. "
+                                                  "Set to 1 for the safest behaviour."}),
+            },
+            "optional": {
+                "multiplier_list": ("STRING", {"multiline": True, "default": "",
+                                               "tooltip": "Per-pair multiplier schedule, comma-separated (e.g. '2,2,4,4,2'). "
+                                                          "When connected, overrides 'multiplier'. Missing trailing pairs default to 2."}),
+                "interpolation_states": ("INTERPOLATION_STATES", {"tooltip": "Optional skip/keep list from MakeInterpolationStateList"}),
             },
         }
 
@@ -488,6 +538,9 @@ class AutoRifeTensorrt:
         clear_cache_after_n_frames=100,
         multiplier=2,
         keep_model_loaded=False,
+        batch_size=1,
+        multiplier_list="",
+        interpolation_states=None,
     ):
         B, H, W, C = frames.shape
         shape_dict = {
@@ -497,7 +550,7 @@ class AutoRifeTensorrt:
         }
 
         cudaStream = cuda.Stream()
-        
+
         # Auto-detect CUDA graph usage based on CUDA version
         use_cuda_graph = False
         try:
@@ -517,23 +570,124 @@ class AutoRifeTensorrt:
         engine = rife_trt_model
         logger(f"Using loaded TensorRT engine")
 
-        # Activate and allocate buffers for the engine
+        # Activate engine and retrieve profile bounds
         engine.activate()
-        engine.allocate_buffers(shape_dict=shape_dict)
+        bounds = engine.get_input_profile_bounds()
+
+        # Determine if tiling/padding is needed
+        needs_tiling = False
+        min_hw = (1, 1)
+        max_hw = (4096, 4096)
+        if bounds and 'img0' in bounds:
+            min_shape = bounds['img0'][0]
+            max_shape = bounds['img0'][2]
+            min_hw = (min_shape[2], min_shape[3])
+            max_hw = (max_shape[2], max_shape[3])
+            needs_tiling = H > max_hw[0] or W > max_hw[1] or H < min_hw[0] or W < min_hw[1]
+            if needs_tiling:
+                logger(f"Frame {H}x{W} outside profile range "
+                       f"[{min_hw[0]}-{max_hw[0]}]x[{min_hw[1]}-{max_hw[1]}] — enabling tiling/padding")
 
         frames = preprocess_frames(frames)
 
-        def return_middle_frame(frame_0, frame_1, timestep):
-            timestep_t = torch.tensor([timestep], dtype=torch.float32).to(get_torch_device())
-            # s = time.time()
-            output = engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream, use_cuda_graph)
-            # e = time.time()
-            # print(f"Time taken to infer: {(e-s)*1000} ms")
+        if needs_tiling:
+            # Disable CUDA graph when tiling (buffers are reallocated per tile)
+            use_cuda_graph = False
+            logger("Tiling mode: CUDA graph disabled")
 
-            result = output['output']
-            return result
+            def _raw_infer(f0, f1, t):
+                """Per-tile inference: reallocate buffers for the tile shape."""
+                _, _, th, tw = f0.shape
+                tile_shape_dict = {
+                    "img0": {"shape": (1, 3, th, tw)},
+                    "img1": {"shape": (1, 3, th, tw)},
+                    "output": {"shape": (1, 3, th, tw)},
+                }
+                engine.allocate_buffers(shape_dict=tile_shape_dict)
+                timestep_t = torch.tensor([t], dtype=torch.float32).to(get_torch_device())
+                output = engine.infer({"img0": f0, "img1": f1, "timestep": timestep_t}, cudaStream, use_cuda_graph)
+                return output['output']
 
-        result = generate_frames_rife(frames, clear_cache_after_n_frames, multiplier, return_middle_frame)
+            def return_middle_frame(frame_0, frame_1, timestep):
+                return infer_tiled(frame_0, frame_1, timestep, _raw_infer, min_hw, max_hw)
+
+            # Batch inference is not compatible with tiling (each tile has different shape)
+            batch_infer_fn = None
+            effective_batch_size = 1
+        else:
+            # Fast path: single allocation, original behaviour
+            engine.allocate_buffers(shape_dict=shape_dict)
+
+            def return_middle_frame(frame_0, frame_1, timestep):
+                timestep_t = torch.tensor([timestep], dtype=torch.float32).to(get_torch_device())
+                output = engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream, use_cuda_graph)
+                result = output['output']
+                return result
+
+            # Batch inference: allocate buffers for the requested batch size
+            if batch_size > 1:
+                batch_shape_dict = {
+                    "img0": {"shape": (batch_size, 3, H, W)},
+                    "img1": {"shape": (batch_size, 3, H, W)},
+                    "output": {"shape": (batch_size, 3, H, W)},
+                }
+                # Check if the engine profile supports this batch size
+                engine_max_batch = max_hw[0]  # placeholder, real check below
+                if bounds and 'img0' in bounds:
+                    engine_max_batch = bounds['img0'][2][0]  # max batch dim
+                if engine_max_batch >= batch_size:
+                    logger(f"Batch inference enabled: batch_size={batch_size}")
+                    # Pre-allocate buffers for the full batch
+                    engine.allocate_buffers(shape_dict=batch_shape_dict)
+
+                    def batch_infer_fn(frame0_batch, frame1_batch, timestep_batch):
+                        B_actual = frame0_batch.shape[0]
+                        # If the actual batch is smaller than the allocated batch,
+                        # we still pass the full buffer but only use the first B_actual rows.
+                        # TensorRT requires the input shape to match what was set in allocate_buffers.
+                        # So we pad to the full batch size if needed.
+                        if B_actual < batch_size:
+                            pad_shape = (batch_size - B_actual,) + frame0_batch.shape[1:]
+                            pad0 = torch.zeros(pad_shape, dtype=frame0_batch.dtype, device=frame0_batch.device)
+                            pad1 = torch.zeros(pad_shape, dtype=frame1_batch.dtype, device=frame1_batch.device)
+                            frame0_batch = torch.cat([frame0_batch, pad0], dim=0)
+                            frame1_batch = torch.cat([frame1_batch, pad1], dim=0)
+                            timestep_batch = torch.cat([timestep_batch, torch.zeros(batch_size - B_actual, dtype=timestep_batch.dtype, device=timestep_batch.device)])
+
+                        timestep_t = timestep_batch.to(get_torch_device())
+                        output = engine.infer({"img0": frame0_batch, "img1": frame1_batch, "timestep": timestep_t}, cudaStream, use_cuda_graph)
+                        return output['output'][:B_actual]
+
+                    effective_batch_size = batch_size
+                else:
+                    logger(f"⚠️  Engine max_batch={engine_max_batch} < requested batch_size={batch_size}. "
+                           "Falling back to sequential inference. Rebuild engine with batch_size >= "
+                           f"{batch_size} to enable batched inference.")
+                    batch_infer_fn = None
+                    effective_batch_size = 1
+            else:
+                batch_infer_fn = None
+                effective_batch_size = 1
+
+        # Resolve multiplier: per-pair schedule (list) overrides uniform int
+        effective_multiplier = multiplier
+        if multiplier_list and multiplier_list.strip():
+            try:
+                effective_multiplier = [int(x.strip()) for x in multiplier_list.split(',') if x.strip() != '']
+                logger(f"Using per-pair multiplier schedule: {effective_multiplier}")
+            except ValueError as e:
+                logger(f"Invalid multiplier_list '{multiplier_list}': {e}. Falling back to uniform multiplier={multiplier}")
+                effective_multiplier = multiplier
+
+        result = generate_frames_rife(
+            frames,
+            clear_cache_after_n_frames,
+            effective_multiplier,
+            return_middle_frame,
+            interpolation_states=interpolation_states,
+            batch_size=effective_batch_size,
+            batch_infer_function=batch_infer_fn,
+        )
         out = postprocess_frames(result)
 
         if not keep_model_loaded:
@@ -546,12 +700,14 @@ NODE_CLASS_MAPPINGS = {
     "AutoRifeTensorrt": AutoRifeTensorrt,
     "AutoLoadRifeTensorrtModel": AutoLoadRifeTensorrtModel,
     "CustomResolutionConfig": CustomResolutionConfig,
+    "MakeInterpolationStateList": MakeInterpolationStateList,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "AutoRifeTensorrt": "Auto RIFE TensorRT",
     "AutoLoadRifeTensorrtModel": "(Down)load RIFE TensorRT Model",
     "CustomResolutionConfig": "RIFE Custom Resolution Config",
+    "MakeInterpolationStateList": "RIFE Interpolation State List",
 }
 
 __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

export_onnx.py

@@ -102,9 +102,10 @@ def export_onnx(ckpt_name, ensemble, scale_factor):
     onnx_save_path = os.path.join("models",onnx_save_name)
 
     dynamic_axes = {
-        'img0': {2: 'height', 3: 'width'},
-        'img1': {2: 'height', 3: 'width'},
-        'output': {2: 'height', 3: 'width'},
+        'img0': {0: 'batch', 2: 'height', 3: 'width'},
+        'img1': {0: 'batch', 2: 'height', 3: 'width'},
+        'timestep': {0: 'batch'},
+        'output': {0: 'batch', 2: 'height', 3: 'width'},
     }
 
     with torch.no_grad():  # Disable gradients for efficiency

trt_utilities.py

@@ -423,6 +423,29 @@ class Engine:
         else:
             self.context = self.engine.create_execution_context()
 
+    def get_input_profile_bounds(self):
+        """Return min/opt/max shapes for the first optimization profile.
+
+        Returns a dict: {tensor_name: (min_shape, opt_shape, max_shape)}.
+        Uses engine.get_tensor_profile_shape (TensorRT 10+). Returns empty
+        dict if the API is unavailable (older TensorRT builds).
+        """
+        bounds = {}
+        try:
+            trt_instance = get_trt()
+            for idx in range(self.engine.num_io_tensors):
+                name = self.engine.get_tensor_name(idx)
+                if self.engine.get_tensor_mode(name) == trt_instance.TensorIOMode.INPUT:
+                    # get_tensor_profile_shape returns [min, opt, max] as trt.Dims
+                    profile_shapes = self.engine.get_tensor_profile_shape(name, 0)
+                    min_shape = tuple(profile_shapes[0])
+                    opt_shape = tuple(profile_shapes[1])
+                    max_shape = tuple(profile_shapes[2])
+                    bounds[name] = (min_shape, opt_shape, max_shape)
+        except Exception as e:
+            print(f"⚠️  Could not retrieve profile bounds: {e}")
+        return bounds
+
     def allocate_buffers(self, shape_dict=None, device="cuda"):
         # Clean up CUDA graph resources since tensors will be recreated
         if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None:

vfi_utilities.py

@@ -1,6 +1,7 @@
 # https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/blob/main/vfi_utils.py
 
 import os
+import gc
 import torch
 import typing
 import einops
@@ -11,6 +12,20 @@ from colored import Fore, Back, Style
 
 DEVICE = get_torch_device()
 
+
+class InterpolationStateList():
+    """Carries a list of frame indices and a flag indicating whether the list
+    is a skip-list (frames to skip) or a keep-list (frames to interpolate).
+    Ported from Fannovel16/ComfyUI-Frame-Interpolation.
+    """
+    def __init__(self, frame_indices: typing.List[int], is_skip_list: bool):
+        self.frame_indices = frame_indices
+        self.is_skip_list = is_skip_list
+
+    def is_frame_skipped(self, frame_index):
+        is_frame_in_list = frame_index in self.frame_indices
+        return self.is_skip_list and is_frame_in_list or not self.is_skip_list and not is_frame_in_list
+
 def load_file_from_github_release(model_type, ckpt_name):
     error_strs = []
     for i, base_model_download_url in enumerate(BASE_MODEL_DOWNLOAD_URLS):
@@ -32,55 +47,238 @@ def preprocess_frames(frames):
     return einops.rearrange(frames[..., :3], "n h w c -> n c h w")
 
 def postprocess_frames(frames):
-    return einops.rearrange(frames, "n c h w -> n h w c")[..., :3].cpu()
+    # Always return float32 — numpy and all downstream ComfyUI nodes require it.
+    # Engine output may be fp16 when built with precision="fp16".
+    return einops.rearrange(frames, "n c h w -> n h w c")[..., :3].to(torch.float32).cpu()
+
+
+def _make_feather_window(tile_h, tile_w, overlap):
+    """Build a 2D cosine-feathered weight window for blending tiles."""
+    def _ramp(size):
+        win = torch.ones(size)
+        ramp = min(overlap, size // 2)
+        if ramp > 0:
+            t = torch.linspace(0, 1, ramp)
+            fade = 0.5 * (1 - torch.cos(torch.pi * t))
+            win[:ramp] = fade
+            win[-ramp:] = fade.flip(0)
+        return win
+    win_h = _ramp(tile_h)
+    win_w = _ramp(tile_w)
+    return win_h.view(1, 1, -1, 1) * win_w.view(1, 1, 1, -1)
+
+
+def _compute_tile_positions(total, tile_size, overlap):
+    """Return list of start positions covering [0, total) with given tile size and overlap."""
+    if total <= tile_size:
+        return [0]
+    stride = tile_size - overlap
+    positions = list(range(0, total - tile_size + 1, stride))
+    # Ensure the last tile reaches the end
+    if positions[-1] + tile_size < total:
+        positions.append(total - tile_size)
+    # Deduplicate (can happen when total is only slightly larger than tile_size)
+    return sorted(set(positions))
+
+
+def infer_tiled(frame_0, frame_1, timestep, infer_fn, min_hw, max_hw, overlap=128):
+    """Run inference with automatic tiling/padding when the frame dimensions
+    fall outside the TensorRT engine's optimization profile.
+
+    Args:
+        frame_0, frame_1: [1, C, H, W] tensors on device.
+        timestep: float.
+        infer_fn: callable(f0, f1, t) -> [1, C, H, W] (the raw engine.infer wrapper).
+        min_hw: (min_h, min_w) from the engine profile.
+        max_hw: (max_h, max_w) from the engine profile.
+        overlap: feathering overlap in pixels between adjacent tiles.
+
+    Returns:
+        [1, C, H, W] interpolated frame at the original resolution.
+    """
+    _, C, H, W = frame_0.shape
+    min_h, min_w = min_hw
+    max_h, max_w = max_hw
+
+    needs_tile = H > max_h or W > max_w
+    needs_pad = H < min_h or W < min_w
+
+    # Fast path: dimensions are within profile
+    if not needs_tile and not needs_pad:
+        return infer_fn(frame_0, frame_1, timestep)
+
+    # --- Padding (under-min) ---
+    pad_h_before = pad_h_after = pad_w_before = pad_w_after = 0
+    if needs_pad:
+        pad_h_before = max(0, (min_h - H + 1) // 2)
+        pad_h_after = max(0, min_h - H - pad_h_before)
+        pad_w_before = max(0, (min_w - W + 1) // 2)
+        pad_w_after = max(0, min_w - W - pad_w_before)
+        frame_0 = torch.nn.functional.pad(frame_0, (pad_w_before, pad_w_after, pad_h_before, pad_h_after), mode='replicate')
+        frame_1 = torch.nn.functional.pad(frame_1, (pad_w_before, pad_w_after, pad_h_before, pad_h_after), mode='replicate')
+        _, _, H_pad, W_pad = frame_0.shape
+    else:
+        H_pad, W_pad = H, W
+
+    # --- Tiling (over-max) ---
+    if H_pad > max_h or W_pad > max_w:
+        tile_h = min(max_h, H_pad)
+        tile_w = min(max_w, W_pad)
+        h_positions = _compute_tile_positions(H_pad, tile_h, overlap)
+        w_positions = _compute_tile_positions(W_pad, tile_w, overlap)
+
+        window = _make_feather_window(tile_h, tile_w, overlap).to(frame_0.device)
+
+        output = torch.zeros(1, C, H_pad, W_pad, device=frame_0.device, dtype=frame_0.dtype)
+        weight = torch.zeros(1, 1, H_pad, W_pad, device=frame_0.device, dtype=frame_0.dtype)
+
+        for h0 in h_positions:
+            for w0 in w_positions:
+                h1, w1 = h0 + tile_h, w0 + tile_w
+                t0 = frame_0[:, :, h0:h1, w0:w1]
+                t1 = frame_1[:, :, h0:h1, w0:w1]
+                tile_out = infer_fn(t0, t1, timestep)
+                output[:, :, h0:h1, w0:w1] += tile_out * window
+                weight[:, :, h0:h1, w0:w1] += window
+
+        result = output / weight.clamp(min=1e-8)
+    else:
+        result = infer_fn(frame_0, frame_1, timestep)
+
+    # --- Crop back if we padded ---
+    if needs_pad:
+        result = result[:, :, pad_h_before:pad_h_before + H, pad_w_before:pad_w_before + W]
+
+    return result
 
 def generate_frames_rife(
         frames,
         clear_cache_after_n_frames,
         multiplier,
-        return_middle_frame_function
+        return_middle_frame_function,
+        interpolation_states: InterpolationStateList = None,
+        batch_size: int = 1,
+        batch_infer_function=None,
         ):
+    """Interpolate frames.
 
-    output_frames = torch.zeros(multiplier*frames.shape[0], *frames.shape[1:], device="cpu")
+    Args:
+        frames: input tensor [N, C, H, W] (already preprocessed).
+        clear_cache_after_n_frames: clear CUDA cache every N interpolated frames.
+        multiplier: int (uniform) or list[int] (per-pair schedule). A list is
+            padded with 2 for any missing trailing pairs, matching Fannovel16.
+        return_middle_frame_function: callable(frame0, frame1, timestep) -> middle.
+            Used when batch_size == 1 (or when batch_infer_function is None).
+        interpolation_states: optional InterpolationStateList to skip specific pairs.
+        batch_size: number of (pair, timestep) tasks per GPU call. >1 enables
+            batched inference (requires batch_infer_function and an engine built
+            with max_batch >= batch_size).
+        batch_infer_function: optional callable(frame0_batch, frame1_batch, timestep_batch)
+            -> middle_batch. frame0_batch/frame1_batch are [B,C,H,W], timestep_batch
+            is [B]. Returns [B,C,H,W]. Required when batch_size > 1.
+    """
+
+    n_pairs = len(frames) - 1
+
+    # Normalise multiplier to a per-pair list
+    if isinstance(multiplier, int):
+        multipliers = [int(multiplier)] * n_pairs
+    else:
+        multipliers = list(map(int, multiplier))
+        multipliers += [2] * (n_pairs - len(multipliers))
+
+    # Upper bound on output length (sum of per-pair multipliers + final frame)
+    total_out = sum(multipliers) + 1
+    output_frames = torch.zeros(total_out, *frames.shape[1:], device="cpu")
     out_len = 0
 
     number_of_frames_processed_since_last_cleared_cuda_cache = 0
-    # pbar = ProgressBar(len(frames))
 
-    for frame_itr in range(len(frames) - 1): # Skip the final frame since there are no frames after it
+    # Decide whether to use batched inference
+    use_batch = batch_size > 1 and batch_infer_function is not None
 
-        frame_0 = frames[frame_itr:frame_itr+1]
-        frame_1 = frames[frame_itr+1:frame_itr+2]
-        output_frames[out_len] = frame_0 # Start with first frame
-        out_len += 1
+    if use_batch:
+        # Build a flat list of all (pair_idx, timestep) tasks, skipping excluded pairs.
+        # Each task produces exactly one intermediate frame.
+        tasks = []
+        for pair_idx in range(n_pairs):
+            if interpolation_states is not None and interpolation_states.is_frame_skipped(pair_idx):
+                continue
+            m = multipliers[pair_idx]
+            for step in range(1, m):
+                tasks.append((pair_idx, step / m))
 
-        for middle_i in range(1, multiplier):
-            timestep = middle_i/multiplier
-            middle_frame = return_middle_frame_function(frame_0, frame_1, timestep).detach().cpu()
+        # Storage for intermediate frames, keyed by pair index
+        results = {i: [] for i in range(n_pairs)}
 
-            # Copy middle frames to output
-            output_frames[out_len] = middle_frame
-            out_len +=1
+        pos = 0
+        while pos < len(tasks):
+            batch_tasks = tasks[pos : pos + batch_size]
+            B = len(batch_tasks)
 
-            # Try to avoid a memory overflow by clearing cuda cache regularly
-            number_of_frames_processed_since_last_cleared_cuda_cache += 1
-            if number_of_frames_processed_since_last_cleared_cuda_cache >= clear_cache_after_n_frames:
-                soft_empty_cache()
-                number_of_frames_processed_since_last_cleared_cuda_cache = 0
-                logger("Clearing cache...")
+            frame0_list = [frames[pair_idx : pair_idx + 1] for pair_idx, _ in batch_tasks]
+            frame1_list = [frames[pair_idx + 1 : pair_idx + 2] for pair_idx, _ in batch_tasks]
+            timestep_list = [dt for _, dt in batch_tasks]
 
-            # pbar.update(1)
+            frame0_batch = torch.cat(frame0_list, dim=0)
+            frame1_batch = torch.cat(frame1_list, dim=0)
+            timestep_batch = torch.tensor(timestep_list, dtype=torch.float32)
 
+            middle_batch = batch_infer_function(frame0_batch, frame1_batch, timestep_batch).detach().cpu()
+
+            for idx, (pair_idx, _) in enumerate(batch_tasks):
+                results[pair_idx].append(middle_batch[idx : idx + 1])
+                number_of_frames_processed_since_last_cleared_cuda_cache += 1
+                if number_of_frames_processed_since_last_cleared_cuda_cache >= clear_cache_after_n_frames:
+                    soft_empty_cache()
+                    gc.collect()
+                    number_of_frames_processed_since_last_cleared_cuda_cache = 0
+                    logger("Clearing cache...")
+
+            pos += B
+
+        # Assemble output: each original frame followed by its interpolated frames
+        for pair_idx in range(n_pairs):
+            output_frames[out_len] = frames[pair_idx : pair_idx + 1]
+            out_len += 1
+            for mid in results[pair_idx]:
+                output_frames[out_len] = mid
+                out_len += 1
+    else:
+        # Sequential path (original behaviour, batch_size == 1)
+        for frame_itr in range(n_pairs):
+            frame_0 = frames[frame_itr:frame_itr+1]
+            frame_1 = frames[frame_itr+1:frame_itr+2]
+            output_frames[out_len] = frame_0
+            out_len += 1
+
+            if interpolation_states is not None and interpolation_states.is_frame_skipped(frame_itr):
+                continue
+
+            m = multipliers[frame_itr]
+            for middle_i in range(1, m):
+                timestep = middle_i / m
+                middle_frame = return_middle_frame_function(frame_0, frame_1, timestep).detach().cpu()
+
+                output_frames[out_len] = middle_frame
+                out_len += 1
+
+                number_of_frames_processed_since_last_cleared_cuda_cache += 1
+                if number_of_frames_processed_since_last_cleared_cuda_cache >= clear_cache_after_n_frames:
+                    soft_empty_cache()
+                    gc.collect()
+                    number_of_frames_processed_since_last_cleared_cuda_cache = 0
+                    logger("Clearing cache...")
 
     # Append final frame
     output_frames[out_len] = frames[-1:]
-    logger(f"done! - {(len(frames) -1) * (multiplier-1)} new frames generated at resolution: {output_frames[0].shape}")
+    new_frames_count = sum(max(m - 1, 0) for m in multipliers)
+    logger(f"done! - {new_frames_count} new frames generated at resolution: {output_frames[0].shape}")
     out_len += 1
 
-    # clear cache for courtesy
     soft_empty_cache()
+    gc.collect()
     logger("Final clearing cache done ...")
-#
-    res = output_frames[:out_len]
 
+    res = output_frames[:out_len]
     return res