more extensions

This commit is contained in:
salt
2024-02-20 08:30:59 -08:00
parent 8048a05d12
commit 7053eec79c
145 changed files with 7747 additions and 2775 deletions
@@ -25,8 +25,6 @@ s_param = '-s' if "python_embeded" in sys.executable else ''
def get_cuda_home_path():
if "CUDA_HOME" in os.environ:
return os.environ["CUDA_HOME"]
if platform.system() == "Windows":
return str(Path(__file__).parent / "nvrtc_dlls") #https://github.com/cupy/cupy/issues/7776
import torch
torch_lib_path = Path(torch.__file__).parent / "lib"
torch_lib_path = str(torch_lib_path.resolve())
@@ -41,8 +41,7 @@ class AMT_VFI:
"multiplier": ("INT", {"default": 2, "min": 2, "max": 1000})
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -57,7 +56,7 @@ class AMT_VFI:
clear_cache_after_n_frames: typing.SupportsInt = 1,
multiplier: typing.SupportsInt = 2,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
model_path = load_file_from_direct_url(MODEL_TYPE, f"https://huggingface.co/lalala125/AMT/resolve/main/{ckpt_name}")
ckpt_config = CKPT_CONFIGS[ckpt_name]
@@ -81,7 +80,7 @@ class AMT_VFI:
args = [interpolation_model]
out = generic_frame_loop(frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args,
interpolation_states=optional_interpolation_states, dtype=torch.float16 if cache_in_fp16 else torch.float32)
interpolation_states=optional_interpolation_states, dtype=torch.float32)
out = padder.unpad(out)
out = postprocess_frames(out)
return (out,)
@@ -20,8 +20,7 @@ class CAIN_VFI:
"multiplier": ("INT", {"default": 2, "min": 2, "max": 1000})
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -36,7 +35,7 @@ class CAIN_VFI:
clear_cache_after_n_frames: typing.SupportsInt = 1,
multiplier: typing.SupportsInt = 2,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
from .cain_arch import CAIN
model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name)
@@ -60,6 +59,6 @@ class CAIN_VFI:
args = [interpolation_model]
out = postprocess_frames(
generic_frame_loop(frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args,
interpolation_states=optional_interpolation_states, use_timestep=False, dtype=torch.float16 if cache_in_fp16 else torch.float32)
interpolation_states=optional_interpolation_states, use_timestep=False, dtype=torch.float32)
)
return (out,)
@@ -50,8 +50,7 @@ class EISAI_VFI:
"multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}),
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -66,7 +65,7 @@ class EISAI_VFI:
clear_cache_after_n_frames = 10,
multiplier: typing.SupportsInt = 2,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
interpolation_model = EISAI(MODEL_FILE_NAMES)
interpolation_model.eval().to(get_torch_device())
@@ -80,6 +79,6 @@ class EISAI_VFI:
args = [interpolation_model, scale]
out = postprocess_frames(
generic_frame_loop(frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args,
interpolation_states=optional_interpolation_states, dtype=torch.float16 if cache_in_fp16 else torch.float32)
interpolation_states=optional_interpolation_states, dtype=torch.float32)
)
return (out,)
@@ -52,8 +52,7 @@ class FILM_VFI:
"multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}),
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -68,14 +67,14 @@ class FILM_VFI:
clear_cache_after_n_frames = 10,
multiplier: typing.SupportsInt = 2,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
interpolation_states = optional_interpolation_states
model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name)
model = torch.jit.load(model_path, map_location='cpu')
model.eval()
model = model.to(DEVICE)
dtype = torch.float16 if cache_in_fp16 else torch.float32
dtype = torch.float32
frames = preprocess_frames(frames)
number_of_frames_processed_since_last_cleared_cuda_cache = 0
@@ -37,8 +37,7 @@ class FLAVR_VFI:
"duplicate_first_last_frames": ("BOOLEAN", {"default": False})
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -55,7 +54,7 @@ class FLAVR_VFI:
multiplier: typing.SupportsInt = 2,
duplicate_first_last_frames: bool = False,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
if multiplier != 2:
warnings.warn("Currently, FLAVR only supports 2x interpolation. The process will continue but please set multiplier=2 afterward")
@@ -105,7 +104,7 @@ class FLAVR_VFI:
print("Done cache clearing")
gc.collect()
dtype = torch.float16 if cache_in_fp16 else torch.float32
dtype = torch.float32
output_frames = [frame.cpu().to(dtype=dtype) for frame in output_frames] #Ensure all frames are in cpu
out = torch.cat(output_frames, dim=0)
out = padder.unpad(out)
@@ -87,8 +87,7 @@ class GMFSS_Fortuna_VFI:
"multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}),
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -103,7 +102,7 @@ class GMFSS_Fortuna_VFI:
clear_cache_after_n_frames = 10,
multiplier: typing.SupportsInt = 2,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
"""
Perform video frame interpolation using a given checkpoint model.
@@ -139,6 +138,6 @@ class GMFSS_Fortuna_VFI:
args = [interpolation_model, scale]
out = postprocess_frames(
generic_frame_loop(frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args,
interpolation_states=optional_interpolation_states, dtype=torch.float16 if cache_in_fp16 else torch.float32)
interpolation_states=optional_interpolation_states, dtype=torch.float32)
)
return (out,)
@@ -20,8 +20,7 @@ class IFRNet_VFI:
"scale_factor": ([0.25, 0.5, 1.0, 2.0, 4.0], {"default": 1.0}),
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -37,7 +36,7 @@ class IFRNet_VFI:
multiplier: typing.SupportsInt = 2,
scale_factor: typing.SupportsFloat = 1.0,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
from .IFRNet_S_arch import IRFNet_S
from .IFRNet_L_arch import IRFNet_L
@@ -53,6 +52,6 @@ class IFRNet_VFI:
args = [interpolation_model, scale_factor]
out = postprocess_frames(
generic_frame_loop(frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args,
interpolation_states=optional_interpolation_states, dtype=torch.float16 if cache_in_fp16 else torch.float32)
interpolation_states=optional_interpolation_states, dtype=torch.float32)
)
return (out,)
@@ -21,8 +21,7 @@ class IFUnet_VFI:
"ensemble": ("BOOLEAN", {"default":True})
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -39,7 +38,7 @@ class IFUnet_VFI:
scale_factor: typing.SupportsFloat = 1.0,
ensemble: bool = True,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
from .IFUNet_arch import IFUNetModel
model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name)
@@ -54,7 +53,7 @@ class IFUnet_VFI:
args = [interpolation_model, scale_factor, ensemble]
out = postprocess_frames(
generic_frame_loop(frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args,
interpolation_states=optional_interpolation_states, dtype=torch.float16 if cache_in_fp16 else torch.float32)
interpolation_states=optional_interpolation_states, dtype=torch.float32)
)
return (out,)
@@ -22,8 +22,7 @@ class M2M_VFI:
"multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}),
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -38,7 +37,7 @@ class M2M_VFI:
clear_cache_after_n_frames: typing.SupportsInt = 1,
multiplier: typing.SupportsInt = 2,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
from .M2M_arch import M2M_PWC
model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name)
@@ -56,6 +55,6 @@ class M2M_VFI:
args = [interpolation_model]
out = postprocess_frames(
generic_frame_loop(frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args,
interpolation_states=optional_interpolation_states, dtype=torch.float16 if cache_in_fp16 else torch.float32)
interpolation_states=optional_interpolation_states, dtype=torch.float32)
)
return (out,)
@@ -217,8 +217,6 @@ def cuda_kernel(strFunction: str, strKernel: str, objVariables: typing.Dict, **r
def get_cuda_home_path():
if "CUDA_HOME" in os.environ:
return os.environ["CUDA_HOME"]
if platform.system() == "Windows":
return str(Path(__file__).parent.parent.parent.parent / "nvrtc_dlls") #https://github.com/cupy/cupy/issues/7776
import torch
torch_lib_path = Path(torch.__file__).parent / "lib"
torch_lib_path = str(torch_lib_path.resolve())
@@ -6,6 +6,7 @@ import typing
from comfy.model_management import get_torch_device
import re
from functools import cmp_to_key
from packaging import version
MODEL_TYPE = pathlib.Path(__file__).parent.name
CKPT_NAME_VER_DICT = {
@@ -19,14 +20,13 @@ CKPT_NAME_VER_DICT = {
"rife47.pth": "4.7",
"rife48.pth": "4.7",
"rife49.pth": "4.7",
"sudo_rife4_269.662_testV1_scale1.pth": "4.0"
#Arch 4.10 doesn't work due to state dict mismatch
#TODO: Investigating and fix it
#"rife410.pth": "4.10",
#"rife411.pth": "4.10",
#"rife412.pth": "4.10"
}
ver_re = re.compile(r'\d+')
ver_cmp_key = cmp_to_key(lambda a,b: int(ver_re.search(a)[0]) > int(ver_re.search(b)[0]))
class RIFE_VFI:
@classmethod
@@ -34,7 +34,7 @@ class RIFE_VFI:
return {
"required": {
"ckpt_name": (
sorted(list(CKPT_NAME_VER_DICT.keys()), key=ver_cmp_key),
sorted(list(CKPT_NAME_VER_DICT.keys()), key=lambda ckpt_name: version.parse(CKPT_NAME_VER_DICT[ckpt_name])),
{"default": "rife47.pth"}
),
"frames": ("IMAGE", ),
@@ -45,8 +45,7 @@ class RIFE_VFI:
"scale_factor": ([0.25, 0.5, 1.0, 2.0, 4.0], {"default": 1.0})
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -64,7 +63,7 @@ class RIFE_VFI:
ensemble = False,
scale_factor = 1.0,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
"""
Perform video frame interpolation using a given checkpoint model.
@@ -103,6 +102,6 @@ class RIFE_VFI:
args = [interpolation_model, scale_list, fast_mode, ensemble]
out = postprocess_frames(
generic_frame_loop(frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args,
interpolation_states=optional_interpolation_states, dtype=torch.float16 if cache_in_fp16 else torch.float32)
interpolation_states=optional_interpolation_states, dtype=torch.float32)
)
return (out,)
@@ -21,8 +21,7 @@ class SepconvVFI:
"multiplier": ("INT", {"default": 2, "min": 2, "max": 1000})
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -37,7 +36,7 @@ class SepconvVFI:
clear_cache_after_n_frames = 10,
multiplier: typing.SupportsInt = 2,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
from .sepconv_enhanced import Network
model_path = load_file_from_github_release(MODEL_TYPE, ckpt_name)
@@ -52,6 +51,6 @@ class SepconvVFI:
args = [interpolation_model]
out = postprocess_frames(
generic_frame_loop(frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args,
interpolation_states=optional_interpolation_states, use_timestep=False, dtype=torch.float16 if cache_in_fp16 else torch.float32)
interpolation_states=optional_interpolation_states, use_timestep=False, dtype=torch.float32)
)
return (out,)
@@ -22,8 +22,7 @@ class STMFNet_VFI:
"duplicate_first_last_frames": ("BOOLEAN", {"default": False})
},
"optional": {
"optional_interpolation_states": ("INTERPOLATION_STATES", ),
"cache_in_fp16": ("BOOLEAN", {"default": True})
"optional_interpolation_states": ("INTERPOLATION_STATES", )
}
}
@@ -40,7 +39,7 @@ class STMFNet_VFI:
multiplier: typing.SupportsInt = 2,
duplicate_first_last_frames: bool = False,
optional_interpolation_states: InterpolationStateList = None,
cache_in_fp16: bool = True
**kwargs
):
from .stmfnet_arch import STMFNet_Model
if multiplier != 2:
@@ -91,7 +90,7 @@ class STMFNet_VFI:
print("Comfy-VFI: Done cache clearing")
gc.collect()
dtype = torch.float16 if cache_in_fp16 else torch.float32
dtype = torch.float32
output_frames = [frame.cpu().to(dtype=dtype) for frame in output_frames] #Ensure all frames are in cpu
out = torch.cat(output_frames, dim=0)
# clear cache for courtesy
@@ -144,15 +144,18 @@ def generic_frame_loop(
return [*first_half, *second_half]
assert_batch_size(frames) # Too lazy to include model name lol
output_frames = [] # List to store processed frames in the correct order
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
#Ensure that input frames are in fp32 - the same dtype as model
frame0 = frames[frame_itr:frame_itr+1].float()
frame1 = frames[frame_itr+1:frame_itr+2].float()
output_frames.append(frame0.to(dtype=dtype)) # Start with first frame
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
@@ -175,8 +178,10 @@ def generic_frame_loop(
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))
# Extend output array by batch
output_frames.extend(middle_frame_batches)
# 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
@@ -189,14 +194,14 @@ def generic_frame_loop(
gc.collect()
print(f"Comfy-VFI done! {len(output_frames)} frames generated at resolution: {output_frames[0].shape}")
output_frames.append(frames[-1:].to(dtype=dtype)) # Append final frame
output_frames = [frame.cpu() for frame in output_frames] #Ensure all frames are in cpu
out = torch.cat(output_frames, dim=0)
# Append final frame
output_frames[out_len] = frames[-1:]
out_len += 1
# clear cache for courtesy
print("Comfy-VFI: Final clearing cache...", end = ' ')
soft_empty_cache()
print("Done cache clearing")
return out
return output_frames[:out_len]
""" def generic_4frame_loop(
frames,