359 lines
12 KiB
Python
359 lines
12 KiB
Python
import yaml
|
|
import os
|
|
import sys
|
|
import warnings
|
|
import logging
|
|
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
|
|
import time
|
|
|
|
# -------------------------------
|
|
# Debug / Logging Setup
|
|
# -------------------------------
|
|
DEBUG = False
|
|
|
|
os.environ["ORT_LOG_LEVEL"] = "3"
|
|
os.environ["ORT_CUDA_MEMORY_LIMIT"] = "0"
|
|
|
|
warnings.simplefilter("ignore", category=FutureWarning)
|
|
warnings.simplefilter("ignore", category=DeprecationWarning)
|
|
warnings.simplefilter("ignore", category=UserWarning)
|
|
warnings.filterwarnings("ignore")
|
|
|
|
logging.getLogger("insightface").setLevel(logging.ERROR)
|
|
logging.getLogger("urllib3").setLevel(logging.ERROR)
|
|
logging.getLogger("onnxruntime").setLevel(logging.ERROR)
|
|
logging.basicConfig(level=logging.ERROR)
|
|
|
|
class suppress_stdout_stderr:
|
|
"""Suppress all console output temporarily"""
|
|
def __enter__(self):
|
|
self._stdout = sys.stdout
|
|
self._stderr = sys.stderr
|
|
sys.stdout = open(os.devnull, "w")
|
|
sys.stderr = open(os.devnull, "w")
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
sys.stdout.close()
|
|
sys.stderr.close()
|
|
sys.stdout = self._stdout
|
|
sys.stderr = self._stderr
|
|
|
|
def log(msg, progress=None):
|
|
if DEBUG:
|
|
print(msg)
|
|
if progress is not None:
|
|
pass
|
|
|
|
def format_eta(seconds: float):
|
|
if seconds < 60:
|
|
return f"{seconds:.1f}s"
|
|
minutes = int(seconds // 60)
|
|
seconds = seconds % 60
|
|
return f"{minutes}m {seconds:.0f}s"
|
|
|
|
# -------------------------------
|
|
# Constants
|
|
# -------------------------------
|
|
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/"
|
|
]
|
|
|
|
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)
|
|
else:
|
|
raise Exception("config.yaml file is necessary. Recreate it from the original repository.")
|
|
|
|
DEVICE = get_torch_device()
|
|
|
|
# -------------------------------
|
|
# Core Classes
|
|
# -------------------------------
|
|
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(',')]
|
|
return (InterpolationStateList(frame_indices_list, is_skip_list),)
|
|
|
|
# -------------------------------
|
|
# Download Helpers
|
|
# -------------------------------
|
|
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):
|
|
if model_dir is None:
|
|
hub_dir = get_dir()
|
|
model_dir = os.path.join(hub_dir, 'checkpoints')
|
|
|
|
os.makedirs(model_dir, exist_ok=True)
|
|
|
|
file_name = os.path.basename(urlparse(url).path)
|
|
cached_file = os.path.abspath(os.path.join(model_dir, file_name))
|
|
|
|
if not os.path.exists(cached_file):
|
|
log(f'Downloading: "{url}" -> {cached_file}')
|
|
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):
|
|
error_strs = []
|
|
for i, base in enumerate(BASE_MODEL_DOWNLOAD_URLS):
|
|
try:
|
|
return load_file_from_url(base + ckpt_name, get_ckpt_container_path(model_type))
|
|
except Exception:
|
|
if i < len(BASE_MODEL_DOWNLOAD_URLS) - 1:
|
|
log("Failed endpoint, trying next…")
|
|
error_strs.append(traceback.format_exc())
|
|
|
|
raise Exception(f"Failed to download {ckpt_name} from all sources:\n\n" + "\n".join(error_strs))
|
|
|
|
def load_file_from_direct_url(model_type, url):
|
|
return load_file_from_url(url, get_ckpt_container_path(model_type))
|
|
|
|
# -------------------------------
|
|
# Frame Helpers
|
|
# -------------------------------
|
|
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):
|
|
label = "Most VFI models require" if vfi_name is None else f"VFI model {vfi_name} requires"
|
|
assert len(frames) >= batch_size, f"{label} at least {batch_size} frames, got {frames.shape[0]}."
|
|
|
|
# -------------------------------
|
|
# Core Frame Loop
|
|
# -------------------------------
|
|
def _generic_frame_loop(
|
|
frames,
|
|
clear_cache_after_n_frames,
|
|
multiplier,
|
|
return_middle_frame_function,
|
|
*return_middle_frame_function_args,
|
|
interpolation_states=None,
|
|
use_timestep=True,
|
|
dtype=torch.float16,
|
|
final_logging=True
|
|
):
|
|
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 = non_timestep_inference(frame0, middle, n=n // 2)
|
|
second = non_timestep_inference(middle, frame1, n=n // 2)
|
|
return [*first, middle, *second] if n % 2 else [*first, *second]
|
|
|
|
output_frames = torch.zeros(multiplier * frames.shape[0], *frames.shape[1:], dtype=dtype, device="cpu")
|
|
out_len = 0
|
|
cache_counter = 0
|
|
|
|
# ----- Progress bar setup -----
|
|
total_steps = (len(frames) - 1) * (multiplier - 1)
|
|
bar_length = 30
|
|
start_time = time.time()
|
|
last_print = 0
|
|
|
|
for frame_itr in range(len(frames) - 1):
|
|
frame0 = frames[frame_itr:frame_itr + 1]
|
|
output_frames[out_len] = frame0
|
|
out_len += 1
|
|
|
|
frame0 = frame0.to(dtype=torch.float32)
|
|
frame1 = frames[frame_itr + 1:frame_itr + 2].to(dtype=torch.float32)
|
|
|
|
if interpolation_states and interpolation_states.is_frame_skipped(frame_itr):
|
|
continue
|
|
|
|
middle_frame_batches = []
|
|
|
|
if use_timestep:
|
|
for i in range(1, multiplier):
|
|
timestep = i / multiplier
|
|
mid = return_middle_frame_function(
|
|
frame0.to(DEVICE),
|
|
frame1.to(DEVICE),
|
|
timestep,
|
|
*return_middle_frame_function_args
|
|
).detach().cpu()
|
|
middle_frame_batches.append(mid.to(dtype=dtype))
|
|
else:
|
|
mids = non_timestep_inference(frame0.to(DEVICE), frame1.to(DEVICE), multiplier - 1)
|
|
middle_frame_batches.extend(torch.cat(mids, dim=0).detach().cpu().to(dtype=dtype))
|
|
|
|
for mf in middle_frame_batches:
|
|
output_frames[out_len] = mf
|
|
out_len += 1
|
|
|
|
# ----- Progress bar + ETA -----
|
|
processed_steps = (frame_itr * (multiplier - 1)) + len(middle_frame_batches)
|
|
now = time.time()
|
|
|
|
# throttle prints to avoid spam
|
|
if now - last_print > 0.05 or processed_steps == total_steps:
|
|
elapsed = now - start_time
|
|
avg = elapsed / max(processed_steps, 1)
|
|
remaining = avg * (total_steps - processed_steps)
|
|
|
|
filled = int(bar_length * processed_steps // max(total_steps, 1))
|
|
bar = "█" * filled + "-" * (bar_length - filled)
|
|
|
|
sys.stdout.write(
|
|
f"\r[Rife] Progress: |{bar}| "
|
|
f"{processed_steps}/{total_steps} "
|
|
f"ETA: {format_eta(remaining)}"
|
|
)
|
|
sys.stdout.flush()
|
|
last_print = now
|
|
|
|
cache_counter += 1
|
|
if cache_counter >= clear_cache_after_n_frames:
|
|
soft_empty_cache()
|
|
cache_counter = 0
|
|
|
|
gc.collect()
|
|
|
|
output_frames[out_len] = frames[-1:]
|
|
out_len += 1
|
|
|
|
soft_empty_cache()
|
|
|
|
# Finalize bar
|
|
sys.stdout.write("\r[Rife] Progress: |██████████████████████████████| 100% ETA: 0s\n")
|
|
sys.stdout.flush()
|
|
|
|
return output_frames[:out_len]
|
|
|
|
# -------------------------------
|
|
# Public Frame Loop
|
|
# -------------------------------
|
|
def generic_frame_loop(
|
|
model_name,
|
|
frames,
|
|
clear_cache_after_n_frames,
|
|
multiplier,
|
|
return_middle_frame_function,
|
|
*return_middle_frame_function_args,
|
|
interpolation_states=None,
|
|
use_timestep=True,
|
|
dtype=torch.float32
|
|
):
|
|
assert_batch_size(frames, vfi_name=model_name.replace('_', ' ').replace('VFI', ''))
|
|
|
|
if isinstance(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 isinstance(multiplier, list):
|
|
multipliers = list(map(int, multiplier))
|
|
multipliers += [2] * (len(frames) - len(multipliers) - 1)
|
|
frame_batches = []
|
|
|
|
for i in range(len(frames) - 1):
|
|
m = multipliers[i]
|
|
if m == 0:
|
|
continue
|
|
|
|
batch = _generic_frame_loop(
|
|
frames[i:i + 2],
|
|
clear_cache_after_n_frames,
|
|
m,
|
|
return_middle_frame_function,
|
|
*return_middle_frame_function_args,
|
|
interpolation_states=interpolation_states,
|
|
use_timestep=use_timestep,
|
|
dtype=dtype,
|
|
final_logging=False
|
|
)
|
|
|
|
if i != len(frames) - 2:
|
|
batch = batch[:-1]
|
|
|
|
frame_batches.append(batch)
|
|
|
|
output = torch.cat(frame_batches)
|
|
log(f"Done! {len(output)} frames generated.", 1.0)
|
|
return output
|
|
|
|
raise NotImplementedError(f"Unsupported multiplier: {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] """ |