nodes.py

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",
    
}