246 lines
8.4 KiB
Python
246 lines
8.4 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_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,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"StreamDiffusion_Loader": StreamDiffusion_Loader,
|
|
"StreamDiffusion_Sampler":StreamDiffusion_Sampler,
|
|
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"StreamDiffusion_Loader": "StreamDiffusion_Loader",
|
|
"StreamDiffusion_Sampler":"StreamDiffusion_Sampler",
|
|
|
|
} |