# !/usr/bin/env python # -*- coding: UTF-8 -*- import time from torchvision.io import read_video from tqdm import tqdm from .utils.wrapper import StreamV2VWrapper import torch import os from PIL import Image import numpy as np import sys 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, EulerDiscreteScheduler, HeunDiscreteScheduler, KDPM2DiscreteScheduler, 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: x = tensor_to_image(x) 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] modified_value = phi2narry(value) 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") samples = img.movedim(1, -1) 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) if sys.platform == 'win32': model_path = model_path.replace('\\', "/") return model_path def get_instance_path(path): instance_path = os.path.normpath(path) if sys.platform == 'win32': instance_path = instance_path.replace('\\', "/") return instance_path def get_lora(lora): lora_path = folder_paths.get_full_path("loras", lora) lora_path = get_instance_path(lora_path) return lora_path def instance_path(path, repo): if repo == "": if path == "none": repo = "none" else: model_path = get_local_path(file_path, path) 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", ] def get_sheduler(name): scheduler = False if name == "Euler": scheduler = EulerDiscreteScheduler() elif name == "Euler a": scheduler = EulerAncestralDiscreteScheduler() elif name == "DDIM": scheduler = DDIMScheduler() elif name == "DDPM": scheduler = DDPMScheduler() elif name == "DPM++ 2M": scheduler = DPMSolverMultistepScheduler() elif name == "DPM++ 2M Karras": scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True) elif name == "DPM++ 2M SDE": scheduler = DPMSolverMultistepScheduler(algorithm_type="sde-dpmsolver++") elif name == "DPM++ 2M SDE Karras": scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True, algorithm_type="sde-dpmsolver++") elif name == "DPM++ SDE": scheduler = DPMSolverSinglestepScheduler() elif name == "DPM++ SDE Karras": scheduler = DPMSolverSinglestepScheduler(use_karras_sigmas=True) elif name == "DPM2": scheduler = KDPM2DiscreteScheduler() elif name == "DPM2 Karras": scheduler = KDPM2DiscreteScheduler(use_karras_sigmas=True) elif name == "DPM2 a": scheduler = KDPM2AncestralDiscreteScheduler() elif name == "DPM2 a Karras": scheduler = KDPM2AncestralDiscreteScheduler(use_karras_sigmas=True) elif name == "Heun": scheduler = HeunDiscreteScheduler() elif name == "LCM": scheduler = LCMScheduler() elif name == "LMS": scheduler = LMSDiscreteScheduler() elif name == "LMS Karras": scheduler = LMSDiscreteScheduler(use_karras_sigmas=True) elif name == "UniPC": scheduler = UniPCMultistepScheduler() return scheduler class Stream_Model_Loader: def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"),), "scheduler": (scheduler_list,), "use_sdxl": ("BOOLEAN", {"default": False},), "vae_id": ("STRING", {"default": ""}), "lcm_lora": (["none"] + folder_paths.get_filename_list("loras"),), } } 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, ): # 加载pipe ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) scheduler_used = get_sheduler(scheduler) if use_sdxl: model_config = "stabilityai/stable-diffusion-xl-base-1.0" original_config_file = os.path.join(CURRENT_DIR, "models", "sd_xl_base.yaml") if dif_version_int >= 28: pipe = StableDiffusionXLPipeline.from_single_file( ckpt_path, config=model_config,original_config=original_config_file, torch_dtype=torch.float16).to("cuda") else: pipe = StableDiffusionXLPipeline.from_single_file( ckpt_path,config=model_config, original_config_file=original_config_file, torch_dtype=torch.float16).to("cuda") info = "sdxl" else: model_config = "Lykon/dreamshaper-8" original_config_file = os.path.join(folder_paths.models_dir, "configs", "v1-inference.yaml") if dif_version_int >= 28: pipe = StableDiffusionPipeline.from_single_file( ckpt_path,config=model_config, original_config=original_config_file, torch_dtype=torch.float16).to("cuda") else: pipe = StableDiffusionPipeline.from_single_file( ckpt_path, config=model_config,original_config_file=original_config_file, torch_dtype=torch.float16).to("cuda") info = "sd15" 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", ) 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",), "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",) FUNCTION = "main" CATEGORY = "StreamV2V_Plus" def main(self, pipe, info, lora, lora_scale, trigger_word): info, ckpt_name = info.split(";") lora_path = get_lora(lora) 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, ckpt_name]) return (pipe, info,) class Stream_Sampler: def __init__(self): pass @classmethod def INPUT_TYPES(cls): video_files = [f for f in os.listdir(input_path) if 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"],), "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}), "noise_strength": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}), "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}), "acceleration": (["xformers", "tensorrt", ],), }, "optional": {"image": ("IMAGE",), } } 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, 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": 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] height_v = int(video.shape[1]) width_v = int(video.shape[2]) 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"] print("get audio from video") audio = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate} init_step = int(50 * (1 - noise_strength)) interval = int(50 * noise_strength) // diffusion_steps t_index_list = [init_step + i * interval for i in range(diffusion_steps)] stream = StreamV2VWrapper( pipe, mode="img2img", t_index_list=t_index_list, frame_buffer_size=1, width=width_v, height=height_v, warmup=10, do_add_noise=True, output_type="pt", enable_similar_image_filter=False, similar_image_filter_threshold=0.98, use_denoising_batch=True, use_cached_attn=True, use_feature_injection=True, feature_injection_strength=0.8, feature_similarity_threshold=0.98, cache_interval=4, cache_maxframes=1, use_tome_cache=True, seed=seed, using_sdxl=using_sdxl ) prompt = prompt + " " + trigger_word + "style" stream.stream.prepare( prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, ) video_result = torch.zeros(video.shape[0], height_v, width_v, 3) for _ in range(stream.batch_size): stream(image=video[0].permute(2, 0, 1)) inference_time = [] for i in tqdm(range(video.shape[0])): iteration_start_time = time.time() output_image = stream(video[i].permute(2, 0, 1)) video_result[i] = output_image.permute(1, 2, 0) iteration_end_time = time.time() 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) images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height_v, width_v, 3))))) return (images, audio, fps) elif sampler_type == "txt2img": stream = StreamV2VWrapper( pipe, t_index_list=[0, 16, 32, 45], output_type="pil", mode="txt2img", device="cuda", dtype=torch.float16, frame_buffer_size=1, warmup=10, width=width, height=height, do_add_noise=True, use_denoising_batch=False, use_cached_attn=True, use_feature_injection=True, feature_injection_strength=0.8, feature_similarity_threshold=0.98, cache_interval=1, cache_maxframes=4, use_tome_cache=True, seed=seed, use_safety_checker=False, using_sdxl=using_sdxl ) prompt = prompt + " " + trigger_word + "style" stream.stream.prepare( prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, ) 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=1, width=width, height=height, warmup=10, dtype=torch.float16, device="cuda", do_add_noise=True, output_type="pil", use_denoising_batch=True, use_cached_attn=False, use_feature_injection=True, feature_injection_strength=0.8, feature_similarity_threshold=0.98, cache_interval=4, cache_maxframes=1, use_tome_cache=True, seed=seed, 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 image = nomarl_upscale(image, width, height) for _ in range(2): # 预处理prompt+img stream(image) images = stream(image) images = phi2narry(images) return (images,) NODE_CLASS_MAPPINGS = { "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", "Stream_Sampler": "Stream_Sampler", }