From 149194598b3af8eadc54953a14117421e9da146c Mon Sep 17 00:00:00 2001 From: Limitex <76650151+Limitex@users.noreply.github.com> Date: Wed, 27 Dec 2023 08:14:45 +0900 Subject: [PATCH] Added Stream Diffusion --- nodes.py | 111 ++++++++++++++++++++++++++++++++++++++++++++--- requirements.txt | 5 ++- 2 files changed, 109 insertions(+), 7 deletions(-) diff --git a/nodes.py b/nodes.py index 1b80ec1..fdc7ee9 100644 --- a/nodes.py +++ b/nodes.py @@ -9,7 +9,8 @@ from diffusers import StableDiffusionPipeline, AutoencoderKL from comfy.cli_args import args from PIL import Image, ImageOps, ImageSequence from PIL.PngImagePlugin import PngInfo - +from streamdiffusion import StreamDiffusion +from streamdiffusion.image_utils import postprocess_image class DiffusersPipelineLoader: def __init__(self): @@ -136,8 +137,8 @@ class DiffusersClipTextEncode: "negative": ("STRING", {"multiline": True}), }} - RETURN_TYPES = ("EMBEDS", "EMBEDS", ) - RETURN_NAMES = ("positive_embeds", "negative_embeds", ) + RETURN_TYPES = ("EMBEDS", "EMBEDS", "STRING", "STRING", ) + RETURN_NAMES = ("positive_embeds", "negative_embeds", "positive", "negative", ) FUNCTION = "concat_embeds" @@ -146,7 +147,7 @@ class DiffusersClipTextEncode: def concat_embeds(self, maked_pipeline, positive, negative): positive_embeds, negative_embeds = token_auto_concat_embeds(maked_pipeline, positive,negative) - return (positive_embeds, negative_embeds, ) + return (positive_embeds, negative_embeds, positive, negative, ) class DiffusersSampler: def __init__(self): @@ -231,6 +232,100 @@ class DiffusersSaveImage: return { "ui": { "images": results } } +class StreamDiffusionCreateStream: + def __init__(self): + self.dtype = torch.float32 + self.torch_device = get_torch_device() + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "maked_pipeline": ("MAKED_PIPELINE", ), + "autoencoder": ("AUTOENCODER", ), + "t_index_list_type": (["txt2image", "image2image"], {"default": "txt2image"}), + "width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), + "height": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), + "do_add_noise": ("BOOLEAN", {"default": True}), + "use_denoising_batch": ("BOOLEAN", {"default": True}), + "frame_buffer_size": ("INT", {"default": 1, "min": 1, "max": 10000}), + "cfg_type": (["none", "full", "self", "initialize"], {"default": "self"}), + }, + } + + RETURN_TYPES = ("STREAM",) + FUNCTION = "load_stream" + + CATEGORY = "Diffusers/StreamDiffusion" + + def load_stream(self, maked_pipeline, autoencoder, t_index_list_type, width, height, do_add_noise, use_denoising_batch, frame_buffer_size, cfg_type): + if t_index_list_type == "txt2image": + t_index_list = [0, 16, 32, 45] + elif t_index_list_type == "image2image": + t_index_list = [32, 45] + + stream = StreamDiffusion( + pipe = maked_pipeline, + t_index_list = t_index_list, + torch_dtype = self.dtype, + cfg_type = cfg_type, + width = width, + height = height, + do_add_noise = do_add_noise, + use_denoising_batch = use_denoising_batch, + frame_buffer_size = frame_buffer_size, + ) + stream.load_lcm_lora() + stream.fuse_lora() + stream.vae = autoencoder.to(self.torch_device) + return ((stream, t_index_list), ) + +class StreamDiffusionSampler: + def __init__(self): + self.torch_device = get_torch_device() + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "stream": ("STREAM", ), + "positive": ("STRING", {"multiline": True}), + "negative": ("STRING", {"multiline": True}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "cfg": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 100.0}), + "delta": ("FLOAT", {"default": 1, "min": 0.0, "max": 1.0}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "num": ("INT", {"default": 1, "min": 1, "max": 10000}), + }, + } + + RETURN_TYPES = ("PIL_IMAGE",) + + FUNCTION = "sample" + + CATEGORY = "Diffusers/StreamDiffusion" + + def sample(self, stream: StreamDiffusion, positive, negative, steps, cfg, delta, seed, num): + t_index_list = stream[1] + stream = stream[0] + stream.prepare( + prompt = positive, + negative_prompt = negative, + num_inference_steps = steps, + guidance_scale = cfg, + delta = delta, + seed = seed + ) + + for _ in t_index_list: + stream() + + result = [] + for _ in range(num): + x_output = stream.txt2img() + result.append(postprocess_image(x_output, output_type="pil")[0]) + return (result,) + NODE_CLASS_MAPPINGS = { "DiffusersPipelineLoader": DiffusersPipelineLoader, "DiffusersVaeLoader": DiffusersVaeLoader, @@ -238,7 +333,9 @@ NODE_CLASS_MAPPINGS = { "DiffusersModelMakeup": DiffusersModelMakeup, "DiffusersClipTextEncode": DiffusersClipTextEncode, "DiffusersSampler": DiffusersSampler, - "DiffusersSaveImage": DiffusersSaveImage + "DiffusersSaveImage": DiffusersSaveImage, + "StreamDiffusionCreateStream": StreamDiffusionCreateStream, + "StreamDiffusionSampler": StreamDiffusionSampler, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -248,5 +345,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "DiffusersModelMakeup": "Diffusers Model Makeup", "DiffusersClipTextEncode": "Diffusers Clip Text Encode", "DiffusersSampler": "Diffusers Sampler", - "DiffusersSaveImage": "Diffusers Save Image" + "DiffusersSaveImage": "Diffusers Save Image", + "StreamDiffusionCreateStream": "StreamDiffusion Create Stream", + "StreamDiffusionSampler": "StreamDiffusion Sampler", } diff --git a/requirements.txt b/requirements.txt index d1b11fa..6c4f8b3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,7 @@ accelerate transformers safetensors omegaconf -pytorch_lightning \ No newline at end of file +pytorch_lightning +xformers +gitpython +git+https://github.com/cumulo-autumn/StreamDiffusion.git@main#egg=streamdiffusion[tensorrt] \ No newline at end of file