Spaces:
Running
on
L40S
Running
on
L40S
File size: 3,812 Bytes
4450790 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
from pathlib import Path
import comfy
import comfy.model_management as model_management
import comfy.utils
import numpy as np
import tensorflow as tf
import torch
from frame_interpolation.eval import interpolator, util
from ..errors import ModelNotFound
from ..log import log
from ..utils import get_model_path
class MTB_LoadFilmModel:
"""Loads a FILM model
[DEPRECATED] Use ComfyUI-FrameInterpolation instead
"""
@staticmethod
def get_models() -> list[Path]:
models_paths = get_model_path("FILM").iterdir()
return [x for x in models_paths if x.suffix in [".onnx", ".pth"]]
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"film_model": (
["L1", "Style", "VGG"],
{"default": "Style"},
),
},
}
RETURN_TYPES = ("FILM_MODEL",)
FUNCTION = "load_model"
CATEGORY = "mtb/frame iterpolation"
DEPRECATED = True
def load_model(self, film_model: str):
model_path = get_model_path("FILM", film_model)
if not model_path or not model_path.exists():
raise ModelNotFound(f"FILM ({model_path})")
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 MTB_FilmInterpolation:
"""Google Research FILM frame interpolation for large motion
[DEPRECATED] Use ComfyUI-FrameInterpolation instead
"""
@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 = "mtb/frame iterpolation"
DEPRECATED = True
def do_interpolation(
self,
images: torch.Tensor,
interpolate: int,
film_model: interpolator.Interpolator,
):
n = images.size(0)
# check if images is an empty tensor and return it...
if n == 0:
return (images,)
# 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,)
__nodes__ = [MTB_LoadFilmModel, MTB_FilmInterpolation]
|