Upload vfi_utils.py

vfi_utils.py

@@ -0,0 +1,397 @@
+import yaml
+import os
+from torch.hub import download_url_to_file, get_dir
+from urllib.parse import urlparse
+import torch
+import typing
+import traceback
+import einops
+import gc
+import torchvision.transforms.functional as transform
+from comfy.model_management import soft_empty_cache, get_torch_device
+import numpy as np
+
+
+# -----------------------------
+# Config
+# -----------------------------
+config_path = os.path.join(os.path.dirname(__file__), "./config.yaml")
+if os.path.exists(config_path):
+    config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader) or {}
+else:
+    raise Exception(
+        "config.yaml file is neccessary, plz recreate the config file by downloading it from "
+        "https://github.com/Fannovel16/ComfyUI-Frame-Interpolation"
+    )
+
+DEVICE = get_torch_device()
+
+
+# -----------------------------
+# Model download sources
+# -----------------------------
+# Original GitHub release bases (some RIFE checkpoints are no longer hosted there -> 404)
+DEFAULT_BASE_MODEL_DOWNLOAD_URLS = [
+    "https://github.com/styler00dollar/VSGAN-tensorrt-docker/releases/download/models/",
+    "https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/releases/download/models/",
+    "https://github.com/dajes/frame-interpolation-pytorch/releases/download/v1.0.0/",
+]
+
+# Optional: override via config.yaml:
+# base_model_download_urls:
+#   - "https://..."
+BASE_MODEL_DOWNLOAD_URLS = config.get("base_model_download_urls", DEFAULT_BASE_MODEL_DOWNLOAD_URLS)
+
+# Optional: add extra base URLs via env var (comma-separated):
+#   COMFY_VFI_EXTRA_MODEL_DOWNLOAD_BASE_URLS="https://mirror1/.../,https://mirror2/.../"
+_env_extra = os.getenv("COMFY_VFI_EXTRA_MODEL_DOWNLOAD_BASE_URLS", "")
+if _env_extra.strip():
+    for u in [x.strip() for x in _env_extra.split(",") if x.strip()]:
+        if not u.endswith("/"):
+            u += "/"
+        if u not in BASE_MODEL_DOWNLOAD_URLS:
+            BASE_MODEL_DOWNLOAD_URLS.append(u)
+
+# Optional: last-resort direct URL overrides per checkpoint.
+# You asked for this HuggingFace mirror:
+#   https://huggingface.co/saliacoel/x/resolve/main/rife47.pth
+#
+# You can override it without editing code by setting:
+#   COMFY_VFI_RIFE47_URL="https://huggingface.co/.../rife47.pth"
+#
+# Or in config.yaml:
+# ckpt_url_mirrors:
+#   rife47.pth: "https://huggingface.co/.../rife47.pth"
+DEFAULT_CKPT_URL_MIRRORS = {
+    "rife47.pth": os.getenv(
+        "COMFY_VFI_RIFE47_URL",
+        "https://huggingface.co/saliacoel/x/resolve/main/rife47.pth",
+    ),
+}
+_ckpt_overrides = config.get("ckpt_url_mirrors", {}) or {}
+CKPT_URL_MIRRORS = {**DEFAULT_CKPT_URL_MIRRORS, **_ckpt_overrides}
+
+
+def _is_http_url(value: str) -> bool:
+    try:
+        parts = urlparse(value)
+        return parts.scheme in ("http", "https") and bool(parts.netloc)
+    except Exception:
+        return False
+
+
+class InterpolationStateList:
+    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)
+
+
+class MakeInterpolationStateList:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "frame_indices": ("STRING", {"multiline": True, "default": "1,2,3"}),
+                "is_skip_list": ("BOOLEAN", {"default": True}),
+            },
+        }
+
+    RETURN_TYPES = ("INTERPOLATION_STATES",)
+    FUNCTION = "create_options"
+    CATEGORY = "ComfyUI-Frame-Interpolation/VFI"
+
+    def create_options(self, frame_indices: str, is_skip_list: bool):
+        frame_indices_list = [int(item) for item in frame_indices.split(",")]
+
+        interpolation_state_list = InterpolationStateList(
+            frame_indices=frame_indices_list,
+            is_skip_list=is_skip_list,
+        )
+        return (interpolation_state_list,)
+
+
+def get_ckpt_container_path(model_type):
+    return os.path.abspath(os.path.join(os.path.dirname(__file__), config["ckpts_path"], model_type))
+
+
+def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
+    """Load file from http url, will download models if necessary.
+
+    Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
+
+    Args:
+        url (str): URL to be downloaded.
+        model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
+        progress (bool): Whether to show the download progress.
+        file_name (str): The downloaded file name. If None, use the file name in the url.
+
+    Returns:
+        str: The path to the downloaded file.
+    """
+    if model_dir is None:  # use the pytorch hub_dir
+        hub_dir = get_dir()
+        model_dir = os.path.join(hub_dir, "checkpoints")
+
+    os.makedirs(model_dir, exist_ok=True)
+
+    if file_name is None:
+        parts = urlparse(url)
+        file_name = os.path.basename(parts.path)
+
+    cached_file = os.path.abspath(os.path.join(model_dir, file_name))
+    if not os.path.exists(cached_file):
+        print(f'Downloading: "{url}" to {cached_file}\n')
+        download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
+    return cached_file
+
+
+def load_file_from_github_release(model_type, ckpt_name):
+    """
+    Backwards-compatible name, but now supports:
+      - direct URL passed as ckpt_name
+      - a per-ckpt mirror URL fallback (e.g. HuggingFace)
+    """
+    # Allow passing a direct URL (future-proofing / manual overrides)
+    if isinstance(ckpt_name, str) and _is_http_url(ckpt_name):
+        return load_file_from_url(ckpt_name, get_ckpt_container_path(model_type))
+
+    error_strs = []
+
+    urls_to_try = [base + ckpt_name for base in BASE_MODEL_DOWNLOAD_URLS]
+
+    # Add last-resort mirror(s) if configured for this ckpt
+    mirror = CKPT_URL_MIRRORS.get(ckpt_name)
+    if mirror:
+        if isinstance(mirror, (list, tuple)):
+            urls_to_try.extend(list(mirror))
+        else:
+            urls_to_try.append(str(mirror))
+
+    # De-duplicate while preserving order
+    seen = set()
+    deduped = []
+    for u in urls_to_try:
+        if u not in seen:
+            seen.add(u)
+            deduped.append(u)
+    urls_to_try = deduped
+
+    for i, url in enumerate(urls_to_try):
+        try:
+            return load_file_from_url(url, get_ckpt_container_path(model_type))
+        except Exception:
+            traceback_str = traceback.format_exc()
+            if i < len(urls_to_try) - 1:
+                print("Failed! Trying another endpoint.")
+            error_strs.append(f"Error when downloading from: {url}\n\n{traceback_str}")
+
+    error_str = "\n\n".join(error_strs)
+    raise Exception(
+        f"Tried all endpoints to download {ckpt_name} but no success. Below is the error log:\n\n{error_str}"
+    )
+
+
+def load_file_from_direct_url(model_type, url):
+    return load_file_from_url(url, get_ckpt_container_path(model_type))
+
+
+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()
+
+
+def assert_batch_size(frames, batch_size=2, vfi_name=None):
+    subject_verb = "Most VFI models require" if vfi_name is None else f"VFI model {vfi_name} requires"
+    assert len(frames) >= batch_size, (
+        f"{subject_verb} at least {batch_size} frames to work with, only found {frames.shape[0]}. "
+        f"Please check the frame input using PreviewImage."
+    )
+
+
+def _generic_frame_loop(
+    frames,
+    clear_cache_after_n_frames,
+    multiplier: typing.Union[typing.SupportsInt, typing.List],
+    return_middle_frame_function,
+    *return_middle_frame_function_args,
+    interpolation_states: InterpolationStateList = None,
+    use_timestep=True,
+    dtype=torch.float16,
+    final_logging=True,
+):
+    # https://github.com/hzwer/Practical-RIFE/blob/main/inference_video.py#L169
+    def non_timestep_inference(frame0, frame1, n):
+        middle = return_middle_frame_function(frame0, frame1, None, *return_middle_frame_function_args)
+        if n == 1:
+            return [middle]
+        first_half = non_timestep_inference(frame0, middle, n=n // 2)
+        second_half = non_timestep_inference(middle, frame1, n=n // 2)
+        if n % 2:
+            return [*first_half, middle, *second_half]
+        else:
+            return [*first_half, *second_half]
+
+    output_frames = torch.zeros(multiplier * frames.shape[0], *frames.shape[1:], dtype=dtype, device="cpu")
+    out_len = 0
+
+    number_of_frames_processed_since_last_cleared_cuda_cache = 0
+
+    for frame_itr in range(len(frames) - 1):  # Skip the final frame since there are no frames after it
+        frame0 = frames[frame_itr : frame_itr + 1]
+        output_frames[out_len] = frame0  # Start with first frame
+        out_len += 1
+
+        # Ensure that input frames are in fp32 - the same dtype as model
+        frame0 = frame0.to(dtype=torch.float32)
+        frame1 = frames[frame_itr + 1 : frame_itr + 2].to(dtype=torch.float32)
+
+        if interpolation_states is not None and interpolation_states.is_frame_skipped(frame_itr):
+            continue
+
+        # Generate and append a batch of middle frames
+        middle_frame_batches = []
+
+        if use_timestep:
+            for middle_i in range(1, multiplier):
+                timestep = middle_i / multiplier
+
+                middle_frame = (
+                    return_middle_frame_function(
+                        frame0.to(DEVICE),
+                        frame1.to(DEVICE),
+                        timestep,
+                        *return_middle_frame_function_args,
+                    )
+                    .detach()
+                    .cpu()
+                )
+                middle_frame_batches.append(middle_frame.to(dtype=dtype))
+        else:
+            middle_frames = non_timestep_inference(frame0.to(DEVICE), frame1.to(DEVICE), multiplier - 1)
+            middle_frame_batches.extend(torch.cat(middle_frames, dim=0).detach().cpu().to(dtype=dtype))
+
+        # Copy middle frames to output
+        for middle_frame in middle_frame_batches:
+            output_frames[out_len] = middle_frame
+            out_len += 1
+
+        number_of_frames_processed_since_last_cleared_cuda_cache += 1
+        # Try to avoid a memory overflow by clearing cuda cache regularly
+        if number_of_frames_processed_since_last_cleared_cuda_cache >= clear_cache_after_n_frames:
+            print("Comfy-VFI: Clearing cache...", end=" ")
+            soft_empty_cache()
+            number_of_frames_processed_since_last_cleared_cuda_cache = 0
+            print("Done cache clearing")
+
+        gc.collect()
+
+    if final_logging:
+        print(f"Comfy-VFI done! {len(output_frames)} frames generated at resolution: {output_frames[0].shape}")
+
+    # Append final frame
+    output_frames[out_len] = frames[-1:]
+    out_len += 1
+
+    # clear cache for courtesy
+    if final_logging:
+        print("Comfy-VFI: Final clearing cache...", end=" ")
+    soft_empty_cache()
+    if final_logging:
+        print("Done cache clearing")
+
+    return output_frames[:out_len]
+
+
+def generic_frame_loop(
+    model_name,
+    frames,
+    clear_cache_after_n_frames,
+    multiplier: typing.Union[typing.SupportsInt, typing.List],
+    return_middle_frame_function,
+    *return_middle_frame_function_args,
+    interpolation_states: InterpolationStateList = None,
+    use_timestep=True,
+    dtype=torch.float32,
+):
+    assert_batch_size(frames, vfi_name=model_name.replace("_", " ").replace("VFI", ""))
+    if type(multiplier) == int:
+        return _generic_frame_loop(
+            frames,
+            clear_cache_after_n_frames,
+            multiplier,
+            return_middle_frame_function,
+            *return_middle_frame_function_args,
+            interpolation_states=interpolation_states,
+            use_timestep=use_timestep,
+            dtype=dtype,
+        )
+    if type(multiplier) == list:
+        multipliers = list(map(int, multiplier))
+        multipliers += [2] * (len(frames) - len(multipliers) - 1)
+        frame_batches = []
+        for frame_itr in range(len(frames) - 1):
+            multiplier = multipliers[frame_itr]
+            if multiplier == 0:
+                continue
+            frame_batch = _generic_frame_loop(
+                frames[frame_itr : frame_itr + 2],
+                clear_cache_after_n_frames,
+                multiplier,
+                return_middle_frame_function,
+                *return_middle_frame_function_args,
+                interpolation_states=interpolation_states,
+                use_timestep=use_timestep,
+                dtype=dtype,
+                final_logging=False,
+            )
+            if frame_itr != len(frames) - 2:  # Not append last frame unless this batch is the last one
+                frame_batch = frame_batch[:-1]
+            frame_batches.append(frame_batch)
+        output_frames = torch.cat(frame_batches)
+        print(f"Comfy-VFI done! {len(output_frames)} frames generated at resolution: {output_frames[0].shape}")
+        return output_frames
+    raise NotImplementedError(f"multipiler of {type(multiplier)}")
+
+
+class FloatToInt:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"float": ("FLOAT", {"default": 0, "min": 0, "step": 0.01})}}
+
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "convert"
+    CATEGORY = "ComfyUI-Frame-Interpolation"
+
+    def convert(self, float):
+        if hasattr(float, "__iter__"):
+            return (list(map(int, float)),)
+        return (int(float),)
+
+
+""" def generic_4frame_loop(
+        frames,
+        clear_cache_after_n_frames,
+        multiplier: typing.SupportsInt,
+        return_middle_frame_function,
+        *return_middle_frame_function_args,
+        interpolation_states: InterpolationStateList = None,
+        use_timestep=False):
+
+    if use_timestep: raise NotImplementedError("Timestep 4 frame VFI model")
+    def non_timestep_inference(frame_0, frame_1, frame_2, frame_3, n):
+        middle = return_middle_frame_function(frame_0, frame_1, None, *return_middle_frame_function_args)
+        if n == 1:
+            return [middle]
+        first_half = non_timestep_inference(frame_0, middle, n=n//2)
+        second_half = non_timestep_inference(middle, frame_1, n=n//2)
+        if n%2:
+            return [*first_half, middle, *second_half]
+        else:
+            return [*first_half, *second_half] """