291 lines
10 KiB
Python
291 lines
10 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
from diffusers import ControlNetModel
|
|
from diffusers import (
|
|
DDIMScheduler,
|
|
DPMSolverMultistepScheduler,
|
|
EulerDiscreteScheduler,
|
|
EulerAncestralDiscreteScheduler,
|
|
AutoencoderKL,
|
|
LCMScheduler,
|
|
DDPMScheduler,
|
|
DEISMultistepScheduler,
|
|
PNDMScheduler
|
|
)
|
|
|
|
from .video_editing_pipeline import VideoEditingPipeline
|
|
from contextlib import nullcontext
|
|
|
|
import folder_paths
|
|
import comfy.model_management as mm
|
|
|
|
|
|
class LoadVideo2Images:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
video_extensions = ['webm', 'mp4', 'mkv', 'gif']
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = []
|
|
for f in os.listdir(input_dir):
|
|
if os.path.isfile(os.path.join(input_dir, f)):
|
|
file_parts = f.split('.')
|
|
if len(file_parts) > 1 and (file_parts[-1] in video_extensions):
|
|
files.append(f)
|
|
return {"required": {
|
|
"video": (sorted(files), {"video_upload": True}),
|
|
"input_interval": ("INT", {"default": 10, "min": -1, "step": 1, "max": 50, "display": "slider"}),
|
|
}, }
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "load_video"
|
|
CATEGORY = "vdeiting"
|
|
|
|
# OUTPUT_IS_LIST = (True,)
|
|
|
|
def load_video(self, video, input_interval):
|
|
# import pdb;pdb.set_trace()
|
|
video_path = folder_paths.get_annotated_filepath(video)
|
|
frames = self.video_to_frame(video_path, input_interval)
|
|
# images = {"frames": frames}
|
|
return (frames,)
|
|
|
|
def video_to_frame(self, video_path: str, interval: int):
|
|
vidcap = cv2.VideoCapture(video_path)
|
|
success = True
|
|
|
|
count = 0
|
|
res = []
|
|
while success:
|
|
count += 1
|
|
success, image = vidcap.read()
|
|
if count % interval != 1:
|
|
continue
|
|
if image is not None:
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
image = cv2.resize(image[:, 100:800], (512, 512))
|
|
res.append(image)
|
|
|
|
vidcap.release()
|
|
return res
|
|
|
|
|
|
class VEdit_image2canny:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"low_threshold": ("INT", {"default": 50, "min": 0, "max": 255, "step": 1, "display": "slider"}),
|
|
"high_threshold": ("INT", {"default": 100, "min": 0, "max": 255, "step": 1, "display": "slider"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONTROL_IMAGE",)
|
|
RETURN_NAMES = ("control_image",)
|
|
FUNCTION = "image2canny"
|
|
CATEGORY = "vdeiting"
|
|
|
|
def image2canny(self, image, low_threshold, high_threshold):
|
|
# import pdb;pdb.set_trace()
|
|
# frames = image["frames"]
|
|
frames = image
|
|
control_images = []
|
|
if not isinstance(frames, list):
|
|
frames = [frames]
|
|
for frame in frames:
|
|
np_image = cv2.Canny(frame, low_threshold, high_threshold)
|
|
np_image = np_image[:, :, None]
|
|
np_image = np.concatenate([np_image, np_image, np_image], axis=2)
|
|
canny_image = Image.fromarray(np_image)
|
|
control_images.append(canny_image)
|
|
return (control_images,)
|
|
|
|
|
|
class VEdit_ControlNet_ModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ([
|
|
"lllyasviel/sd-controlnet-canny"
|
|
],
|
|
{"default": "lllyasviel/sd-controlnet-canny"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_controlnet"
|
|
CATEGORY = "vdeiting"
|
|
|
|
def load_controlnet(self, model):
|
|
device = mm.get_torch_device()
|
|
controlnet = ControlNetModel.from_pretrained(model).to(device)
|
|
return (controlnet,)
|
|
|
|
|
|
class VEdit_ModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"base_model": (
|
|
[
|
|
"runwayml/stable-diffusion-v1-5",
|
|
"stabilityai/stable-diffusion-2-1",
|
|
],
|
|
{
|
|
"default": "runwayml/stable-diffusion-v1-5",
|
|
}),
|
|
"controlnet": ("MODEL",)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "vdeiting"
|
|
|
|
def load_model(self, base_model, controlnet):
|
|
device = mm.get_torch_device()
|
|
pipeline = VideoEditingPipeline.from_pretrained(base_model,
|
|
controlnet=controlnet,
|
|
torch_dtype=torch.float16).to(device)
|
|
pipeline.safety_checker = None
|
|
|
|
return (pipeline,)
|
|
|
|
|
|
class VEdit_Sampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"image": ("IMAGE",),
|
|
"control_image": ("IMAGE",),
|
|
"prompt": ("STRING", {"multiline": True, "default": "a beautiful woman with red hair", }),
|
|
"n_prompt": ("STRING", {"multiline": True,
|
|
"default": "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality", }),
|
|
"steps": ("INT", {"default": 25, "min": 1, "max": 200, "step": 1, "display": "slider"}),
|
|
"controlnet_conditioning_scale": (
|
|
"FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01, "display": "slider"}),
|
|
"strength": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01, "display": "slider"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "display": "slider"}),
|
|
"scheduler": (
|
|
[
|
|
'DPMSolverMultistepScheduler',
|
|
'DPMSolverMultistepScheduler_SDE_karras',
|
|
'DDPMScheduler',
|
|
'DDIMScheduler',
|
|
'LCMScheduler',
|
|
'PNDMScheduler',
|
|
'DEISMultistepScheduler',
|
|
'EulerDiscreteScheduler',
|
|
'EulerAncestralDiscreteScheduler'
|
|
], {
|
|
"default": 'DPMSolverMultistepScheduler'
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "sampler"
|
|
CATEGORY = "vdeiting"
|
|
|
|
def sampler(self,
|
|
model,
|
|
image,
|
|
control_image,
|
|
prompt,
|
|
n_prompt,
|
|
steps,
|
|
controlnet_conditioning_scale,
|
|
strength,
|
|
seed,
|
|
scheduler):
|
|
device = mm.get_torch_device()
|
|
mm.unload_all_models()
|
|
mm.soft_empty_cache()
|
|
dtype = mm.unet_dtype()
|
|
offload_device = mm.unet_offload_device()
|
|
|
|
pipe = model
|
|
scheduler_config = {
|
|
'num_train_timesteps': 1000,
|
|
'beta_start': 0.00085,
|
|
'beta_end': 0.012,
|
|
'beta_schedule': "linear",
|
|
'steps_offset': 1,
|
|
}
|
|
if scheduler == 'DPMSolverMultistepScheduler':
|
|
noise_scheduler = DPMSolverMultistepScheduler(**scheduler_config)
|
|
elif scheduler == 'DDIMScheduler':
|
|
noise_scheduler = DDIMScheduler(**scheduler_config)
|
|
elif scheduler == 'DPMSolverMultistepScheduler_SDE_karras':
|
|
scheduler_config.update({"algorithm_type": "sde-dpmsolver++"})
|
|
scheduler_config.update({"use_karras_sigmas": True})
|
|
noise_scheduler = DPMSolverMultistepScheduler(**scheduler_config)
|
|
elif scheduler == 'DDPMScheduler':
|
|
noise_scheduler = DDPMScheduler(**scheduler_config)
|
|
elif scheduler == 'LCMScheduler':
|
|
noise_scheduler = LCMScheduler(**scheduler_config)
|
|
elif scheduler == 'PNDMScheduler':
|
|
scheduler_config.update({"set_alpha_to_one": False})
|
|
scheduler_config.update({"trained_betas": None})
|
|
noise_scheduler = PNDMScheduler(**scheduler_config)
|
|
elif scheduler == 'DEISMultistepScheduler':
|
|
noise_scheduler = DEISMultistepScheduler(**scheduler_config)
|
|
elif scheduler == 'EulerDiscreteScheduler':
|
|
noise_scheduler = EulerDiscreteScheduler(**scheduler_config)
|
|
elif scheduler == 'EulerAncestralDiscreteScheduler':
|
|
noise_scheduler = EulerAncestralDiscreteScheduler(**scheduler_config)
|
|
else:
|
|
raise TypeError(f"not support {scheduler}!!!")
|
|
|
|
# import pdb;pdb.set_trace()
|
|
pipe.scheduler = noise_scheduler
|
|
autocast_condition = (dtype != torch.float32) and not mm.is_device_mps(device)
|
|
image = [Image.fromarray(frame) for frame in image]
|
|
with torch.autocast(mm.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext():
|
|
generator = torch.Generator(device=device)
|
|
generator.manual_seed(seed)
|
|
|
|
out_frames = pipe(
|
|
images=image,
|
|
control_images=control_image,
|
|
prompt=prompt,
|
|
negative_prompt=n_prompt,
|
|
num_inference_steps=steps,
|
|
strength=strength,
|
|
generator=generator,
|
|
controlnet_conditioning_scale=controlnet_conditioning_scale,
|
|
)
|
|
to_tensor = transforms.ToTensor()
|
|
# out_frames = list(map(lambda x: to_tensor(x), out_frames))
|
|
tensors = [to_tensor(img) for img in out_frames]
|
|
transposed_tensors = [tensor.permute(1, 2, 0) for tensor in tensors]
|
|
return (transposed_tensors,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LoadVideo2Images": LoadVideo2Images,
|
|
"VEdit_image2canny": VEdit_image2canny,
|
|
"VEdit_ControlNet_ModelLoader": VEdit_ControlNet_ModelLoader,
|
|
"VEdit_ModelLoader": VEdit_ModelLoader,
|
|
"VEdit_Sampler": VEdit_Sampler
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LoadVideo2Images": "Load Video to Images",
|
|
"VEdit_image2canny": "Image to Canny",
|
|
"VEdit_ControlNet_ModelLoader": "ControlNet Model Loader",
|
|
"VEdit_ModelLoader": "VEDit Model Loader",
|
|
"VEdit_Sampler": "VEdit Sampler"
|
|
}
|