From ba535c91ddbc2c8e41ac3b208ede740afe05d15e Mon Sep 17 00:00:00 2001 From: smthemex <138738845+smthemex@users.noreply.github.com> Date: Sun, 28 Jul 2024 12:13:20 +0800 Subject: [PATCH] Update streamv2v_node.py --- streamv2v_node.py | 181 +++++++++++++++++++++++++--------------------- 1 file changed, 99 insertions(+), 82 deletions(-) diff --git a/streamv2v_node.py b/streamv2v_node.py index 302802a..d2a577a 100644 --- a/streamv2v_node.py +++ b/streamv2v_node.py @@ -13,26 +13,30 @@ import folder_paths from comfy.utils import common_upscale import cv2 import diffusers + dif_version = str(diffusers.__version__) dif_version_int = int(dif_version.split(".")[1]) from diffusers import (DDIMScheduler, KDPM2AncestralDiscreteScheduler, LMSDiscreteScheduler, DPMSolverMultistepScheduler, - DPMSolverSinglestepScheduler,AutoencoderTiny, + DPMSolverSinglestepScheduler, AutoencoderTiny, EulerDiscreteScheduler, HeunDiscreteScheduler, KDPM2DiscreteScheduler, - EulerAncestralDiscreteScheduler, UniPCMultistepScheduler,DiffusionPipeline, + EulerAncestralDiscreteScheduler, UniPCMultistepScheduler, DiffusionPipeline, DDPMScheduler, LCMScheduler, StableDiffusionPipeline, StableDiffusionXLPipeline, AutoencoderKL, ) + CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) path_dir = os.path.dirname(CURRENT_DIR) file_path = os.path.dirname(path_dir) BIGMAX = (2 ** 53 - 1) input_path = folder_paths.get_input_directory() + def tensor_to_image(tensor): # tensor = tensor.cpu() image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() image = Image.fromarray(image_np, mode='RGB') return image + def get_video_img(tensor): outputs = [] for x in tensor: @@ -40,10 +44,12 @@ def get_video_img(tensor): outputs.append(x) return outputs + def phi2narry(img): img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0) return img + def narry_list(list_in): for i in range(len(list_in)): value = list_in[i] @@ -51,6 +57,7 @@ def narry_list(list_in): list_in[i] = modified_value return list_in + def nomarl_upscale(img_tensor, width, height): samples = img_tensor.movedim(-1, 1) img = common_upscale(samples, width, height, "nearest-exact", "center") @@ -58,12 +65,14 @@ def nomarl_upscale(img_tensor, width, height): img_pil = tensor_to_image(samples) return img_pil + def tensor_upscale(img_tensor, width, height): samples = img_tensor.movedim(-1, 1) img = common_upscale(samples, width, height, "nearest-exact", "center") samples = img.movedim(1, -1) return samples + def get_local_path(file_path, model_path): path = os.path.join(file_path, "models", "diffusers", model_path) model_path = os.path.normpath(path) @@ -94,25 +103,27 @@ def instance_path(path, repo): repo = get_instance_path(model_path) return repo -scheduler_list = ["LCM","DDIM", - "Euler", - "Euler a", - "DDPM", - "DPM++ 2M", - "DPM++ 2M Karras", - "DPM++ 2M SDE", - "DPM++ 2M SDE Karras", - "DPM++ SDE", - "DPM++ SDE Karras", - "DPM2", - "DPM2 Karras", - "DPM2 a", - "DPM2 a Karras", - "Heun", - "LMS", - "LMS Karras", - "UniPC", -] + +scheduler_list = ["LCM", "DDIM", + "Euler", + "Euler a", + "DDPM", + "DPM++ 2M", + "DPM++ 2M Karras", + "DPM++ 2M SDE", + "DPM++ 2M SDE Karras", + "DPM++ SDE", + "DPM++ SDE Karras", + "DPM2", + "DPM2 Karras", + "DPM2 a", + "DPM2 a Karras", + "Heun", + "LMS", + "LMS Karras", + "UniPC", + ] + def get_sheduler(name): scheduler = False @@ -156,10 +167,11 @@ def get_sheduler(name): scheduler = UniPCMultistepScheduler() return scheduler + class Stream_Model_Loader: def __init__(self): pass - + @classmethod def INPUT_TYPES(cls): return { @@ -172,12 +184,12 @@ class Stream_Model_Loader: } } - RETURN_TYPES = ("MODEL","STRING") - RETURN_NAMES = ("pipe","info") + RETURN_TYPES = ("MODEL", "STRING") + RETURN_NAMES = ("pipe", "info") FUNCTION = "main" CATEGORY = "StreamV2V_Plus" - def main(self,ckpt_name,scheduler,use_sdxl,vae_id,lcm_lora,): + def main(self, ckpt_name, scheduler, use_sdxl, vae_id, lcm_lora, ): # 加载pipe ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) scheduler_used = get_sheduler(scheduler) @@ -200,49 +212,50 @@ class Stream_Model_Loader: pipe = StableDiffusionPipeline.from_single_file( ckpt_path, original_config_file=original_config_file, torch_dtype=torch.float16).to("cuda") info = "sd15" - if vae_id!="": + if vae_id != "": pipe.vae = AutoencoderKL.from_pretrained(vae_id, torch_dtype=torch.float16).to("cuda") pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config) pipe.enable_xformers_memory_efficient_attention() pipe.enable_vae_tiling() lcm_lora_path = get_lora(lcm_lora) pipe.load_lora_weights(lcm_lora_path, adapter_name="lcm", ) - return (pipe,info,) + info = ";".join([info, ckpt_name]) + return (pipe, info,) class Stream_Lora_Loader: def __init__(self): pass - + @classmethod def INPUT_TYPES(cls): return { "required": { - "pipe":("MODEL",), + "pipe": ("MODEL",), "info": ("STRING", {"forceInput": True, "default": ""}), "lora": (["none"] + folder_paths.get_filename_list("loras"),), "lora_scale": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.0, "step": 0.1}), "trigger_word": ("STRING", {"default": "best quality"}), } } - - RETURN_TYPES = ("MODEL","STRING",) - RETURN_NAMES = ("pipe","info",) + + RETURN_TYPES = ("MODEL", "STRING",) + RETURN_NAMES = ("pipe", "info",) FUNCTION = "main" CATEGORY = "StreamV2V_Plus" - - def main(self, pipe,info,lora,lora_scale,trigger_word): - + + def main(self, pipe, info, lora, lora_scale, trigger_word): + info, ckpt_name = info.split(";") lora_path = get_lora(lora) - if lora == "none" : + if lora == "none": raise "you need a style lora" pipe.load_lora_weights(lora_path, adapter_name=trigger_word, ) pipe.set_adapters(adapter_names=["lcm", trigger_word], adapter_weights=[1.0, lora_scale]) pipe.fuse_lora() print(f"Use style lora: {trigger_word} in {lora}") - info=";".join([info,trigger_word]) - return (pipe,info,) + info = ";".join([info, trigger_word, ckpt_name]) + return (pipe, info,) class Stream_Sampler: @@ -255,12 +268,12 @@ class Stream_Sampler: os.path.isfile(os.path.join(input_path, f)) and f.split('.')[-1] in ["mp4", "webm", "mkv", "avi"]] return { - "required": { - "pipe":("MODEL",), - "info":("STRING",{"forceInput": True, "default": ""}), - "prompt": ("STRING", {"multiline": True, "default": "Claymation, a man is giving a talk"}), - "video": (["none"]+video_files,), - "sampler_type": (["txt2img", "vdieo2vdieo","WebCam2Video"],), + "required": { + "pipe": ("MODEL",), + "info": ("STRING", {"forceInput": True, "default": ""}), + "prompt": ("STRING", {"multiline": True, "default": "Claymation, a man is giving a talk"}), + "video": (["none"] + video_files,), + "sampler_type": (["txt2img", "vdieo2vdieo", "WebCam2Video"],), "guidance_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.1, "round": 0.01}), "diffusion_steps": ("INT", {"default": 4, "min": 1, "max": 1000}), "num_inference_steps": ("INT", {"default": 50, "min": 1, "max": 1000}), @@ -268,38 +281,40 @@ class Stream_Sampler: "seed": ("INT", {"default": 2, "min": 0, "max": 0xffffffffffffffff}), "width": ("INT", {"default": 512, "min": 256, "max": 4096, "step": 64, }), "height": ("INT", {"default": 512, "min": 256, "max": 4096, "step": 64}), - "batch_size": ("INT", {"default": 1, "min": 1, "max": 1024, "step": 1}), - }, - "optional":{"image": ("IMAGE",), - } + "acceleration": (["xformers", "tensorrt", ],), + + }, + "optional": {"image": ("IMAGE",), + } } - RETURN_TYPES = ("IMAGE","AUDIO","FLOAT",) - RETURN_NAMES = ("image","audio","fps",) + RETURN_TYPES = ("IMAGE", "AUDIO", "FLOAT",) + RETURN_NAMES = ("image", "audio", "fps",) FUNCTION = "main_stream" CATEGORY = "StreamV2V_Plus" - def main_stream(self,pipe,info,prompt,video,sampler_type,guidance_scale,diffusion_steps,num_inference_steps, - noise_strength,seed,width,height,batch_size,**kwargs): - model_type,trigger_word=info.split(";") - if model_type=="sdxl": - using_sdxl=True + def main_stream(self, pipe, info, prompt, video, sampler_type, guidance_scale, diffusion_steps, num_inference_steps, + noise_strength, seed, width, height, acceleration, **kwargs): + model_type, trigger_word, ckpt_name = info.split(";") + if model_type == "sdxl": + using_sdxl = True else: - using_sdxl=False - image=kwargs.get("image") - if sampler_type=="vdieo2vdieo": + using_sdxl = False + image = kwargs.get("image") + if sampler_type == "vdieo2vdieo": if video == "none": raise "need video input" input_video = get_instance_path(os.path.join(input_path, video)) - video_info = read_video(input_video) # input path to the file name "THWC") -> tuple[Tensor, Tensor, dict[str, Any]] - video = video_info[0] / 255 #Tensor THWC - fps = video_info[2]["video_fps"] #dict[str, Any] + video_info = read_video( + input_video) # input path to the file name "THWC") -> tuple[Tensor, Tensor, dict[str, Any]] + video = video_info[0] / 255 # Tensor THWC + fps = video_info[2]["video_fps"] # dict[str, Any] height_v = int(video.shape[1]) width_v = int(video.shape[2]) - waveform=video_info[1] - if waveform.size(1)==0: - audio=None + waveform = video_info[1] + if waveform.size(1) == 0: + audio = None print("no audio in input video!") else: sample_rate = video_info[2]["audio_fps"] @@ -313,7 +328,7 @@ class Stream_Sampler: pipe, mode="img2img", t_index_list=t_index_list, - frame_buffer_size=batch_size, + frame_buffer_size=1, width=width_v, height=height_v, warmup=10, @@ -332,7 +347,7 @@ class Stream_Sampler: seed=seed, using_sdxl=using_sdxl ) - + prompt = prompt + " " + trigger_word + "style" stream.stream.prepare( prompt=prompt, @@ -354,11 +369,11 @@ class Stream_Sampler: inference_time.append(iteration_end_time - iteration_start_time) print(f'Avg time: {sum(inference_time[20:]) / len(inference_time[20:])}') gen = get_video_img(video_result) - gen = narry_list(gen) + gen = narry_list(gen) images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height_v, width_v, 3))))) - return (images,audio, fps) + return (images, audio, fps) - elif sampler_type=="txt2img": + elif sampler_type == "txt2img": stream = StreamV2VWrapper( pipe, t_index_list=[0, 16, 32, 45], @@ -366,7 +381,7 @@ class Stream_Sampler: mode="txt2img", device="cuda", dtype=torch.float16, - frame_buffer_size=batch_size, + frame_buffer_size=1, warmup=10, width=width, height=height, @@ -381,7 +396,7 @@ class Stream_Sampler: use_tome_cache=True, seed=seed, use_safety_checker=False, - using_sdxl = using_sdxl + using_sdxl=using_sdxl ) prompt = prompt + " " + trigger_word + "style" stream.stream.prepare( @@ -389,16 +404,16 @@ class Stream_Sampler: num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, ) - output_img=stream(prompt=prompt) + output_img = stream(prompt=prompt) images = phi2narry(output_img) return (images,) - + else: stream = StreamV2VWrapper( pipe, mode="img2img", t_index_list=[32, 45], - frame_buffer_size=batch_size, + frame_buffer_size=1, width=width, height=height, warmup=10, @@ -407,7 +422,7 @@ class Stream_Sampler: do_add_noise=True, output_type="pil", use_denoising_batch=True, - use_cached_attn=True, + use_cached_attn=False, use_feature_injection=True, feature_injection_strength=0.8, feature_similarity_threshold=0.98, @@ -415,16 +430,18 @@ class Stream_Sampler: cache_maxframes=1, use_tome_cache=True, seed=seed, - using_sdxl=using_sdxl + acceleration=acceleration, + using_sdxl=using_sdxl, + model_id_or_path=ckpt_name ) prompt = prompt + " " + trigger_word + "style" stream.stream.prepare( prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, - ) #预处理prompt + ) # 预处理prompt image = nomarl_upscale(image, width, height) - for _ in range(2): #预处理prompt+img + for _ in range(2): # 预处理prompt+img stream(image) images = stream(image) images = phi2narry(images) @@ -432,13 +449,13 @@ class Stream_Sampler: NODE_CLASS_MAPPINGS = { - "Stream_Model_Loader":Stream_Model_Loader, - "Stream_Lora_Loader":Stream_Lora_Loader, + "Stream_Model_Loader": Stream_Model_Loader, + "Stream_Lora_Loader": Stream_Lora_Loader, "Stream_Sampler": Stream_Sampler, } NODE_DISPLAY_NAME_MAPPINGS = { - "Load_Stream":"Load_Stream", - "Stream_Lora_Loader":"Stream_Lora_Loader", + "Load_Stream": "Load_Stream", + "Stream_Lora_Loader": "Stream_Lora_Loader", "Stream_Sampler": "Stream_Sampler", }