diff --git a/__init__.py b/__init__.py index c4a73c5..c404606 100644 --- a/__init__.py +++ b/__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'] diff --git a/export_onnx.py b/export_onnx.py index 96b77fc..bc34ea6 100644 --- a/export_onnx.py +++ b/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 diff --git a/trt_utilities.py b/trt_utilities.py index bec743c..9e85304 100644 --- a/trt_utilities.py +++ b/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: diff --git a/vfi_utilities.py b/vfi_utilities.py index 3fad020..adafc1f 100644 --- a/vfi_utilities.py +++ b/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