merge: 🎉 pull request #11 from dev/frame_interpolation

.dockerignore

@@ -0,0 +1,34 @@
+# Include any files or directories that you don't want to be copied to your
+# container here (e.g., local build artifacts, temporary files, etc.).
+#
+# For more help, visit the .dockerignore file reference guide at
+# https://docs.docker.com/engine/reference/builder/#dockerignore-file
+
+**/.DS_Store
+**/__pycache__
+**/.venv
+**/.classpath
+**/.dockerignore
+**/.env
+**/.git
+**/.gitignore
+**/.project
+**/.settings
+**/.toolstarget
+**/.vs
+**/.vscode
+**/*.*proj.user
+**/*.dbmdl
+**/*.jfm
+**/bin
+**/charts
+**/docker-compose*
+**/compose*
+**/Dockerfile*
+**/node_modules
+**/npm-debug.log
+**/obj
+**/secrets.dev.yaml
+**/values.dev.yaml
+LICENSE
+README.md

.github/CODEOWNERS

@@ -0,0 +1,5 @@
+* @melMass
+extern/GFPGAN/* @TencentARC
+extern/SadTalker/* @OpenTalker
+nodes/deep_bump.py @HugoTini
+web/imageFeed.js @pythongosssss @melMass

.gitmodules

@@ -1,6 +1,9 @@
-[submodule "extern/SadTalker"]
-	path = extern/SadTalker
-	url = https://github.com/OpenTalker/SadTalker.git
+[submodule "extern/google-FILM"]
+	path = extern/frame_interpolation
+	url = https://github.com/google-research/frame-interpolation
 [submodule "extern/GFPGAN"]
 	path = extern/GFPGAN
 	url = https://github.com/TencentARC/GFPGAN.git
+[submodule "extern/frame_interpolation"]
+	path = extern/frame_interpolation
+	url = https://github.com/google-research/frame-interpolation

nodes/faceswap.py

@@ -16,6 +16,8 @@ from insightface.model_zoo.inswapper import INSwapper
 from ..utils import pil2tensor, tensor2pil
 from ..log import mklog, NullWriter
 import sys
+import comfy.model_management as model_management
+
 
 # endregion
 
@@ -100,6 +102,7 @@ class FaceSwap:
         debug="false",
     ):
         def do_swap(img):
+            model_management.throw_exception_if_processing_interrupted()
             img = tensor2pil(img)
             ref = tensor2pil(reference)
             face_ids = {

nodes/image_interpolation.py

@@ -0,0 +1,242 @@
+from typing import List
+from pathlib import Path
+import os
+import glob
+import folder_paths
+from ..log import log
+import torch
+from frame_interpolation.eval import util, interpolator
+from ..utils import tensor2np
+import uuid
+import numpy as np
+import subprocess
+import comfy
+
+import tensorflow as tf
+import comfy.model_management as model_management
+
+
+class LoadFilmModel:
+    """Loads a FILM model"""
+
+    @staticmethod
+    def get_models() -> List[Path]:
+        models_path = os.path.join(folder_paths.models_dir, "FILM/*")
+        models = glob.glob(models_path)
+        models = [Path(x) for x in models if x.endswith(".onnx") or x.endswith(".pth")]
+        return models
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "film_model": (
+                    ["L1", "Style", "VGG"],
+                    {"default": "Style"},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("FILM_MODEL",)
+    FUNCTION = "load_model"
+    CATEGORY = "face"
+
+    def load_model(self, film_model: str):
+        model_path = Path(folder_paths.models_dir) / "FILM" / film_model
+        if not (model_path / "saved_model.pb").exists():
+            model_path = model_path / "saved_model"
+
+        if not model_path.exists():
+            log.error(f"Model {model_path} does not exist")
+            raise ValueError(f"Model {model_path} does not exist")
+
+        log.info(f"Loading model {model_path}")
+
+        return (interpolator.Interpolator(model_path.as_posix(), None),)
+
+
+class FilmInterpolation:
+    """Google Research FILM frame interpolation for large motion"""
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "interpolate": ("INT", {"default": 2, "min": 1, "max": 50}),
+                "film_model": ("FILM_MODEL",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "do_interpolation"
+    CATEGORY = "animation"
+
+    def do_interpolation(
+        self,
+        images: torch.Tensor,
+        interpolate: int,
+        film_model: interpolator.Interpolator,
+    ):
+        n = images.size(0)
+        # check if tensorflow GPU is available
+        available_gpus = tf.config.list_physical_devices("GPU")
+        if not len(available_gpus):
+            log.warning(
+                "Tensorflow GPU not available, falling back to CPU this will be very slow"
+            )
+        else:
+            log.debug(f"Tensorflow GPU available, using {available_gpus}")
+
+        num_frames = (n - 1) * (2 ** (interpolate) - 1)
+        log.debug(f"Will interpolate into {num_frames} frames")
+
+        in_frames = [images[i] for i in range(n)]
+        out_tensors = []
+
+        pbar = comfy.utils.ProgressBar(num_frames)
+
+        for frame in util.interpolate_recursively_from_memory(
+            in_frames, interpolate, film_model
+        ):
+            out_tensors.append(
+                torch.from_numpy(frame) if isinstance(frame, np.ndarray) else frame
+            )
+            model_management.throw_exception_if_processing_interrupted()
+            pbar.update(1)
+
+        out_tensors = torch.cat([tens.unsqueeze(0) for tens in out_tensors], dim=0)
+
+        log.debug(f"Returning {len(out_tensors)} tensors")
+        log.debug(f"Output shape {out_tensors.shape}")
+        log.debug(f"Output type {out_tensors.dtype}")
+        return (out_tensors,)
+
+
+class ConcatImages:
+    """Add images to batch"""
+
+    def __init__(self):
+        pass
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "concat_images"
+    CATEGORY = "animation"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "imageA": ("IMAGE",),
+                "imageB": ("IMAGE",),
+            },
+        }
+
+    @classmethod
+    def concatenate_tensors(cls, A: torch.Tensor, B: torch.Tensor):
+        # Get the batch sizes of A and B
+        batch_size_A = A.size(0)
+        batch_size_B = B.size(0)
+
+        # Concatenate the tensors along the batch dimension
+        concatenated = torch.cat((A, B), dim=0)
+
+        # Update the batch size in the concatenated tensor
+        concatenated_size = list(concatenated.size())
+        concatenated_size[0] = batch_size_A + batch_size_B
+        concatenated = concatenated.view(*concatenated_size)
+
+        return concatenated
+
+    def concat_images(self, imageA: torch.Tensor, imageB: torch.Tensor):
+        log.debug(f"Concatenating A ({imageA.shape}) and B ({imageB.shape})")
+        return (self.concatenate_tensors(imageA, imageB),)
+
+
+class ExportToProRes:
+    """Export to ProRes 4444 (Experimental)"""
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                # "frames": ("FRAMES",),
+                "fps": ("FLOAT", {"default": 24, "min": 1}),
+                "prefix": ("STRING", {"default": "export"}),
+            }
+        }
+
+    RETURN_TYPES = ("VIDEO",)
+    OUTPUT_NODE = True
+    FUNCTION = "export_prores"
+    CATEGORY = "animation"
+
+    def export_prores(
+        self,
+        images: torch.Tensor,
+        fps: float,
+        prefix: str,
+    ):
+        output_dir = Path(folder_paths.get_output_directory())
+        id = f"{prefix}_{uuid.uuid4()}.mov"
+
+        log.debug(f"Exporting to {output_dir / id}")
+
+        frames = tensor2np(images)
+        log.debug(f"Frames type {type(frames)}")
+        log.debug(f"Exporting {len(frames)} frames")
+
+        frames = [frame.astype(np.uint16) * 257 for frame in frames]
+
+        height, width, _ = frames[0].shape
+
+        out_path = (output_dir / id).as_posix()
+
+        # Prepare the FFmpeg command
+        command = [
+            "ffmpeg",
+            "-y",
+            "-f",
+            "rawvideo",
+            "-vcodec",
+            "rawvideo",
+            "-s",
+            f"{width}x{height}",
+            "-pix_fmt",
+            "rgb48le",
+            "-r",
+            str(fps),
+            "-i",
+            "-",
+            "-c:v",
+            "prores_ks",
+            "-profile:v",
+            "4",
+            "-pix_fmt",
+            "yuva444p10le",
+            "-r",
+            str(fps),
+            "-y",
+            out_path,
+        ]
+
+        process = subprocess.Popen(command, stdin=subprocess.PIPE)
+
+        for frame in frames:
+            model_management.throw_exception_if_processing_interrupted()
+            process.stdin.write(frame.tobytes())
+
+        process.stdin.close()
+        process.wait()
+
+        return (out_path,)
+
+
+__nodes__ = [LoadFilmModel, FilmInterpolation, ExportToProRes, ConcatImages]

nodes/image_processing.py

@@ -14,6 +14,8 @@ import folder_paths
 from PIL.PngImagePlugin import PngInfo
 import json
 import os
+import comfy.model_management as model_management
+
 
 try:
     from cv2.ximgproc import guidedFilter
@@ -514,6 +516,7 @@ class ImageResizeFactor:
         resample="lanczos",
         mask=None,
     ) -> torch.Tensor:
+        model_management.throw_exception_if_processing_interrupted()
         batch_count = 1
         img = tensor2pil(image)
 

requirements.txt

@@ -1,11 +1,17 @@
-onnxruntime-gpu
-imageio
+onnxruntime-gpu==1.15.1
+imageio===2.28.1
 qrcode[pil]
 numpy==1.23.5
-ifnude==0.0.3
 insightface==0.7.3
 mmcv==2.0.0
 mmdet==3.0.0
 rembg==2.0.37
 facexlib==0.3.0
-basicsr==1.4.2
+basicsr==1.4.2
+# on windows non WSL 2.10 is the last version with GPU support
+tensorflow<2.11.0; platform_system == "Windows"
+tb-nightly==2.12.0a20230126; platform_system == "Windows"
+tensorflow; platform_system != "Windows"
+# the old tf version on windows comes with a breaking protobuf version
+protobuf==3.19.6
+gdown @ git+https://github.com/melMass/gdown@main

scripts/download_models.py

@@ -38,12 +38,20 @@ models_to_download = {
         ],
         "destination": "upscale_models",
     },
+    "FILM: Frame Interpolation for Large Motion": {
+        "size": 402,
+        "download_url": [
+            "https://drive.google.com/drive/folders/131_--QrieM4aQbbLWrUtbO2cGbX8-war"
+        ],
+        "destination": "FILMOS",
+    },
 }
 
 console = Console()
 
 from urllib.parse import urlparse
 from pathlib import Path
+import gdown
 
 
 def download_model(download_url, destination):
@@ -53,6 +61,16 @@ def download_model(download_url, destination):
         return
 
     filename = os.path.basename(urlparse(download_url).path)
+    response = None
+    if "drive.google.com" in download_url:
+        if "/folders/" in download_url:
+            # download folder
+            gdown.download_folder(download_url, output=destination, resume=True)
+            return
+        # download from google drive
+        gdown.download(download_url, destination, quiet=False, resume=True)
+        return
+
     response = requests.get(download_url, stream=True)
     total_size = int(response.headers.get("content-length", 0))
 
@@ -150,6 +168,7 @@ def main(models_to_download, skip_input=False):
         for model_name, model_details in models_to_download_selected.items():
             download_url = model_details["download_url"]
             destination = model_details["destination"]
+            console.print(f"Downloading {model_name}...")
             download_model(download_url, destination)
 
     except KeyboardInterrupt:

utils.py

@@ -4,6 +4,8 @@ import torch
 from pathlib import Path
 import sys
 
+from typing import Union, List
+
 
 def add_path(path, prepend=False):
     if isinstance(path, list):
@@ -43,37 +45,36 @@ add_path(comfy_dir)
 add_path((comfy_dir / "custom_nodes"))
 
 
-# Tensor to PIL (grabbed from WAS Suite)
-def tensor2pil(image: torch.Tensor) -> Image.Image:
-    return Image.fromarray(
-        np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
-    )
+def tensor2pil(image: torch.Tensor) -> Union[Image.Image, List[Image.Image]]:
+    batch_count = 1
+    if len(image.shape) > 3:
+        batch_count = image.size(0)
+
+    if batch_count == 1:
+        return Image.fromarray(
+            np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
+        )
+    return [tensor2pil(image[i]) for i in range(batch_count)]
 
 
-# TODO: write pil2tensor counterpart (batch support)
-# def tensor2pil(image: torch.Tensor) -> Union[Image.Image, List[Image.Image]]:
-#     batch_count = 1
-#     if len(image.shape) > 3:
-#         batch_count = image.size(0)
+def pil2tensor(image: Image.Image | List[Image.Image]) -> torch.Tensor:
+    if isinstance(image, list):
+        return torch.cat([pil2tensor(img) for img in image], dim=0)
 
-#     if batch_count == 1:
-#         return Image.fromarray(
-#             np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
-#         )
-#     return [tensor2pil(image[i]) for i in range(batch_count)]
-
-
-# Convert PIL to Tensor (grabbed from WAS Suite)
-def pil2tensor(image: Image.Image) -> torch.Tensor:
     return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
 
 
-def np2tensor(img_np):
+def np2tensor(img_np: np.ndarray | List[np.ndarray]) -> torch.Tensor:
+    if isinstance(img_np, list):
+        return torch.cat([np2tensor(img) for img in img_np], dim=0)
+
     return torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0)
 
 
-def tensor2np(tensor: torch.Tensor) -> np.ndarray:
+def tensor2np(tensor: torch.Tensor) -> Union[np.ndarray, List[np.ndarray]]:
+    batch_count = 1
+    if len(tensor.shape) > 3:
+        batch_count = tensor.size(0)
+    if batch_count > 1:
+        return [tensor2np(tensor[i]) for i in range(batch_count)]
     return np.clip(255.0 * tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
-
-
-