diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..656adce --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..05a6505 --- /dev/null +++ b/nodes.py @@ -0,0 +1,555 @@ +import os +import sys +from .streamdiffusion.pipeline import StreamDiffusion +from .streamdiffusion.wrapper import StreamDiffusionWrapper +from diffusers import AutoencoderTiny, StableDiffusionPipeline +from pathlib import Path +import traceback +from typing import List, Literal, Optional, Union, Dict +import torch +import gc +# Get the absolute path of various directories +my_dir = os.path.dirname(os.path.abspath(__file__)) +custom_nodes_dir = os.path.abspath(os.path.join(my_dir, '..')) +comfy_dir = os.path.abspath(os.path.join(my_dir, '..', '..')) + +# Construct the path to the font file +font_path = os.path.join(my_dir, 'arial.ttf') + +# Append comfy_dir to sys.path & import files +sys.path.append(comfy_dir) +import folder_paths + + +import comfy.sample +import comfy.samplers +import comfy.sd +import comfy.utils +import comfy.latent_formats +import comfy.model_management + +# Append my_dir to sys.path & import files +sys.path.append(my_dir) + +from PIL import Image, ImageOps, ImageSequence +from PIL.PngImagePlugin import PngInfo +import numpy as np + +MAX_RESOLUTION=8192 + +CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) + +def get_async_loop(): + loop = None + try: + loop = asyncio.get_event_loop() + except: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + return loop + + +class StreamDiffusion_LoRA_Stacker: + modes = ["simple", "advanced"] + + @classmethod + def INPUT_TYPES(cls): + loras = ["None"] + folder_paths.get_filename_list("loras") + inputs = { + "required": { + "input_mode": (cls.modes,), + "lora_count": ("INT", {"default": 3, "min": 0, "max": 50, "step": 1}), + } + } + + for i in range(1, 50): + inputs["required"][f"lora_name_{i}"] = (loras,) + inputs["required"][f"lora_wt_{i}"] = ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}) + inputs["required"][f"model_str_{i}"] = ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}) + inputs["required"][f"clip_str_{i}"] = ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}) + + inputs["optional"] = { + "lora_stack": ("LORA_STACK",) + } + return inputs + + RETURN_TYPES = ("LORA_STACK",) + RETURN_NAMES = ("LORA_STACK",) + FUNCTION = "lora_stacker" + CATEGORY = "StreamDiffusion/Stackers" + + def lora_stacker(self, input_mode, lora_count, lora_stack=None, **kwargs): + + # Extract values from kwargs + loras = [kwargs.get(f"lora_name_{i}") for i in range(1, lora_count + 1)] + + # Create a list of tuples using provided parameters, exclude tuples with lora_name as "None" + if input_mode == "simple": + weights = [kwargs.get(f"lora_wt_{i}") for i in range(1, lora_count + 1)] + loras = [(lora_name, lora_weight, lora_weight) for lora_name, lora_weight in zip(loras, weights) if + lora_name != "None"] + else: + model_strs = [kwargs.get(f"model_str_{i}") for i in range(1, lora_count + 1)] + clip_strs = [kwargs.get(f"clip_str_{i}") for i in range(1, lora_count + 1)] + loras = [(lora_name, model_str, clip_str) for lora_name, model_str, clip_str in + zip(loras, model_strs, clip_strs) if lora_name != "None"] + + # If lora_stack is not None, extend the loras list with lora_stack + if lora_stack is not None: + loras.extend([l for l in lora_stack if l[0] != "None"]) + + return (loras,) + + +class StreamDiffusion_Loader: + @classmethod + def INPUT_TYPES(s): + loras = ["None"] + folder_paths.get_filename_list("loras") + return { + "required": { + "ckpt_name": (["Baked ckpt"]+folder_paths.get_filename_list("checkpoints"), ), + "vae_name": (["Baked VAE"] + folder_paths.get_filename_list("vae"),), + "lcm_lora": (loras,), + "acceleration": (["none", "xfomers", "sfast", "tensorrt"],), + + "use_tiny_vae": ("BOOLEAN", { "default": True }), + "use_lcm_lora": ("BOOLEAN", { "default": True }), + }, + "optional": { + "lora_stack": ("LORA_STACK", ), + }, + } + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("MODEL",) + + FUNCTION = "efficientloader" + CATEGORY = "StreamDiffusion/Loader" + + def efficientloader(self,ckpt_name,vae_name,lcm_lora,acceleration,use_tiny_vae,use_lcm_lora,lora_stack=None): + + device = comfy.model_management.get_torch_device() + device_name = comfy.model_management.get_torch_device_name(device) + vae_dtype=comfy.model_management.vae_dtype() + + if ckpt_name =='Baked ckpt': + ckpt_path="KBlueLeaf/kohaku-v2.1" + else: + ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + + if vae_name=='Baked VAE': + vae_id = None + else: + vae_id=vae_name + + if lcm_lora=='None': + lcm_lora_id=None + else: + lcm_lora_id= folder_paths.get_full_path("loras", lcm_lora) + + lora_dict =None + if lora_stack is not None: + lora_dict={} + for lora_name, lora_scale, strength_clip in lora_stack: + full_lora_name=folder_paths.get_full_path("loras", lora_name) + lora_dict[full_lora_name]=lora_scale + + t_index_list=[32,40,45] + + stream = StreamDiffusionWrapper( + model_id_or_path=ckpt_path, + lora_dict=lora_dict, + t_index_list=t_index_list, + frame_buffer_size=1, + width=512, + height=512, + warmup=10, + acceleration=acceleration, + use_tiny_vae =use_tiny_vae, + device=device, + use_lcm_lora = use_lcm_lora, + output_type = 'pt', + dtype = torch.float16, + lcm_lora_id=lcm_lora_id, + vae_id =vae_id, + ) + + return (stream,) + + +class StreamDiffusion_Sampler: + @classmethod + def INPUT_TYPES(s): + return { + "required":{ + "model": ("MODEL",), + "positive": ("STRING", {"default": "CLIP_POSITIVE","multiline": True}), + "negative": ("STRING", {"default": "CLIP_NEGATIVE", "multiline": True}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "steps": ("INT", {"default": 50, "min": 1, "max": 10000}), + "cfg": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 100.0}), + "delta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0}), + + "width": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1}), + "height": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 10000, "step": 1}),#The frame buffer size for denoising batch, by default 1. + + "index_list": ("STRING", {"default": "32,40,45","multiline": False}), + + "cfg_type": (["none", "full", "self", "initialize"],), + + "add_noise": ("BOOLEAN", { "default": True }), + "use_denoising_batch": ("BOOLEAN", { "default": True }), + "enable_similar_image_filter": ("BOOLEAN", { "default": False }), + "use_safety_checker": ("BOOLEAN", { "default": False }), + + }, + "optional": { + "similar_image_filter_threshold": ("FLOAT", {"default": 0.98, "min": 0.0, "max": 100.0,"step": 0.01}), + "similar_image_filter_max_skip_frame": ("INT", {"default": 10, "min": 0, "max": 100}), + "latent": ("LATENT",), + "image": ("IMAGE",), + "lora_stack": ("LORA_STACK", ), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("IMAGE",) + + FUNCTION = "sample" + CATEGORY = "StreamDiffusion/Sampler" + + @torch.no_grad() + def sample(self,model,positive,negative,seed,steps,cfg,delta,width,height,batch_size,index_list,cfg_type,add_noise,use_denoising_batch,enable_similar_image_filter=False,use_safety_checker=False,similar_image_filter_threshold= 0.98,similar_image_filter_max_skip_frame=10,latent=None,image=None,lora_stack=None): + device = comfy.model_management.get_torch_device() + device_name = comfy.model_management.get_torch_device_name(device) + vae_dtype=comfy.model_management.vae_dtype() + + # latent_image = latent["samples"] + # latent_image = latent_image.to(model.device) + # batch_size,channel,latent_height,latent_width,=latent_image.shape + # width=latent_width*8 + # height =latent_height*8 + + t_index_list=[32,40,45] + t_index_list =[int(i) for i in index_list.split(',')] + + if image ==None: + mode = "txt2img" + else: + mode = "img2img" + image = image.movedim(-1,1) + + if cfg <= 1.0: + cfg_type = "none" + + if batch_size>1 and mode=="txt2img": + use_denoising_batch=False + + + # stream.set_sampler_param(t_index_list=t_index_list, + # width=width, + # height=height, + # do_add_noise= add_noise=='enable', + # frame_buffer_size=frame_buffer_size, + # use_denoising_batch=use_denoising_batch, + # cfg_type=cfg_type,) + + model.prepare( + positive, + negative, + steps, + cfg, + delta, + t_index_list, + add_noise, + enable_similar_image_filter, + similar_image_filter_threshold, + similar_image_filter_max_skip_frame, + use_denoising_batch, + cfg_type, + seed, + batch_size, + use_safety_checker, + ) + # latent_image = self.predict_x0_batch( + # torch.randn((stream.batch_size, 4, stream.latent_height, stream.latent_width)).to( + # device=stream.device, dtype=stream.dtype + # ) + # ) + + # if batch_size==1: + # for _ in range(stream.batch_size - 1): + # stream() + + output = model.sample(image).permute(0, 2, 3, 1) + + return (output,) + + + +def original( + output: str = r"G:\Workspace\AI\ComfyUI\output\output.png", + model_id_or_path: str = "KBlueLeaf/kohaku-v2.1", + lora_dict: Optional[Dict[str, float]] = None, + prompt: str = "a dog is running in the garden", + width: int = 512, + height: int = 512, + acceleration: Literal["none", "xformers", "tensorrt"] = "xformers", + use_denoising_batch: bool = False, + seed: int = 2, +): + + """ + Process for generating images based on a prompt using a specified model. + + Parameters + ---------- + output : str, optional + The output image file to save images to. + model_id_or_path : str + The name of the model to use for image generation. + lora_dict : Optional[Dict[str, float]], optional + The lora_dict to load, by default None. + Keys are the LoRA names and values are the LoRA scales. + Example: {'LoRA_1' : 0.5 , 'LoRA_2' : 0.7 ,...} + prompt : str + The prompt to generate images from. + width : int, optional + The width of the image, by default 512. + height : int, optional + The height of the image, by default 512. + acceleration : Literal["none", "xformers", "tensorrt"] + The type of acceleration to use for image generation. + use_denoising_batch : bool, optional + Whether to use denoising batch or not, by default False. + seed : int, optional + The seed, by default 2. if -1, use random seed. + """ + + stream = StreamDiffusionWrapper( + model_id_or_path=model_id_or_path, + lora_dict=lora_dict, + t_index_list=[0, 16, 32, 45], + frame_buffer_size=1, + width=width, + height=height, + warmup=10, + acceleration=acceleration, + mode="txt2img", + use_denoising_batch=use_denoising_batch, + cfg_type="none", + seed=seed, + ) + + stream.prepare( + prompt=prompt, + num_inference_steps=50, + ) + + for _ in range(stream.batch_size - 1): + stream() + + output_image = stream() + output_image.save(output) + +class StreamDiffusion_Wrapper: + @classmethod + def INPUT_TYPES(s): + loras = ["None"] + folder_paths.get_filename_list("loras") + + return { + "required":{ + + "ckpt_name": (["Baked ckpt"]+folder_paths.get_filename_list("checkpoints"), ), + "vae_name": (["Baked VAE"] + folder_paths.get_filename_list("vae"),), + "lcm_lora": (loras,), + "acceleration": (["none", "xfomers", "sfast", "tensorrt"],), + + "use_tiny_vae": ("BOOLEAN", { "default": True }), + "use_lcm_lora": ("BOOLEAN", { "default": True }), + + + "positive": ("STRING", {"default": "CLIP_POSITIVE","multiline": True}), + "negative": ("STRING", {"default": "CLIP_NEGATIVE", "multiline": True}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "steps": ("INT", {"default": 50, "min": 1, "max": 10000}), + "cfg": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 100.0}), + "delta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0}), + + "width": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1}), + "height": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 10000, "step": 1}),#The frame buffer size for denoising batch, by default 1. + + "index_list": ("STRING", {"default": "32,40,45","multiline": False}), + + "cfg_type": (["none", "full", "self", "initialize"],), + + "add_noise": ("BOOLEAN", { "default": True }), + "use_denoising_batch": ("BOOLEAN", { "default": True }), + "enable_similar_image_filter": ("BOOLEAN", { "default": True }), + "use_safety_checker": ("BOOLEAN", { "default": True }), + + }, + "optional": { + "similar_image_filter_threshold": ("FLOAT", {"default": 0.98, "min": 0.0, "max": 100.0,"step": 0.01}), + "similar_image_filter_max_skip_frame": ("INT", {"default": 10, "min": 0, "max": 100}), + "image": ("IMAGE",), + "lora_stack": ("LORA_STACK", ), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("IMAGE",) + + FUNCTION = "sample" + CATEGORY = "StreamDiffusion/Wrapper" + + @torch.no_grad() + def sample(self,ckpt_name,vae_name,lcm_lora,acceleration,use_tiny_vae,use_lcm_lora,positive,negative,seed,steps,cfg,delta,width,height,batch_size,index_list,cfg_type,add_noise,use_denoising_batch,enable_similar_image_filter=False,use_safety_checker=False,similar_image_filter_threshold= 0.98,similar_image_filter_max_skip_frame=10,image=None,lora_stack=None): + + device = comfy.model_management.get_torch_device() + device_name = comfy.model_management.get_torch_device_name(device) + vae_dtype=comfy.model_management.vae_dtype() + unet_dtype=comfy.model_management.unet_dtype + text_encoder_dtype=comfy.model_management.text_encoder_dtype + output_device = comfy.model_management.intermediate_device() + + if ckpt_name =='Baked ckpt': + ckpt_path="KBlueLeaf/kohaku-v2.1" + else: + ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + + if vae_name=='Baked VAE': + vae_id = None + else: + vae_id=vae_name + + if lcm_lora=='None': + lcm_lora_id=None + else: + lcm_lora_id= folder_paths.get_full_path("loras", lcm_lora) + + lora_dict =None + if lora_stack is not None: + lora_dict={} + for lora_name, lora_scale, strength_clip in lora_stack: + full_lora_name=folder_paths.get_full_path("loras", lora_name) + lora_dict[full_lora_name]=lora_scale + + + t_index_list=[32,40,45] + t_index_list =[int(i) for i in index_list.split(',')] + + if image ==None: + mode = "txt2img" + else: + mode = "img2img" + image = image.movedim(-1,1) + + if cfg <= 1.0: + cfg_type = "none" + + if batch_size>1 and mode=="txt2img": + use_denoising_batch=False + + stream = StreamDiffusionWrapper( + model_id_or_path=ckpt_path, + lora_dict=lora_dict, + t_index_list=t_index_list, + frame_buffer_size=batch_size, + width=width, + height=height, + warmup=10, + acceleration=acceleration, + mode=mode, + use_denoising_batch=use_denoising_batch, + cfg_type=cfg_type, + seed=seed, + use_tiny_vae =use_tiny_vae, + do_add_noise = add_noise, + device=device, + use_lcm_lora = use_lcm_lora, + output_type = 'pt', + dtype = torch.float16, + lcm_lora_id=lcm_lora_id, + vae_id =vae_id, + enable_similar_image_filter = enable_similar_image_filter, + similar_image_filter_threshold =similar_image_filter_threshold, + similar_image_filter_max_skip_frame=similar_image_filter_max_skip_frame, + use_safety_checker = use_safety_checker, + ) + + # stream.set_sampler_param(t_index_list=t_index_list, + # width=width, + # height=height, + # do_add_noise= add_noise=='enable', + # frame_buffer_size=frame_buffer_size, + # use_denoising_batch=use_denoising_batch, + # cfg_type=cfg_type,) + + stream.prepare( + positive, + negative, + num_inference_steps=steps, + guidance_scale=cfg, + delta=delta, + # generator=torch.manual_seed(seed), + # seed=seed, + ) + # latent_image = self.predict_x0_batch( + # torch.randn((stream.batch_size, 4, stream.latent_height, stream.latent_width)).to( + # device=stream.device, dtype=stream.dtype + # ) + # ) + + # if batch_size==1: + # for _ in range(stream.batch_size - 1): + # stream() + + output = stream.sample(image).permute(0, 2, 3, 1) + + # if not isinstance(output,list): + # output_pils=[output] + # else: + # output_pils=output + + # output_images=[] + # output_masks=[] + + # for i in output_pils: + # i = ImageOps.exif_transpose(i) + # image = i.convert("RGB") + # image = np.array(image).astype(np.float32) / 255.0 + # image = torch.from_numpy(image)[None,] + # if 'A' in i.getbands(): + # mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + # mask = 1. - torch.from_numpy(mask) + # else: + # mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") + # output_images.append(image) + # output_masks.append(mask.unsqueeze(0)) + + # if len(output_images) > 1: + # output_image = torch.cat(output_images, dim=0) + # output_mask = torch.cat(output_masks, dim=0) + # else: + # output_image = output_images[0] + # output_mask = output_masks[0] + + return (output,) + + +NODE_CLASS_MAPPINGS = { + "StreamDiffusion_Loader": StreamDiffusion_Loader, + "StreamDiffusion_Sampler":StreamDiffusion_Sampler, + "StreamDiffusion_LoRA_Stacker":StreamDiffusion_LoRA_Stacker, + "StreamDiffusion_Wrapper":StreamDiffusion_Wrapper + +} +NODE_DISPLAY_NAME_MAPPINGS = { + "StreamDiffusion_Loader": "StreamDiffusion_Loader", + "StreamDiffusion_Sampler":"StreamDiffusion_Sampler", + "StreamDiffusion_LoRA_Stacker":"StreamDiffusion_LoRA_Stacker", + "StreamDiffusion_Wrapper":"StreamDiffusion_Wrapper", + +} \ No newline at end of file diff --git a/single.py b/single.py new file mode 100644 index 0000000..20b2557 --- /dev/null +++ b/single.py @@ -0,0 +1,70 @@ +import os +import sys +from typing import Literal, Dict, Optional + +CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) + +def main( + output: str = os.path.join(CURRENT_DIR, "..", "..", "images", "outputs", "output.png"), + model_id_or_path: str = "KBlueLeaf/kohaku-v2.1", + lora_dict: Optional[Dict[str, float]] = None, + prompt: str = "a dog is running in the garden", + width: int = 512, + height: int = 512, + acceleration: Literal["none", "xformers", "tensorrt"] = "xformers", + use_denoising_batch: bool = False, + seed: int = 2, +): + + """ + Process for generating images based on a prompt using a specified model. + + Parameters + ---------- + output : str, optional + The output image file to save images to. + model_id_or_path : str + The name of the model to use for image generation. + lora_dict : Optional[Dict[str, float]], optional + The lora_dict to load, by default None. + Keys are the LoRA names and values are the LoRA scales. + Example: {'LoRA_1' : 0.5 , 'LoRA_2' : 0.7 ,...} + prompt : str + The prompt to generate images from. + width : int, optional + The width of the image, by default 512. + height : int, optional + The height of the image, by default 512. + acceleration : Literal["none", "xformers", "tensorrt"] + The type of acceleration to use for image generation. + use_denoising_batch : bool, optional + Whether to use denoising batch or not, by default False. + seed : int, optional + The seed, by default 2. if -1, use random seed. + """ + + stream = StreamDiffusionWrapper( + model_id_or_path=model_id_or_path, + lora_dict=lora_dict, + t_index_list=[0, 16, 32, 45], + frame_buffer_size=1, + width=width, + height=height, + warmup=10, + acceleration=acceleration, + mode="txt2img", + use_denoising_batch=use_denoising_batch, + cfg_type="none", + seed=seed, + ) + + stream.prepare( + prompt=prompt, + num_inference_steps=50, + ) + + for _ in range(stream.batch_size - 1): + stream() + + output_image = stream() + output_image.save(output) diff --git a/streamdiffusion/__init__.py b/streamdiffusion/__init__.py new file mode 100644 index 0000000..3021899 --- /dev/null +++ b/streamdiffusion/__init__.py @@ -0,0 +1 @@ +from .pipeline import StreamDiffusion diff --git a/streamdiffusion/__pycache__/__init__.cpython-310.pyc b/streamdiffusion/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..a00240b Binary files /dev/null and b/streamdiffusion/__pycache__/__init__.cpython-310.pyc differ diff --git a/streamdiffusion/__pycache__/__init__.cpython-311.pyc b/streamdiffusion/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..d83069d Binary files /dev/null and b/streamdiffusion/__pycache__/__init__.cpython-311.pyc differ diff --git a/streamdiffusion/__pycache__/image_filter.cpython-310.pyc b/streamdiffusion/__pycache__/image_filter.cpython-310.pyc new file mode 100644 index 0000000..59e4d71 Binary files /dev/null and b/streamdiffusion/__pycache__/image_filter.cpython-310.pyc differ diff --git a/streamdiffusion/__pycache__/image_filter.cpython-311.pyc b/streamdiffusion/__pycache__/image_filter.cpython-311.pyc new file mode 100644 index 0000000..b4da82c Binary files /dev/null and b/streamdiffusion/__pycache__/image_filter.cpython-311.pyc differ diff --git a/streamdiffusion/__pycache__/image_utils.cpython-310.pyc b/streamdiffusion/__pycache__/image_utils.cpython-310.pyc new file mode 100644 index 0000000..be4f63b Binary files /dev/null and b/streamdiffusion/__pycache__/image_utils.cpython-310.pyc differ diff --git a/streamdiffusion/__pycache__/image_utils.cpython-311.pyc b/streamdiffusion/__pycache__/image_utils.cpython-311.pyc new file mode 100644 index 0000000..2187038 Binary files /dev/null and b/streamdiffusion/__pycache__/image_utils.cpython-311.pyc differ diff --git a/streamdiffusion/__pycache__/pipeline.cpython-310.pyc b/streamdiffusion/__pycache__/pipeline.cpython-310.pyc new file mode 100644 index 0000000..a9720dd Binary files /dev/null and b/streamdiffusion/__pycache__/pipeline.cpython-310.pyc differ diff --git a/streamdiffusion/__pycache__/pipeline.cpython-311.pyc b/streamdiffusion/__pycache__/pipeline.cpython-311.pyc new file mode 100644 index 0000000..fd932b4 Binary files /dev/null and b/streamdiffusion/__pycache__/pipeline.cpython-311.pyc differ diff --git a/streamdiffusion/__pycache__/wrapper.cpython-310.pyc b/streamdiffusion/__pycache__/wrapper.cpython-310.pyc new file mode 100644 index 0000000..871b9b4 Binary files /dev/null and b/streamdiffusion/__pycache__/wrapper.cpython-310.pyc differ diff --git a/streamdiffusion/__pycache__/wrapper.cpython-311.pyc b/streamdiffusion/__pycache__/wrapper.cpython-311.pyc new file mode 100644 index 0000000..769a137 Binary files /dev/null and b/streamdiffusion/__pycache__/wrapper.cpython-311.pyc differ diff --git a/streamdiffusion/acceleration/__init__.py b/streamdiffusion/acceleration/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/streamdiffusion/acceleration/sfast/__init__.py b/streamdiffusion/acceleration/sfast/__init__.py new file mode 100644 index 0000000..962a7f5 --- /dev/null +++ b/streamdiffusion/acceleration/sfast/__init__.py @@ -0,0 +1,33 @@ +from typing import Optional + +from sfast.compilers.stable_diffusion_pipeline_compiler import CompilationConfig, compile + +from ...pipeline import StreamDiffusion + + +def accelerate_with_stable_fast( + stream: StreamDiffusion, + config: Optional[CompilationConfig] = None, +): + if config is None: + config = CompilationConfig.Default() + # xformers and Triton are suggested for achieving best performance. + try: + import xformers + + config.enable_xformers = True + except ImportError: + print("xformers not installed, skip") + try: + import triton + + config.enable_triton = True + except ImportError: + print("Triton not installed, skip") + # CUDA Graph is suggested for small batch sizes and small resolutions to reduce CPU overhead. + config.enable_cuda_graph = True + stream.pipe = compile(stream.pipe, config) + stream.unet = stream.pipe.unet + stream.vae = stream.pipe.vae + stream.text_encoder = stream.pipe.text_encoder + return stream diff --git a/streamdiffusion/acceleration/tensorrt/__init__.py b/streamdiffusion/acceleration/tensorrt/__init__.py new file mode 100644 index 0000000..e629567 --- /dev/null +++ b/streamdiffusion/acceleration/tensorrt/__init__.py @@ -0,0 +1,187 @@ +import gc +import os + +import torch +from diffusers import AutoencoderKL, UNet2DConditionModel +from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import ( + retrieve_latents, +) +from polygraphy import cuda + +from ...pipeline import StreamDiffusion +from .builder import EngineBuilder, create_onnx_path +from .engine import AutoencoderKLEngine, UNet2DConditionModelEngine +from .models import VAE, BaseModel, UNet, VAEEncoder + + +class TorchVAEEncoder(torch.nn.Module): + def __init__(self, vae: AutoencoderKL): + super().__init__() + self.vae = vae + + def forward(self, x: torch.Tensor): + return retrieve_latents(self.vae.encode(x)) + + +def compile_vae_encoder( + vae: TorchVAEEncoder, + model_data: BaseModel, + onnx_path: str, + onnx_opt_path: str, + engine_path: str, + opt_batch_size: int = 1, + engine_build_options: dict = {}, +): + builder = EngineBuilder(model_data, vae, device=torch.device("cuda")) + builder.build( + onnx_path, + onnx_opt_path, + engine_path, + opt_batch_size=opt_batch_size, + **engine_build_options, + ) + + +def compile_vae_decoder( + vae: AutoencoderKL, + model_data: BaseModel, + onnx_path: str, + onnx_opt_path: str, + engine_path: str, + opt_batch_size: int = 1, + engine_build_options: dict = {}, +): + vae = vae.to(torch.device("cuda")) + builder = EngineBuilder(model_data, vae, device=torch.device("cuda")) + builder.build( + onnx_path, + onnx_opt_path, + engine_path, + opt_batch_size=opt_batch_size, + **engine_build_options, + ) + + +def compile_unet( + unet: UNet2DConditionModel, + model_data: BaseModel, + onnx_path: str, + onnx_opt_path: str, + engine_path: str, + opt_batch_size: int = 1, + engine_build_options: dict = {}, +): + unet = unet.to(torch.device("cuda"), dtype=torch.float16) + builder = EngineBuilder(model_data, unet, device=torch.device("cuda")) + builder.build( + onnx_path, + onnx_opt_path, + engine_path, + opt_batch_size=opt_batch_size, + **engine_build_options, + ) + + +def accelerate_with_tensorrt( + stream: StreamDiffusion, + engine_dir: str, + max_batch_size: int = 2, + min_batch_size: int = 1, + use_cuda_graph: bool = False, + engine_build_options: dict = {}, +): + if "opt_batch_size" not in engine_build_options or engine_build_options["opt_batch_size"] is None: + engine_build_options["opt_batch_size"] = max_batch_size + text_encoder = stream.text_encoder + unet = stream.unet + vae = stream.vae + + del stream.unet, stream.vae, stream.pipe.unet, stream.pipe.vae + + vae_config = vae.config + vae_dtype = vae.dtype + + unet.to(torch.device("cpu")) + vae.to(torch.device("cpu")) + + gc.collect() + torch.cuda.empty_cache() + + onnx_dir = os.path.join(engine_dir, "onnx") + os.makedirs(onnx_dir, exist_ok=True) + + unet_engine_path = f"{engine_dir}/unet.engine" + vae_encoder_engine_path = f"{engine_dir}/vae_encoder.engine" + vae_decoder_engine_path = f"{engine_dir}/vae_decoder.engine" + + unet_model = UNet( + fp16=True, + device=stream.device, + max_batch_size=max_batch_size, + min_batch_size=min_batch_size, + embedding_dim=text_encoder.config.hidden_size, + unet_dim=unet.config.in_channels, + ) + vae_decoder_model = VAE( + device=stream.device, + max_batch_size=max_batch_size, + min_batch_size=min_batch_size, + ) + vae_encoder_model = VAEEncoder( + device=stream.device, + max_batch_size=max_batch_size, + min_batch_size=min_batch_size, + ) + + if not os.path.exists(unet_engine_path): + compile_unet( + unet, + unet_model, + create_onnx_path("unet", onnx_dir, opt=False), + create_onnx_path("unet", onnx_dir, opt=True), + opt_batch_size=max_batch_size, + **engine_build_options, + ) + else: + del unet + + if not os.path.exists(vae_decoder_engine_path): + compile_vae_decoder( + vae, + vae_decoder_model, + create_onnx_path("vae_decoder", onnx_dir, opt=False), + create_onnx_path("vae_decoder", onnx_dir, opt=True), + opt_batch_size=max_batch_size, + **engine_build_options, + ) + + if not os.path.exists(vae_encoder_engine_path): + vae_encoder = TorchVAEEncoder(vae).to(torch.device("cuda")) + compile_vae_encoder( + vae_encoder, + vae_encoder_model, + create_onnx_path("vae_encoder", onnx_dir, opt=False), + create_onnx_path("vae_encoder", onnx_dir, opt=True), + opt_batch_size=max_batch_size, + **engine_build_options, + ) + + del vae + + cuda_steram = cuda.Stream() + + stream.unet = UNet2DConditionModelEngine(unet_engine_path, cuda_steram, use_cuda_graph=use_cuda_graph) + stream.vae = AutoencoderKLEngine( + vae_encoder_engine_path, + vae_decoder_engine_path, + cuda_steram, + stream.pipe.vae_scale_factor, + use_cuda_graph=use_cuda_graph, + ) + setattr(stream.vae, "config", vae_config) + setattr(stream.vae, "dtype", vae_dtype) + + gc.collect() + torch.cuda.empty_cache() + + return stream diff --git a/streamdiffusion/acceleration/tensorrt/builder.py b/streamdiffusion/acceleration/tensorrt/builder.py new file mode 100644 index 0000000..69de1bc --- /dev/null +++ b/streamdiffusion/acceleration/tensorrt/builder.py @@ -0,0 +1,94 @@ +import gc +import os +from typing import * + +import torch + +from .models import BaseModel +from .utilities import ( + build_engine, + export_onnx, + optimize_onnx, +) + + +def create_onnx_path(name, onnx_dir, opt=True): + return os.path.join(onnx_dir, name + (".opt" if opt else "") + ".onnx") + + +class EngineBuilder: + def __init__( + self, + model: BaseModel, + network: Any, + device=torch.device("cuda"), + ): + self.device = device + + self.model = model + self.network = network + + def build( + self, + onnx_path: str, + onnx_opt_path: str, + engine_path: str, + opt_image_height: int = 512, + opt_image_width: int = 512, + opt_batch_size: int = 1, + min_image_resolution: int = 256, + max_image_resolution: int = 1024, + build_enable_refit: bool = False, + build_static_batch: bool = False, + build_dynamic_shape: bool = False, + build_all_tactics: bool = False, + onnx_opset: int = 17, + force_engine_build: bool = False, + force_onnx_export: bool = False, + force_onnx_optimize: bool = False, + ): + if not force_onnx_export and os.path.exists(onnx_path): + print(f"Found cached model: {onnx_path}") + else: + print(f"Exporting model: {onnx_path}") + export_onnx( + self.network, + onnx_path=onnx_path, + model_data=self.model, + opt_image_height=opt_image_height, + opt_image_width=opt_image_width, + opt_batch_size=opt_batch_size, + onnx_opset=onnx_opset, + ) + del self.network + gc.collect() + torch.cuda.empty_cache() + if not force_onnx_optimize and os.path.exists(onnx_opt_path): + print(f"Found cached model: {onnx_opt_path}") + else: + print(f"Generating optimizing model: {onnx_opt_path}") + optimize_onnx( + onnx_path=onnx_path, + onnx_opt_path=onnx_opt_path, + model_data=self.model, + ) + self.model.min_latent_shape = min_image_resolution // 8 + self.model.max_latent_shape = max_image_resolution // 8 + if not force_engine_build and os.path.exists(engine_path): + print(f"Found cached engine: {engine_path}") + else: + build_engine( + engine_path=engine_path, + onnx_opt_path=onnx_opt_path, + model_data=self.model, + opt_image_height=opt_image_height, + opt_image_width=opt_image_width, + opt_batch_size=opt_batch_size, + build_static_batch=build_static_batch, + build_dynamic_shape=build_dynamic_shape, + build_all_tactics=build_all_tactics, + build_enable_refit=build_enable_refit, + ) + + gc.collect() + torch.cuda.empty_cache() diff --git a/streamdiffusion/acceleration/tensorrt/engine.py b/streamdiffusion/acceleration/tensorrt/engine.py new file mode 100644 index 0000000..43d6524 --- /dev/null +++ b/streamdiffusion/acceleration/tensorrt/engine.py @@ -0,0 +1,123 @@ +from typing import * + +import torch +from diffusers.models.autoencoder_tiny import AutoencoderTinyOutput +from diffusers.models.unet_2d_condition import UNet2DConditionOutput +from diffusers.models.vae import DecoderOutput +from polygraphy import cuda + +from .utilities import Engine + + +class UNet2DConditionModelEngine: + def __init__(self, filepath: str, stream: cuda.Stream, use_cuda_graph: bool = False): + self.engine = Engine(filepath) + self.stream = stream + self.use_cuda_graph = use_cuda_graph + + self.engine.load() + self.engine.activate() + + def __call__( + self, + latent_model_input: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + **kwargs, + ) -> Any: + if timestep.dtype != torch.float32: + timestep = timestep.float() + + self.engine.allocate_buffers( + shape_dict={ + "sample": latent_model_input.shape, + "timestep": timestep.shape, + "encoder_hidden_states": encoder_hidden_states.shape, + "latent": latent_model_input.shape, + }, + device=latent_model_input.device, + ) + + noise_pred = self.engine.infer( + { + "sample": latent_model_input, + "timestep": timestep, + "encoder_hidden_states": encoder_hidden_states, + }, + self.stream, + use_cuda_graph=self.use_cuda_graph, + )["latent"] + return UNet2DConditionOutput(sample=noise_pred) + + def to(self, *args, **kwargs): + pass + + def forward(self, *args, **kwargs): + pass + + +class AutoencoderKLEngine: + def __init__( + self, + encoder_path: str, + decoder_path: str, + stream: cuda.Stream, + scaling_factor: int, + use_cuda_graph: bool = False, + ): + self.encoder = Engine(encoder_path) + self.decoder = Engine(decoder_path) + self.stream = stream + self.vae_scale_factor = scaling_factor + self.use_cuda_graph = use_cuda_graph + + self.encoder.load() + self.decoder.load() + self.encoder.activate() + self.decoder.activate() + + def encode(self, images: torch.Tensor, **kwargs): + self.encoder.allocate_buffers( + shape_dict={ + "images": images.shape, + "latent": ( + images.shape[0], + 4, + images.shape[2] // self.vae_scale_factor, + images.shape[3] // self.vae_scale_factor, + ), + }, + device=images.device, + ) + latents = self.encoder.infer( + {"images": images}, + self.stream, + use_cuda_graph=self.use_cuda_graph, + )["latent"] + return AutoencoderTinyOutput(latents=latents) + + def decode(self, latent: torch.Tensor, **kwargs): + self.decoder.allocate_buffers( + shape_dict={ + "latent": latent.shape, + "images": ( + latent.shape[0], + 3, + latent.shape[2] * self.vae_scale_factor, + latent.shape[3] * self.vae_scale_factor, + ), + }, + device=latent.device, + ) + images = self.decoder.infer( + {"latent": latent}, + self.stream, + use_cuda_graph=self.use_cuda_graph, + )["images"] + return DecoderOutput(sample=images) + + def to(self, *args, **kwargs): + pass + + def forward(self, *args, **kwargs): + pass diff --git a/streamdiffusion/acceleration/tensorrt/models.py b/streamdiffusion/acceleration/tensorrt/models.py new file mode 100644 index 0000000..3c631ae --- /dev/null +++ b/streamdiffusion/acceleration/tensorrt/models.py @@ -0,0 +1,434 @@ +#! fork: https://github.com/NVIDIA/TensorRT/blob/main/demo/Diffusion/models.py + +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import onnx_graphsurgeon as gs +import torch +from onnx import shape_inference +from polygraphy.backend.onnx.loader import fold_constants + + +class Optimizer: + def __init__(self, onnx_graph, verbose=False): + self.graph = gs.import_onnx(onnx_graph) + self.verbose = verbose + + def info(self, prefix): + if self.verbose: + print( + f"{prefix} .. {len(self.graph.nodes)} nodes, {len(self.graph.tensors().keys())} tensors, {len(self.graph.inputs)} inputs, {len(self.graph.outputs)} outputs" + ) + + def cleanup(self, return_onnx=False): + self.graph.cleanup().toposort() + if return_onnx: + return gs.export_onnx(self.graph) + + def select_outputs(self, keep, names=None): + self.graph.outputs = [self.graph.outputs[o] for o in keep] + if names: + for i, name in enumerate(names): + self.graph.outputs[i].name = name + + def fold_constants(self, return_onnx=False): + onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True) + self.graph = gs.import_onnx(onnx_graph) + if return_onnx: + return onnx_graph + + def infer_shapes(self, return_onnx=False): + onnx_graph = gs.export_onnx(self.graph) + if onnx_graph.ByteSize() > 2147483648: + raise TypeError("ERROR: model size exceeds supported 2GB limit") + else: + onnx_graph = shape_inference.infer_shapes(onnx_graph) + + self.graph = gs.import_onnx(onnx_graph) + if return_onnx: + return onnx_graph + + +class BaseModel: + def __init__( + self, + fp16=False, + device="cuda", + verbose=True, + max_batch_size=16, + min_batch_size=1, + embedding_dim=768, + text_maxlen=77, + ): + self.name = "SD Model" + self.fp16 = fp16 + self.device = device + self.verbose = verbose + + self.min_batch = min_batch_size + self.max_batch = max_batch_size + self.min_image_shape = 256 # min image resolution: 256x256 + self.max_image_shape = 1024 # max image resolution: 1024x1024 + self.min_latent_shape = self.min_image_shape // 8 + self.max_latent_shape = self.max_image_shape // 8 + + self.embedding_dim = embedding_dim + self.text_maxlen = text_maxlen + + def get_model(self): + pass + + def get_input_names(self): + pass + + def get_output_names(self): + pass + + def get_dynamic_axes(self): + return None + + def get_sample_input(self, batch_size, image_height, image_width): + pass + + def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): + return None + + def get_shape_dict(self, batch_size, image_height, image_width): + return None + + def optimize(self, onnx_graph): + opt = Optimizer(onnx_graph, verbose=self.verbose) + opt.info(self.name + ": original") + opt.cleanup() + opt.info(self.name + ": cleanup") + opt.fold_constants() + opt.info(self.name + ": fold constants") + opt.infer_shapes() + opt.info(self.name + ": shape inference") + onnx_opt_graph = opt.cleanup(return_onnx=True) + opt.info(self.name + ": finished") + return onnx_opt_graph + + def check_dims(self, batch_size, image_height, image_width): + assert batch_size >= self.min_batch and batch_size <= self.max_batch + assert image_height % 8 == 0 or image_width % 8 == 0 + latent_height = image_height // 8 + latent_width = image_width // 8 + assert latent_height >= self.min_latent_shape and latent_height <= self.max_latent_shape + assert latent_width >= self.min_latent_shape and latent_width <= self.max_latent_shape + return (latent_height, latent_width) + + def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape): + min_batch = batch_size if static_batch else self.min_batch + max_batch = batch_size if static_batch else self.max_batch + latent_height = image_height // 8 + latent_width = image_width // 8 + min_image_height = image_height if static_shape else self.min_image_shape + max_image_height = image_height if static_shape else self.max_image_shape + min_image_width = image_width if static_shape else self.min_image_shape + max_image_width = image_width if static_shape else self.max_image_shape + min_latent_height = latent_height if static_shape else self.min_latent_shape + max_latent_height = latent_height if static_shape else self.max_latent_shape + min_latent_width = latent_width if static_shape else self.min_latent_shape + max_latent_width = latent_width if static_shape else self.max_latent_shape + return ( + min_batch, + max_batch, + min_image_height, + max_image_height, + min_image_width, + max_image_width, + min_latent_height, + max_latent_height, + min_latent_width, + max_latent_width, + ) + + +class CLIP(BaseModel): + def __init__(self, device, max_batch_size, embedding_dim, min_batch_size=1): + super(CLIP, self).__init__( + device=device, + max_batch_size=max_batch_size, + min_batch_size=min_batch_size, + embedding_dim=embedding_dim, + ) + self.name = "CLIP" + + def get_input_names(self): + return ["input_ids"] + + def get_output_names(self): + return ["text_embeddings", "pooler_output"] + + def get_dynamic_axes(self): + return {"input_ids": {0: "B"}, "text_embeddings": {0: "B"}} + + def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): + self.check_dims(batch_size, image_height, image_width) + min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims( + batch_size, image_height, image_width, static_batch, static_shape + ) + return { + "input_ids": [ + (min_batch, self.text_maxlen), + (batch_size, self.text_maxlen), + (max_batch, self.text_maxlen), + ] + } + + def get_shape_dict(self, batch_size, image_height, image_width): + self.check_dims(batch_size, image_height, image_width) + return { + "input_ids": (batch_size, self.text_maxlen), + "text_embeddings": (batch_size, self.text_maxlen, self.embedding_dim), + } + + def get_sample_input(self, batch_size, image_height, image_width): + self.check_dims(batch_size, image_height, image_width) + return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device) + + def optimize(self, onnx_graph): + opt = Optimizer(onnx_graph) + opt.info(self.name + ": original") + opt.select_outputs([0]) # delete graph output#1 + opt.cleanup() + opt.info(self.name + ": remove output[1]") + opt.fold_constants() + opt.info(self.name + ": fold constants") + opt.infer_shapes() + opt.info(self.name + ": shape inference") + opt.select_outputs([0], names=["text_embeddings"]) # rename network output + opt.info(self.name + ": remove output[0]") + opt_onnx_graph = opt.cleanup(return_onnx=True) + opt.info(self.name + ": finished") + return opt_onnx_graph + + +class UNet(BaseModel): + def __init__( + self, + fp16=False, + device="cuda", + max_batch_size=16, + min_batch_size=1, + embedding_dim=768, + text_maxlen=77, + unet_dim=4, + ): + super(UNet, self).__init__( + fp16=fp16, + device=device, + max_batch_size=max_batch_size, + min_batch_size=min_batch_size, + embedding_dim=embedding_dim, + text_maxlen=text_maxlen, + ) + self.unet_dim = unet_dim + self.name = "UNet" + + def get_input_names(self): + return ["sample", "timestep", "encoder_hidden_states"] + + def get_output_names(self): + return ["latent"] + + def get_dynamic_axes(self): + return { + "sample": {0: "2B", 2: "H", 3: "W"}, + "timestep": {0: "2B"}, + "encoder_hidden_states": {0: "2B"}, + "latent": {0: "2B", 2: "H", 3: "W"}, + } + + def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + ( + min_batch, + max_batch, + _, + _, + _, + _, + min_latent_height, + max_latent_height, + min_latent_width, + max_latent_width, + ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape) + return { + "sample": [ + (min_batch, self.unet_dim, min_latent_height, min_latent_width), + (batch_size, self.unet_dim, latent_height, latent_width), + (max_batch, self.unet_dim, max_latent_height, max_latent_width), + ], + "timestep": [(min_batch,), (batch_size,), (max_batch,)], + "encoder_hidden_states": [ + (min_batch, self.text_maxlen, self.embedding_dim), + (batch_size, self.text_maxlen, self.embedding_dim), + (max_batch, self.text_maxlen, self.embedding_dim), + ], + } + + def get_shape_dict(self, batch_size, image_height, image_width): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + return { + "sample": (2 * batch_size, self.unet_dim, latent_height, latent_width), + "timestep": (2 * batch_size,), + "encoder_hidden_states": (2 * batch_size, self.text_maxlen, self.embedding_dim), + "latent": (2 * batch_size, 4, latent_height, latent_width), + } + + def get_sample_input(self, batch_size, image_height, image_width): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + dtype = torch.float16 if self.fp16 else torch.float32 + return ( + torch.randn( + 2 * batch_size, self.unet_dim, latent_height, latent_width, dtype=torch.float32, device=self.device + ), + torch.ones((2 * batch_size,), dtype=torch.float32, device=self.device), + torch.randn(2 * batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device), + ) + + +class VAE(BaseModel): + def __init__(self, device, max_batch_size, min_batch_size=1): + super(VAE, self).__init__( + device=device, + max_batch_size=max_batch_size, + min_batch_size=min_batch_size, + embedding_dim=None, + ) + self.name = "VAE decoder" + + def get_input_names(self): + return ["latent"] + + def get_output_names(self): + return ["images"] + + def get_dynamic_axes(self): + return { + "latent": {0: "B", 2: "H", 3: "W"}, + "images": {0: "B", 2: "8H", 3: "8W"}, + } + + def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + ( + min_batch, + max_batch, + _, + _, + _, + _, + min_latent_height, + max_latent_height, + min_latent_width, + max_latent_width, + ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape) + return { + "latent": [ + (min_batch, 4, min_latent_height, min_latent_width), + (batch_size, 4, latent_height, latent_width), + (max_batch, 4, max_latent_height, max_latent_width), + ] + } + + def get_shape_dict(self, batch_size, image_height, image_width): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + return { + "latent": (batch_size, 4, latent_height, latent_width), + "images": (batch_size, 3, image_height, image_width), + } + + def get_sample_input(self, batch_size, image_height, image_width): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + return torch.randn( + batch_size, + 4, + latent_height, + latent_width, + dtype=torch.float32, + device=self.device, + ) + + +class VAEEncoder(BaseModel): + def __init__(self, device, max_batch_size, min_batch_size=1): + super(VAEEncoder, self).__init__( + device=device, + max_batch_size=max_batch_size, + min_batch_size=min_batch_size, + embedding_dim=None, + ) + self.name = "VAE encoder" + + def get_input_names(self): + return ["images"] + + def get_output_names(self): + return ["latent"] + + def get_dynamic_axes(self): + return { + "images": {0: "B", 2: "8H", 3: "8W"}, + "latent": {0: "B", 2: "H", 3: "W"}, + } + + def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): + assert batch_size >= self.min_batch and batch_size <= self.max_batch + min_batch = batch_size if static_batch else self.min_batch + max_batch = batch_size if static_batch else self.max_batch + self.check_dims(batch_size, image_height, image_width) + ( + min_batch, + max_batch, + min_image_height, + max_image_height, + min_image_width, + max_image_width, + _, + _, + _, + _, + ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape) + + return { + "images": [ + (min_batch, 3, min_image_height, min_image_width), + (batch_size, 3, image_height, image_width), + (max_batch, 3, max_image_height, max_image_width), + ], + } + + def get_shape_dict(self, batch_size, image_height, image_width): + latent_height, latent_width = self.check_dims(batch_size, image_height, image_width) + return { + "images": (batch_size, 3, image_height, image_width), + "latent": (batch_size, 4, latent_height, latent_width), + } + + def get_sample_input(self, batch_size, image_height, image_width): + self.check_dims(batch_size, image_height, image_width) + return torch.randn( + batch_size, + 3, + image_height, + image_width, + dtype=torch.float32, + device=self.device, + ) diff --git a/streamdiffusion/acceleration/tensorrt/utilities.py b/streamdiffusion/acceleration/tensorrt/utilities.py new file mode 100644 index 0000000..eb1b01e --- /dev/null +++ b/streamdiffusion/acceleration/tensorrt/utilities.py @@ -0,0 +1,441 @@ +#! fork: https://github.com/NVIDIA/TensorRT/blob/main/demo/Diffusion/utilities.py + +# +# Copyright 2022 The HuggingFace Inc. team. +# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import gc +from collections import OrderedDict +from typing import * + +import numpy as np +import onnx +import onnx_graphsurgeon as gs +import tensorrt as trt +import torch +from cuda import cudart +from PIL import Image +from polygraphy import cuda +from polygraphy.backend.common import bytes_from_path +from polygraphy.backend.trt import ( + CreateConfig, + Profile, + engine_from_bytes, + engine_from_network, + network_from_onnx_path, + save_engine, +) +from polygraphy.backend.trt import util as trt_util + +from .models import CLIP, VAE, BaseModel, UNet, VAEEncoder + + +TRT_LOGGER = trt.Logger(trt.Logger.ERROR) + +# Map of numpy dtype -> torch dtype +numpy_to_torch_dtype_dict = { + np.uint8: torch.uint8, + np.int8: torch.int8, + np.int16: torch.int16, + np.int32: torch.int32, + np.int64: torch.int64, + np.float16: torch.float16, + np.float32: torch.float32, + np.float64: torch.float64, + np.complex64: torch.complex64, + np.complex128: torch.complex128, +} +if np.version.full_version >= "1.24.0": + numpy_to_torch_dtype_dict[np.bool_] = torch.bool +else: + numpy_to_torch_dtype_dict[np.bool] = torch.bool + +# Map of torch dtype -> numpy dtype +torch_to_numpy_dtype_dict = {value: key for (key, value) in numpy_to_torch_dtype_dict.items()} + + +def CUASSERT(cuda_ret): + err = cuda_ret[0] + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError( + f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" + ) + if len(cuda_ret) > 1: + return cuda_ret[1] + return None + + +class Engine: + def __init__( + self, + engine_path, + ): + self.engine_path = engine_path + self.engine = None + self.context = None + self.buffers = OrderedDict() + self.tensors = OrderedDict() + self.cuda_graph_instance = None # cuda graph + + def __del__(self): + [buf.free() for buf in self.buffers.values() if isinstance(buf, cuda.DeviceArray)] + del self.engine + del self.context + del self.buffers + del self.tensors + + def refit(self, onnx_path, onnx_refit_path): + def convert_int64(arr): + # TODO: smarter conversion + if len(arr.shape) == 0: + return np.int32(arr) + return arr + + def add_to_map(refit_dict, name, values): + if name in refit_dict: + assert refit_dict[name] is None + if values.dtype == np.int64: + values = convert_int64(values) + refit_dict[name] = values + + print(f"Refitting TensorRT engine with {onnx_refit_path} weights") + refit_nodes = gs.import_onnx(onnx.load(onnx_refit_path)).toposort().nodes + + # Construct mapping from weight names in refit model -> original model + name_map = {} + for n, node in enumerate(gs.import_onnx(onnx.load(onnx_path)).toposort().nodes): + refit_node = refit_nodes[n] + assert node.op == refit_node.op + # Constant nodes in ONNX do not have inputs but have a constant output + if node.op == "Constant": + name_map[refit_node.outputs[0].name] = node.outputs[0].name + # Handle scale and bias weights + elif node.op == "Conv": + if node.inputs[1].__class__ == gs.Constant: + name_map[refit_node.name + "_TRTKERNEL"] = node.name + "_TRTKERNEL" + if node.inputs[2].__class__ == gs.Constant: + name_map[refit_node.name + "_TRTBIAS"] = node.name + "_TRTBIAS" + # For all other nodes: find node inputs that are initializers (gs.Constant) + else: + for i, inp in enumerate(node.inputs): + if inp.__class__ == gs.Constant: + name_map[refit_node.inputs[i].name] = inp.name + + def map_name(name): + if name in name_map: + return name_map[name] + return name + + # Construct refit dictionary + refit_dict = {} + refitter = trt.Refitter(self.engine, TRT_LOGGER) + all_weights = refitter.get_all() + for layer_name, role in zip(all_weights[0], all_weights[1]): + # for speciailized roles, use a unique name in the map: + if role == trt.WeightsRole.KERNEL: + name = layer_name + "_TRTKERNEL" + elif role == trt.WeightsRole.BIAS: + name = layer_name + "_TRTBIAS" + else: + name = layer_name + + assert name not in refit_dict, "Found duplicate layer: " + name + refit_dict[name] = None + + for n in refit_nodes: + # Constant nodes in ONNX do not have inputs but have a constant output + if n.op == "Constant": + name = map_name(n.outputs[0].name) + print(f"Add Constant {name}\n") + add_to_map(refit_dict, name, n.outputs[0].values) + + # Handle scale and bias weights + elif n.op == "Conv": + if n.inputs[1].__class__ == gs.Constant: + name = map_name(n.name + "_TRTKERNEL") + add_to_map(refit_dict, name, n.inputs[1].values) + + if n.inputs[2].__class__ == gs.Constant: + name = map_name(n.name + "_TRTBIAS") + add_to_map(refit_dict, name, n.inputs[2].values) + + # For all other nodes: find node inputs that are initializers (AKA gs.Constant) + else: + for inp in n.inputs: + name = map_name(inp.name) + if inp.__class__ == gs.Constant: + add_to_map(refit_dict, name, inp.values) + + for layer_name, weights_role in zip(all_weights[0], all_weights[1]): + if weights_role == trt.WeightsRole.KERNEL: + custom_name = layer_name + "_TRTKERNEL" + elif weights_role == trt.WeightsRole.BIAS: + custom_name = layer_name + "_TRTBIAS" + else: + custom_name = layer_name + + # Skip refitting Trilu for now; scalar weights of type int64 value 1 - for clip model + if layer_name.startswith("onnx::Trilu"): + continue + + if refit_dict[custom_name] is not None: + refitter.set_weights(layer_name, weights_role, refit_dict[custom_name]) + else: + print(f"[W] No refit weights for layer: {layer_name}") + + if not refitter.refit_cuda_engine(): + print("Failed to refit!") + exit(0) + + def build( + self, + onnx_path, + fp16, + input_profile=None, + enable_refit=False, + enable_all_tactics=False, + timing_cache=None, + workspace_size=0, + ): + print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}") + p = Profile() + if input_profile: + for name, dims in input_profile.items(): + assert len(dims) == 3 + p.add(name, min=dims[0], opt=dims[1], max=dims[2]) + + config_kwargs = {} + + if workspace_size > 0: + config_kwargs["memory_pool_limits"] = {trt.MemoryPoolType.WORKSPACE: workspace_size} + if not enable_all_tactics: + config_kwargs["tactic_sources"] = [] + + engine = engine_from_network( + network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]), + config=CreateConfig( + fp16=fp16, refittable=enable_refit, profiles=[p], load_timing_cache=timing_cache, **config_kwargs + ), + save_timing_cache=timing_cache, + ) + save_engine(engine, path=self.engine_path) + + def load(self): + print(f"Loading TensorRT engine: {self.engine_path}") + self.engine = engine_from_bytes(bytes_from_path(self.engine_path)) + + def activate(self, reuse_device_memory=None): + if reuse_device_memory: + self.context = self.engine.create_execution_context_without_device_memory() + self.context.device_memory = reuse_device_memory + else: + self.context = self.engine.create_execution_context() + + def allocate_buffers(self, shape_dict=None, device="cuda"): + for idx in range(trt_util.get_bindings_per_profile(self.engine)): + binding = self.engine[idx] + if shape_dict and binding in shape_dict: + shape = shape_dict[binding] + else: + shape = self.engine.get_binding_shape(binding) + dtype = trt.nptype(self.engine.get_binding_dtype(binding)) + if self.engine.binding_is_input(binding): + self.context.set_binding_shape(idx, shape) + tensor = torch.empty(tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]).to(device=device) + self.tensors[binding] = tensor + + def infer(self, feed_dict, stream, use_cuda_graph=False): + for name, buf in feed_dict.items(): + self.tensors[name].copy_(buf) + + for name, tensor in self.tensors.items(): + self.context.set_tensor_address(name, tensor.data_ptr()) + + if use_cuda_graph: + if self.cuda_graph_instance is not None: + CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream.ptr)) + CUASSERT(cudart.cudaStreamSynchronize(stream.ptr)) + else: + # do inference before CUDA graph capture + noerror = self.context.execute_async_v3(stream.ptr) + if not noerror: + raise ValueError("ERROR: inference failed.") + # capture cuda graph + CUASSERT( + cudart.cudaStreamBeginCapture(stream.ptr, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal) + ) + self.context.execute_async_v3(stream.ptr) + self.graph = CUASSERT(cudart.cudaStreamEndCapture(stream.ptr)) + self.cuda_graph_instance = CUASSERT(cudart.cudaGraphInstantiate(self.graph, 0)) + else: + noerror = self.context.execute_async_v3(stream.ptr) + if not noerror: + raise ValueError("ERROR: inference failed.") + + return self.tensors + + +def decode_images(images: torch.Tensor): + images = ( + ((images + 1) * 255 / 2).clamp(0, 255).detach().permute(0, 2, 3, 1).round().type(torch.uint8).cpu().numpy() + ) + return [Image.fromarray(x) for x in images] + + +def preprocess_image(image: Image.Image): + w, h = image.size + w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32 + image = image.resize((w, h)) + init_image = np.array(image).astype(np.float32) / 255.0 + init_image = init_image[None].transpose(0, 3, 1, 2) + init_image = torch.from_numpy(init_image).contiguous() + return 2.0 * init_image - 1.0 + + +def prepare_mask_and_masked_image(image: Image.Image, mask: Image.Image): + if isinstance(image, Image.Image): + image = np.array(image.convert("RGB")) + image = image[None].transpose(0, 3, 1, 2) + image = torch.from_numpy(image).to(dtype=torch.float32).contiguous() / 127.5 - 1.0 + if isinstance(mask, Image.Image): + mask = np.array(mask.convert("L")) + mask = mask.astype(np.float32) / 255.0 + mask = mask[None, None] + mask[mask < 0.5] = 0 + mask[mask >= 0.5] = 1 + mask = torch.from_numpy(mask).to(dtype=torch.float32).contiguous() + + masked_image = image * (mask < 0.5) + + return mask, masked_image + + +def create_models( + model_id: str, + use_auth_token: Optional[str], + device: Union[str, torch.device], + max_batch_size: int, + unet_in_channels: int = 4, + embedding_dim: int = 768, +): + models = { + "clip": CLIP( + hf_token=use_auth_token, + device=device, + max_batch_size=max_batch_size, + embedding_dim=embedding_dim, + ), + "unet": UNet( + hf_token=use_auth_token, + fp16=True, + device=device, + max_batch_size=max_batch_size, + embedding_dim=embedding_dim, + unet_dim=unet_in_channels, + ), + "vae": VAE( + hf_token=use_auth_token, + device=device, + max_batch_size=max_batch_size, + embedding_dim=embedding_dim, + ), + "vae_encoder": VAEEncoder( + hf_token=use_auth_token, + device=device, + max_batch_size=max_batch_size, + embedding_dim=embedding_dim, + ), + } + return models + + +def build_engine( + engine_path: str, + onnx_opt_path: str, + model_data: BaseModel, + opt_image_height: int, + opt_image_width: int, + opt_batch_size: int, + build_static_batch: bool = False, + build_dynamic_shape: bool = False, + build_all_tactics: bool = False, + build_enable_refit: bool = False, +): + _, free_mem, _ = cudart.cudaMemGetInfo() + GiB = 2**30 + if free_mem > 6 * GiB: + activation_carveout = 4 * GiB + max_workspace_size = free_mem - activation_carveout + else: + max_workspace_size = 0 + engine = Engine(engine_path) + input_profile = model_data.get_input_profile( + opt_batch_size, + opt_image_height, + opt_image_width, + static_batch=build_static_batch, + static_shape=not build_dynamic_shape, + ) + engine.build( + onnx_opt_path, + fp16=True, + input_profile=input_profile, + enable_refit=build_enable_refit, + enable_all_tactics=build_all_tactics, + workspace_size=max_workspace_size, + ) + + return engine + + +def export_onnx( + model, + onnx_path: str, + model_data: BaseModel, + opt_image_height: int, + opt_image_width: int, + opt_batch_size: int, + onnx_opset: int, +): + with torch.inference_mode(), torch.autocast("cuda"): + inputs = model_data.get_sample_input(opt_batch_size, opt_image_height, opt_image_width) + torch.onnx.export( + model, + inputs, + onnx_path, + export_params=True, + opset_version=onnx_opset, + do_constant_folding=True, + input_names=model_data.get_input_names(), + output_names=model_data.get_output_names(), + dynamic_axes=model_data.get_dynamic_axes(), + ) + del model + gc.collect() + torch.cuda.empty_cache() + + +def optimize_onnx( + onnx_path: str, + onnx_opt_path: str, + model_data: BaseModel, +): + onnx_opt_graph = model_data.optimize(onnx.load(onnx_path)) + onnx.save(onnx_opt_graph, onnx_opt_path) + del onnx_opt_graph + gc.collect() + torch.cuda.empty_cache() diff --git a/streamdiffusion/image_filter.py b/streamdiffusion/image_filter.py new file mode 100644 index 0000000..3c6a1d4 --- /dev/null +++ b/streamdiffusion/image_filter.py @@ -0,0 +1,45 @@ +from typing import Optional +import random + +import torch + + +class SimilarImageFilter: + def __init__(self, threshold: float = 0.98, max_skip_frame: float = 10) -> None: + self.threshold = threshold + self.prev_tensor = None + self.cos = torch.nn.CosineSimilarity(dim=0, eps=1e-6) + self.max_skip_frame = max_skip_frame + self.skip_count = 0 + + def __call__(self, x: torch.Tensor) -> Optional[torch.Tensor]: + if self.prev_tensor is None: + self.prev_tensor = x.detach().clone() + return x + else: + cos_sim = self.cos(self.prev_tensor.reshape(-1), x.reshape(-1)).item() + sample = random.uniform(0, 1) + if self.threshold >= 1: + skip_prob = 0 + else: + skip_prob = max(0, 1 - (1 - cos_sim) / (1 - self.threshold)) + + # not skip frame + if skip_prob < sample: + self.prev_tensor = x.detach().clone() + return x + # skip frame + else: + if self.skip_count > self.max_skip_frame: + self.skip_count = 0 + self.prev_tensor = x.detach().clone() + return x + else: + self.skip_count += 1 + return None + + def set_threshold(self, threshold: float) -> None: + self.threshold = threshold + + def set_max_skip_frame(self, max_skip_frame: float) -> None: + self.max_skip_frame = max_skip_frame diff --git a/streamdiffusion/image_utils.py b/streamdiffusion/image_utils.py new file mode 100644 index 0000000..200295b --- /dev/null +++ b/streamdiffusion/image_utils.py @@ -0,0 +1,98 @@ +from typing import List, Optional, Tuple, Union + +import numpy as np +import PIL.Image +import torch +import torchvision + + +def denormalize(images: Union[torch.Tensor, np.ndarray]) -> torch.Tensor: + """ + Denormalize an image array to [0,1]. + """ + return (images / 2 + 0.5).clamp(0, 1) + + +def pt_to_numpy(images: torch.Tensor) -> np.ndarray: + """ + Convert a PyTorch tensor to a NumPy image. + """ + images = images.cpu().permute(0, 2, 3, 1).float().numpy() + return images + + +def numpy_to_pil(images: np.ndarray) -> PIL.Image.Image: + """ + Convert a NumPy image or a batch of images to a PIL image. + """ + if images.ndim == 3: + images = images[None, ...] + images = (images * 255).round().astype("uint8") + if images.shape[-1] == 1: + # special case for grayscale (single channel) images + pil_images = [ + PIL.Image.fromarray(image.squeeze(), mode="L") for image in images + ] + else: + pil_images = [PIL.Image.fromarray(image) for image in images] + + return pil_images + + +def postprocess_image( + image: torch.Tensor, + output_type: str = "pil", + do_denormalize: Optional[List[bool]] = None, +) -> Union[torch.Tensor, np.ndarray, PIL.Image.Image]: + if not isinstance(image, torch.Tensor): + raise ValueError( + f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor" + ) + + if output_type == "latent": + return image + + do_normalize_flg = True + if do_denormalize is None: + do_denormalize = [do_normalize_flg] * image.shape[0] + + image = torch.stack( + [ + denormalize(image[i]) if do_denormalize[i] else image[i] + for i in range(image.shape[0]) + ] + ) + + if output_type == "pt": + return image + + image = pt_to_numpy(image) + + if output_type == "np": + return image + + if output_type == "pil": + return numpy_to_pil(image) + + +def process_image( + image_pil: PIL.Image.Image, range: Tuple[int, int] = (-1, 1) +) -> Tuple[torch.Tensor, PIL.Image.Image]: + image = torchvision.transforms.ToTensor()(image_pil) + r_min, r_max = range[0], range[1] + image = image * (r_max - r_min) + r_min + return image[None, ...], image_pil + + +def pil2tensor(image_pil: PIL.Image.Image) -> torch.Tensor: + height = image_pil.height + width = image_pil.width + imgs = [] + img, _ = process_image(image_pil) + imgs.append(img) + imgs = torch.vstack(imgs) + images = torch.nn.functional.interpolate( + imgs, size=(height, width), mode="bilinear" + ) + image_tensors = images.to(torch.float16) + return image_tensors diff --git a/streamdiffusion/pip_utils.py b/streamdiffusion/pip_utils.py new file mode 100644 index 0000000..200b74c --- /dev/null +++ b/streamdiffusion/pip_utils.py @@ -0,0 +1,52 @@ +import importlib +import importlib.util +import os +import subprocess +import sys +from typing import Dict, Optional + +from packaging.version import Version + + +python = sys.executable +index_url = os.environ.get("INDEX_URL", "") + + +def version(package: str) -> Optional[Version]: + try: + return Version(importlib.import_module(package).__version__) + except ModuleNotFoundError: + return None + + +def is_installed(package: str) -> bool: + try: + spec = importlib.util.find_spec(package) + except ModuleNotFoundError: + return False + + return spec is not None + + +def run_python(command: str, env: Dict[str, str] = None) -> str: + run_kwargs = { + "args": f"{python} {command}", + "shell": True, + "env": os.environ if env is None else env, + "encoding": "utf8", + "errors": "ignore", + } + + print(run_kwargs["args"]) + + result = subprocess.run(**run_kwargs) + + if result.returncode != 0: + print(f"Error running command: {command}", file=sys.stderr) + raise RuntimeError(f"Error running command: {command}") + + return result.stdout or "" + + +def run_pip(command: str, env: Dict[str, str] = None) -> str: + return run_python(f"-m pip {command}", env) diff --git a/streamdiffusion/pipeline.py b/streamdiffusion/pipeline.py new file mode 100644 index 0000000..f4fb61d --- /dev/null +++ b/streamdiffusion/pipeline.py @@ -0,0 +1,575 @@ +import time +from typing import List, Optional, Union, Any, Dict, Tuple, Literal + +import numpy as np +import PIL.Image +import torch +from diffusers import LCMScheduler, StableDiffusionPipeline +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import ( + retrieve_latents, +) + +from .image_filter import SimilarImageFilter +from .image_utils import postprocess_image + + +class StreamDiffusion: + def __init__( + self, + pipe: StableDiffusionPipeline, + t_index_list: List[int], + torch_dtype: torch.dtype = torch.float16, + width: int = 512, + height: int = 512, + do_add_noise: bool = True, + use_denoising_batch: bool = True, + frame_buffer_size: int = 1, + cfg_type: Literal["none", "full", "self", "initialize"] = "self", + ) -> None: + self.device = pipe.device + self.dtype = torch_dtype + self.generator = None + + self.height = height + self.width = width + + self.latent_height = int(height // pipe.vae_scale_factor) + self.latent_width = int(width // pipe.vae_scale_factor) + + self.frame_bff_size = frame_buffer_size + self.denoising_steps_num = len(t_index_list) + + self.cfg_type = cfg_type + + if use_denoising_batch: + self.batch_size = self.denoising_steps_num * frame_buffer_size + if self.cfg_type == "initialize": + self.trt_unet_batch_size = ( + self.denoising_steps_num + 1 + ) * self.frame_bff_size + elif self.cfg_type == "full": + self.trt_unet_batch_size = ( + 2 * self.denoising_steps_num * self.frame_bff_size + ) + else: + self.trt_unet_batch_size = self.denoising_steps_num * frame_buffer_size + else: + self.trt_unet_batch_size = self.frame_bff_size + self.batch_size = frame_buffer_size + + self.t_list = t_index_list + + self.do_add_noise = do_add_noise + self.use_denoising_batch = use_denoising_batch + + self.similar_image_filter = False + self.similar_filter = SimilarImageFilter() + self.prev_image_result = None + + self.pipe = pipe + self.image_processor = VaeImageProcessor(pipe.vae_scale_factor) + + self.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config) + self.text_encoder = pipe.text_encoder + self.unet = pipe.unet + self.vae = pipe.vae + + self.inference_time_ema = 0 + + def set_sampler_param(self, + t_index_list: List[int], + width: int = 512, + height: int = 512, + do_add_noise: bool = True, + use_denoising_batch: bool = True, + frame_buffer_size: int = 1, + cfg_type: Literal["none", "full", "self", "initialize"] = "self",): + self.height = height + self.width = width + + self.latent_height = int(height // self.pipe.vae_scale_factor) + self.latent_width = int(width // self.pipe.vae_scale_factor) + + self.frame_bff_size = frame_buffer_size + + self.cfg_type = cfg_type + self.t_list = t_index_list + + self.do_add_noise = do_add_noise + self.use_denoising_batch = use_denoising_batch + + self.inference_time_ema = 0 + + self.denoising_steps_num = len(self.t_list) + if self.use_denoising_batch: + self.batch_size = self.denoising_steps_num * self.frame_bff_size + if self.cfg_type == "initialize": + self.trt_unet_batch_size = ( + self.denoising_steps_num + 1 + ) * self.frame_bff_size + elif self.cfg_type == "full": + self.trt_unet_batch_size = ( + 2 * self.denoising_steps_num * self.frame_bff_size + ) + else: + self.trt_unet_batch_size = self.denoising_steps_num * self.frame_bff_size + else: + self.trt_unet_batch_size = self.frame_bff_size + self.batch_size = self.frame_bff_size + + def load_lcm_lora( + self, + pretrained_model_name_or_path_or_dict: Union[ + str, Dict[str, torch.Tensor] + ] = "latent-consistency/lcm-lora-sdv1-5", + adapter_name: Optional[Any] = None, + **kwargs, + ) -> None: + self.pipe.load_lora_weights( + pretrained_model_name_or_path_or_dict, adapter_name, **kwargs + ) + + def load_lora( + self, + pretrained_lora_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], + adapter_name: Optional[Any] = None, + **kwargs, + ) -> None: + self.pipe.load_lora_weights( + pretrained_lora_model_name_or_path_or_dict, adapter_name, **kwargs + ) + + def fuse_lora( + self, + fuse_unet: bool = True, + fuse_text_encoder: bool = True, + lora_scale: float = 1.0, + safe_fusing: bool = False, + ) -> None: + self.pipe.fuse_lora( + fuse_unet=fuse_unet, + fuse_text_encoder=fuse_text_encoder, + lora_scale=lora_scale, + safe_fusing=safe_fusing, + ) + + def enable_similar_image_filter(self, threshold: float = 0.98, max_skip_frame: float = 10) -> None: + self.similar_image_filter = True + self.similar_filter.set_threshold(threshold) + self.similar_filter.set_max_skip_frame(max_skip_frame) + + def disable_similar_image_filter(self) -> None: + self.similar_image_filter = False + + @torch.no_grad() + def prepare( + self, + prompt: str, + negative_prompt: str = "", + num_inference_steps: int = 50, + guidance_scale: float = 1.2, + delta: float = 1.0, + generator: Optional[torch.Generator] = torch.Generator(), + seed: int = 2, + ) -> None: + self.generator = generator + self.generator.manual_seed(seed) + # initialize x_t_latent (it can be any random tensor) + if self.denoising_steps_num > 1: + self.x_t_latent_buffer = torch.zeros( + ( + (self.denoising_steps_num - 1) * self.frame_bff_size, + 4, + self.latent_height, + self.latent_width, + ), + dtype=self.dtype, + device=self.device, + ) + else: + self.x_t_latent_buffer = None + + if self.cfg_type == "none": + self.guidance_scale = 1.0 + else: + self.guidance_scale = guidance_scale + self.delta = delta + + do_classifier_free_guidance = False + if self.guidance_scale > 1.0: + do_classifier_free_guidance = True + + encoder_output = self.pipe.encode_prompt( + prompt=prompt, + device=self.device, + num_images_per_prompt=1, + do_classifier_free_guidance=do_classifier_free_guidance, + negative_prompt=negative_prompt, + ) + self.prompt_embeds = encoder_output[0].repeat(self.batch_size, 1, 1) + + if self.use_denoising_batch and self.cfg_type == "full": + uncond_prompt_embeds = encoder_output[1].repeat(self.batch_size, 1, 1) + elif self.cfg_type == "initialize": + uncond_prompt_embeds = encoder_output[1].repeat(self.frame_bff_size, 1, 1) + + if self.guidance_scale > 1.0 and ( + self.cfg_type == "initialize" or self.cfg_type == "full" + ): + self.prompt_embeds = torch.cat( + [uncond_prompt_embeds, self.prompt_embeds], dim=0 + ) + + self.scheduler.set_timesteps(num_inference_steps, self.device) + self.timesteps = self.scheduler.timesteps.to(self.device) + + # make sub timesteps list based on the indices in the t_list list and the values in the timesteps list + self.sub_timesteps = [] + for t in self.t_list: + self.sub_timesteps.append(self.timesteps[t]) + + sub_timesteps_tensor = torch.tensor( + self.sub_timesteps, dtype=torch.long, device=self.device + ) + self.sub_timesteps_tensor = torch.repeat_interleave( + sub_timesteps_tensor, + repeats=self.frame_bff_size if self.use_denoising_batch else 1, + dim=0, + ) + + self.init_noise = torch.randn( + (self.batch_size, 4, self.latent_height, self.latent_width), + generator=generator, + ).to(device=self.device, dtype=self.dtype) + + self.stock_noise = torch.zeros_like(self.init_noise) + + c_skip_list = [] + c_out_list = [] + for timestep in self.sub_timesteps: + c_skip, c_out = self.scheduler.get_scalings_for_boundary_condition_discrete( + timestep + ) + c_skip_list.append(c_skip) + c_out_list.append(c_out) + + self.c_skip = ( + torch.stack(c_skip_list) + .view(len(self.t_list), 1, 1, 1) + .to(dtype=self.dtype, device=self.device) + ) + self.c_out = ( + torch.stack(c_out_list) + .view(len(self.t_list), 1, 1, 1) + .to(dtype=self.dtype, device=self.device) + ) + + alpha_prod_t_sqrt_list = [] + beta_prod_t_sqrt_list = [] + for timestep in self.sub_timesteps: + alpha_prod_t_sqrt = self.scheduler.alphas_cumprod[timestep].sqrt() + beta_prod_t_sqrt = (1 - self.scheduler.alphas_cumprod[timestep]).sqrt() + alpha_prod_t_sqrt_list.append(alpha_prod_t_sqrt) + beta_prod_t_sqrt_list.append(beta_prod_t_sqrt) + alpha_prod_t_sqrt = ( + torch.stack(alpha_prod_t_sqrt_list) + .view(len(self.t_list), 1, 1, 1) + .to(dtype=self.dtype, device=self.device) + ) + beta_prod_t_sqrt = ( + torch.stack(beta_prod_t_sqrt_list) + .view(len(self.t_list), 1, 1, 1) + .to(dtype=self.dtype, device=self.device) + ) + self.alpha_prod_t_sqrt = torch.repeat_interleave( + alpha_prod_t_sqrt, + repeats=self.frame_bff_size if self.use_denoising_batch else 1, + dim=0, + ) + self.beta_prod_t_sqrt = torch.repeat_interleave( + beta_prod_t_sqrt, + repeats=self.frame_bff_size if self.use_denoising_batch else 1, + dim=0, + ) + + @torch.no_grad() + def update_prompt(self, prompt: str,negative_prompt: Optional[str] = None) -> None: + encoder_output = self.pipe.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + device=self.device, + num_images_per_prompt=1, + do_classifier_free_guidance=False, + ) + self.prompt_embeds = encoder_output[0].repeat(self.batch_size, 1, 1) + + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + t_index: int, + ) -> torch.Tensor: + noisy_samples = ( + self.alpha_prod_t_sqrt[t_index] * original_samples + + self.beta_prod_t_sqrt[t_index] * noise + ) + return noisy_samples + + def scheduler_step_batch( + self, + model_pred_batch: torch.Tensor, + x_t_latent_batch: torch.Tensor, + idx: Optional[int] = None, + ) -> torch.Tensor: + # TODO: use t_list to select beta_prod_t_sqrt + if idx is None: + F_theta = ( + x_t_latent_batch - self.beta_prod_t_sqrt * model_pred_batch + ) / self.alpha_prod_t_sqrt + denoised_batch = self.c_out * F_theta + self.c_skip * x_t_latent_batch + else: + F_theta = ( + x_t_latent_batch - self.beta_prod_t_sqrt[idx] * model_pred_batch + ) / self.alpha_prod_t_sqrt[idx] + denoised_batch = ( + self.c_out[idx] * F_theta + self.c_skip[idx] * x_t_latent_batch + ) + + return denoised_batch + + def unet_step( + self, + x_t_latent: torch.Tensor, + t_list: Union[torch.Tensor, list[int]], + idx: Optional[int] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if self.guidance_scale > 1.0 and (self.cfg_type == "initialize"): + x_t_latent_plus_uc = torch.concat([x_t_latent[0:1], x_t_latent], dim=0) + t_list = torch.concat([t_list[0:1], t_list], dim=0) + elif self.guidance_scale > 1.0 and (self.cfg_type == "full"): + x_t_latent_plus_uc = torch.concat([x_t_latent, x_t_latent], dim=0) + t_list = torch.concat([t_list, t_list], dim=0) + else: + x_t_latent_plus_uc = x_t_latent + + model_pred = self.unet( + x_t_latent_plus_uc, + t_list, + encoder_hidden_states=self.prompt_embeds, + return_dict=False, + )[0] + + if self.guidance_scale > 1.0 and (self.cfg_type == "initialize"): + noise_pred_text = model_pred[1:] + self.stock_noise = torch.concat( + [model_pred[0:1], self.stock_noise[1:]], dim=0 + ) # ここコメントアウトでself out cfg + elif self.guidance_scale > 1.0 and (self.cfg_type == "full"): + noise_pred_uncond, noise_pred_text = model_pred.chunk(2) + else: + noise_pred_text = model_pred + if self.guidance_scale > 1.0 and ( + self.cfg_type == "self" or self.cfg_type == "initialize" + ): + noise_pred_uncond = self.stock_noise * self.delta + if self.guidance_scale > 1.0 and self.cfg_type != "none": + model_pred = noise_pred_uncond + self.guidance_scale * ( + noise_pred_text - noise_pred_uncond + ) + else: + model_pred = noise_pred_text + + # compute the previous noisy sample x_t -> x_t-1 + if self.use_denoising_batch: + denoised_batch = self.scheduler_step_batch(model_pred, x_t_latent, idx) + if self.cfg_type == "self" or self.cfg_type == "initialize": + scaled_noise = self.beta_prod_t_sqrt * self.stock_noise + delta_x = self.scheduler_step_batch(model_pred, scaled_noise, idx) + alpha_next = torch.concat( + [ + self.alpha_prod_t_sqrt[1:], + torch.ones_like(self.alpha_prod_t_sqrt[0:1]), + ], + dim=0, + ) + delta_x = alpha_next * delta_x + beta_next = torch.concat( + [ + self.beta_prod_t_sqrt[1:], + torch.ones_like(self.beta_prod_t_sqrt[0:1]), + ], + dim=0, + ) + delta_x = delta_x / beta_next + init_noise = torch.concat( + [self.init_noise[1:], self.init_noise[0:1]], dim=0 + ) + self.stock_noise = init_noise + delta_x + + else: + # denoised_batch = self.scheduler.step(model_pred, t_list[0], x_t_latent).denoised + denoised_batch = self.scheduler_step_batch(model_pred, x_t_latent, idx) + + return denoised_batch, model_pred + + def encode_image(self, image_tensors: torch.Tensor) -> torch.Tensor: + image_tensors = image_tensors.to( + device=self.device, + dtype=self.vae.dtype, + ) + img_latent = retrieve_latents(self.vae.encode(image_tensors), self.generator) + img_latent = img_latent * self.vae.config.scaling_factor + x_t_latent = self.add_noise(img_latent, self.init_noise[0], 0) + return x_t_latent + + def decode_image(self, x_0_pred_out: torch.Tensor) -> torch.Tensor: + output_latents = self.vae.decode( + x_0_pred_out / self.vae.config.scaling_factor, return_dict=False + ) + output_latent =output_latents[0] + return output_latent + + def predict_x0_batch(self, x_t_latent: torch.Tensor) -> torch.Tensor: + prev_latent_batch = self.x_t_latent_buffer + + if self.use_denoising_batch: + t_list = self.sub_timesteps_tensor + if self.denoising_steps_num > 1: + x_t_latent = torch.cat((x_t_latent, prev_latent_batch), dim=0) + self.stock_noise = torch.cat( + (self.init_noise[0:1], self.stock_noise[:-1]), dim=0 + ) + x_0_pred_batch, model_pred = self.unet_step(x_t_latent, t_list) + + if self.denoising_steps_num > 1: + x_0_pred_out = x_0_pred_batch[-1].unsqueeze(0) + if self.do_add_noise: + self.x_t_latent_buffer = ( + self.alpha_prod_t_sqrt[1:] * x_0_pred_batch[:-1] + + self.beta_prod_t_sqrt[1:] * self.init_noise[1:] + ) + else: + self.x_t_latent_buffer = ( + self.alpha_prod_t_sqrt[1:] * x_0_pred_batch[:-1] + ) + else: + x_0_pred_out = x_0_pred_batch + self.x_t_latent_buffer = None + else: + self.init_noise = x_t_latent + for idx, t in enumerate(self.sub_timesteps_tensor): + t = t.view( + 1, + ).repeat( + self.frame_bff_size, + ) + x_0_pred, model_pred = self.unet_step(x_t_latent, t, idx) + if idx < len(self.sub_timesteps_tensor) - 1: + if self.do_add_noise: + x_t_latent = self.alpha_prod_t_sqrt[ + idx + 1 + ] * x_0_pred + self.beta_prod_t_sqrt[ + idx + 1 + ] * torch.randn_like( + x_0_pred, device=self.device, dtype=self.dtype + ) + else: + x_t_latent = self.alpha_prod_t_sqrt[idx + 1] * x_0_pred + x_0_pred_out = x_0_pred + + return x_0_pred_out + + @torch.no_grad() + def __call__( + self, x: Union[torch.Tensor, PIL.Image.Image, np.ndarray] = None + ) -> torch.Tensor: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + if x is not None: + x = self.image_processor.preprocess(x, self.height, self.width).to( + device=self.device, dtype=self.dtype + ) + if self.similar_image_filter: + x = self.similar_filter(x) + if x is None: + time.sleep(self.inference_time_ema) + return self.prev_image_result + x_t_latent = self.encode_image(x) + else: + # TODO: check the dimension of x_t_latent + x_t_latent = torch.randn((1, 4, self.latent_height, self.latent_width)).to( + device=self.device, dtype=self.dtype + ) + x_0_pred_out = self.predict_x0_batch(x_t_latent) + x_output = self.decode_image(x_0_pred_out).detach().clone() + + self.prev_image_result = x_output + end.record() + torch.cuda.synchronize() + inference_time = start.elapsed_time(end) / 1000 + self.inference_time_ema = 0.9 * self.inference_time_ema + 0.1 * inference_time + return x_output + + @torch.no_grad() + def sample( + self, x: Union[torch.Tensor, PIL.Image.Image, np.ndarray] = None + ) -> torch.Tensor: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + if x is not None: + x = self.image_processor.preprocess(x, self.height, self.width).to( + device=self.device, dtype=self.dtype + ) + if self.similar_image_filter: + x = self.similar_filter(x) + if x is None: + time.sleep(self.inference_time_ema) + return self.prev_image_result + x_t_latent = self.encode_image(x) + b,c,h,w=x_t_latent.shape + + # x_t_latent=x_t_latent.repeat((2, 1,1,1)) + else: + # TODO: check the dimension of x_t_latent + x_t_latent = torch.randn((self.frame_bff_size, 4, self.latent_height, self.latent_width)).to( + device=self.device, dtype=self.dtype + ) + x_0_pred_out = self.predict_x0_batch(x_t_latent) + x_output = self.decode_image(x_0_pred_out).detach().clone() + self.prev_image_result = x_output + end.record() + torch.cuda.synchronize() + inference_time = start.elapsed_time(end) / 1000 + self.inference_time_ema = 0.9 * self.inference_time_ema + 0.1 * inference_time + return x_output + + + @torch.no_grad() + def txt2img(self, batch_size: int = 1) -> torch.Tensor: + x_0_pred_out = self.predict_x0_batch( + torch.randn((batch_size, 4, self.latent_height, self.latent_width)).to( + device=self.device, dtype=self.dtype + ) + ) + x_output = self.decode_image(x_0_pred_out).detach().clone() + return x_output + + def txt2img_sd_turbo(self, batch_size: int = 1) -> torch.Tensor: + x_t_latent = torch.randn( + (batch_size, 4, self.latent_height, self.latent_width), + device=self.device, + dtype=self.dtype, + ) + model_pred = self.unet( + x_t_latent, + self.sub_timesteps_tensor, + encoder_hidden_states=self.prompt_embeds, + return_dict=False, + )[0] + x_0_pred_out = ( + x_t_latent - self.beta_prod_t_sqrt * model_pred + ) / self.alpha_prod_t_sqrt + return self.decode_image(x_0_pred_out) diff --git a/streamdiffusion/tools/__init__.py b/streamdiffusion/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/streamdiffusion/tools/install-tensorrt.py b/streamdiffusion/tools/install-tensorrt.py new file mode 100644 index 0000000..6389c07 --- /dev/null +++ b/streamdiffusion/tools/install-tensorrt.py @@ -0,0 +1,49 @@ +from typing import Literal, Optional + +import fire +from packaging.version import Version + +from ..pip_utils import is_installed, run_pip, version + + +def get_cuda_version_from_torch() -> Optional[Literal["11", "12"]]: + try: + import torch + except ImportError: + return None + + return torch.version.cuda.split(".")[0] + + +def install(cu: Optional[Literal["11", "12"]] = get_cuda_version_from_torch()): + if cu is None or cu not in ["11", "12"]: + print("Could not detect CUDA version. Please specify manually.") + return + print("Installing TensorRT requirements...") + + if is_installed("tensorrt"): + if version("tensorrt") < Version("9.0.0"): + run_pip("uninstall -y tensorrt") + + cudnn_name = f"nvidia-cudnn-cu{cu}==8.9.4.25" + + if not is_installed("tensorrt"): + run_pip(f"install {cudnn_name} --no-cache-dir") + run_pip( + "install --pre --extra-index-url https://pypi.nvidia.com tensorrt==9.0.1.post11.dev4 --no-cache-dir" + ) + + if not is_installed("polygraphy"): + run_pip( + "install polygraphy==0.47.1 --extra-index-url https://pypi.ngc.nvidia.com" + ) + if not is_installed("onnx_graphsurgeon"): + run_pip( + "install onnx-graphsurgeon==0.3.26 --extra-index-url https://pypi.ngc.nvidia.com" + ) + + pass + + +if __name__ == "__main__": + fire.Fire(install) diff --git a/streamdiffusion/wrapper.py b/streamdiffusion/wrapper.py new file mode 100644 index 0000000..f3d714d --- /dev/null +++ b/streamdiffusion/wrapper.py @@ -0,0 +1,773 @@ +import gc +import os +from pathlib import Path +import traceback +from typing import List, Literal, Optional, Union, Dict + +import numpy as np +import torch +from diffusers import AutoencoderTiny, StableDiffusionPipeline +from PIL import Image + +from .pipeline import StreamDiffusion +from .image_utils import postprocess_image + + +torch.set_grad_enabled(False) +torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.allow_tf32 = True + + +class StreamDiffusionWrapper: + def __init__( + self, + model_id_or_path: str, + t_index_list: List[int], + lora_dict: Optional[Dict[str, float]] = None, + mode: Literal["img2img", "txt2img"] = "img2img", + output_type: Literal["pil", "pt", "np", "latent"] = "pil", + lcm_lora_id: Optional[str] = None, + vae_id: Optional[str] = None, + device: Literal["cpu", "cuda"] = "cuda", + dtype: torch.dtype = torch.float16, + frame_buffer_size: int = 1, + width: int = 512, + height: int = 512, + warmup: int = 10, + acceleration: Literal["none", "xformers", "tensorrt"] = "tensorrt", + do_add_noise: bool = True, + device_ids: Optional[List[int]] = None, + use_lcm_lora: bool = True, + use_tiny_vae: bool = True, + enable_similar_image_filter: bool = False, + similar_image_filter_threshold: float = 0.98, + similar_image_filter_max_skip_frame: int = 10, + use_denoising_batch: bool = True, + cfg_type: Literal["none", "full", "self", "initialize"] = "self", + seed: int = 2, + use_safety_checker: bool = False, + ): + """ + Initializes the StreamDiffusionWrapper. + + Parameters + ---------- + model_id_or_path : str + The model id or path to load. + t_index_list : List[int] + The t_index_list to use for inference. + lora_dict : Optional[Dict[str, float]], optional + The lora_dict to load, by default None. + Keys are the LoRA names and values are the LoRA scales. + Example: {'LoRA_1' : 0.5 , 'LoRA_2' : 0.7 ,...} + mode : Literal["img2img", "txt2img"], optional + txt2img or img2img, by default "img2img". + output_type : Literal["pil", "pt", "np", "latent"], optional + The output type of image, by default "pil". + lcm_lora_id : Optional[str], optional + The lcm_lora_id to load, by default None. + If None, the default LCM-LoRA + ("latent-consistency/lcm-lora-sdv1-5") will be used. + vae_id : Optional[str], optional + The vae_id to load, by default None. + If None, the default TinyVAE + ("madebyollin/taesd") will be used. + device : Literal["cpu", "cuda"], optional + The device to use for inference, by default "cuda". + dtype : torch.dtype, optional + The dtype for inference, by default torch.float16. + frame_buffer_size : int, optional + The frame buffer size for denoising batch, by default 1. + width : int, optional + The width of the image, by default 512. + height : int, optional + The height of the image, by default 512. + warmup : int, optional + The number of warmup steps to perform, by default 10. + acceleration : Literal["none", "xformers", "tensorrt"], optional + The acceleration method, by default "tensorrt". + do_add_noise : bool, optional + Whether to add noise for following denoising steps or not, + by default True. + device_ids : Optional[List[int]], optional + The device ids to use for DataParallel, by default None. + use_lcm_lora : bool, optional + Whether to use LCM-LoRA or not, by default True. + use_tiny_vae : bool, optional + Whether to use TinyVAE or not, by default True. + enable_similar_image_filter : bool, optional + Whether to enable similar image filter or not, + by default False. + similar_image_filter_threshold : float, optional + The threshold for similar image filter, by default 0.98. + similar_image_filter_max_skip_frame : int, optional + The max skip frame for similar image filter, by default 10. + use_denoising_batch : bool, optional + Whether to use denoising batch or not, by default True. + cfg_type : Literal["none", "full", "self", "initialize"], + optional + The cfg_type for img2img mode, by default "self". + You cannot use anything other than "none" for txt2img mode. + seed : int, optional + The seed, by default 2. + use_safety_checker : bool, optional + Whether to use safety checker or not, by default False. + """ + self.sd_turbo = "turbo" in model_id_or_path + + if mode == "txt2img": + if cfg_type != "none": + raise ValueError( + f"txt2img mode accepts only cfg_type = 'none', but got {cfg_type}" + ) + if use_denoising_batch and frame_buffer_size > 1: + if not self.sd_turbo: + raise ValueError( + "txt2img mode cannot use denoising batch with frame_buffer_size > 1." + ) + + if mode == "img2img": + if not use_denoising_batch: + raise NotImplementedError( + "img2img mode must use denoising batch for now." + ) + + self.device = device + self.dtype = dtype + self.width = width + self.height = height + self.mode = mode + self.output_type = output_type + self.frame_buffer_size = frame_buffer_size + self.batch_size = ( + len(t_index_list) * frame_buffer_size + if use_denoising_batch + else frame_buffer_size + ) + self.t_index_list =t_index_list + self.cfg_type =cfg_type + self.use_denoising_batch = use_denoising_batch + self.use_safety_checker = use_safety_checker + self.do_add_noise =do_add_noise + self.seed=seed + + self.stream: StreamDiffusion = self._load_model( + model_id_or_path=model_id_or_path, + lora_dict=lora_dict, + lcm_lora_id=lcm_lora_id, + vae_id=vae_id, + t_index_list=t_index_list, + acceleration=acceleration, + warmup=warmup, + do_add_noise=do_add_noise, + use_lcm_lora=use_lcm_lora, + use_tiny_vae=use_tiny_vae, + cfg_type=cfg_type, + seed=seed, + ) + + if device_ids is not None: + self.stream.unet = torch.nn.DataParallel( + self.stream.unet, device_ids=device_ids + ) + + if enable_similar_image_filter: + self.stream.enable_similar_image_filter(similar_image_filter_threshold, similar_image_filter_max_skip_frame) + + def prepare( + self, + prompt: str, + negative_prompt: str = "", + num_inference_steps: int = 50, + guidance_scale: float = 1.2, + delta: float = 1.0, + t_index_list: List[int]=[16,32,45], + do_add_noise: bool = True, + enable_similar_image_filter: bool = False, + similar_image_filter_threshold: float = 0.98, + similar_image_filter_max_skip_frame: int = 10, + use_denoising_batch: bool = True, + cfg_type: Literal["none", "full", "self", "initialize"] = "self", + seed: int = 2, + frame_buffer_size:int=1, + use_safety_checker: bool = False, + ) -> None: + """ + Prepares the model for inference. + + Parameters + ---------- + prompt : str + The prompt to generate images from. + num_inference_steps : int, optional + The number of inference steps to perform, by default 50. + guidance_scale : float, optional + The guidance scale to use, by default 1.2. + delta : float, optional + The delta multiplier of virtual residual noise, + by default 1.0. + """ + # self.stream.prepare( + # prompt, + # negative_prompt, + # num_inference_steps=num_inference_steps, + # guidance_scale=guidance_scale, + # delta=delta, + # ) + self.prompt =prompt + self.negative_prompt=negative_prompt + self.num_inference_steps=num_inference_steps + self.guidance_scale=guidance_scale + self.delta=delta + + self.frame_buffer_size = frame_buffer_size + self.batch_size = ( + len(t_index_list) * frame_buffer_size + if use_denoising_batch + else frame_buffer_size + ) + self.t_index_list =t_index_list + self.cfg_type =cfg_type + self.use_denoising_batch = use_denoising_batch + self.use_safety_checker = use_safety_checker + self.do_add_noise =do_add_noise + self.seed=seed + + if enable_similar_image_filter: + self.stream.enable_similar_image_filter(similar_image_filter_threshold, similar_image_filter_max_skip_frame) + else: + self.stream.disable_similar_image_filter() + + if self.use_safety_checker: + if self.safety_checker==None or self.feature_extractor==None: + from transformers import CLIPFeatureExtractor + from diffusers.pipelines.stable_diffusion.safety_checker import ( + StableDiffusionSafetyChecker, + ) + + self.safety_checker = StableDiffusionSafetyChecker.from_pretrained( + "CompVis/stable-diffusion-safety-checker" + ).to(self.stream.device) + self.feature_extractor = CLIPFeatureExtractor.from_pretrained( + "openai/clip-vit-base-patch32" + ) + self.nsfw_fallback_img = Image.new("RGB", (512, 512), (0, 0, 0)) + + def __call__( + self, + image: Optional[Union[str, Image.Image, torch.Tensor]] = None, + prompt: Optional[str] = None, + ) -> Union[Image.Image, List[Image.Image]]: + """ + Performs img2img or txt2img based on the mode. + + Parameters + ---------- + image : Optional[Union[str, Image.Image, torch.Tensor]] + The image to generate from. + prompt : Optional[str] + The prompt to generate images from. + + Returns + ------- + Union[Image.Image, List[Image.Image]] + The generated image. + """ + if self.mode == "img2img": + return self.img2img(image) + else: + return self.txt2img(prompt) + + def sample(self, image: Optional[Union[str, Image.Image, torch.Tensor]] = None, + prompt: Optional[str] = None,negative_prompt: Optional[str] = None)-> List[Image.Image]: + + use_denoising_batch=self.use_denoising_batch + if not image == None: + #图生图 + if isinstance(image, str) or isinstance(image, Image.Image): + image = self.preprocess_image(image) + + use_denoising_batch = True + self.stream.set_sampler_param(t_index_list=self.t_index_list, + width=self.width, + height=self.height, + do_add_noise=self.do_add_noise, + use_denoising_batch=use_denoising_batch, + frame_buffer_size=self.frame_buffer_size, + cfg_type=self.cfg_type) + else: + #文生图 + if self.frame_buffer_size >1 and self.use_denoising_batch: + use_denoising_batch = False + self.stream.set_sampler_param(t_index_list=self.t_index_list, + width=self.width, + height=self.height, + do_add_noise=self.do_add_noise, + use_denoising_batch=use_denoising_batch, + frame_buffer_size=self.frame_buffer_size, + cfg_type='none') + + self.stream.prepare( + prompt=self.prompt, + negative_prompt=self.negative_prompt, + num_inference_steps=self.num_inference_steps, + guidance_scale=self.guidance_scale, + delta=self.delta, + seed=self.seed, + ) + + if prompt is not None: + self.stream.update_prompt(prompt,negative_prompt) + + self.batch_size = ( + len(self.t_index_list) * self.frame_buffer_size + if use_denoising_batch + else self.frame_buffer_size + ) + + if self.frame_buffer_size==1: + for _ in range(self.batch_size): + self.stream.sample(image) + + image_tensor = self.stream.sample(image) + image = postprocess_image(image_tensor.cpu(), output_type=self.output_type) + + if self.use_safety_checker: + safety_checker_input = self.feature_extractor( + image, return_tensors="pt" + ).to(self.device) + _, has_nsfw_concept = self.safety_checker( + images=image_tensor.to(self.dtype), + clip_input=safety_checker_input.pixel_values.to(self.dtype), + ) + image = self.nsfw_fallback_img if has_nsfw_concept[0] else image + return image + + def txt2img( + self, prompt: Optional[str] = None + ) -> Union[Image.Image, List[Image.Image], torch.Tensor, np.ndarray]: + """ + Performs txt2img. + + Parameters + ---------- + prompt : Optional[str] + The prompt to generate images from. + + Returns + ------- + Union[Image.Image, List[Image.Image]] + The generated image. + """ + if prompt is not None: + self.stream.update_prompt(prompt) + + if self.sd_turbo: + image_tensor = self.stream.txt2img_sd_turbo(self.batch_size) + else: + image_tensor = self.stream.txt2img(self.frame_buffer_size) + image = self.postprocess_image(image_tensor, output_type=self.output_type) + + if self.use_safety_checker: + safety_checker_input = self.feature_extractor( + image, return_tensors="pt" + ).to(self.device) + _, has_nsfw_concept = self.safety_checker( + images=image_tensor.to(self.dtype), + clip_input=safety_checker_input.pixel_values.to(self.dtype), + ) + image = self.nsfw_fallback_img if has_nsfw_concept[0] else image + + return image + + def img2img( + self, image: Union[str, Image.Image, torch.Tensor] + ) -> Union[Image.Image, List[Image.Image], torch.Tensor, np.ndarray]: + """ + Performs img2img. + + Parameters + ---------- + image : Union[str, Image.Image, torch.Tensor] + The image to generate from. + + Returns + ------- + Image.Image + The generated image. + """ + if isinstance(image, str) or isinstance(image, Image.Image): + image = self.preprocess_image(image) + + image_tensor = self.stream(image) + image = self.postprocess_image(image_tensor, output_type=self.output_type) + + if self.use_safety_checker: + safety_checker_input = self.feature_extractor( + image, return_tensors="pt" + ).to(self.device) + _, has_nsfw_concept = self.safety_checker( + images=image_tensor.to(self.dtype), + clip_input=safety_checker_input.pixel_values.to(self.dtype), + ) + image = self.nsfw_fallback_img if has_nsfw_concept[0] else image + + return image + + def preprocess_image(self, image: Union[str, Image.Image]) -> torch.Tensor: + """ + Preprocesses the image. + + Parameters + ---------- + image : Union[str, Image.Image, torch.Tensor] + The image to preprocess. + + Returns + ------- + torch.Tensor + The preprocessed image. + """ + if isinstance(image, str): + image = Image.open(image).convert("RGB").resize((self.width, self.height)) + if isinstance(image, Image.Image): + image = image.convert("RGB").resize((self.width, self.height)) + + return self.stream.image_processor.preprocess( + image, self.height, self.width + ).to(device=self.device, dtype=self.dtype) + + def postprocess_image( + self, image_tensor: torch.Tensor, output_type: str = "pil" + ) -> Union[Image.Image, List[Image.Image], torch.Tensor, np.ndarray]: + """ + Postprocesses the image. + + Parameters + ---------- + image_tensor : torch.Tensor + The image tensor to postprocess. + + Returns + ------- + Union[Image.Image, List[Image.Image]] + The postprocessed image. + """ + if self.frame_buffer_size > 1: + return postprocess_image(image_tensor.cpu(), output_type=output_type) + else: + return postprocess_image(image_tensor.cpu(), output_type=output_type)[0] + + def _load_model( + self, + model_id_or_path: str, + t_index_list: List[int], + lora_dict: Optional[Dict[str, float]] = None, + lcm_lora_id: Optional[str] = None, + vae_id: Optional[str] = None, + acceleration: Literal["none", "xformers", "tensorrt"] = "tensorrt", + warmup: int = 10, + do_add_noise: bool = True, + use_lcm_lora: bool = True, + use_tiny_vae: bool = True, + cfg_type: Literal["none", "full", "self", "initialize"] = "self", + seed: int = 2, + ) -> StreamDiffusion: + """ + Loads the model. + + This method does the following: + + 1. Loads the model from the model_id_or_path. + 2. Loads and fuses the LCM-LoRA model from the lcm_lora_id if needed. + 3. Loads the VAE model from the vae_id if needed. + 4. Enables acceleration if needed. + 5. Prepares the model for inference. + 6. Load the safety checker if needed. + + Parameters + ---------- + model_id_or_path : str + The model id or path to load. + t_index_list : List[int] + The t_index_list to use for inference. + lora_dict : Optional[Dict[str, float]], optional + The lora_dict to load, by default None. + Keys are the LoRA names and values are the LoRA scales. + Example: {'LoRA_1' : 0.5 , 'LoRA_2' : 0.7 ,...} + lcm_lora_id : Optional[str], optional + The lcm_lora_id to load, by default None. + vae_id : Optional[str], optional + The vae_id to load, by default None. + acceleration : Literal["none", "xfomers", "sfast", "tensorrt"], optional + The acceleration method, by default "tensorrt". + warmup : int, optional + The number of warmup steps to perform, by default 10. + do_add_noise : bool, optional + Whether to add noise for following denoising steps or not, + by default True. + use_lcm_lora : bool, optional + Whether to use LCM-LoRA or not, by default True. + use_tiny_vae : bool, optional + Whether to use TinyVAE or not, by default True. + cfg_type : Literal["none", "full", "self", "initialize"], + optional + The cfg_type for img2img mode, by default "self". + You cannot use anything other than "none" for txt2img mode. + seed : int, optional + The seed, by default 2. + + Returns + ------- + StreamDiffusion + The loaded model. + """ + + try: # Load from local directory + pipe: StableDiffusionPipeline = StableDiffusionPipeline.from_pretrained( + model_id_or_path, + ).to(device=self.device, dtype=self.dtype) + + except ValueError: # Load from huggingface + pipe: StableDiffusionPipeline = StableDiffusionPipeline.from_single_file( + model_id_or_path, + ).to(device=self.device, dtype=self.dtype) + except Exception: # No model found + traceback.print_exc() + print("Model load has failed. Doesn't exist.") + exit() + + stream = StreamDiffusion( + pipe=pipe, + t_index_list=t_index_list, + torch_dtype=self.dtype, + width=self.width, + height=self.height, + do_add_noise=do_add_noise, + frame_buffer_size=self.frame_buffer_size, + use_denoising_batch=self.use_denoising_batch, + cfg_type=cfg_type, + ) + if not self.sd_turbo: + if use_lcm_lora: + if lcm_lora_id is not None: + stream.load_lcm_lora( + pretrained_model_name_or_path_or_dict=lcm_lora_id + ) + else: + stream.load_lcm_lora() + stream.fuse_lora() + + if lora_dict is not None: + for lora_name, lora_scale in lora_dict.items(): + stream.load_lora(lora_name) + stream.fuse_lora(lora_scale=lora_scale) + print(f"Use LoRA: {lora_name} in weights {lora_scale}") + + if use_tiny_vae: + if vae_id is not None: + stream.vae = AutoencoderTiny.from_pretrained(vae_id).to( + device=pipe.device, dtype=pipe.dtype + ) + else: + stream.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd").to( + device=pipe.device, dtype=pipe.dtype + ) + + try: + if acceleration == "xformers": + stream.pipe.enable_xformers_memory_efficient_attention() + if acceleration == "tensorrt": + from polygraphy import cuda + from streamdiffusion.acceleration.tensorrt import ( + TorchVAEEncoder, + compile_unet, + compile_vae_decoder, + compile_vae_encoder, + ) + from streamdiffusion.acceleration.tensorrt.engine import ( + AutoencoderKLEngine, + UNet2DConditionModelEngine, + ) + from streamdiffusion.acceleration.tensorrt.models import ( + VAE, + UNet, + VAEEncoder, + ) + + def create_prefix( + model_id_or_path: str, + max_batch_size: int, + min_batch_size: int, + ): + maybe_path = Path(model_id_or_path) + if maybe_path.exists(): + return f"{maybe_path.stem}--lcm_lora-{use_lcm_lora}--tiny_vae-{use_tiny_vae}--max_batch-{max_batch_size}--min_batch-{min_batch_size}--mode-{self.mode}" + else: + return f"{model_id_or_path}--lcm_lora-{use_lcm_lora}--tiny_vae-{use_tiny_vae}--max_batch-{max_batch_size}--min_batch-{min_batch_size}--mode-{self.mode}" + + engine_dir = os.path.join("engines") + unet_path = os.path.join( + engine_dir, + create_prefix( + model_id_or_path=model_id_or_path, + max_batch_size=stream.trt_unet_batch_size, + min_batch_size=stream.trt_unet_batch_size, + ), + "unet.engine", + ) + vae_encoder_path = os.path.join( + engine_dir, + create_prefix( + model_id_or_path=model_id_or_path, + max_batch_size=self.batch_size + if self.mode == "txt2img" + else stream.frame_bff_size, + min_batch_size=self.batch_size + if self.mode == "txt2img" + else stream.frame_bff_size, + ), + "vae_encoder.engine", + ) + vae_decoder_path = os.path.join( + engine_dir, + create_prefix( + model_id_or_path=model_id_or_path, + max_batch_size=self.batch_size + if self.mode == "txt2img" + else stream.frame_bff_size, + min_batch_size=self.batch_size + if self.mode == "txt2img" + else stream.frame_bff_size, + ), + "vae_decoder.engine", + ) + + if not os.path.exists(unet_path): + os.makedirs(os.path.dirname(unet_path), exist_ok=True) + unet_model = UNet( + fp16=True, + device=stream.device, + max_batch_size=stream.trt_unet_batch_size, + min_batch_size=stream.trt_unet_batch_size, + embedding_dim=stream.text_encoder.config.hidden_size, + unet_dim=stream.unet.config.in_channels, + ) + compile_unet( + stream.unet, + unet_model, + unet_path + ".onnx", + unet_path + ".opt.onnx", + unet_path, + opt_batch_size=stream.trt_unet_batch_size, + ) + + if not os.path.exists(vae_decoder_path): + os.makedirs(os.path.dirname(vae_decoder_path), exist_ok=True) + stream.vae.forward = stream.vae.decode + vae_decoder_model = VAE( + device=stream.device, + max_batch_size=self.batch_size + if self.mode == "txt2img" + else stream.frame_bff_size, + min_batch_size=self.batch_size + if self.mode == "txt2img" + else stream.frame_bff_size, + ) + compile_vae_decoder( + stream.vae, + vae_decoder_model, + vae_decoder_path + ".onnx", + vae_decoder_path + ".opt.onnx", + vae_decoder_path, + opt_batch_size=self.batch_size + if self.mode == "txt2img" + else stream.frame_bff_size, + ) + delattr(stream.vae, "forward") + + if not os.path.exists(vae_encoder_path): + os.makedirs(os.path.dirname(vae_encoder_path), exist_ok=True) + vae_encoder = TorchVAEEncoder(stream.vae).to(torch.device("cuda")) + vae_encoder_model = VAEEncoder( + device=stream.device, + max_batch_size=self.batch_size + if self.mode == "txt2img" + else stream.frame_bff_size, + min_batch_size=self.batch_size + if self.mode == "txt2img" + else stream.frame_bff_size, + ) + compile_vae_encoder( + vae_encoder, + vae_encoder_model, + vae_encoder_path + ".onnx", + vae_encoder_path + ".opt.onnx", + vae_encoder_path, + opt_batch_size=self.batch_size + if self.mode == "txt2img" + else stream.frame_bff_size, + ) + + cuda_steram = cuda.Stream() + + vae_config = stream.vae.config + vae_dtype = stream.vae.dtype + + stream.unet = UNet2DConditionModelEngine( + unet_path, cuda_steram, use_cuda_graph=False + ) + stream.vae = AutoencoderKLEngine( + vae_encoder_path, + vae_decoder_path, + cuda_steram, + stream.pipe.vae_scale_factor, + use_cuda_graph=False, + ) + setattr(stream.vae, "config", vae_config) + setattr(stream.vae, "dtype", vae_dtype) + + gc.collect() + torch.cuda.empty_cache() + + print("TensorRT acceleration enabled.") + if acceleration == "sfast": + from streamdiffusion.acceleration.sfast import ( + accelerate_with_stable_fast, + ) + + stream = accelerate_with_stable_fast(stream) + print("StableFast acceleration enabled.") + except Exception: + traceback.print_exc() + print("Acceleration has failed. Falling back to normal mode.") + + if seed < 0: # Random seed + seed = np.random.randint(0, 1000000) + + stream.prepare( + "", + "", + num_inference_steps=50, + guidance_scale=1.1 + if stream.cfg_type in ["full", "self", "initialize"] + else 1.0, + generator=torch.manual_seed(seed), + seed=seed, + ) + + if self.use_safety_checker: + from transformers import CLIPFeatureExtractor + from diffusers.pipelines.stable_diffusion.safety_checker import ( + StableDiffusionSafetyChecker, + ) + + self.safety_checker = StableDiffusionSafetyChecker.from_pretrained( + "CompVis/stable-diffusion-safety-checker" + ).to(pipe.device) + self.feature_extractor = CLIPFeatureExtractor.from_pretrained( + "openai/clip-vit-base-patch32" + ) + self.nsfw_fallback_img = Image.new("RGB", (512, 512), (0, 0, 0)) + + return stream diff --git a/workflow.png b/workflow.png new file mode 100644 index 0000000..dcf9454 Binary files /dev/null and b/workflow.png differ