update
This commit is contained in:
KUN
@@ -0,0 +1,33 @@
|
||||
# ComfyUI nodes to use M3net
|
||||
|
||||
https://github.com/SingleZombie/DiffusersExample/tree/main/ReplaceAttn
|
||||
|
||||
|
||||
|
||||
## Install
|
||||
|
||||
- 推荐使用管理器 ComfyUI Manager 安装（On the Way）
|
||||
|
||||
- 手动安装：
|
||||
1. `cd custom_nodes`
|
||||
2. `git clone https://github.com/leeguandong/ComfyUI_VideoEditing.git`
|
||||
3. 重启 ComfyUI
|
||||
|
||||
|
||||
## Use
|
||||
|
||||
- 节点：
|
||||
- "Load Video to Images": 将视频转成图片
|
||||
- "Image to Canny": controlnet使用的canny算子
|
||||
- "ControlNet Model Loader": controlnet模型加载
|
||||
- "VEDit Model Loader": SD 模型加载
|
||||
- "VEdit Sampler": 推理
|
||||
|
||||
|
||||
|
||||
## Stars
|
||||
|
||||
[](https://star-history.com/#leeguandong/ComfyUI_VideoEditing&Date)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,292 @@
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
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,
|
||||
)
|
||||
# pipe.to(offload_device)
|
||||
# pipe.to(offload_device)
|
||||
# image_out = out_frames.squeeze(0).permute(1, 2, 3, 0).cpu().float()
|
||||
# return (image_out,)
|
||||
out_frames = torch.cat(out_frames)
|
||||
out_frames = torch.from_numpy(out_frames)
|
||||
return (out_frames,)
|
||||
|
||||
|
||||
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"
|
||||
}
|
||||
BIN
Binary file not shown.
|
After Width: | Height: | Size: 974 KiB |
@@ -0,0 +1,2 @@
|
||||
opencv-python
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
from diffusers import StableDiffusionControlNetImg2ImgPipeline, AutoencoderKL, UNet2DConditionModel
|
||||
from diffusers.schedulers import KarrasDiffusionSchedulers
|
||||
from diffusers.models.attention_processor import Attention, AttnProcessor
|
||||
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
|
||||
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
|
||||
import torch
|
||||
|
||||
|
||||
class AttnState:
|
||||
STORE = 0
|
||||
LOAD = 1
|
||||
|
||||
def __init__(self):
|
||||
self.reset()
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
return self.__state
|
||||
|
||||
@property
|
||||
def timestep(self):
|
||||
return self.__timestep
|
||||
|
||||
def set_timestep(self, t):
|
||||
self.__timestep = t
|
||||
|
||||
def reset(self):
|
||||
self.__state = AttnState.STORE
|
||||
self.__timestep = 0
|
||||
|
||||
def to_load(self):
|
||||
self.__state = AttnState.LOAD
|
||||
|
||||
|
||||
class CrossFrameAttnProcessor(AttnProcessor):
|
||||
"""
|
||||
Cross frame attention processor. Each frame attends the first frame and previous frame.
|
||||
|
||||
Args:
|
||||
attn_state: Whether the model is processing the first frame or an intermediate frame
|
||||
"""
|
||||
|
||||
def __init__(self, attn_state: AttnState):
|
||||
super().__init__()
|
||||
self.attn_state = attn_state
|
||||
self.cur_timestep = 0
|
||||
self.first_maps = {}
|
||||
self.prev_maps = {}
|
||||
|
||||
def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, **kwargs):
|
||||
|
||||
if encoder_hidden_states is None: # kv
|
||||
# Is self attention
|
||||
|
||||
tot_timestep = self.attn_state.timestep
|
||||
if self.attn_state.state == AttnState.STORE:
|
||||
self.first_maps[self.cur_timestep] = hidden_states.detach()
|
||||
self.prev_maps[self.cur_timestep] = hidden_states.detach()
|
||||
res = super().__call__(attn, hidden_states, encoder_hidden_states, **kwargs)
|
||||
else:
|
||||
tmp = hidden_states.detach()
|
||||
cross_map = torch.cat(
|
||||
(self.first_maps[self.cur_timestep], self.prev_maps[self.cur_timestep]), dim=1)
|
||||
res = super().__call__(attn, hidden_states, cross_map, **kwargs)
|
||||
self.prev_maps[self.cur_timestep] = tmp
|
||||
|
||||
self.cur_timestep += 1
|
||||
if self.cur_timestep == tot_timestep:
|
||||
self.cur_timestep = 0
|
||||
else:
|
||||
# Is cross attention
|
||||
res = super().__call__(attn, hidden_states, encoder_hidden_states, **kwargs)
|
||||
|
||||
return res
|
||||
|
||||
|
||||
class VideoEditingPipeline(StableDiffusionControlNetImg2ImgPipeline):
|
||||
def __init__(
|
||||
self,
|
||||
vae: AutoencoderKL,
|
||||
text_encoder: CLIPTextModel,
|
||||
tokenizer: CLIPTokenizer,
|
||||
unet: UNet2DConditionModel,
|
||||
controlnet,
|
||||
scheduler: KarrasDiffusionSchedulers,
|
||||
safety_checker: StableDiffusionSafetyChecker,
|
||||
feature_extractor: CLIPImageProcessor,
|
||||
image_encoder: CLIPVisionModelWithProjection = None,
|
||||
requires_safety_checker: bool = True,
|
||||
):
|
||||
super().__init__(vae, text_encoder, tokenizer, unet, controlnet, scheduler,
|
||||
safety_checker, feature_extractor, image_encoder, requires_safety_checker)
|
||||
self.attn_state = AttnState()
|
||||
attn_processor_dict = {}
|
||||
for k in unet.attn_processors.keys():
|
||||
if k.startswith("up"):
|
||||
attn_processor_dict[k] = CrossFrameAttnProcessor(
|
||||
self.attn_state)
|
||||
else:
|
||||
attn_processor_dict[k] = AttnProcessor()
|
||||
|
||||
self.unet.set_attn_processor(attn_processor_dict)
|
||||
|
||||
def __call__(self, *args, images=None, control_images=None, **kwargs):
|
||||
self.attn_state.reset()
|
||||
self.attn_state.set_timestep(
|
||||
int(kwargs['num_inference_steps'] * kwargs['strength']))
|
||||
outputs = [super().__call__(
|
||||
*args, **kwargs, image=images[0], control_image=control_images[0]).images[0]]
|
||||
self.attn_state.to_load()
|
||||
for i in range(1, len(images)):
|
||||
image = images[i]
|
||||
control_image = control_images[i]
|
||||
outputs.append(super().__call__(
|
||||
*args, **kwargs, image=image, control_image=control_image).images[0])
|
||||
return outputs
|
||||
Reference in New Issue
Block a user