555 lines
20 KiB
Python
555 lines
20 KiB
Python
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",
|
|
|
|
} |