commit 97698dbef41edff0b76942ea9728fcc1a61aeaea Author: KUN Date: Tue Aug 13 08:03:27 2024 +0800 update diff --git a/README.md b/README.md new file mode 100644 index 0000000..c2ea6dd --- /dev/null +++ b/README.md @@ -0,0 +1,33 @@ +# ComfyUI nodes to use M3net + +https://github.com/SingleZombie/DiffusersExample/tree/main/ReplaceAttn + +![image](preview.png) + +## 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 + +[![Star History Chart](https://api.star-history.com/svg?repos=leeguandong/ComfyUI_VideoEditing&type=Date)](https://star-history.com/#leeguandong/ComfyUI_VideoEditing&Date) + + + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..b3700cf --- /dev/null +++ b/nodes.py @@ -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" +} diff --git a/preview.png b/preview.png new file mode 100644 index 0000000..bb33fca Binary files /dev/null and b/preview.png differ diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a512a8a --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +opencv-python + diff --git a/video_editing_pipeline.py b/video_editing_pipeline.py new file mode 100644 index 0000000..80174ed --- /dev/null +++ b/video_editing_pipeline.py @@ -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