init commit
upload
This commit is contained in:
jesenzhang
@@ -0,0 +1,4 @@
|
||||
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -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",
|
||||
|
||||
}
|
||||
@@ -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)
|
||||
@@ -0,0 +1 @@
|
||||
from .pipeline import StreamDiffusion
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -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
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 582 KiB |
Reference in New Issue
Block a user