Upload folder using huggingface_hub (part 6)

.gitattributes

@@ -46,3 +46,28 @@ ComfyUI/comfy_api_nodes/__pycache__/nodes_bytedance.cpython-312.pyc filter=lfs d
 ComfyUI/comfy_api_nodes/__pycache__/nodes_kling.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text
 ComfyUI/comfy_api_nodes/apis/__pycache__/__init__.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text
 ComfyUI/comfy_extras/__pycache__/nodes_gaussian_splat.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/brushnet/__pycache__/unet_2d_blocks.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/kolors/chatglm/tokenizer/vocab.txt filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-Easy-Use/resources/OpenSans-Medium.ttf filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-Easy-Use/resources/wenquan.ttf filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-KJNodes/docs/images/2024-04-03_20_49_29-ComfyUI.png filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-KJNodes/fonts/FreeMono.ttf filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-KJNodes/fonts/FreeMonoBoldOblique.otf filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-KJNodes/fonts/TTNorms-Black.otf filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-KJNodes/nodes/__pycache__/image_nodes.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-KJNodes/nodes/__pycache__/ltxv_nodes.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-KJNodes/nodes/__pycache__/model_optimization_nodes.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-KJNodes/nodes/__pycache__/nodes.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/base[[:space:]]model[[:space:]]image.png filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/buildings[[:space:]]ff.png filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/buildings.mp4 filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/distilled[[:space:]]image.png filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/hdr_input_video.mp4 filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/ingredients_input.jpg filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/inpainting_input.mp4 filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/instant_shave_input.mp4 filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/lipdub_input.mp4 filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/motion_track_input.jpg filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/outpainting_input.mp4 filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/assets/spatial_upscale_4x_input.mp4 filter=lfs diff=lfs merge=lfs -text
+ComfyUI/custom_nodes/ComfyUI-LTXVideo/gemma_configs/tokenizer.json filter=lfs diff=lfs merge=lfs -text

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/api/bizyair.py

@@ -0,0 +1,372 @@
+import yaml
+import pathlib
+import base64
+import io
+import json
+import os
+import pickle
+import zlib
+import urllib.parse
+import urllib.request
+import urllib.error
+from enum import Enum
+from functools import singledispatch
+from typing import Any, List, Union
+
+import numpy as np
+import torch
+from PIL import Image
+
+root_path = pathlib.Path(__file__).parent.parent.parent.parent
+config_path = os.path.join(root_path, 'config.yaml')
+
+class BizyAIRAPI:
+    def __init__(self):
+        self.base_url = 'https://bizyair-api.siliconflow.cn/x/v1'
+        self.api_key = None
+
+
+    def getAPIKey(self):
+        if self.api_key is None:
+            if os.path.isfile(config_path):
+                with open(config_path, 'r') as f:
+                    data = yaml.load(f, Loader=yaml.FullLoader)
+                    if 'BIZYAIR_API_KEY' not in data:
+                        raise Exception("Please add BIZYAIR_API_KEY to config.yaml")
+                    self.api_key = data['BIZYAIR_API_KEY']
+            else:
+                raise Exception("Please add config.yaml to root path")
+        return self.api_key
+
+    def send_post_request(self, url, payload, headers):
+        try:
+            data = json.dumps(payload).encode("utf-8")
+            req = urllib.request.Request(url, data=data, headers=headers, method="POST")
+            with urllib.request.urlopen(req) as response:
+                response_data = response.read().decode("utf-8")
+            return response_data
+        except urllib.error.URLError as e:
+            if "Unauthorized" in str(e):
+                raise Exception(
+                    "Key is invalid, please refer to https://cloud.siliconflow.cn to get the API key.\n"
+                    "If you have the key, please click the 'BizyAir Key' button at the bottom right to set the key."
+                )
+            else:
+                raise Exception(
+                    f"Failed to connect to the server: {e}, if you have no key, "
+                )
+
+    # joycaption
+    def joyCaption(self, payload, image, apikey_override=None, API_URL='/supernode/joycaption2'):
+        if apikey_override is not None:
+            api_key = apikey_override
+        else:
+            api_key = self.getAPIKey()
+        url = f"{self.base_url}{API_URL}"
+        print('Sending request to:', url)
+        auth = f"Bearer {api_key}"
+        headers = {
+            "accept": "application/json",
+            "content-type": "application/json",
+            "authorization": auth,
+        }
+        input_image = encode_data(image, disable_image_marker=True)
+        payload["image"] = input_image
+
+        ret: str = self.send_post_request(url=url, payload=payload, headers=headers)
+        ret = json.loads(ret)
+
+        try:
+            if "result" in ret:
+                ret = json.loads(ret["result"])
+        except Exception as e:
+            raise Exception(f"Unexpected response: {ret} {e=}")
+
+        if ret["type"] == "error":
+            raise Exception(ret["message"])
+
+        msg = ret["data"]
+        if msg["type"] not in ("comfyair", "bizyair",):
+            raise Exception(f"Unexpected response type: {msg}")
+
+        caption = msg["data"]
+
+        return caption
+
+bizyairAPI = BizyAIRAPI()
+
+
+
+BIZYAIR_DEBUG = True
+# Marker to identify base64-encoded tensors
+TENSOR_MARKER = "TENSOR:"
+IMAGE_MARKER = "IMAGE:"
+
+
+class TaskStatus(Enum):
+    PENDING = "pending"
+    PROCESSING = "processing"
+    COMPLETED = "completed"
+
+
+def convert_image_to_rgb(image: Image.Image) -> Image.Image:
+    if image.mode != "RGB":
+        return image.convert("RGB")
+    return image
+
+
+def encode_image_to_base64(
+    image: Image.Image, format: str = "png", quality: int = 100, lossless=False
+) -> str:
+    image = convert_image_to_rgb(image)
+    with io.BytesIO() as output:
+        image.save(output, format=format, quality=quality, lossless=lossless)
+        output.seek(0)
+        img_bytes = output.getvalue()
+        if BIZYAIR_DEBUG:
+            print(f"encode_image_to_base64: {format_bytes(len(img_bytes))}")
+    return base64.b64encode(img_bytes).decode("utf-8")
+
+
+def decode_base64_to_np(img_data: str, format: str = "png") -> np.ndarray:
+    img_bytes = base64.b64decode(img_data)
+    if BIZYAIR_DEBUG:
+        print(f"decode_base64_to_np: {format_bytes(len(img_bytes))}")
+    with io.BytesIO(img_bytes) as input_buffer:
+        img = Image.open(input_buffer)
+        # https://github.com/comfyanonymous/ComfyUI/blob/a178e25912b01abf436eba1cfaab316ba02d272d/nodes.py#L1511
+        img = img.convert("RGB")
+        return np.array(img)
+
+
+def decode_base64_to_image(img_data: str) -> Image.Image:
+    img_bytes = base64.b64decode(img_data)
+    with io.BytesIO(img_bytes) as input_buffer:
+        img = Image.open(input_buffer)
+        if BIZYAIR_DEBUG:
+            format_info = img.format.upper() if img.format else "Unknown"
+            print(f"decode image format: {format_info}")
+        return img
+
+
+def format_bytes(num_bytes: int) -> str:
+    """
+    Converts a number of bytes to a human-readable string with units (B, KB, or MB).
+
+    :param num_bytes: The number of bytes to convert.
+    :return: A string representing the number of bytes in a human-readable format.
+    """
+    if num_bytes < 1024:
+        return f"{num_bytes} B"
+    elif num_bytes < 1024 * 1024:
+        return f"{num_bytes / 1024:.2f} KB"
+    else:
+        return f"{num_bytes / (1024 * 1024):.2f} MB"
+
+
+def _legacy_encode_comfy_image(image: torch.Tensor, image_format="png") -> str:
+    input_image = image.cpu().detach().numpy()
+    i = 255.0 * input_image[0]
+    input_image = np.clip(i, 0, 255).astype(np.uint8)
+    base64ed_image = encode_image_to_base64(
+        Image.fromarray(input_image), format=image_format
+    )
+    return base64ed_image
+
+
+def _legacy_decode_comfy_image(
+    img_data: Union[List, str], image_format="png"
+) -> torch.tensor:
+    if isinstance(img_data, List):
+        decoded_imgs = [decode_comfy_image(x, old_version=True) for x in img_data]
+
+        combined_imgs = torch.cat(decoded_imgs, dim=0)
+        return combined_imgs
+
+    out = decode_base64_to_np(img_data, format=image_format)
+    out = np.array(out).astype(np.float32) / 255.0
+    output = torch.from_numpy(out)[None,]
+    return output
+
+
+def _new_encode_comfy_image(images: torch.Tensor, image_format="WEBP", **kwargs) -> str:
+    """https://docs.comfy.org/essentials/custom_node_snippets#save-an-image-batch
+    Encode a batch of images to base64 strings.
+
+    Args:
+        images (torch.Tensor): A batch of images.
+        image_format (str, optional): The format of the images. Defaults to "WEBP".
+
+    Returns:
+        str: A JSON string containing the base64-encoded images.
+    """
+    results = {}
+    for batch_number, image in enumerate(images):
+        i = 255.0 * image.cpu().numpy()
+        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+        base64ed_image = encode_image_to_base64(img, format=image_format, **kwargs)
+        results[batch_number] = base64ed_image
+
+    return json.dumps(results)
+
+
+def _new_decode_comfy_image(img_datas: str, image_format="WEBP") -> torch.tensor:
+    """
+    Decode a batch of base64-encoded images.
+
+    Args:
+        img_datas (str): A JSON string containing the base64-encoded images.
+        image_format (str, optional): The format of the images. Defaults to "WEBP".
+
+    Returns:
+        torch.Tensor: A tensor containing the decoded images.
+    """
+    img_datas = json.loads(img_datas)
+
+    decoded_imgs = []
+    for img_data in img_datas.values():
+        decoded_image = decode_base64_to_np(img_data, format=image_format)
+        decoded_image = np.array(decoded_image).astype(np.float32) / 255.0
+        decoded_imgs.append(torch.from_numpy(decoded_image)[None,])
+
+    return torch.cat(decoded_imgs, dim=0)
+
+
+def encode_comfy_image(
+    image: torch.Tensor, image_format="WEBP", old_version=False, lossless=False
+) -> str:
+    if old_version:
+        return _legacy_encode_comfy_image(image, image_format)
+    return _new_encode_comfy_image(image, image_format, lossless=lossless)
+
+
+def decode_comfy_image(
+    img_data: Union[List, str], image_format="WEBP", old_version=False
+) -> torch.tensor:
+    if old_version:
+        return _legacy_decode_comfy_image(img_data, image_format)
+    return _new_decode_comfy_image(img_data, image_format)
+
+
+def tensor_to_base64(tensor: torch.Tensor, compress=True) -> str:
+    tensor_np = tensor.cpu().detach().numpy()
+
+    tensor_bytes = pickle.dumps(tensor_np)
+    if compress:
+        tensor_bytes = zlib.compress(tensor_bytes)
+
+    tensor_b64 = base64.b64encode(tensor_bytes).decode("utf-8")
+    return tensor_b64
+
+
+def base64_to_tensor(tensor_b64: str, compress=True) -> torch.Tensor:
+    tensor_bytes = base64.b64decode(tensor_b64)
+
+    if compress:
+        tensor_bytes = zlib.decompress(tensor_bytes)
+
+    tensor_np = pickle.loads(tensor_bytes)
+
+    tensor = torch.from_numpy(tensor_np)
+    return tensor
+
+
+@singledispatch
+def decode_data(input, old_version=False):
+    raise NotImplementedError(f"Unsupported type: {type(input)}")
+
+
+@decode_data.register(int)
+@decode_data.register(float)
+@decode_data.register(bool)
+@decode_data.register(type(None))
+def _(input, **kwargs):
+    return input
+
+
+@decode_data.register(dict)
+def _(input, **kwargs):
+    return {k: decode_data(v, **kwargs) for k, v in input.items()}
+
+
+@decode_data.register(list)
+def _(input, **kwargs):
+    return [decode_data(x, **kwargs) for x in input]
+
+
+@decode_data.register(str)
+def _(input: str, **kwargs):
+    if input.startswith(TENSOR_MARKER):
+        tensor_b64 = input[len(TENSOR_MARKER) :]
+        return base64_to_tensor(tensor_b64)
+    elif input.startswith(IMAGE_MARKER):
+        tensor_b64 = input[len(IMAGE_MARKER) :]
+        old_version = kwargs.get("old_version", False)
+        return decode_comfy_image(tensor_b64, old_version=old_version)
+    return input
+
+
+@singledispatch
+def encode_data(output, disable_image_marker=False, old_version=False):
+    raise NotImplementedError(f"Unsupported type: {type(output)}")
+
+
+@encode_data.register(dict)
+def _(output, **kwargs):
+    return {k: encode_data(v, **kwargs) for k, v in output.items()}
+
+
+@encode_data.register(list)
+def _(output, **kwargs):
+    return [encode_data(x, **kwargs) for x in output]
+
+
+def is_image_tensor(tensor) -> bool:
+    """https://docs.comfy.org/essentials/custom_node_datatypes#image
+
+    Check if the given tensor is in the format of an IMAGE (shape [B, H, W, C] where C=3).
+
+    `Args`:
+        tensor (torch.Tensor): The tensor to check.
+
+    `Returns`:
+        bool: True if the tensor is in the IMAGE format, False otherwise.
+    """
+    try:
+        if not isinstance(tensor, torch.Tensor):
+            return False
+
+        if len(tensor.shape) != 4:
+            return False
+
+        B, H, W, C = tensor.shape
+        if C != 3:
+            return False
+
+        return True
+    except:
+        return False
+
+
+@encode_data.register(torch.Tensor)
+def _(output, **kwargs):
+    if is_image_tensor(output) and not kwargs.get("disable_image_marker", False):
+        old_version = kwargs.get("old_version", False)
+        lossless = kwargs.get("lossless", True)
+        return IMAGE_MARKER + encode_comfy_image(
+            output, image_format="WEBP", old_version=old_version, lossless=lossless
+        )
+    return TENSOR_MARKER + tensor_to_base64(output)
+
+
+@encode_data.register(int)
+@encode_data.register(float)
+@encode_data.register(bool)
+@encode_data.register(type(None))
+def _(output, **kwargs):
+    return output
+
+
+@encode_data.register(str)
+def _(output, **kwargs):
+    return output

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/api/fluxai.py

@@ -0,0 +1,51 @@
+import json
+import os
+import yaml
+import requests
+import pathlib
+from aiohttp import web
+
+root_path = pathlib.Path(__file__).parent.parent.parent.parent
+config_path = os.path.join(root_path,'config.yaml')
+class FluxAIAPI:
+    def __init__(self):
+        self.api_url = "https://fluxaiimagegenerator.com/api"
+        self.origin = "https://fluxaiimagegenerator.com"
+        self.user_agent = None
+        self.cookie = None
+
+    def promptGenerate(self, text, cookies=None):
+        cookie = self.cookie if cookies is None else cookies
+        if cookie is None:
+            if os.path.isfile(config_path):
+                with open(config_path, 'r') as f:
+                    data = yaml.load(f, Loader=yaml.FullLoader)
+                    if 'FLUXAI_COOKIE' not in data:
+                        raise Exception("Please add FLUXAI_COOKIE to config.yaml")
+                    if "FLUXAI_USER_AGENT" in data:
+                        self.user_agent = data["FLUXAI_USER_AGENT"]
+                    self.cookie = cookie = data['FLUXAI_COOKIE']
+
+        headers = {
+            "Cookie": cookie,
+            "Referer": "https://fluxaiimagegenerator.com/flux-prompt-generator",
+            "Origin": self.origin,
+            "Content-Type": "application/json",
+        }
+        if self.user_agent is not None:
+            headers['User-Agent'] = self.user_agent
+
+        url = self.api_url + '/prompt'
+        json = {
+            "prompt": text
+        }
+
+        response = requests.post(url, json=json, headers=headers)
+        res = response.json()
+        if "error" in res:
+            return res['error']
+        elif "data" in res and "prompt" in res['data']:
+            return res['data']['prompt']
+
+fluxaiAPI = FluxAIAPI()
+

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/api/stability.py

@@ -0,0 +1,200 @@
+import json
+import os
+import yaml
+import requests
+import pathlib
+from aiohttp import web
+from server import PromptServer
+from ..image import tensor2pil, pil2tensor, image2base64, pil2byte
+from ..log import log_node_error
+
+
+root_path = pathlib.Path(__file__).parent.parent.parent.parent
+config_path = os.path.join(root_path,'config.yaml')
+default_key = [{'name':'Default', 'key':''}]
+
+
+class StabilityAPI:
+    def __init__(self):
+        self.api_url = "https://api.stability.ai"
+        self.api_keys = None
+        self.api_current = 0
+        self.user_info = {}
+
+    def getErrors(self, code):
+        errors = {
+            400: "Bad Request",
+            403: "ApiKey Forbidden",
+            413: "Your request was larger than 10MiB.",
+            429: "You have made more than 150 requests in 10 seconds.",
+            500: "Internal Server Error",
+        }
+        return errors.get(code, "Unknown Error")
+
+    def getAPIKeys(self):
+        if os.path.isfile(config_path):
+            with open(config_path, 'r') as f:
+                data = yaml.load(f, Loader=yaml.FullLoader)
+                if not data:
+                    data = {'STABILITY_API_KEY': default_key, 'STABILITY_API_DEFAULT':0}
+                    with open(config_path, 'w') as f:
+                        yaml.dump(data, f)
+                if 'STABILITY_API_KEY' not in data:
+                    data['STABILITY_API_KEY'] = default_key
+                    data['STABILITY_API_DEFAULT'] = 0
+                    with open(config_path, 'w') as f:
+                        yaml.dump(data, f)
+                api_keys = data['STABILITY_API_KEY']
+                self.api_current = data['STABILITY_API_DEFAULT']
+                self.api_keys = api_keys
+                return api_keys
+        else:
+            # create a yaml file
+            with open(config_path, 'w') as f:
+                data = {'STABILITY_API_KEY': default_key, 'STABILITY_API_DEFAULT':0}
+                yaml.dump(data, f)
+                return data['STABILITY_API_KEY']
+        pass
+
+    def setAPIKeys(self, api_keys):
+        if len(api_keys) > 0:
+            self.api_keys = api_keys
+            # load and save the yaml file
+            with open(config_path, 'r') as f:
+                data = yaml.load(f, Loader=yaml.FullLoader)
+                data['STABILITY_API_KEY'] = api_keys
+                with open(config_path, 'w') as f:
+                    yaml.dump(data, f)
+        return True
+
+    def setAPIDefault(self, current):
+        if current is not None:
+            self.api_current = current
+            # load and save the yaml file
+            with open(config_path, 'r') as f:
+                data = yaml.load(f, Loader=yaml.FullLoader)
+                data['STABILITY_API_DEFAULT'] = current
+                with open(config_path, 'w') as f:
+                    yaml.dump(data, f)
+        return True
+
+    def generate_sd3_image(self, prompt, negative_prompt, aspect_ratio,  model, seed, mode='text-to-image', image=None, strength=1, output_format='png', node_name='easy stableDiffusion3API'):
+        url = f"{self.api_url}/v2beta/stable-image/generate/sd3"
+        api_key = self.api_keys[self.api_current]['key']
+        files = None
+        data = {
+            "prompt": prompt,
+            "mode": mode,
+            "model": model,
+            "seed": seed,
+            "output_format": output_format,
+        }
+        if model == 'sd3':
+            data['negative_prompt'] = negative_prompt
+
+        if mode == 'text-to-image':
+            files = {"none": ''}
+            data['aspect_ratio'] = aspect_ratio
+        elif mode == 'image-to-image':
+            pil_image = tensor2pil(image)
+            image_byte = pil2byte(pil_image)
+            files = {"image": ("output.png", image_byte, 'image/png')}
+            data['strength'] = strength
+
+        response = requests.post(url,
+            headers={"authorization": f"{api_key}", "accept": "application/json"},
+            files=files,
+            data=data,
+        )
+        if response.status_code == 200:
+            PromptServer.instance.send_sync('stable-diffusion-api-generate-succeed',{"model":model})
+            json_data = response.json()
+            image_base64 = json_data['image']
+            image_data = image2base64(image_base64)
+            output_t = pil2tensor(image_data)
+            return output_t
+        else:
+            if 'application/json' in response.headers['Content-Type']:
+                error_info = response.json()
+                log_node_error(node_name,  error_info.get('name', 'No name provided'))
+                log_node_error(node_name, error_info.get('errors', ['No details provided']))
+            error_status_text = self.getErrors(response.status_code)
+            PromptServer.instance.send_sync('easyuse-toast',{"type": "error", "content": error_status_text})
+            raise Exception(f"Failed to generate image: {error_status_text}")
+
+    # get user account
+    async def getUserAccount(self, cache=True):
+        url = f"{self.api_url}/v1/user/account"
+        api_key = self.api_keys[self.api_current]['key']
+        name = self.api_keys[self.api_current]['name']
+        if cache and name in self.user_info:
+            return self.user_info[name]
+        else:
+            response = requests.get(url, headers={"Authorization": f"Bearer {api_key}"})
+            if response.status_code == 200:
+                user_info = response.json()
+                self.user_info[name] = user_info
+                return user_info
+            else:
+                PromptServer.instance.send_sync('easyuse-toast',{'type': 'error', 'content': self.getErrors(response.status_code)})
+                return None
+
+    # get user balance
+    async def getUserBalance(self):
+        url = f"{self.api_url}/v1/user/balance"
+        api_key = self.api_keys[self.api_current]['key']
+        response = requests.get(url, headers={
+            "Authorization": f"Bearer {api_key}"
+        })
+        if response.status_code == 200:
+            return response.json()
+        else:
+            PromptServer.instance.send_sync('easyuse-toast', {'type': 'error', 'content': self.getErrors(response.status_code)})
+            return None
+
+stableAPI = StabilityAPI()
+
+@PromptServer.instance.routes.get("/easyuse/stability/api_keys")
+async def get_stability_api_keys(request):
+    stableAPI.getAPIKeys()
+    return web.json_response({"keys": stableAPI.api_keys, "current": stableAPI.api_current})
+
+@PromptServer.instance.routes.post("/easyuse/stability/set_api_keys")
+async def set_stability_api_keys(request):
+    post = await request.post()
+    api_keys = post.get("api_keys")
+    current = post.get('current')
+    if api_keys is not None:
+        api_keys = json.loads(api_keys)
+        stableAPI.setAPIKeys(api_keys)
+        if current is not None:
+            print(current)
+            stableAPI.setAPIDefault(int(current))
+            account = await stableAPI.getUserAccount()
+            balance = await stableAPI.getUserBalance()
+            return web.json_response({'account': account, 'balance': balance})
+        else:
+            return web.json_response({'status': 'ok'})
+    else:
+        return web.Response(status=400)
+
+@PromptServer.instance.routes.post("/easyuse/stability/set_apikey_default")
+async def set_stability_api_default(request):
+    post = await request.post()
+    current = post.get("current")
+    if current is not None and current < len(stableAPI.api_keys):
+        stableAPI.api_current = current
+        return web.json_response({'status': 'ok'})
+    else:
+        return web.Response(status=400)
+
+@PromptServer.instance.routes.get("/easyuse/stability/user_info")
+async def get_account_info(request):
+    account = await stableAPI.getUserAccount()
+    balance = await stableAPI.getUserBalance()
+    return web.json_response({'account': account, 'balance': balance})
+
+@PromptServer.instance.routes.get("/easyuse/stability/balance")
+async def get_balance_info(request):
+    balance = await stableAPI.getUserBalance()
+    return web.json_response({'balance': balance})

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/cache.py

@@ -0,0 +1,87 @@
+import itertools
+from typing import Optional
+
+class TaggedCache:
+    def __init__(self, tag_settings: Optional[dict]=None):
+        self._tag_settings = tag_settings or {}  # tag cache size
+        self._data = {}
+
+    def __getitem__(self, key):
+        for tag_data in self._data.values():
+            if key in tag_data:
+                return tag_data[key]
+        raise KeyError(f'Key `{key}` does not exist')
+
+    def __setitem__(self, key, value: tuple):
+        # value: (tag: str, (islist: bool, data: *))
+
+        # if key already exists, pop old value
+        for tag_data in self._data.values():
+            if key in tag_data:
+                tag_data.pop(key, None)
+                break
+
+        tag = value[0]
+        if tag not in self._data:
+
+            try:
+                from cachetools import LRUCache
+
+                default_size = 20
+                if 'ckpt' in tag:
+                    default_size = 5
+                elif tag in ['latent', 'image']:
+                    default_size = 100
+
+                self._data[tag] = LRUCache(maxsize=self._tag_settings.get(tag, default_size))
+
+            except (ImportError, ModuleNotFoundError):
+                # TODO: implement a simple lru dict
+                self._data[tag] = {}
+        self._data[tag][key] = value
+
+    def __delitem__(self, key):
+        for tag_data in self._data.values():
+            if key in tag_data:
+                del tag_data[key]
+                return
+        raise KeyError(f'Key `{key}` does not exist')
+
+    def __contains__(self, key):
+        return any(key in tag_data for tag_data in self._data.values())
+
+    def items(self):
+        yield from itertools.chain(*map(lambda x :x.items(), self._data.values()))
+
+    def get(self, key, default=None):
+        """D.get(k[,d]) -> D[k] if k in D, else d.  d defaults to None."""
+        for tag_data in self._data.values():
+            if key in tag_data:
+                return tag_data[key]
+        return default
+
+    def clear(self):
+        # clear all cache
+        self._data = {}
+
+cache_settings = {}
+cache = TaggedCache(cache_settings)
+cache_count = {}
+
+def update_cache(k, tag, v):
+    cache[k] = (tag, v)
+    cnt = cache_count.get(k)
+    if cnt is None:
+        cnt = 0
+        cache_count[k] = cnt
+    else:
+        cache_count[k] += 1
+def remove_cache(key):
+    if key == '*':
+        cache.clear()
+        cache_count.clear()
+    elif key in cache:
+        del cache[key]
+        cache_count.pop(key, None)
+    else:
+        print(f"invalid {key}")

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/chooser.py

@@ -0,0 +1,153 @@
+from threading import Event
+
+import torch
+
+from server import PromptServer
+from aiohttp import web
+from comfy import model_management as mm
+from comfy_execution.graph import ExecutionBlocker
+import time
+
+class ChooserCancelled(Exception):
+    pass
+
+def get_chooser_cache():
+    """获取选择器缓存"""
+    if not hasattr(PromptServer.instance, '_easyuse_chooser_node'):
+        PromptServer.instance._easyuse_chooser_node = {}
+    return PromptServer.instance._easyuse_chooser_node
+
+def cleanup_session_data(node_id):
+    """清理会话数据"""
+    node_data = get_chooser_cache()
+    if node_id in node_data:
+        session_keys = ["event", "selected", "images", "total_count", "cancelled"]
+        for key in session_keys:
+            if key in node_data[node_id]:
+                del node_data[node_id][key]
+
+def wait_for_chooser(id, images, mode, period=0.1):
+    try:
+        node_data = get_chooser_cache()
+        images = [images[i:i + 1, ...] for i in range(images.shape[0])]
+        if mode == "Keep Last Selection":
+            if id in node_data and "last_selection" in node_data[id]:
+                last_selection = node_data[id]["last_selection"]
+                if last_selection and len(last_selection) > 0:
+                    valid_indices = [idx for idx in last_selection if 0 <= idx < len(images)]
+                    if valid_indices:
+                        try:
+                            PromptServer.instance.send_sync("easyuse-image-keep-selection", {
+                                "id": id,
+                                "selected": valid_indices
+                            })
+                        except Exception as e:
+                            pass
+                        cleanup_session_data(id)
+                        indices_str = ','.join(str(i) for i in valid_indices)
+                        images = [images[idx] for idx in valid_indices]
+                        images = torch.cat(images, dim=0)
+                        return {"result": (images,)}
+
+        if id in node_data:
+            del node_data[id]
+
+        event = Event()
+        node_data[id] = {
+            "event": event,
+            "images": images,
+            "selected": None,
+            "total_count": len(images),
+            "cancelled": False,
+        }
+
+        while id in node_data:
+            node_info = node_data[id]
+            if node_info.get("cancelled", False):
+                cleanup_session_data(id)
+                raise ChooserCancelled("Manual selection cancelled")
+
+            if "selected" in node_info and node_info["selected"] is not None:
+                break
+
+            time.sleep(period)
+
+        if id in node_data:
+            node_info = node_data[id]
+            selected_indices = node_info.get("selected")
+
+            if selected_indices is not None and len(selected_indices) > 0:
+                valid_indices = [idx for idx in selected_indices if 0 <= idx < len(images)]
+                if valid_indices:
+                    selected_images = [images[idx] for idx in valid_indices]
+
+                    if id not in node_data:
+                        node_data[id] = {}
+                    node_data[id]["last_selection"] = valid_indices
+                    cleanup_session_data(id)
+                    selected_images = torch.cat(selected_images, dim=0)
+                    return {"result": (selected_images,)}
+                else:
+                    cleanup_session_data(id)
+                    return {"result": (images[0] if len(images) > 0 else ExecutionBlocker(None),)}
+            else:
+                cleanup_session_data(id)
+                return {
+                    "result": (images[0] if len(images) > 0 else ExecutionBlocker(None),)}
+        else:
+            return {"result": (images[0] if len(images) > 0 else ExecutionBlocker(None),)}
+
+    except ChooserCancelled:
+        raise mm.InterruptProcessingException()
+    except Exception as e:
+        node_data = get_chooser_cache()
+        if id in node_data:
+            cleanup_session_data(id)
+        if 'image_list' in locals() and len(images) > 0:
+            return {"result": (images[0])}
+        else:
+            return {"result": (ExecutionBlocker(None),)}
+
+
+@PromptServer.instance.routes.post('/easyuse/image_chooser_message')
+async def handle_image_selection(request):
+    try:
+        data = await request.json()
+        node_id = data.get("node_id")
+        selected = data.get("selected", [])
+        action = data.get("action")
+
+        node_data = get_chooser_cache()
+
+        if node_id not in node_data:
+            return web.json_response({"code": -1, "error": "Node data does not exist"})
+
+        try:
+            node_info = node_data[node_id]
+
+            if "total_count" not in node_info:
+                return web.json_response({"code": -1, "error": "The node has been processed"})
+
+            if action == "cancel":
+                node_info["cancelled"] = True
+                node_info["selected"] = []
+            elif action == "select" and isinstance(selected, list):
+                valid_indices = [idx for idx in selected if isinstance(idx, int) and 0 <= idx < node_info["total_count"]]
+                if valid_indices:
+                    node_info["selected"] = valid_indices
+                    node_info["cancelled"] = False
+                else:
+                    return web.json_response({"code": -1, "error": "Invalid Selection Index"})
+            else:
+                return web.json_response({"code": -1, "error": "Invalid operation"})
+
+            node_info["event"].set()
+            return web.json_response({"code": 1})
+
+        except Exception as e:
+            if node_id in node_data and "event" in node_data[node_id]:
+                node_data[node_id]["event"].set()
+            return web.json_response({"code": -1, "message": "Processing Failed"})
+
+    except Exception as e:
+        return web.json_response({"code": -1, "message": "Request Failed"})

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/colorfix.py

@@ -0,0 +1,115 @@
+import torch
+from PIL import Image
+from torch import Tensor
+from torch.nn import functional as F
+
+from torchvision.transforms import ToTensor, ToPILImage
+
+def adain_color_fix(target: Image, source: Image):
+    # Convert images to tensors
+    to_tensor = ToTensor()
+    target_tensor = to_tensor(target).unsqueeze(0)
+    source_tensor = to_tensor(source).unsqueeze(0)
+
+    # Apply adaptive instance normalization
+    result_tensor = adaptive_instance_normalization(target_tensor, source_tensor)
+
+    # Convert tensor back to image
+    to_image = ToPILImage()
+    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
+
+    return result_image
+
+def wavelet_color_fix(target: Image, source: Image):
+    source = source.resize(target.size, resample=Image.Resampling.LANCZOS)
+
+    # Convert images to tensors
+    to_tensor = ToTensor()
+    target_tensor = to_tensor(target).unsqueeze(0)
+    source_tensor = to_tensor(source).unsqueeze(0)
+
+    # Apply wavelet reconstruction
+    result_tensor = wavelet_reconstruction(target_tensor, source_tensor)
+
+    # Convert tensor back to image
+    to_image = ToPILImage()
+    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
+
+    return result_image
+
+def calc_mean_std(feat: Tensor, eps=1e-5):
+    """Calculate mean and std for adaptive_instance_normalization.
+    Args:
+        feat (Tensor): 4D tensor.
+        eps (float): A small value added to the variance to avoid
+            divide-by-zero. Default: 1e-5.
+    """
+    size = feat.size()
+    assert len(size) == 4, 'The input feature should be 4D tensor.'
+    b, c = size[:2]
+    feat_var = feat.view(b, c, -1).var(dim=2) + eps
+    feat_std = feat_var.sqrt().view(b, c, 1, 1)
+    feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1)
+    return feat_mean, feat_std
+
+def adaptive_instance_normalization(content_feat:Tensor, style_feat:Tensor):
+    """Adaptive instance normalization.
+    Adjust the reference features to have the similar color and illuminations
+    as those in the degradate features.
+    Args:
+        content_feat (Tensor): The reference feature.
+        style_feat (Tensor): The degradate features.
+    """
+    size = content_feat.size()
+    style_mean, style_std = calc_mean_std(style_feat)
+    content_mean, content_std = calc_mean_std(content_feat)
+    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
+    return normalized_feat * style_std.expand(size) + style_mean.expand(size)
+
+def wavelet_blur(image: Tensor, radius: int):
+    """
+    Apply wavelet blur to the input tensor.
+    """
+    # input shape: (1, 3, H, W)
+    # convolution kernel
+    kernel_vals = [
+        [0.0625, 0.125, 0.0625],
+        [0.125, 0.25, 0.125],
+        [0.0625, 0.125, 0.0625],
+    ]
+    kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device)
+    # add channel dimensions to the kernel to make it a 4D tensor
+    kernel = kernel[None, None]
+    # repeat the kernel across all input channels
+    kernel = kernel.repeat(3, 1, 1, 1)
+    image = F.pad(image, (radius, radius, radius, radius), mode='replicate')
+    # apply convolution
+    output = F.conv2d(image, kernel, groups=3, dilation=radius)
+    return output
+
+def wavelet_decomposition(image: Tensor, levels=5):
+    """
+    Apply wavelet decomposition to the input tensor.
+    This function only returns the low frequency & the high frequency.
+    """
+    high_freq = torch.zeros_like(image)
+    for i in range(levels):
+        radius = 2 ** i
+        low_freq = wavelet_blur(image, radius)
+        high_freq += (image - low_freq)
+        image = low_freq
+
+    return high_freq, low_freq
+
+def wavelet_reconstruction(content_feat:Tensor, style_feat:Tensor):
+    """
+    Apply wavelet decomposition, so that the content will have the same color as the style.
+    """
+    # calculate the wavelet decomposition of the content feature
+    content_high_freq, content_low_freq = wavelet_decomposition(content_feat)
+    del content_low_freq
+    # calculate the wavelet decomposition of the style feature
+    style_high_freq, style_low_freq = wavelet_decomposition(style_feat)
+    del style_high_freq
+    # reconstruct the content feature with the style's high frequency
+    return content_high_freq + style_low_freq

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/conditioning.py

@@ -0,0 +1,57 @@
+from .utils import find_wildcards_seed, find_nearest_steps, is_linked_styles_selector
+from .log import log_node_warn
+from .translate import zh_to_en, has_chinese
+from .wildcards import process_with_loras
+from .adv_encode import advanced_encode
+
+from nodes import ConditioningConcat, ConditioningCombine, ConditioningAverage, ConditioningSetTimestepRange, CLIPTextEncode
+
+def prompt_to_cond(type, model, clip, clip_skip, lora_stack, text, prompt_token_normalization, prompt_weight_interpretation, a1111_prompt_style ,my_unique_id, prompt, easyCache, can_load_lora=True, steps=None, model_type=None):
+    styles_selector = is_linked_styles_selector(prompt, my_unique_id, type)
+    title = "Positive encoding" if type == 'positive' else "Negative encoding"
+
+    # Translate cn to en
+    if model_type not in ['hydit'] and text is not None and has_chinese(text):
+        text = zh_to_en([text])[0]
+
+    if model_type in ['hydit', 'flux', 'mochi']:
+        log_node_warn(title + "...")
+        embeddings_final, = CLIPTextEncode().encode(clip, text) if text is not None else (None,)
+
+        return (embeddings_final, "", model, clip)
+
+    log_node_warn(title + "...")
+
+    positive_seed = find_wildcards_seed(my_unique_id, text, prompt)
+    model, clip, text, cond_decode, show_prompt, pipe_lora_stack = process_with_loras(
+        text, model, clip, type, positive_seed, can_load_lora, lora_stack, easyCache)
+    wildcard_prompt = cond_decode if show_prompt or styles_selector else ""
+
+    clipped = clip.clone()
+    # 当clip模型不存在t5xxl时，可执行跳过层
+    if not hasattr(clip.cond_stage_model, 't5xxl'):
+        if clip_skip != 0:
+            clipped.clip_layer(clip_skip)
+
+    steps = steps if steps is not None else find_nearest_steps(my_unique_id, prompt)
+    return (advanced_encode(clipped, text, prompt_token_normalization,
+                            prompt_weight_interpretation, w_max=1.0,
+                            apply_to_pooled='enable',
+                            a1111_prompt_style=a1111_prompt_style, steps=steps) if text is not None else None, wildcard_prompt, model, clipped)
+
+def set_cond(old_cond, new_cond, mode, average_strength, old_cond_start, old_cond_end, new_cond_start, new_cond_end):
+    if not old_cond:
+        return new_cond
+    else:
+        if mode == "replace":
+            return new_cond
+        elif mode == "concat":
+            return ConditioningConcat().concat(new_cond, old_cond)[0]
+        elif mode == "combine":
+            return ConditioningCombine().combine(old_cond, new_cond)[0]
+        elif mode == 'average':
+            return ConditioningAverage().addWeighted(new_cond, old_cond, average_strength)[0]
+        elif mode == 'timestep':
+            cond_1 = ConditioningSetTimestepRange().set_range(old_cond, old_cond_start, old_cond_end)[0]
+            cond_2 = ConditioningSetTimestepRange().set_range(new_cond, new_cond_start, new_cond_end)[0]
+            return ConditioningCombine().combine(cond_1, cond_2)[0]

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/controlnet.py

@@ -0,0 +1,93 @@
+import folder_paths
+import comfy.controlnet
+import comfy.model_management
+from nodes import NODE_CLASS_MAPPINGS
+
+union_controlnet_types = {"auto": -1, "openpose": 0, "depth": 1, "hed/pidi/scribble/ted": 2, "canny/lineart/anime_lineart/mlsd": 3, "normal": 4, "segment": 5, "tile": 6, "repaint": 7}
+
+class easyControlnet:
+    def __init__(self):
+        pass
+
+    def apply(self, control_net_name, image, positive, negative, strength, start_percent=0, end_percent=1, control_net=None, scale_soft_weights=1, mask=None, union_type=None, easyCache=None, use_cache=True, model=None, vae=None):
+        if strength == 0:
+            return (positive, negative)
+
+        # kolors controlnet patch
+        from ..modules.kolors.loader import is_kolors_model, applyKolorsUnet
+        if is_kolors_model(model):
+            from ..modules.kolors.model_patch import patch_controlnet
+            if control_net is None:
+                with applyKolorsUnet():
+                    control_net = easyCache.load_controlnet(control_net_name, scale_soft_weights, use_cache)
+                control_net = patch_controlnet(model, control_net)
+        else:
+            if control_net is None:
+                if easyCache is not None:
+                    control_net = easyCache.load_controlnet(control_net_name, scale_soft_weights, use_cache)
+                else:
+                    controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
+                    control_net = comfy.controlnet.load_controlnet(controlnet_path)
+
+        # union controlnet
+        if union_type is not None:
+            control_net = control_net.copy()
+            type_number = union_controlnet_types[union_type]
+            if type_number >= 0:
+                control_net.set_extra_arg("control_type", [type_number])
+            else:
+                control_net.set_extra_arg("control_type", [])
+
+        if mask is not None:
+            mask = mask.to(self.device)
+
+        if mask is not None and len(mask.shape) < 3:
+            mask = mask.unsqueeze(0)
+
+        control_hint = image.movedim(-1, 1)
+
+        is_cond = True
+        if negative is None:
+            p = []
+            for t in positive:
+                n = [t[0], t[1].copy()]
+                c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent))
+                if 'control' in t[1]:
+                    c_net.set_previous_controlnet(t[1]['control'])
+                n[1]['control'] = c_net
+                n[1]['control_apply_to_uncond'] = True
+                if mask is not None:
+                    n[1]['mask'] = mask
+                    n[1]['set_area_to_bounds'] = False
+                p.append(n)
+            positive = p
+        else:
+            cnets = {}
+            out = []
+            for conditioning in [positive, negative]:
+                c = []
+                for t in conditioning:
+                    d = t[1].copy()
+
+                    prev_cnet = d.get('control', None)
+                    if prev_cnet in cnets:
+                        c_net = cnets[prev_cnet]
+                    else:
+                        c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae)
+                        c_net.set_previous_controlnet(prev_cnet)
+                        cnets[prev_cnet] = c_net
+
+                    d['control'] = c_net
+                    d['control_apply_to_uncond'] = False
+
+                    if mask is not None:
+                        d['mask'] = mask
+                        d['set_area_to_bounds'] = False
+
+                    n = [t[0], d]
+                    c.append(n)
+                out.append(c)
+            positive = out[0]
+            negative = out[1]
+
+        return (positive, negative)

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/dynthres_core.py

@@ -0,0 +1,167 @@
+import torch, math
+
+######################### DynThresh Core #########################
+
+class DynThresh:
+
+    Modes = ["Constant", "Linear Down", "Cosine Down", "Half Cosine Down", "Linear Up", "Cosine Up", "Half Cosine Up", "Power Up", "Power Down", "Linear Repeating", "Cosine Repeating", "Sawtooth"]
+    Startpoints = ["MEAN", "ZERO"]
+    Variabilities = ["AD", "STD"]
+
+    def __init__(self, mimic_scale, threshold_percentile, mimic_mode, mimic_scale_min, cfg_mode, cfg_scale_min, sched_val, experiment_mode, max_steps, separate_feature_channels, scaling_startpoint, variability_measure, interpolate_phi):
+        self.mimic_scale = mimic_scale
+        self.threshold_percentile = threshold_percentile
+        self.mimic_mode = mimic_mode
+        self.cfg_mode = cfg_mode
+        self.max_steps = max_steps
+        self.cfg_scale_min = cfg_scale_min
+        self.mimic_scale_min = mimic_scale_min
+        self.experiment_mode = experiment_mode
+        self.sched_val = sched_val
+        self.sep_feat_channels = separate_feature_channels
+        self.scaling_startpoint = scaling_startpoint
+        self.variability_measure = variability_measure
+        self.interpolate_phi = interpolate_phi
+
+    def interpret_scale(self, scale, mode, min):
+        scale -= min
+        max = self.max_steps - 1
+        frac = self.step / max
+        if mode == "Constant":
+            pass
+        elif mode == "Linear Down":
+            scale *= 1.0 - frac
+        elif mode == "Half Cosine Down":
+            scale *= math.cos(frac)
+        elif mode == "Cosine Down":
+            scale *= math.cos(frac * 1.5707)
+        elif mode == "Linear Up":
+            scale *= frac
+        elif mode == "Half Cosine Up":
+            scale *= 1.0 - math.cos(frac)
+        elif mode == "Cosine Up":
+            scale *= 1.0 - math.cos(frac * 1.5707)
+        elif mode == "Power Up":
+            scale *= math.pow(frac, self.sched_val)
+        elif mode == "Power Down":
+            scale *= 1.0 - math.pow(frac, self.sched_val)
+        elif mode == "Linear Repeating":
+            portion = (frac * self.sched_val) % 1.0
+            scale *= (0.5 - portion) * 2 if portion < 0.5 else (portion - 0.5) * 2
+        elif mode == "Cosine Repeating":
+            scale *= math.cos(frac * 6.28318 * self.sched_val) * 0.5 + 0.5
+        elif mode == "Sawtooth":
+            scale *= (frac * self.sched_val) % 1.0
+        scale += min
+        return scale
+
+    def dynthresh(self, cond, uncond, cfg_scale, weights):
+        mimic_scale = self.interpret_scale(self.mimic_scale, self.mimic_mode, self.mimic_scale_min)
+        cfg_scale = self.interpret_scale(cfg_scale, self.cfg_mode, self.cfg_scale_min)
+        # uncond shape is (batch, 4, height, width)
+        conds_per_batch = cond.shape[0] / uncond.shape[0]
+        assert conds_per_batch == int(conds_per_batch), "Expected # of conds per batch to be constant across batches"
+        cond_stacked = cond.reshape((-1, int(conds_per_batch)) + uncond.shape[1:])
+
+        ### Normal first part of the CFG Scale logic, basically
+        diff = cond_stacked - uncond.unsqueeze(1)
+        if weights is not None:
+            diff = diff * weights
+        relative = diff.sum(1)
+
+        ### Get the normal result for both mimic and normal scale
+        mim_target = uncond + relative * mimic_scale
+        cfg_target = uncond + relative * cfg_scale
+        ### If we weren't doing mimic scale, we'd just return cfg_target here
+
+        ### Now recenter the values relative to their average rather than absolute, to allow scaling from average
+        mim_flattened = mim_target.flatten(2)
+        cfg_flattened = cfg_target.flatten(2)
+        mim_means = mim_flattened.mean(dim=2).unsqueeze(2)
+        cfg_means = cfg_flattened.mean(dim=2).unsqueeze(2)
+        mim_centered = mim_flattened - mim_means
+        cfg_centered = cfg_flattened - cfg_means
+
+        if self.sep_feat_channels:
+            if self.variability_measure == 'STD':
+                mim_scaleref = mim_centered.std(dim=2).unsqueeze(2)
+                cfg_scaleref = cfg_centered.std(dim=2).unsqueeze(2)
+            else: # 'AD'
+                mim_scaleref = mim_centered.abs().max(dim=2).values.unsqueeze(2)
+                cfg_scaleref = torch.quantile(cfg_centered.abs(), self.threshold_percentile, dim=2).unsqueeze(2)
+
+        else:
+            if self.variability_measure == 'STD':
+                mim_scaleref = mim_centered.std()
+                cfg_scaleref = cfg_centered.std()
+            else: # 'AD'
+                mim_scaleref = mim_centered.abs().max()
+                cfg_scaleref = torch.quantile(cfg_centered.abs(), self.threshold_percentile)
+
+        if self.scaling_startpoint == 'ZERO':
+            scaling_factor = mim_scaleref / cfg_scaleref
+            result = cfg_flattened * scaling_factor
+
+        else: # 'MEAN'
+            if self.variability_measure == 'STD':
+                cfg_renormalized = (cfg_centered / cfg_scaleref) * mim_scaleref
+            else: # 'AD'
+                ### Get the maximum value of all datapoints (with an optional threshold percentile on the uncond)
+                max_scaleref = torch.maximum(mim_scaleref, cfg_scaleref)
+                ### Clamp to the max
+                cfg_clamped = cfg_centered.clamp(-max_scaleref, max_scaleref)
+                ### Now shrink from the max to normalize and grow to the mimic scale (instead of the CFG scale)
+                cfg_renormalized = (cfg_clamped / max_scaleref) * mim_scaleref
+
+            ### Now add it back onto the averages to get into real scale again and return
+            result = cfg_renormalized + cfg_means
+
+        actual_res = result.unflatten(2, mim_target.shape[2:])
+
+        if self.interpolate_phi != 1.0:
+            actual_res = actual_res * self.interpolate_phi + cfg_target * (1.0 - self.interpolate_phi)
+
+        if self.experiment_mode == 1:
+            num = actual_res.cpu().numpy()
+            for y in range(0, 64):
+                for x in range (0, 64):
+                    if num[0][0][y][x] > 1.0:
+                        num[0][1][y][x] *= 0.5
+                    if num[0][1][y][x] > 1.0:
+                        num[0][1][y][x] *= 0.5
+                    if num[0][2][y][x] > 1.5:
+                        num[0][2][y][x] *= 0.5
+            actual_res = torch.from_numpy(num).to(device=uncond.device)
+        elif self.experiment_mode == 2:
+            num = actual_res.cpu().numpy()
+            for y in range(0, 64):
+                for x in range (0, 64):
+                    over_scale = False
+                    for z in range(0, 4):
+                        if abs(num[0][z][y][x]) > 1.5:
+                            over_scale = True
+                    if over_scale:
+                        for z in range(0, 4):
+                            num[0][z][y][x] *= 0.7
+            actual_res = torch.from_numpy(num).to(device=uncond.device)
+        elif self.experiment_mode == 3:
+            coefs = torch.tensor([
+                #  R       G        B      W
+                [0.298,   0.207,  0.208, 0.0], # L1
+                [0.187,   0.286,  0.173, 0.0], # L2
+                [-0.158,  0.189,  0.264, 0.0], # L3
+                [-0.184, -0.271, -0.473, 1.0], # L4
+            ], device=uncond.device)
+            res_rgb = torch.einsum("laxy,ab -> lbxy", actual_res, coefs)
+            max_r, max_g, max_b, max_w = res_rgb[0][0].max(), res_rgb[0][1].max(), res_rgb[0][2].max(), res_rgb[0][3].max()
+            max_rgb = max(max_r, max_g, max_b)
+            print(f"test max = r={max_r}, g={max_g}, b={max_b}, w={max_w}, rgb={max_rgb}")
+            if self.step / (self.max_steps - 1) > 0.2:
+                if max_rgb < 2.0 and max_w < 3.0:
+                    res_rgb /= max_rgb / 2.4
+            else:
+                if max_rgb > 2.4 and max_w > 3.0:
+                    res_rgb /= max_rgb / 2.4
+            actual_res = torch.einsum("laxy,ab -> lbxy", res_rgb, coefs.inverse())
+
+        return actual_res

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/easing.py

@@ -0,0 +1,27 @@
+@staticmethod
+def easyIn(t: float)-> float:
+    return t*t
+@staticmethod
+def easyOut(t: float)-> float:
+    return -(t * (t - 2))
+@staticmethod
+def easyInOut(t: float)-> float:
+    if t < 0.5:
+        return 2*t*t
+    else:
+        return (-2*t*t) + (4*t) - 1
+
+class EasingBase:
+
+    def easing(self, t: float, function='linear') -> float:
+        if function == 'easyIn':
+            return easyIn(t)
+        elif function == 'easyOut':
+            return easyOut(t)
+        elif function == 'easyInOut':
+            return easyInOut(t)
+        else:
+            return t
+
+    def ease(self, start, end, t) -> float:
+        return end * t + start * (1 - t)

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/gradual_latent_hires_fix.py

@@ -0,0 +1,273 @@
+import torch
+from torchvision.transforms.functional import gaussian_blur
+from comfy.k_diffusion.sampling import default_noise_sampler, get_ancestral_step, to_d, BrownianTreeNoiseSampler
+from tqdm.auto import trange
+
+@torch.no_grad()
+def sample_euler_ancestral(
+    model,
+    x,
+    sigmas,
+    extra_args=None,
+    callback=None,
+    disable=None,
+    eta=1.0,
+    s_noise=1.0,
+    noise_sampler=None,
+    upscale_ratio=2.0,
+    start_step=5,
+    end_step=15,
+    upscale_n_step=3,
+    unsharp_kernel_size=3,
+    unsharp_sigma=0.5,
+    unsharp_strength=0.0,
+):
+    """Ancestral sampling with Euler method steps."""
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    s_in = x.new_ones([x.shape[0]])
+
+    # make upscale info
+    upscale_steps = []
+    step = start_step - 1
+    while step < end_step - 1:
+        upscale_steps.append(step)
+        step += upscale_n_step
+    height, width = x.shape[2:]
+    upscale_shapes = [
+        (int(height * (((upscale_ratio - 1) / i) + 1)), int(width * (((upscale_ratio - 1) / i) + 1)))
+        for i in reversed(range(1, len(upscale_steps) + 1))
+    ]
+    upscale_info = {k: v for k, v in zip(upscale_steps, upscale_shapes)}
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
+        if callback is not None:
+            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised})
+        d = to_d(x, sigmas[i], denoised)
+        # Euler method
+        dt = sigma_down - sigmas[i]
+        x = x + d * dt
+        if sigmas[i + 1] > 0:
+            # Resize
+            if i in upscale_info:
+                x = torch.nn.functional.interpolate(x, size=upscale_info[i], mode="bicubic", align_corners=False)
+                if unsharp_strength > 0:
+                    blurred = gaussian_blur(x, kernel_size=unsharp_kernel_size, sigma=unsharp_sigma)
+                    x = x + unsharp_strength * (x - blurred)
+
+            noise_sampler = default_noise_sampler(x)
+            noise = noise_sampler(sigmas[i], sigmas[i + 1])
+            x = x + noise * sigma_up * s_noise
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_2s_ancestral(
+    model,
+    x,
+    sigmas,
+    extra_args=None,
+    callback=None,
+    disable=None,
+    eta=1.0,
+    s_noise=1.0,
+    noise_sampler=None,
+    upscale_ratio=2.0,
+    start_step=5,
+    end_step=15,
+    upscale_n_step=3,
+    unsharp_kernel_size=3,
+    unsharp_sigma=0.5,
+    unsharp_strength=0.0,
+):
+    """Ancestral sampling with DPM-Solver++(2S) second-order steps."""
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    sigma_fn = lambda t: t.neg().exp()
+    t_fn = lambda sigma: sigma.log().neg()
+
+    # make upscale info
+    upscale_steps = []
+    step = start_step - 1
+    while step < end_step - 1:
+        upscale_steps.append(step)
+        step += upscale_n_step
+    height, width = x.shape[2:]
+    upscale_shapes = [
+        (int(height * (((upscale_ratio - 1) / i) + 1)), int(width * (((upscale_ratio - 1) / i) + 1)))
+        for i in reversed(range(1, len(upscale_steps) + 1))
+    ]
+    upscale_info = {k: v for k, v in zip(upscale_steps, upscale_shapes)}
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
+        if callback is not None:
+            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised})
+        if sigma_down == 0:
+            # Euler method
+            d = to_d(x, sigmas[i], denoised)
+            dt = sigma_down - sigmas[i]
+            x = x + d * dt
+        else:
+            # DPM-Solver++(2S)
+            t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)
+            r = 1 / 2
+            h = t_next - t
+            s = t + r * h
+            x_2 = (sigma_fn(s) / sigma_fn(t)) * x - (-h * r).expm1() * denoised
+            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
+            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_2
+        # Noise addition
+        if sigmas[i + 1] > 0:
+            # Resize
+            if i in upscale_info:
+                x = torch.nn.functional.interpolate(x, size=upscale_info[i], mode="bicubic", align_corners=False)
+                if unsharp_strength > 0:
+                    blurred = gaussian_blur(x, kernel_size=unsharp_kernel_size, sigma=unsharp_sigma)
+                    x = x + unsharp_strength * (x - blurred)
+            noise_sampler = default_noise_sampler(x)
+            noise = noise_sampler(sigmas[i], sigmas[i + 1])
+            x = x + noise * sigma_up * s_noise
+    return x
+
+
+@torch.no_grad()
+def sample_dpmpp_2m_sde(
+    model,
+    x,
+    sigmas,
+    extra_args=None,
+    callback=None,
+    disable=None,
+    eta=1.0,
+    s_noise=1.0,
+    noise_sampler=None,
+    solver_type="midpoint",
+    upscale_ratio=2.0,
+    start_step=5,
+    end_step=15,
+    upscale_n_step=3,
+    unsharp_kernel_size=3,
+    unsharp_sigma=0.5,
+    unsharp_strength=0.0,
+):
+    """DPM-Solver++(2M) SDE."""
+
+    if solver_type not in {"heun", "midpoint"}:
+        raise ValueError("solver_type must be 'heun' or 'midpoint'")
+
+    seed = extra_args.get("seed", None)
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    old_denoised = None
+    h_last = None
+    h = None
+
+    # make upscale info
+    upscale_steps = []
+    step = start_step - 1
+    while step < end_step - 1:
+        upscale_steps.append(step)
+        step += upscale_n_step
+    height, width = x.shape[2:]
+    upscale_shapes = [
+        (int(height * (((upscale_ratio - 1) / i) + 1)), int(width * (((upscale_ratio - 1) / i) + 1)))
+        for i in reversed(range(1, len(upscale_steps) + 1))
+    ]
+    upscale_info = {k: v for k, v in zip(upscale_steps, upscale_shapes)}
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised})
+        if sigmas[i + 1] == 0:
+            # Denoising step
+            x = denoised
+        else:
+            # DPM-Solver++(2M) SDE
+            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = s - t
+            eta_h = eta * h
+
+            x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised
+
+            if old_denoised is not None:
+                r = h_last / h
+                if solver_type == "heun":
+                    x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised)
+                elif solver_type == "midpoint":
+                    x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised)
+
+            if eta:
+                # Resize
+                if i in upscale_info:
+                    x = torch.nn.functional.interpolate(x, size=upscale_info[i], mode="bicubic", align_corners=False)
+                    if unsharp_strength > 0:
+                        blurred = gaussian_blur(x, kernel_size=unsharp_kernel_size, sigma=unsharp_sigma)
+                        x = x + unsharp_strength * (x - blurred)
+                    denoised = None  # 次ステップとサイズがあわないのでとりあえずNoneにしておく。
+                noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True)
+                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
+
+        old_denoised = denoised
+        h_last = h
+    return x
+
+
+@torch.no_grad()
+def sample_lcm(
+    model,
+    x,
+    sigmas,
+    extra_args=None,
+    callback=None,
+    disable=None,
+    noise_sampler=None,
+    eta=None,
+    s_noise=None,
+    upscale_ratio=2.0,
+    start_step=5,
+    end_step=15,
+    upscale_n_step=3,
+    unsharp_kernel_size=3,
+    unsharp_sigma=0.5,
+    unsharp_strength=0.0,
+):
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    # make upscale info
+    upscale_steps = []
+    step = start_step - 1
+    while step < end_step - 1:
+        upscale_steps.append(step)
+        step += upscale_n_step
+    height, width = x.shape[2:]
+    upscale_shapes = [
+        (int(height * (((upscale_ratio - 1) / i) + 1)), int(width * (((upscale_ratio - 1) / i) + 1)))
+        for i in reversed(range(1, len(upscale_steps) + 1))
+    ]
+    upscale_info = {k: v for k, v in zip(upscale_steps, upscale_shapes)}
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised})
+
+        x = denoised
+        if sigmas[i + 1] > 0:
+            # Resize
+            if i in upscale_info:
+                x = torch.nn.functional.interpolate(x, size=upscale_info[i], mode="bicubic", align_corners=False)
+                if unsharp_strength > 0:
+                    blurred = gaussian_blur(x, kernel_size=unsharp_kernel_size, sigma=unsharp_sigma)
+                    x = x + unsharp_strength * (x - blurred)
+            noise_sampler = default_noise_sampler(x)
+            x += sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1])
+
+    return x

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/image.py

@@ -0,0 +1,227 @@
+import os
+import base64
+import torch
+import numpy as np
+from enum import Enum
+from PIL import Image
+from io import BytesIO
+from typing import List, Union
+
+import folder_paths
+from .utils import install_package
+
+# PIL to Tensor
+def pil2tensor(image):
+  return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+# Tensor to PIL
+def tensor2pil(image):
+    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+# np to Tensor
+def np2tensor(img_np: Union[np.ndarray, List[np.ndarray]]) -> torch.Tensor:
+  if isinstance(img_np, list):
+    return torch.cat([np2tensor(img) for img in img_np], dim=0)
+  return torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0)
+# Tensor to np
+def tensor2np(tensor: torch.Tensor) -> List[np.ndarray]:
+  if len(tensor.shape) == 3:  # Single image
+    return np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8)
+  else:  # Batch of images
+    return [np.clip(255.0 * t.cpu().numpy(), 0, 255).astype(np.uint8) for t in tensor]
+
+def pil2byte(pil_image, format='PNG'):
+  byte_arr = BytesIO()
+  pil_image.save(byte_arr, format=format)
+  byte_arr.seek(0)
+  return byte_arr
+
+def image2base64(image_base64):
+    image_bytes = base64.b64decode(image_base64)
+    image_data = Image.open(BytesIO(image_bytes))
+    return image_data
+
+# Get new bounds
+def get_new_bounds(width, height, left, right, top, bottom):
+  """Returns the new bounds for an image with inset crop data."""
+  left = 0 + left
+  right = width - right
+  top = 0 + top
+  bottom = height - bottom
+  return (left, right, top, bottom)
+
+def RGB2RGBA(image: Image, mask: Image) -> Image:
+  (R, G, B) = image.convert('RGB').split()
+  return Image.merge('RGBA', (R, G, B, mask.convert('L')))
+
+def image2mask(image: Image) -> torch.Tensor:
+  _image = image.convert('RGBA')
+  alpha = _image.split()[0]
+  bg = Image.new("L", _image.size)
+  _image = Image.merge('RGBA', (bg, bg, bg, alpha))
+  ret_mask = torch.tensor([pil2tensor(_image)[0, :, :, 3].tolist()])
+  return ret_mask
+
+def mask2image(mask: torch.Tensor) -> Image:
+  masks = tensor2np(mask)
+  for m in masks:
+    _mask = Image.fromarray(m).convert("L")
+    _image = Image.new("RGBA", _mask.size, color='white')
+    _image = Image.composite(
+      _image, Image.new("RGBA", _mask.size, color='black'), _mask)
+  return _image
+
+# 图像融合
+class blendImage:
+  def g(self, x):
+    return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x))
+
+  def blend_mode(self, img1, img2, mode):
+    if mode == "normal":
+      return img2
+    elif mode == "multiply":
+      return img1 * img2
+    elif mode == "screen":
+      return 1 - (1 - img1) * (1 - img2)
+    elif mode == "overlay":
+      return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
+    elif mode == "soft_light":
+      return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1),
+                         img1 + (2 * img2 - 1) * (self.g(img1) - img1))
+    elif mode == "difference":
+      return img1 - img2
+    else:
+      raise ValueError(f"Unsupported blend mode: {mode}")
+
+  def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str = 'normal'):
+    image2 = image2.to(image1.device)
+    if image1.shape != image2.shape:
+      image2 = image2.permute(0, 3, 1, 2)
+      image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic',
+                                          crop='center')
+      image2 = image2.permute(0, 2, 3, 1)
+
+    blended_image = self.blend_mode(image1, image2, blend_mode)
+    blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor
+    blended_image = torch.clamp(blended_image, 0, 1)
+    return blended_image
+
+
+def empty_image(width, height, batch_size=1, color=0):
+    r = torch.full([batch_size, height, width, 1], ((color >> 16) & 0xFF) / 0xFF)
+    g = torch.full([batch_size, height, width, 1], ((color >> 8) & 0xFF) / 0xFF)
+    b = torch.full([batch_size, height, width, 1], ((color) & 0xFF) / 0xFF)
+    return torch.cat((r, g, b), dim=-1)
+
+
+class ResizeMode(Enum):
+  RESIZE = "Just Resize"
+  INNER_FIT = "Crop and Resize"
+  OUTER_FIT = "Resize and Fill"
+  def int_value(self):
+    if self == ResizeMode.RESIZE:
+      return 0
+    elif self == ResizeMode.INNER_FIT:
+      return 1
+    elif self == ResizeMode.OUTER_FIT:
+      return 2
+    assert False, "NOTREACHED"
+
+# credit by https://github.com/chflame163/ComfyUI_LayerStyle/blob/main/py/imagefunc.py#L591C1-L617C22
+def fit_resize_image(image: Image, target_width: int, target_height: int, fit: str, resize_sampler: str,
+                     background_color: str = '#000000') -> Image:
+  image = image.convert('RGB')
+  orig_width, orig_height = image.size
+  if image is not None:
+    if fit == 'letterbox':
+      if orig_width / orig_height > target_width / target_height:  # 更宽，上下留黑
+        fit_width = target_width
+        fit_height = int(target_width / orig_width * orig_height)
+      else:  # 更瘦，左右留黑
+        fit_height = target_height
+        fit_width = int(target_height / orig_height * orig_width)
+      fit_image = image.resize((fit_width, fit_height), resize_sampler)
+      ret_image = Image.new('RGB', size=(target_width, target_height), color=background_color)
+      ret_image.paste(fit_image, box=((target_width - fit_width) // 2, (target_height - fit_height) // 2))
+    elif fit == 'crop':
+      if orig_width / orig_height > target_width / target_height:  # 更宽，裁左右
+        fit_width = int(orig_height * target_width / target_height)
+        fit_image = image.crop(
+          ((orig_width - fit_width) // 2, 0, (orig_width - fit_width) // 2 + fit_width, orig_height))
+      else:  # 更瘦，裁上下
+        fit_height = int(orig_width * target_height / target_width)
+        fit_image = image.crop(
+          (0, (orig_height - fit_height) // 2, orig_width, (orig_height - fit_height) // 2 + fit_height))
+      ret_image = fit_image.resize((target_width, target_height), resize_sampler)
+    else:
+      ret_image = image.resize((target_width, target_height), resize_sampler)
+  return ret_image
+
+# CLIP反推
+import comfy.utils
+from torchvision import transforms
+Config, Interrogator = None, None
+class CI_Inference:
+  ci_model = None
+  cache_path: str
+
+  def __init__(self):
+    self.ci_model = None
+    self.low_vram = False
+    self.cache_path = os.path.join(folder_paths.models_dir, "clip_interrogator")
+
+  def _load_model(self, model_name, low_vram=False):
+    if not (self.ci_model and model_name == self.ci_model.config.clip_model_name and self.low_vram == low_vram):
+      self.low_vram = low_vram
+      print(f"Load model: {model_name}")
+
+      config = Config(
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        download_cache=True,
+        clip_model_name=model_name,
+        clip_model_path=self.cache_path,
+        cache_path=self.cache_path,
+        caption_model_name='blip-large'
+      )
+
+      if low_vram:
+        config.apply_low_vram_defaults()
+
+      self.ci_model = Interrogator(config)
+
+  def _interrogate(self, image, mode, caption=None):
+    if mode == 'best':
+      prompt = self.ci_model.interrogate(image, caption=caption)
+    elif mode == 'classic':
+      prompt = self.ci_model.interrogate_classic(image, caption=caption)
+    elif mode == 'fast':
+      prompt = self.ci_model.interrogate_fast(image, caption=caption)
+    elif mode == 'negative':
+      prompt = self.ci_model.interrogate_negative(image)
+    else:
+      raise Exception(f"Unknown mode {mode}")
+    return prompt
+
+  def image_to_prompt(self, image, mode, model_name='ViT-L-14/openai', low_vram=False):
+    try:
+      from clip_interrogator import Config, Interrogator
+      global Config, Interrogator
+    except:
+      install_package("clip_interrogator", "0.6.0")
+      from clip_interrogator import Config, Interrogator
+
+    pbar = comfy.utils.ProgressBar(len(image))
+
+    self._load_model(model_name, low_vram)
+    prompt = []
+    for i in range(len(image)):
+      im = image[i]
+
+      im = tensor2pil(im)
+      im = im.convert('RGB')
+
+      _prompt = self._interrogate(im, mode)
+      pbar.update(1)
+      prompt.append(_prompt)
+
+    return prompt
+
+ci = CI_Inference()

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/lllite.py

@@ -0,0 +1,237 @@
+import math
+import torch
+import comfy
+
+
+def extra_options_to_module_prefix(extra_options):
+    # extra_options = {'transformer_index': 2, 'block_index': 8, 'original_shape': [2, 4, 128, 128], 'block': ('input', 7), 'n_heads': 20, 'dim_head': 64}
+
+    # block is: [('input', 4), ('input', 5), ('input', 7), ('input', 8), ('middle', 0),
+    #   ('output', 0), ('output', 1), ('output', 2), ('output', 3), ('output', 4), ('output', 5)]
+    # transformer_index is: [0, 1, 2, 3, 4, 5, 6, 7, 8], for each block
+    # block_index is: 0-1 or 0-9, depends on the block
+    # input 7 and 8, middle has 10 blocks
+
+    # make module name from extra_options
+    block = extra_options["block"]
+    block_index = extra_options["block_index"]
+    if block[0] == "input":
+        module_pfx = f"lllite_unet_input_blocks_{block[1]}_1_transformer_blocks_{block_index}"
+    elif block[0] == "middle":
+        module_pfx = f"lllite_unet_middle_block_1_transformer_blocks_{block_index}"
+    elif block[0] == "output":
+        module_pfx = f"lllite_unet_output_blocks_{block[1]}_1_transformer_blocks_{block_index}"
+    else:
+        raise Exception("invalid block name")
+    return module_pfx
+
+
+def load_control_net_lllite_patch(path, cond_image, multiplier, num_steps, start_percent, end_percent):
+    # calculate start and end step
+    start_step = math.floor(num_steps * start_percent * 0.01) if start_percent > 0 else 0
+    end_step = math.floor(num_steps * end_percent * 0.01) if end_percent > 0 else num_steps
+
+    # load weights
+    ctrl_sd = comfy.utils.load_torch_file(path, safe_load=True)
+
+    # split each weights for each module
+    module_weights = {}
+    for key, value in ctrl_sd.items():
+        fragments = key.split(".")
+        module_name = fragments[0]
+        weight_name = ".".join(fragments[1:])
+
+        if module_name not in module_weights:
+            module_weights[module_name] = {}
+        module_weights[module_name][weight_name] = value
+
+    # load each module
+    modules = {}
+    for module_name, weights in module_weights.items():
+        # ここの自動判定を何とかしたい
+        if "conditioning1.4.weight" in weights:
+            depth = 3
+        elif weights["conditioning1.2.weight"].shape[-1] == 4:
+            depth = 2
+        else:
+            depth = 1
+
+        module = LLLiteModule(
+            name=module_name,
+            is_conv2d=weights["down.0.weight"].ndim == 4,
+            in_dim=weights["down.0.weight"].shape[1],
+            depth=depth,
+            cond_emb_dim=weights["conditioning1.0.weight"].shape[0] * 2,
+            mlp_dim=weights["down.0.weight"].shape[0],
+            multiplier=multiplier,
+            num_steps=num_steps,
+            start_step=start_step,
+            end_step=end_step,
+        )
+        info = module.load_state_dict(weights)
+        modules[module_name] = module
+        if len(modules) == 1:
+            module.is_first = True
+
+    print(f"loaded {path} successfully, {len(modules)} modules")
+
+    # cond imageをセットする
+    cond_image = cond_image.permute(0, 3, 1, 2)  # b,h,w,3 -> b,3,h,w
+    cond_image = cond_image * 2.0 - 1.0  # 0-1 -> -1-+1
+
+    for module in modules.values():
+        module.set_cond_image(cond_image)
+
+    class control_net_lllite_patch:
+        def __init__(self, modules):
+            self.modules = modules
+
+        def __call__(self, q, k, v, extra_options):
+            module_pfx = extra_options_to_module_prefix(extra_options)
+
+            is_attn1 = q.shape[-1] == k.shape[-1]  # self attention
+            if is_attn1:
+                module_pfx = module_pfx + "_attn1"
+            else:
+                module_pfx = module_pfx + "_attn2"
+
+            module_pfx_to_q = module_pfx + "_to_q"
+            module_pfx_to_k = module_pfx + "_to_k"
+            module_pfx_to_v = module_pfx + "_to_v"
+
+            if module_pfx_to_q in self.modules:
+                q = q + self.modules[module_pfx_to_q](q)
+            if module_pfx_to_k in self.modules:
+                k = k + self.modules[module_pfx_to_k](k)
+            if module_pfx_to_v in self.modules:
+                v = v + self.modules[module_pfx_to_v](v)
+
+            return q, k, v
+
+        def to(self, device):
+            for d in self.modules.keys():
+                self.modules[d] = self.modules[d].to(device)
+            return self
+
+    return control_net_lllite_patch(modules)
+
+class LLLiteModule(torch.nn.Module):
+    def __init__(
+        self,
+        name: str,
+        is_conv2d: bool,
+        in_dim: int,
+        depth: int,
+        cond_emb_dim: int,
+        mlp_dim: int,
+        multiplier: int,
+        num_steps: int,
+        start_step: int,
+        end_step: int,
+    ):
+        super().__init__()
+        self.name = name
+        self.is_conv2d = is_conv2d
+        self.multiplier = multiplier
+        self.num_steps = num_steps
+        self.start_step = start_step
+        self.end_step = end_step
+        self.is_first = False
+
+        modules = []
+        modules.append(torch.nn.Conv2d(3, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0))  # to latent (from VAE) size*2
+        if depth == 1:
+            modules.append(torch.nn.ReLU(inplace=True))
+            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0))
+        elif depth == 2:
+            modules.append(torch.nn.ReLU(inplace=True))
+            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=4, stride=4, padding=0))
+        elif depth == 3:
+            # kernel size 8は大きすぎるので、4にする / kernel size 8 is too large, so set it to 4
+            modules.append(torch.nn.ReLU(inplace=True))
+            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0))
+            modules.append(torch.nn.ReLU(inplace=True))
+            modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0))
+
+        self.conditioning1 = torch.nn.Sequential(*modules)
+
+        if self.is_conv2d:
+            self.down = torch.nn.Sequential(
+                torch.nn.Conv2d(in_dim, mlp_dim, kernel_size=1, stride=1, padding=0),
+                torch.nn.ReLU(inplace=True),
+            )
+            self.mid = torch.nn.Sequential(
+                torch.nn.Conv2d(mlp_dim + cond_emb_dim, mlp_dim, kernel_size=1, stride=1, padding=0),
+                torch.nn.ReLU(inplace=True),
+            )
+            self.up = torch.nn.Sequential(
+                torch.nn.Conv2d(mlp_dim, in_dim, kernel_size=1, stride=1, padding=0),
+            )
+        else:
+            self.down = torch.nn.Sequential(
+                torch.nn.Linear(in_dim, mlp_dim),
+                torch.nn.ReLU(inplace=True),
+            )
+            self.mid = torch.nn.Sequential(
+                torch.nn.Linear(mlp_dim + cond_emb_dim, mlp_dim),
+                torch.nn.ReLU(inplace=True),
+            )
+            self.up = torch.nn.Sequential(
+                torch.nn.Linear(mlp_dim, in_dim),
+            )
+
+        self.depth = depth
+        self.cond_image = None
+        self.cond_emb = None
+        self.current_step = 0
+
+    # @torch.inference_mode()
+    def set_cond_image(self, cond_image):
+        # print("set_cond_image", self.name)
+        self.cond_image = cond_image
+        self.cond_emb = None
+        self.current_step = 0
+
+    def forward(self, x):
+        if self.num_steps > 0:
+            if self.current_step < self.start_step:
+                self.current_step += 1
+                return torch.zeros_like(x)
+            elif self.current_step >= self.end_step:
+                if self.is_first and self.current_step == self.end_step:
+                    print(f"end LLLite: step {self.current_step}")
+                self.current_step += 1
+                if self.current_step >= self.num_steps:
+                    self.current_step = 0  # reset
+                return torch.zeros_like(x)
+            else:
+                if self.is_first and self.current_step == self.start_step:
+                    print(f"start LLLite: step {self.current_step}")
+                self.current_step += 1
+                if self.current_step >= self.num_steps:
+                    self.current_step = 0  # reset
+
+        if self.cond_emb is None:
+            # print(f"cond_emb is None, {self.name}")
+            cx = self.conditioning1(self.cond_image.to(x.device, dtype=x.dtype))
+            if not self.is_conv2d:
+                # reshape / b,c,h,w -> b,h*w,c
+                n, c, h, w = cx.shape
+                cx = cx.view(n, c, h * w).permute(0, 2, 1)
+            self.cond_emb = cx
+
+        cx = self.cond_emb
+        # print(f"forward {self.name}, {cx.shape}, {x.shape}")
+
+        # uncond/condでxはバッチサイズが2倍
+        if x.shape[0] != cx.shape[0]:
+            if self.is_conv2d:
+                cx = cx.repeat(x.shape[0] // cx.shape[0], 1, 1, 1)
+            else:
+                # print("x.shape[0] != cx.shape[0]", x.shape[0], cx.shape[0])
+                cx = cx.repeat(x.shape[0] // cx.shape[0], 1, 1)
+
+        cx = torch.cat([cx, self.down(x)], dim=1 if self.is_conv2d else 2)
+        cx = self.mid(cx)
+        cx = self.up(cx)
+        return cx * self.multiplier

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/loader.py

@@ -0,0 +1,570 @@
+import re, time, os, psutil
+import folder_paths
+import comfy.utils
+import comfy.sd
+import comfy.controlnet
+
+from comfy.model_patcher import ModelPatcher
+from nodes import NODE_CLASS_MAPPINGS
+from collections import defaultdict
+from .log import log_node_info, log_node_error
+from ..modules.dit.pixArt.loader import load_pixart
+
+diffusion_loaders = ["easy fullLoader", "easy a1111Loader", "easy fluxLoader", "easy comfyLoader", "easy hunyuanDiTLoader", "easy zero123Loader", "easy svdLoader"]
+stable_cascade_loaders = ["easy cascadeLoader"]
+dit_loaders = ['easy pixArtLoader']
+controlnet_loaders = ["easy controlnetLoader", "easy controlnetLoaderADV", "easy controlnetLoader++"]
+instant_loaders = ["easy instantIDApply", "easy instantIDApplyADV"]
+cascade_vae_node = ["easy preSamplingCascade", "easy fullCascadeKSampler"]
+model_merge_node = ["easy XYInputs: ModelMergeBlocks"]
+lora_widget = ["easy fullLoader", "easy a1111Loader", "easy comfyLoader", "easy fluxLoader"]
+
+class easyLoader:
+    def __init__(self):
+        self.loaded_objects = {
+            "ckpt": defaultdict(tuple),  # {ckpt_name: (model, ...)}
+            "unet": defaultdict(tuple),
+            "clip": defaultdict(tuple),
+            "clip_vision": defaultdict(tuple),
+            "bvae": defaultdict(tuple),
+            "vae": defaultdict(object),
+            "lora": defaultdict(dict),  # {lora_name: {UID: (model_lora, clip_lora)}}
+            "controlnet": defaultdict(dict),
+            "t5": defaultdict(tuple),
+            "chatglm3": defaultdict(tuple),
+        }
+        self.memory_threshold = self.determine_memory_threshold(1)
+        self.lora_name_cache = []
+
+    def clean_values(self, values: str):
+        original_values = values.split("; ")
+        cleaned_values = []
+
+        for value in original_values:
+            cleaned_value = value.strip(';').strip()
+            if cleaned_value == "":
+                continue
+            try:
+                cleaned_value = int(cleaned_value)
+            except ValueError:
+                try:
+                    cleaned_value = float(cleaned_value)
+                except ValueError:
+                    pass
+            cleaned_values.append(cleaned_value)
+
+        return cleaned_values
+
+    def clear_unused_objects(self, desired_names: set, object_type: str):
+        keys = set(self.loaded_objects[object_type].keys())
+        for key in keys - desired_names:
+            del self.loaded_objects[object_type][key]
+
+    def get_input_value(self, entry, key, prompt=None):
+        val = entry["inputs"][key]
+        if isinstance(val, str):
+            return val
+        elif isinstance(val, list):
+            if prompt is not None and val[0]:
+                return prompt[val[0]]['inputs'][key]
+            else:
+                return val[0]
+        else:
+            return str(val)
+
+    def process_pipe_loader(self, entry, desired_ckpt_names, desired_vae_names, desired_lora_names, desired_lora_settings, num_loras=3, suffix=""):
+        for idx in range(1, num_loras + 1):
+            lora_name_key = f"{suffix}lora{idx}_name"
+            desired_lora_names.add(self.get_input_value(entry, lora_name_key))
+            setting = f'{self.get_input_value(entry, lora_name_key)};{entry["inputs"][f"{suffix}lora{idx}_model_strength"]};{entry["inputs"][f"{suffix}lora{idx}_clip_strength"]}'
+            desired_lora_settings.add(setting)
+
+        desired_ckpt_names.add(self.get_input_value(entry, f"{suffix}ckpt_name"))
+        desired_vae_names.add(self.get_input_value(entry, f"{suffix}vae_name"))
+
+    def update_loaded_objects(self, prompt):
+        desired_ckpt_names = set()
+        desired_unet_names = set()
+        desired_clip_names = set()
+        desired_vae_names = set()
+        desired_lora_names = set()
+        desired_lora_settings = set()
+        desired_controlnet_names = set()
+        desired_t5_names = set()
+        desired_glm3_names = set()
+
+        for entry in prompt.values():
+            class_type = entry["class_type"]
+            if class_type in lora_widget:
+                lora_name = self.get_input_value(entry, "lora_name")
+                desired_lora_names.add(lora_name)
+                setting = f'{lora_name};{entry["inputs"]["lora_model_strength"]};{entry["inputs"]["lora_clip_strength"]}'
+                desired_lora_settings.add(setting)
+
+            if class_type in diffusion_loaders:
+                desired_ckpt_names.add(self.get_input_value(entry, "ckpt_name", prompt))
+                desired_vae_names.add(self.get_input_value(entry, "vae_name"))
+
+            elif class_type in ['easy kolorsLoader']:
+                desired_unet_names.add(self.get_input_value(entry, "unet_name"))
+                desired_vae_names.add(self.get_input_value(entry, "vae_name"))
+                desired_glm3_names.add(self.get_input_value(entry, "chatglm3_name"))
+
+            elif class_type in dit_loaders:
+                t5_name = self.get_input_value(entry, "mt5_name") if "mt5_name" in entry["inputs"] else  None
+                clip_name = self.get_input_value(entry, "clip_name") if "clip_name" in entry["inputs"] else None
+                model_name = self.get_input_value(entry, "model_name")
+                ckpt_name = self.get_input_value(entry, "ckpt_name", prompt)
+                if t5_name:
+                    desired_t5_names.add(t5_name)
+                if clip_name:
+                    desired_clip_names.add(clip_name)
+                desired_ckpt_names.add(ckpt_name+'_'+model_name)
+
+            elif class_type in stable_cascade_loaders:
+                desired_unet_names.add(self.get_input_value(entry, "stage_c"))
+                desired_unet_names.add(self.get_input_value(entry, "stage_b"))
+                desired_clip_names.add(self.get_input_value(entry, "clip_name"))
+                desired_vae_names.add(self.get_input_value(entry, "stage_a"))
+
+            elif class_type in cascade_vae_node:
+                encode_vae_name = self.get_input_value(entry, "encode_vae_name")
+                decode_vae_name = self.get_input_value(entry, "decode_vae_name")
+                if encode_vae_name and encode_vae_name != 'None':
+                    desired_vae_names.add(encode_vae_name)
+                if decode_vae_name and decode_vae_name != 'None':
+                    desired_vae_names.add(decode_vae_name)
+
+            elif class_type in controlnet_loaders:
+                control_net_name = self.get_input_value(entry, "control_net_name", prompt)
+                scale_soft_weights = self.get_input_value(entry, "scale_soft_weights")
+                desired_controlnet_names.add(f'{control_net_name};{scale_soft_weights}')
+
+            elif class_type in instant_loaders:
+                control_net_name = self.get_input_value(entry, "control_net_name", prompt)
+                scale_soft_weights = self.get_input_value(entry, "cn_soft_weights")
+                desired_controlnet_names.add(f'{control_net_name};{scale_soft_weights}')
+
+            elif class_type in model_merge_node:
+                desired_ckpt_names.add(self.get_input_value(entry, "ckpt_name_1"))
+                desired_ckpt_names.add(self.get_input_value(entry, "ckpt_name_2"))
+                vae_use = self.get_input_value(entry, "vae_use")
+                if vae_use != 'Use Model 1' and vae_use != 'Use Model 2':
+                    desired_vae_names.add(vae_use)
+
+        object_types = ["ckpt", "unet", "clip", "bvae", "vae", "lora", "controlnet", "t5"]
+        for object_type in object_types:
+            if object_type == 'unet':
+                desired_names = desired_unet_names
+            elif object_type in ["ckpt", "clip", "bvae"]:
+                if object_type == 'clip':
+                    desired_names = desired_ckpt_names.union(desired_clip_names)
+                else:
+                    desired_names = desired_ckpt_names
+            elif object_type == "vae":
+                desired_names = desired_vae_names
+            elif object_type == "controlnet":
+                desired_names = desired_controlnet_names
+            elif object_type == "t5":
+                desired_names = desired_t5_names
+            elif object_type == "chatglm3":
+                desired_names = desired_glm3_names
+            else:
+                desired_names = desired_lora_names
+            self.clear_unused_objects(desired_names, object_type)
+
+    def add_to_cache(self, obj_type, key, value):
+        """
+        Add an item to the cache with the current timestamp.
+        """
+        timestamped_value = (value, time.time())
+        self.loaded_objects[obj_type][key] = timestamped_value
+
+    def determine_memory_threshold(self, percentage=0.8):
+        """
+        Determines the memory threshold as a percentage of the total available memory.
+        Args:
+        - percentage (float): The fraction of total memory to use as the threshold.
+                              Should be a value between 0 and 1. Default is 0.8 (80%).
+        Returns:
+        - memory_threshold (int): Memory threshold in bytes.
+        """
+        total_memory = psutil.virtual_memory().total
+        memory_threshold = total_memory * percentage
+        return memory_threshold
+
+    def get_memory_usage(self):
+        """
+        Returns the memory usage of the current process in bytes.
+        """
+        process = psutil.Process(os.getpid())
+        return process.memory_info().rss
+
+    def eviction_based_on_memory(self):
+        """
+        Evicts objects from cache based on memory usage and priority.
+        """
+        current_memory = self.get_memory_usage()
+        if current_memory < self.memory_threshold:
+            return
+        eviction_order = ["vae", "lora", "bvae", "clip", "ckpt", "controlnet", "unet", "t5", "chatglm3"]
+        for obj_type in eviction_order:
+            if current_memory < self.memory_threshold:
+                break
+            # Sort items based on age (using the timestamp)
+            items = list(self.loaded_objects[obj_type].items())
+            items.sort(key=lambda x: x[1][1])  # Sorting by timestamp
+
+            for item in items:
+                if current_memory < self.memory_threshold:
+                    break
+                del self.loaded_objects[obj_type][item[0]]
+                current_memory = self.get_memory_usage()
+
+    def load_checkpoint(self, ckpt_name, config_name=None, load_vision=False):
+        cache_name = ckpt_name
+        if config_name not in [None, "Default"]:
+            cache_name = ckpt_name + "_" + config_name
+        if cache_name in self.loaded_objects["ckpt"]:
+            clip_vision = self.loaded_objects["clip_vision"][cache_name][0] if load_vision else None
+            clip = self.loaded_objects["clip"][cache_name][0] if not load_vision else None
+            return self.loaded_objects["ckpt"][cache_name][0], clip, self.loaded_objects["bvae"][cache_name][0], clip_vision
+
+        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
+
+        output_clip = False if load_vision else True
+        output_clipvision = True if load_vision else False
+        if config_name not in [None, "Default"]:
+            config_path = folder_paths.get_full_path("configs", config_name)
+            loaded_ckpt = comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings"))
+        else:
+            model_options = {}
+            if re.search("nf4", ckpt_name):
+                from ..modules.bitsandbytes_NF4 import OPS
+                model_options = {"custom_operations": OPS}
+            loaded_ckpt = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=output_clip, output_clipvision=output_clipvision, embedding_directory=folder_paths.get_folder_paths("embeddings"), model_options=model_options)
+
+        self.add_to_cache("ckpt", cache_name, loaded_ckpt[0])
+        self.add_to_cache("bvae", cache_name, loaded_ckpt[2])
+
+        clip = loaded_ckpt[1]
+        clip_vision = loaded_ckpt[3]
+        if clip:
+            self.add_to_cache("clip", cache_name, clip)
+        if clip_vision:
+            self.add_to_cache("clip_vision", cache_name, clip_vision)
+
+        self.eviction_based_on_memory()
+
+        return loaded_ckpt[0], clip, loaded_ckpt[2], clip_vision
+
+    def load_vae(self, vae_name):
+        if vae_name in self.loaded_objects["vae"]:
+            return self.loaded_objects["vae"][vae_name][0]
+
+        vae_path = folder_paths.get_full_path("vae", vae_name)
+        sd = comfy.utils.load_torch_file(vae_path)
+        loaded_vae = comfy.sd.VAE(sd=sd)
+        self.add_to_cache("vae", vae_name, loaded_vae)
+        self.eviction_based_on_memory()
+
+        return loaded_vae
+
+    def load_unet(self, unet_name):
+        if unet_name in self.loaded_objects["unet"]:
+            log_node_info("Load UNet", f"{unet_name} cached")
+            return self.loaded_objects["unet"][unet_name][0]
+
+        unet_path = folder_paths.get_full_path("unet", unet_name)
+        model = comfy.sd.load_unet(unet_path)
+        self.add_to_cache("unet", unet_name, model)
+        self.eviction_based_on_memory()
+
+        return model
+
+    def load_controlnet(self, control_net_name, scale_soft_weights=1, use_cache=True):
+        unique_id = f'{control_net_name};{str(scale_soft_weights)}'
+        if use_cache and unique_id in self.loaded_objects["controlnet"]:
+            return self.loaded_objects["controlnet"][unique_id][0]
+        if scale_soft_weights < 1:
+            if "ScaledSoftControlNetWeights" in NODE_CLASS_MAPPINGS:
+                soft_weight_cls = NODE_CLASS_MAPPINGS['ScaledSoftControlNetWeights']
+                (weights, timestep_keyframe) = soft_weight_cls().load_weights(scale_soft_weights, False)
+                cn_adv_cls = NODE_CLASS_MAPPINGS['ControlNetLoaderAdvanced']
+                control_net, = cn_adv_cls().load_controlnet(control_net_name, timestep_keyframe)
+            else:
+                raise Exception(f"[Advanced-ControlNet Not Found] you need to install 'COMFYUI-Advanced-ControlNet'")
+        else:
+            controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
+            control_net = comfy.controlnet.load_controlnet(controlnet_path)
+        if use_cache:
+            self.add_to_cache("controlnet", unique_id, control_net)
+            self.eviction_based_on_memory()
+
+        return control_net
+    def load_clip(self, clip_name, type='stable_diffusion', load_clip=None):
+        if clip_name in self.loaded_objects["clip"]:
+            return self.loaded_objects["clip"][clip_name][0]
+
+        if type == 'stable_diffusion':
+            clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION
+        elif type == 'stable_cascade':
+            clip_type = comfy.sd.CLIPType.STABLE_CASCADE
+        elif type == 'sd3':
+            clip_type = comfy.sd.CLIPType.SD3
+        elif type == 'flux':
+            clip_type = comfy.sd.CLIPType.FLUX
+        elif type == 'stable_audio':
+            clip_type = comfy.sd.CLIPType.STABLE_AUDIO
+        clip_path = folder_paths.get_full_path("clip", clip_name)
+        load_clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type)
+        self.add_to_cache("clip", clip_name, load_clip)
+        self.eviction_based_on_memory()
+
+        return load_clip
+
+    def load_lora(self, lora, model=None, clip=None, type=None , use_cache=True):
+        lora_name = lora["lora_name"]
+        model = model if model is not None else lora["model"]
+        clip = clip if clip is not None else lora["clip"]
+        model_strength = lora["model_strength"]
+        clip_strength = lora["clip_strength"]
+        lbw = lora["lbw"] if "lbw" in lora else None
+        lbw_a = lora["lbw_a"] if "lbw_a" in lora else None
+        lbw_b = lora["lbw_b"] if "lbw_b" in lora else None
+
+        model_hash = str(model)[44:-1]
+        clip_hash = str(clip)[25:-1] if clip else ''
+
+        unique_id = f'{model_hash};{clip_hash};{lora_name};{model_strength};{clip_strength}'
+
+        if use_cache and unique_id in self.loaded_objects["lora"]:
+            log_node_info("Load LORA",f"{lora_name} cached")
+            return self.loaded_objects["lora"][unique_id][0]
+
+        orig_lora_name = lora_name
+        lora_name = self.resolve_lora_name(lora_name)
+
+        if lora_name is not None:
+            lora_path = folder_paths.get_full_path("loras", lora_name)
+        else:
+            lora_path = None
+
+        if lora_path is not None:
+            log_node_info("Load LORA",f"{lora_name}: model={model_strength:.3f}, clip={clip_strength:.3f}, LBW={lbw}, A={lbw_a}, B={lbw_b}")
+            if lbw:
+                lbw = lora["lbw"]
+                lbw_a = lora["lbw_a"]
+                lbw_b = lora["lbw_b"]
+                if 'LoraLoaderBlockWeight //Inspire' not in NODE_CLASS_MAPPINGS:
+                    raise Exception('[InspirePack Not Found] you need to install ComfyUI-Inspire-Pack')
+                cls = NODE_CLASS_MAPPINGS['LoraLoaderBlockWeight //Inspire']
+                model, clip, _ = cls().doit(model, clip, lora_name, model_strength, clip_strength, False, 0,
+                                            lbw_a, lbw_b, "", lbw)
+            else:
+                _lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
+                keys = _lora.keys()
+                if "down_blocks.0.resnets.0.norm1.bias" in keys:
+                    print('Using LORA for Resadapter')
+                    key_map = {}
+                    key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)
+                    mapping_norm = {}
+
+                    for key in keys:
+                        if ".weight" in key:
+                            key_name_in_ori_sd = key_map[key.replace(".weight", "")]
+                            mapping_norm[key_name_in_ori_sd] = _lora[key]
+                        elif ".bias" in key:
+                            key_name_in_ori_sd = key_map[key.replace(".bias", "")]
+                            mapping_norm[key_name_in_ori_sd.replace(".weight", ".bias")] = _lora[
+                                key
+                            ]
+                        else:
+                            print("===>Unexpected key", key)
+                            mapping_norm[key] = _lora[key]
+
+                    for k in mapping_norm.keys():
+                        if k not in model.model.state_dict():
+                            print("===>Missing key:", k)
+                    model.model.load_state_dict(mapping_norm, strict=False)
+                    return (model, clip)
+
+                # PixArt
+                if type is not None and type == 'PixArt':
+                    from ..modules.dit.pixArt.loader import load_pixart_lora
+                    model = load_pixart_lora(model, _lora, lora_path, model_strength)
+                else:
+                    model, clip = comfy.sd.load_lora_for_models(model, clip, _lora, model_strength, clip_strength)
+
+            if use_cache:
+                self.add_to_cache("lora", unique_id, (model, clip))
+                self.eviction_based_on_memory()
+        else:
+            log_node_error(f"LORA NOT FOUND", orig_lora_name)
+
+        return model, clip
+
+    def resolve_lora_name(self, name):
+        if os.path.exists(name):
+            return name
+        else:
+            if len(self.lora_name_cache) == 0:
+                loras = folder_paths.get_filename_list("loras")
+                self.lora_name_cache.extend(loras)
+            for x in self.lora_name_cache:
+                if x.endswith(name):
+                    return x
+
+            # 如果刷新网页后新添加的lora走这个逻辑
+            log_node_info("LORA NOT IN CACHE", f"{name}")
+            loras = folder_paths.get_filename_list("loras")
+            for x in loras:
+                if x.endswith(name):
+                    self.lora_name_cache.append(x)
+                    return x
+
+            return None
+
+    def load_main(self, ckpt_name, config_name, vae_name, lora_name, lora_model_strength, lora_clip_strength, optional_lora_stack, model_override, clip_override, vae_override, prompt, nf4=False):
+        model: ModelPatcher | None = None
+        clip: comfy.sd.CLIP | None = None
+        vae: comfy.sd.VAE | None = None
+        clip_vision = None
+        lora_stack = []
+
+        # Check for model override
+        can_load_lora = True
+        # 判断是否存在 模型或Lora叠加xyplot, 若存在优先缓存第一个模型
+        # Determine whether there is a model or Lora overlapping xyplot, and if there is, prioritize caching the first model.
+        xy_model_id = next((x for x in prompt if str(prompt[x]["class_type"]) in ["easy XYInputs: ModelMergeBlocks",
+                                                                                  "easy XYInputs: Checkpoint"]), None)
+        # This will find nodes that aren't actively connected to anything, and skip loading lora's for them.
+        xy_lora_id = next((x for x in prompt if str(prompt[x]["class_type"]) == "easy XYInputs: Lora"), None)
+        if xy_lora_id is not None:
+            can_load_lora = False
+        if xy_model_id is not None:
+            node = prompt[xy_model_id]
+            if "ckpt_name_1" in node["inputs"]:
+                ckpt_name_1 = node["inputs"]["ckpt_name_1"]
+                model, clip, vae, clip_vision = self.load_checkpoint(ckpt_name_1)
+                can_load_lora = False
+        elif model_override is not None and clip_override is not None and vae_override is not None:
+            model = model_override
+            clip = clip_override
+            vae = vae_override
+        else:
+            model, clip, vae, clip_vision = self.load_checkpoint(ckpt_name, config_name)
+            if model_override is not None:
+                model = model_override
+            if vae_override is not None:
+                vae = vae_override
+            elif clip_override is not None:
+                clip = clip_override
+
+
+        if optional_lora_stack is not None and can_load_lora:
+            for lora in optional_lora_stack:
+                # This is a subtle bit of code because it uses the model created by the last call, and passes it to the next call.
+                lora = {"lora_name": lora[0], "model": model, "clip": clip, "model_strength": lora[1],
+                        "clip_strength": lora[2]}
+                model, clip = self.load_lora(lora)
+                lora['model'] = model
+                lora['clip'] = clip
+                lora_stack.append(lora)
+
+        if lora_name != "None" and can_load_lora:
+            lora = {"lora_name": lora_name, "model": model, "clip": clip, "model_strength": lora_model_strength,
+                    "clip_strength": lora_clip_strength}
+            model, clip = self.load_lora(lora)
+            lora_stack.append(lora)
+
+        # Check for custom VAE
+        if vae_name not in ["Baked VAE", "Baked-VAE"]:
+            vae = self.load_vae(vae_name)
+        # CLIP skip
+        if not clip:
+            raise Exception("No CLIP found")
+
+        return model, clip, vae, clip_vision, lora_stack
+
+    # Kolors
+    def load_kolors_unet(self, unet_name):
+        if unet_name in self.loaded_objects["unet"]:
+            log_node_info("Load Kolors UNet", f"{unet_name} cached")
+            return self.loaded_objects["unet"][unet_name][0]
+        else:
+            from ..modules.kolors.loader import applyKolorsUnet
+            with applyKolorsUnet():
+                unet_path = folder_paths.get_full_path("unet", unet_name)
+                sd = comfy.utils.load_torch_file(unet_path)
+                model = comfy.sd.load_unet_state_dict(sd)
+                if model is None:
+                    raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
+
+                self.add_to_cache("unet", unet_name, model)
+                self.eviction_based_on_memory()
+
+                return model
+
+    def load_chatglm3(self, chatglm3_name):
+        from ..modules.kolors.loader import load_chatglm3
+        if chatglm3_name in self.loaded_objects["chatglm3"]:
+            log_node_info("Load ChatGLM3", f"{chatglm3_name} cached")
+            return self.loaded_objects["chatglm3"][chatglm3_name][0]
+
+        chatglm_model = load_chatglm3(model_path=folder_paths.get_full_path("llm", chatglm3_name))
+        self.add_to_cache("chatglm3", chatglm3_name, chatglm_model)
+        self.eviction_based_on_memory()
+
+        return chatglm_model
+
+
+    # DiT
+    def load_dit_ckpt(self, ckpt_name, model_name, **kwargs):
+        if (ckpt_name+'_'+model_name) in self.loaded_objects["ckpt"]:
+            return self.loaded_objects["ckpt"][ckpt_name+'_'+model_name][0]
+        model = None
+        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
+        model_type = kwargs['model_type'] if "model_type" in kwargs else 'PixArt'
+        if model_type == 'PixArt':
+            pixart_conf = kwargs['pixart_conf']
+            model_conf = pixart_conf[model_name]
+            model = load_pixart(ckpt_path, model_conf)
+        if model:
+            self.add_to_cache("ckpt", ckpt_name + '_' + model_name, model)
+            self.eviction_based_on_memory()
+        return model
+
+    def load_t5_from_sd3_clip(self, sd3_clip, padding):
+        try:
+            from comfy.text_encoders.sd3_clip import SD3Tokenizer, SD3ClipModel
+        except:
+            from comfy.sd3_clip import SD3Tokenizer, SD3ClipModel
+        import copy
+
+        clip = sd3_clip.clone()
+        assert clip.cond_stage_model.t5xxl is not None, "CLIP must have T5 loaded!"
+
+        # remove transformer
+        transformer = clip.cond_stage_model.t5xxl.transformer
+        clip.cond_stage_model.t5xxl.transformer = None
+
+        # clone object
+        tmp = SD3ClipModel(clip_l=False, clip_g=False, t5=False)
+        tmp.t5xxl = copy.deepcopy(clip.cond_stage_model.t5xxl)
+        # put transformer back
+        clip.cond_stage_model.t5xxl.transformer = transformer
+        tmp.t5xxl.transformer = transformer
+
+        # override special tokens
+        tmp.t5xxl.special_tokens = copy.deepcopy(clip.cond_stage_model.t5xxl.special_tokens)
+        tmp.t5xxl.special_tokens.pop("end")  # make sure empty tokens match
+
+        # tokenizer
+        tok = SD3Tokenizer()
+        tok.t5xxl.min_length = padding
+
+        clip.cond_stage_model = tmp
+        clip.tokenizer = tok
+
+        return clip

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/log.py

@@ -0,0 +1,77 @@
+COLORS_FG = {
+  'BLACK': '\33[30m',
+  'RED': '\33[31m',
+  'GREEN': '\33[32m',
+  'YELLOW': '\33[33m',
+  'BLUE': '\33[34m',
+  'MAGENTA': '\33[35m',
+  'CYAN': '\33[36m',
+  'WHITE': '\33[37m',
+  'GREY': '\33[90m',
+  'BRIGHT_RED': '\33[91m',
+  'BRIGHT_GREEN': '\33[92m',
+  'BRIGHT_YELLOW': '\33[93m',
+  'BRIGHT_BLUE': '\33[94m',
+  'BRIGHT_MAGENTA': '\33[95m',
+  'BRIGHT_CYAN': '\33[96m',
+  'BRIGHT_WHITE': '\33[97m',
+}
+COLORS_STYLE = {
+  'RESET': '\33[0m',
+  'BOLD': '\33[1m',
+  'NORMAL': '\33[22m',
+  'ITALIC': '\33[3m',
+  'UNDERLINE': '\33[4m',
+  'BLINK': '\33[5m',
+  'BLINK2': '\33[6m',
+  'SELECTED': '\33[7m',
+}
+COLORS_BG = {
+  'BLACK': '\33[40m',
+  'RED': '\33[41m',
+  'GREEN': '\33[42m',
+  'YELLOW': '\33[43m',
+  'BLUE': '\33[44m',
+  'MAGENTA': '\33[45m',
+  'CYAN': '\33[46m',
+  'WHITE': '\33[47m',
+  'GREY': '\33[100m',
+  'BRIGHT_RED': '\33[101m',
+  'BRIGHT_GREEN': '\33[102m',
+  'BRIGHT_YELLOW': '\33[103m',
+  'BRIGHT_BLUE': '\33[104m',
+  'BRIGHT_MAGENTA': '\33[105m',
+  'BRIGHT_CYAN': '\33[106m',
+  'BRIGHT_WHITE': '\33[107m',
+}
+
+def log_node_success(node_name, message=None):
+  """Logs a success message."""
+  _log_node(COLORS_FG["GREEN"], node_name, message)
+
+def log_node_info(node_name, message=None):
+  """Logs an info message."""
+  _log_node(COLORS_FG["CYAN"], node_name, message)
+
+
+def log_node_warn(node_name, message=None):
+  """Logs an warn message."""
+  _log_node(COLORS_FG["YELLOW"], node_name, message)
+
+def log_node_error(node_name, message=None):
+  """Logs an warn message."""
+  _log_node(COLORS_FG["RED"], node_name, message)
+
+def log_node(node_name, message=None):
+  """Logs a message."""
+  _log_node(COLORS_FG["CYAN"], node_name, message)
+
+
+def _log_node(color, node_name, message=None, prefix=''):
+  print(_get_log_msg(color, node_name, message, prefix=prefix))
+
+def _get_log_msg(color, node_name, message=None, prefix=''):
+  msg = f'{COLORS_STYLE["BOLD"]}{color}{prefix}[EasyUse] {node_name.replace(" (EasyUse)", "")}'
+  msg += f':{COLORS_STYLE["RESET"]} {message}' if message is not None else f'{COLORS_STYLE["RESET"]}'
+  return msg
+

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/math.py

@@ -0,0 +1,148 @@
+"""
+Math utility functions for formula evaluation
+"""
+import math
+import re
+
+def evaluate_formula(formula: str, a=0, b=0, c=0, d=0):
+    """
+    计算字符串数学公式
+    
+    支持的运算符和函数：
+    - 基本运算：+, -, *, /, //, %, **
+    - 比较运算：>, <, >=, <=, ==, !=
+    - 数学函数：abs, pow, round, ceil, floor, sqrt, exp, log, log10
+    - 三角函数：sin, cos, tan, asin, acos, atan
+    - 常量：pi, e
+    
+    Args:
+        formula: 数学公式字符串，可以使用变量a、b、c、d
+        a: 变量a的值
+        b: 变量b的值
+        c: 变量c的值
+        d: 变量d的值
+    
+    Returns:
+        如果任意输入为list则返回list[float]，否则返回float
+    
+    Examples:
+        >>> evaluate_formula("a + b", 1, 2)
+        3.0
+        >>> evaluate_formula("pow(a, 2)", 5)
+        25.0
+        >>> evaluate_formula("ceil(a / b)", 5, 2)
+        3.0
+        >>> evaluate_formula("(a>b)*b+(a<=b)*a", 5, 3)
+        3.0
+        >>> evaluate_formula("(a>b)*b+(a<=b)*a", 2, 3)
+        2.0
+    """
+    # 安全的数学函数白名单
+    safe_dict = {
+        # 基本运算
+        'abs': abs,
+        'pow': pow,
+        'round': round,
+        # 数学函数
+        'ceil': math.ceil,
+        'floor': math.floor,
+        'sqrt': math.sqrt,
+        'exp': math.exp,
+        'log': math.log,
+        'log10': math.log10,
+        # 三角函数
+        'sin': math.sin,
+        'cos': math.cos,
+        'tan': math.tan,
+        'asin': math.asin,
+        'acos': math.acos,
+        'atan': math.atan,
+        # 常量
+        'pi': math.pi,
+        'e': math.e,
+    }
+
+    # 判断是否有 list 输入
+    list_inputs = {k: v for k, v in {'a': a, 'b': b, 'c': c, 'd': d}.items() if isinstance(v, (list, tuple))}
+    scalar_inputs = {k: v for k, v in {'a': a, 'b': b, 'c': c, 'd': d}.items() if not isinstance(v, (list, tuple))}
+
+    def _eval_single(vals: dict) -> float:
+        env = dict(safe_dict)
+        env.update({k: float(v) for k, v in vals.items()})
+        try:
+            result = eval(formula, {"__builtins__": {}}, env)
+            return float(result)
+        except Exception as e:
+            raise ValueError(f"公式计算错误: {str(e)}")
+
+    if not list_inputs:
+        # 全是标量
+        return _eval_single({k: v for k, v in {'a': a, 'b': b, 'c': c, 'd': d}.items()})
+
+    # 有 list 输入，逐元素计算
+    max_len = max(len(v) for v in list_inputs.values())
+    results = []
+    for i in range(max_len):
+        vals = {k: float(v) for k, v in scalar_inputs.items()}
+        for k, v in list_inputs.items():
+            vals[k] = float(v[i] if i < len(v) else v[-1])
+        results.append(_eval_single(vals))
+    return results
+
+
+def ceil_value(value: float) -> int:
+    """向上取整"""
+    return math.ceil(value)
+
+
+def floor_value(value: float) -> int:
+    """向下取整"""
+    return math.floor(value)
+
+
+def round_value(value: float, decimals: int = 0) -> float:
+    """
+    四舍五入
+    
+    Args:
+        value: 要取整的值
+        decimals: 保留小数位数
+    
+    Returns:
+        四舍五入后的值
+    """
+    return round(value, decimals)
+
+
+def power(base: float, exponent: float) -> float:
+    """计算幂运算"""
+    return math.pow(base, exponent)
+
+
+def sqrt_value(value: float) -> float:
+    """计算平方根"""
+    if value < 0:
+        raise ValueError("不能对负数求平方根")
+    return math.sqrt(value)
+
+
+def add(a: float, b: float) -> float:
+    """加法"""
+    return a + b
+
+
+def subtract(a: float, b: float) -> float:
+    """减法"""
+    return a - b
+
+
+def multiply(a: float, b: float) -> float:
+    """乘法"""
+    return a * b
+
+
+def divide(a: float, b: float) -> float:
+    """除法"""
+    if b == 0:
+        raise ValueError("除数不能为零")
+    return a / b

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/messages.py

@@ -0,0 +1,55 @@
+from server import PromptServer
+from aiohttp import web
+import time
+import json
+
+class MessageCancelled(Exception):
+    pass
+
+class Message:
+    stash = {}
+    messages = {}
+    cancelled = False
+
+    @classmethod
+    def addMessage(cls, id, message):
+        if message == '__cancel__':
+            cls.messages = {}
+            cls.cancelled = True
+        elif message == '__start__':
+            cls.messages = {}
+            cls.stash = {}
+            cls.cancelled = False
+        else:
+            cls.messages[str(id)] = message
+
+    @classmethod
+    def waitForMessage(cls, id, period=0.1, asList=False):
+        sid = str(id)
+        while not (sid in cls.messages) and not ("-1" in cls.messages):
+            if cls.cancelled:
+                cls.cancelled = False
+                raise MessageCancelled()
+            time.sleep(period)
+        if cls.cancelled:
+            cls.cancelled = False
+            raise MessageCancelled()
+        message = cls.messages.pop(str(id), None) or cls.messages.pop("-1")
+        try:
+            if asList:
+                return [str(x.strip()) for x in message.split(",")]
+            else:
+                try:
+                    return json.loads(message)
+                except ValueError:
+                    return message
+        except ValueError:
+            print( f"ERROR IN MESSAGE - failed to parse '${message}' as ${'comma separated list of strings' if asList else 'string'}")
+            return [message] if asList else message
+
+
+@PromptServer.instance.routes.post('/easyuse/message_callback')
+async def message_callback(request):
+    post = await request.post()
+    Message.addMessage(post.get("id"), post.get("message"))
+    return web.json_response({})

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/model.py

@@ -0,0 +1,58 @@
+import json
+import os
+import folder_paths
+import server
+from .utils import find_tags
+
+class easyModelManager:
+
+    def __init__(self):
+        self.img_suffixes = [".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff", ".svg", ".tif", ".tiff"]
+        self.default_suffixes = [".ckpt", ".pt", ".bin", ".pth", ".safetensors"]
+        self.models_config = {
+            "checkpoints": {"suffix": self.default_suffixes},
+            "loras": {"suffix": self.default_suffixes},
+            "unet": {"suffix": self.default_suffixes},
+        }
+        self.model_lists = {}
+
+    def find_thumbnail(self, model_type, name):
+        file_no_ext = os.path.splitext(name)[0]
+        for ext in self.img_suffixes:
+            full_path = folder_paths.get_full_path(model_type, file_no_ext + ext)
+            if os.path.isfile(str(full_path)):
+                return full_path
+        return None
+
+    def get_model_lists(self, model_type):
+        if model_type not in self.models_config:
+            return []
+        filenames = folder_paths.get_filename_list(model_type)
+        model_lists = []
+        for name in filenames:
+            model_suffix = os.path.splitext(name)[-1]
+            if model_suffix not in self.models_config[model_type]["suffix"]:
+                continue
+            else:
+                cfg = {
+                    "name": os.path.basename(os.path.splitext(name)[0]),
+                    "full_name": name,
+                    "remark": '',
+                    "file_path": folder_paths.get_full_path(model_type, name),
+                    "type": model_type,
+                    "suffix": model_suffix,
+                    "dir_tags": find_tags(name),
+                    "cover": self.find_thumbnail(model_type, name),
+                    "metadata": None,
+                    "sha256": None
+                }
+                model_lists.append(cfg)
+
+        return model_lists
+
+    def get_model_info(self, model_type, model_name):
+        pass
+
+# if __name__ == "__main__":
+#     manager = easyModelManager()
+#     print(manager.get_model_lists("checkpoints"))

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/sampler.py

@@ -0,0 +1,1053 @@
+import comfy
+import comfy.model_management
+import comfy.samplers
+import torch
+import numpy as np
+import latent_preview
+from nodes import MAX_RESOLUTION
+from PIL import Image
+from typing import Dict, List, Optional, Tuple, Union, Any
+from ..modules.brushnet.model_patch import add_model_patch
+
+class easySampler:
+    def __init__(self):
+        self.last_helds: dict[str, list] = {
+            "results": [],
+            "pipe_line": [],
+        }
+        self.device = comfy.model_management.intermediate_device()
+
+    @staticmethod
+    def tensor2pil(image: torch.Tensor) -> Image.Image:
+        """Convert a torch tensor to a PIL image."""
+        return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+    @staticmethod
+    def pil2tensor(image: Image.Image) -> torch.Tensor:
+        """Convert a PIL image to a torch tensor."""
+        return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+    @staticmethod
+    def enforce_mul_of_64(d):
+        d = int(d)
+        if d <= 7:
+            d = 8
+        leftover = d % 8  # 8 is the number of pixels per byte
+        if leftover != 0:  # if the number of pixels is not a multiple of 8
+            if (leftover < 4):  # if the number of pixels is less than 4
+                d -= leftover  # remove the leftover pixels
+            else:  # if the number of pixels is more than 4
+                d += 8 - leftover  # add the leftover pixels
+
+        return int(d)
+
+    @staticmethod
+    def safe_split(to_split: str, delimiter: str) -> List[str]:
+        """Split the input string and return a list of non-empty parts."""
+        parts = to_split.split(delimiter)
+        parts = [part for part in parts if part not in ('', ' ', '  ')]
+
+        while len(parts) < 2:
+            parts.append('None')
+        return parts
+
+    def emptyLatent(self, resolution, empty_latent_width, empty_latent_height, batch_size=1, compression=0, model_type='sd', video_length=25):
+        if resolution not in ["自定义 x 自定义", 'width x height (custom)']:
+            try:
+                width, height = map(int, resolution.split(' x '))
+                empty_latent_width = width
+                empty_latent_height = height
+            except ValueError:
+                raise ValueError("Invalid base_resolution format.")
+        if model_type == 'sd3':
+            latent = torch.ones([batch_size, 16, empty_latent_height // 8, empty_latent_width // 8], device=self.device) * 0.0609
+            samples = {"samples": latent}
+        elif model_type == 'mochi':
+            latent = torch.zeros([batch_size, 12, ((video_length - 1) // 6) + 1, empty_latent_height // 8, empty_latent_width // 8], device=self.device)
+            samples = {"samples": latent}
+        elif compression == 0:
+            latent = torch.zeros([batch_size, 4, empty_latent_height // 8, empty_latent_width // 8], device=self.device)
+            samples = {"samples": latent}
+        else:
+            latent_c = torch.zeros(
+                [batch_size, 16, empty_latent_height // compression, empty_latent_width // compression])
+            latent_b = torch.zeros([batch_size, 4, empty_latent_height // 4, empty_latent_width // 4])
+
+            samples = ({"samples": latent_c}, {"samples": latent_b})
+        return samples
+
+    def prepare_noise(self, latent_image, seed, noise_inds=None, noise_device="cpu", incremental_seed_mode="comfy",
+                      variation_seed=None, variation_strength=None):
+        """
+        creates random noise given a latent image and a seed.
+        optional arg skip can be used to skip and discard x number of noise generations for a given seed
+        """
+
+        latent_size = latent_image.size()
+        latent_size_1batch = [1, latent_size[1], latent_size[2], latent_size[3]]
+
+        if variation_strength is not None and variation_strength > 0 or incremental_seed_mode.startswith(
+                "variation str inc"):
+            if noise_device == "cpu":
+                variation_generator = torch.manual_seed(variation_seed)
+            else:
+                torch.cuda.manual_seed(variation_seed)
+                variation_generator = None
+
+            variation_latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout,
+                                           generator=variation_generator, device=noise_device)
+        else:
+            variation_latent = None
+
+        def apply_variation(input_latent, strength_up=None):
+            if variation_latent is None:
+                return input_latent
+            else:
+                strength = variation_strength
+
+                if strength_up is not None:
+                    strength += strength_up
+
+                variation_noise = variation_latent.expand(input_latent.size()[0], -1, -1, -1)
+                result = (1 - strength) * input_latent + strength * variation_noise
+                return result
+
+        # method: incremental seed batch noise
+        if noise_inds is None and incremental_seed_mode == "incremental":
+            batch_cnt = latent_size[0]
+
+            latents = None
+            for i in range(batch_cnt):
+                if noise_device == "cpu":
+                    generator = torch.manual_seed(seed + i)
+                else:
+                    torch.cuda.manual_seed(seed + i)
+                    generator = None
+
+                latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout,
+                                     generator=generator, device=noise_device)
+
+                latent = apply_variation(latent)
+
+                if latents is None:
+                    latents = latent
+                else:
+                    latents = torch.cat((latents, latent), dim=0)
+
+            return latents
+
+        # method: incremental variation batch noise
+        elif noise_inds is None and incremental_seed_mode.startswith("variation str inc"):
+            batch_cnt = latent_size[0]
+
+            latents = None
+            for i in range(batch_cnt):
+                if noise_device == "cpu":
+                    generator = torch.manual_seed(seed)
+                else:
+                    torch.cuda.manual_seed(seed)
+                    generator = None
+
+                latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout,
+                                     generator=generator, device=noise_device)
+
+                step = float(incremental_seed_mode[18:])
+                latent = apply_variation(latent, step * i)
+
+                if latents is None:
+                    latents = latent
+                else:
+                    latents = torch.cat((latents, latent), dim=0)
+
+            return latents
+
+        # method: comfy batch noise
+        if noise_device == "cpu":
+            generator = torch.manual_seed(seed)
+        else:
+            torch.cuda.manual_seed(seed)
+            generator = None
+
+        if noise_inds is None:
+            latents = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout,
+                                  generator=generator, device=noise_device)
+            latents = apply_variation(latents)
+            return latents
+
+        unique_inds, inverse = np.unique(noise_inds, return_inverse=True)
+        noises = []
+        for i in range(unique_inds[-1] + 1):
+            noise = torch.randn([1] + list(latent_image.size())[1:], dtype=latent_image.dtype,
+                                layout=latent_image.layout,
+                                generator=generator, device=noise_device)
+            if i in unique_inds:
+                noises.append(noise)
+        noises = [noises[i] for i in inverse]
+        noises = torch.cat(noises, axis=0)
+        return noises
+
+    def common_ksampler(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0,
+                        disable_noise=False, start_step=None, last_step=None, force_full_denoise=False,
+                        preview_latent=True, disable_pbar=False, noise_device='CPU'):
+        device = comfy.model_management.get_torch_device()
+        noise_device = 'cpu' if noise_device == 'CPU' else device
+        latent_image = latent["samples"]
+        latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image)
+
+        noise_mask = None
+        if "noise_mask" in latent:
+            noise_mask = latent["noise_mask"]
+
+        preview_format = "JPEG"
+        if preview_format not in ["JPEG", "PNG"]:
+            preview_format = "JPEG"
+
+        previewer = False
+
+        if preview_latent:
+            previewer = latent_preview.get_previewer(device, model.model.latent_format)
+
+        pbar = comfy.utils.ProgressBar(steps)
+
+        def callback(step, x0, x, total_steps):
+            preview_bytes = None
+            if previewer:
+                preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
+            pbar.update_absolute(step + 1, total_steps, preview_bytes)
+
+        if disable_noise:
+            noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout,
+                                device=noise_device)
+        else:
+            batch_inds = latent["batch_index"] if "batch_index" in latent else None
+            noise = self.prepare_noise(latent_image, seed, batch_inds, noise_device=noise_device)
+
+        #######################################################################################
+        # add model patch
+        # brushnet
+        add_model_patch(model)
+        #######################################################################################
+        samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative,
+                                      latent_image,
+                                      denoise=denoise, disable_noise=disable_noise, start_step=start_step,
+                                      last_step=last_step,
+                                      force_full_denoise=force_full_denoise, noise_mask=noise_mask,
+                                      callback=callback,
+                                      disable_pbar=disable_pbar, seed=seed)
+        out = latent.copy()
+        out["samples"] = samples
+        return out
+
+    def custom_ksampler(self, model, seed, steps, cfg, _sampler, sigmas, positive, negative, latent,
+                        disable_noise=False, preview_latent=True,  disable_pbar=False, noise_device='CPU'):
+
+        device = comfy.model_management.get_torch_device()
+        noise_device = 'cpu' if noise_device == 'CPU' else device
+
+        latent_image = latent["samples"]
+
+        if disable_noise:
+            noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device=noise_device)
+        else:
+            batch_inds = latent["batch_index"] if "batch_index" in latent else None
+            noise = self.prepare_noise(latent_image, seed, batch_inds, noise_device=noise_device)
+
+        noise_mask = None
+        if "noise_mask" in latent:
+            noise_mask = latent["noise_mask"]
+
+        preview_format = "JPEG"
+        if preview_format not in ["JPEG", "PNG"]:
+            preview_format = "JPEG"
+
+        previewer = False
+
+        if preview_latent:
+            previewer = latent_preview.get_previewer(device, model.model.latent_format)
+
+        pbar = comfy.utils.ProgressBar(steps)
+
+        def callback(step, x0, x, total_steps):
+            preview_bytes = None
+            if previewer:
+                preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
+            pbar.update_absolute(step + 1, total_steps, preview_bytes)
+
+        samples = comfy.samplers.sample(model, noise, positive, negative, cfg, device, _sampler, sigmas, latent_image=latent_image, model_options=model.model_options,
+               denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
+
+        out = latent.copy()
+        out["samples"] = samples
+        return out
+
+    def custom_advanced_ksampler(self, guider, sampler, sigmas, latent_image, add_noise='enable', seed=0, preview_latent=False):
+        latent = latent_image
+        latent_image = latent["samples"]
+        latent = latent.copy()
+        latent_image = comfy.sample.fix_empty_latent_channels(guider.model_patcher, latent_image)
+        latent["samples"] = latent_image
+
+        device = comfy.model_management.get_torch_device()
+        noise_device = device if add_noise == 'enable (GPU=A1111)' else 'cpu'
+
+        if add_noise == 'disable':
+            noise = torch.zeros(latent_image.shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
+        else:
+            batch_inds = latent["batch_index"] if "batch_index" in latent else None
+            noise = self.prepare_noise(latent_image, seed, batch_inds, noise_device=noise_device)
+
+        noise_mask = None
+        if "noise_mask" in latent:
+            noise_mask = latent["noise_mask"]
+
+        x0_output = {}
+        previewer = False
+
+        model = guider.model_patcher
+        steps = sigmas.shape[-1] - 1
+        if preview_latent:
+            previewer = latent_preview.get_previewer(model.load_device, model.model.latent_format)
+
+        pbar = comfy.utils.ProgressBar(steps)
+
+        preview_format = "JPEG"
+        if preview_format not in ["JPEG", "PNG"]:
+            preview_format = "JPEG"
+        def callback(step, x0, x, total_steps):
+            if x0_output is not None:
+                x0_output["x0"] = x0
+
+            preview_bytes = None
+            if previewer:
+                preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
+            pbar.update_absolute(step + 1, total_steps, preview_bytes)
+
+        disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
+        samples = guider.sample(noise, latent_image, sampler, sigmas, denoise_mask=noise_mask,
+                                callback=callback, disable_pbar=disable_pbar, seed=seed)
+        samples = samples.to(comfy.model_management.intermediate_device())
+
+        out = latent.copy()
+        out["samples"] = samples
+        if "x0" in x0_output:
+            out_denoised = latent.copy()
+            out_denoised["samples"] = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())
+        else:
+            out_denoised = out
+        return (out, out_denoised)
+
+    def get_value_by_id(self, key: str, my_unique_id: Any) -> Optional[Any]:
+        """Retrieve value by its associated ID."""
+        try:
+            for value, id_ in self.last_helds[key]:
+                if id_ == my_unique_id:
+                    return value
+        except KeyError:
+            return None
+
+    def update_value_by_id(self, key: str, my_unique_id: Any, new_value: Any) -> Union[bool, None]:
+        """Update the value associated with a given ID. Return True if updated, False if appended, None if key doesn't exist."""
+        try:
+            for i, (value, id_) in enumerate(self.last_helds[key]):
+                if id_ == my_unique_id:
+                    self.last_helds[key][i] = (new_value, id_)
+                    return True
+            self.last_helds[key].append((new_value, my_unique_id))
+            return False
+        except KeyError:
+            return False
+
+    def upscale(self, samples, upscale_method, scale_by, crop):
+        s = samples.copy()
+        width = self.enforce_mul_of_64(round(samples["samples"].shape[3] * scale_by))
+        height = self.enforce_mul_of_64(round(samples["samples"].shape[2] * scale_by))
+
+        if (width > MAX_RESOLUTION):
+            width = MAX_RESOLUTION
+        if (height > MAX_RESOLUTION):
+            height = MAX_RESOLUTION
+
+        s["samples"] = comfy.utils.common_upscale(samples["samples"], width, height, upscale_method, crop)
+        return (s,)
+
+    def handle_upscale(self, samples: dict, upscale_method: str, factor: float, crop: bool) -> dict:
+        """Upscale the samples if the upscale_method is not set to 'None'."""
+        if upscale_method != "None":
+            samples = self.upscale(samples, upscale_method, factor, crop)[0]
+        return samples
+
+    def init_state(self, my_unique_id: Any, key: str, default: Any) -> Any:
+        """Initialize the state by either fetching the stored value or setting a default."""
+        value = self.get_value_by_id(key, my_unique_id)
+        if value is not None:
+            return value
+        return default
+
+    def get_output(self, pipe: dict,) -> Tuple:
+        """Return a tuple of various elements fetched from the input pipe dictionary."""
+        return (
+            pipe,
+            pipe.get("images"),
+            pipe.get("model"),
+            pipe.get("positive"),
+            pipe.get("negative"),
+            pipe.get("samples"),
+            pipe.get("vae"),
+            pipe.get("clip"),
+            pipe.get("seed"),
+        )
+
+    def get_output_sdxl(self, sdxl_pipe: dict) -> Tuple:
+        """Return a tuple of various elements fetched from the input sdxl_pipe dictionary."""
+        return (
+            sdxl_pipe,
+            sdxl_pipe.get("model"),
+            sdxl_pipe.get("positive"),
+            sdxl_pipe.get("negative"),
+            sdxl_pipe.get("vae"),
+            sdxl_pipe.get("refiner_model"),
+            sdxl_pipe.get("refiner_positive"),
+            sdxl_pipe.get("refiner_negative"),
+            sdxl_pipe.get("refiner_vae"),
+            sdxl_pipe.get("samples"),
+            sdxl_pipe.get("clip"),
+            sdxl_pipe.get("images"),
+            sdxl_pipe.get("seed")
+        )
+
+def loglinear_interp(t_steps, num_steps):
+    """
+    Performs log-linear interpolation of a given array of decreasing numbers.
+    """
+    xs = np.linspace(0, 1, len(t_steps))
+    ys = np.log(t_steps[::-1])
+
+    new_xs = np.linspace(0, 1, num_steps)
+    new_ys = np.interp(new_xs, xs, ys)
+
+    interped_ys = np.exp(new_ys)[::-1].copy()
+    return interped_ys
+
+class alignYourStepsScheduler:
+
+    NOISE_LEVELS = {
+        "SD1": [14.6146412293, 6.4745760956, 3.8636745985, 2.6946151520, 1.8841921177, 1.3943805092, 0.9642583904,
+                0.6523686016, 0.3977456272, 0.1515232662, 0.0291671582],
+        "SDXL": [14.6146412293, 6.3184485287, 3.7681790315, 2.1811480769, 1.3405244945, 0.8620721141, 0.5550693289,
+                 0.3798540708, 0.2332364134, 0.1114188177, 0.0291671582],
+        "SVD": [700.00, 54.5, 15.886, 7.977, 4.248, 1.789, 0.981, 0.403, 0.173, 0.034, 0.002]}
+
+    def get_sigmas(self, model_type, steps, denoise):
+
+        total_steps = steps
+        if denoise < 1.0:
+            if denoise <= 0.0:
+                return (torch.FloatTensor([]),)
+            total_steps = round(steps * denoise)
+
+        sigmas = self.NOISE_LEVELS[model_type][:]
+        if (steps + 1) != len(sigmas):
+            sigmas = loglinear_interp(sigmas, steps + 1)
+
+        sigmas = sigmas[-(total_steps + 1):]
+        sigmas[-1] = 0
+        return (torch.FloatTensor(sigmas),)
+
+
+class gitsScheduler:
+
+    NOISE_LEVELS = {
+        0.80: [
+            [14.61464119, 7.49001646, 0.02916753],
+            [14.61464119, 11.54541874, 6.77309084, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 3.07277966, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 2.05039096, 0.02916753],
+            [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 2.05039096, 0.02916753],
+            [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 3.07277966, 1.56271636, 0.02916753],
+            [14.61464119, 12.96784878, 11.54541874, 8.75849152, 7.49001646, 5.85520077, 3.07277966, 1.56271636,
+             0.02916753],
+            [14.61464119, 13.76078796, 12.2308979, 10.90732002, 8.75849152, 7.49001646, 5.85520077, 3.07277966,
+             1.56271636, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 8.75849152, 7.49001646, 5.85520077,
+             3.07277966, 1.56271636, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646,
+             5.85520077, 3.07277966, 1.56271636, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646,
+             6.14220476, 4.86714602, 3.07277966, 1.56271636, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.30717278,
+             7.49001646, 6.14220476, 4.86714602, 3.07277966, 1.56271636, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.24142551,
+             8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.07277966, 1.56271636, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.24142551,
+             8.75849152, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.07277966, 1.56271636, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.24142551,
+             8.75849152, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.1956799, 1.98035145, 0.86115354, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547,
+             9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.1956799, 1.98035145, 0.86115354,
+             0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547,
+             9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.07277966, 1.84880662,
+             0.83188516, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547,
+             9.24142551, 8.75849152, 8.30717278, 7.88507891, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.07277966,
+             1.84880662, 0.83188516, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547,
+             9.24142551, 8.75849152, 8.30717278, 7.88507891, 7.49001646, 6.77309084, 5.85520077, 4.86714602, 3.75677586,
+             2.84484982, 1.78698075, 0.803307, 0.02916753],
+        ],
+        0.85: [
+            [14.61464119, 7.49001646, 0.02916753],
+            [14.61464119, 7.49001646, 1.84880662, 0.02916753],
+            [14.61464119, 11.54541874, 6.77309084, 1.56271636, 0.02916753],
+            [14.61464119, 11.54541874, 7.11996698, 3.07277966, 1.24153244, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.09240818, 2.84484982, 0.95350921, 0.02916753],
+            [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.09240818, 2.84484982, 0.95350921, 0.02916753],
+            [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.58536053, 3.1956799, 1.84880662, 0.803307, 0.02916753],
+            [14.61464119, 12.96784878, 11.54541874, 8.75849152, 7.49001646, 5.58536053, 3.1956799, 1.84880662, 0.803307,
+             0.02916753],
+            [14.61464119, 12.96784878, 11.54541874, 8.75849152, 7.49001646, 6.14220476, 4.65472794, 3.07277966,
+             1.84880662, 0.803307, 0.02916753],
+            [14.61464119, 13.76078796, 12.2308979, 10.90732002, 8.75849152, 7.49001646, 6.14220476, 4.65472794,
+             3.07277966, 1.84880662, 0.803307, 0.02916753],
+            [14.61464119, 13.76078796, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.14220476,
+             4.65472794, 3.07277966, 1.84880662, 0.803307, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646,
+             6.14220476, 4.65472794, 3.07277966, 1.84880662, 0.803307, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.30717278,
+             7.49001646, 6.14220476, 4.65472794, 3.07277966, 1.84880662, 0.803307, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.30717278,
+             7.49001646, 6.14220476, 4.86714602, 3.60512662, 2.6383388, 1.56271636, 0.72133851, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.30717278,
+             7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.46139455, 2.45070267, 1.56271636, 0.72133851,
+             0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.75849152,
+             8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.46139455, 2.45070267, 1.56271636, 0.72133851,
+             0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.24142551,
+             8.75849152, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.46139455, 2.45070267, 1.56271636,
+             0.72133851, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547,
+             9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.46139455, 2.45070267,
+             1.56271636, 0.72133851, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547,
+             9.24142551, 8.75849152, 8.30717278, 7.88507891, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.46139455,
+             2.45070267, 1.56271636, 0.72133851, 0.02916753],
+        ],
+        0.90: [
+            [14.61464119, 6.77309084, 0.02916753],
+            [14.61464119, 7.49001646, 1.56271636, 0.02916753],
+            [14.61464119, 7.49001646, 3.07277966, 0.95350921, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 2.54230714, 0.89115214, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 4.86714602, 2.54230714, 0.89115214, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.09240818, 3.07277966, 1.61558151, 0.69515091, 0.02916753],
+            [14.61464119, 12.2308979, 8.75849152, 7.11996698, 4.86714602, 3.07277966, 1.61558151, 0.69515091,
+             0.02916753],
+            [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 2.95596409, 1.61558151,
+             0.69515091, 0.02916753],
+            [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.19988537, 1.24153244,
+             0.57119018, 0.02916753],
+            [14.61464119, 12.96784878, 10.90732002, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 3.1956799,
+             2.19988537, 1.24153244, 0.57119018, 0.02916753],
+            [14.61464119, 12.96784878, 11.54541874, 9.24142551, 8.30717278, 7.49001646, 5.85520077, 4.45427561,
+             3.1956799, 2.19988537, 1.24153244, 0.57119018, 0.02916753],
+            [14.61464119, 12.96784878, 11.54541874, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 4.86714602,
+             3.75677586, 2.84484982, 1.84880662, 1.08895338, 0.52423614, 0.02916753],
+            [14.61464119, 13.76078796, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.14220476,
+             4.86714602, 3.75677586, 2.84484982, 1.84880662, 1.08895338, 0.52423614, 0.02916753],
+            [14.61464119, 13.76078796, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.44769001,
+             5.58536053, 4.45427561, 3.32507086, 2.45070267, 1.61558151, 0.95350921, 0.45573691, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646,
+             6.44769001, 5.58536053, 4.45427561, 3.32507086, 2.45070267, 1.61558151, 0.95350921, 0.45573691,
+             0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646,
+             6.77309084, 5.85520077, 4.86714602, 3.91689563, 3.07277966, 2.27973175, 1.56271636, 0.95350921, 0.45573691,
+             0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.30717278,
+             7.49001646, 6.77309084, 5.85520077, 4.86714602, 3.91689563, 3.07277966, 2.27973175, 1.56271636, 0.95350921,
+             0.45573691, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.75849152,
+             8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.86714602, 3.91689563, 3.07277966, 2.27973175, 1.56271636,
+             0.95350921, 0.45573691, 0.02916753],
+            [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.75849152,
+             8.30717278, 7.49001646, 6.77309084, 5.85520077, 5.09240818, 4.45427561, 3.60512662, 2.95596409, 2.19988537,
+             1.51179266, 0.89115214, 0.43325692, 0.02916753],
+        ],
+        0.95: [
+            [14.61464119, 6.77309084, 0.02916753],
+            [14.61464119, 6.77309084, 1.56271636, 0.02916753],
+            [14.61464119, 7.49001646, 2.84484982, 0.89115214, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 2.36326075, 0.803307, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 2.95596409, 1.56271636, 0.64427125, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 4.86714602, 2.95596409, 1.56271636, 0.64427125, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 4.86714602, 3.07277966, 1.91321158, 1.08895338, 0.50118381,
+             0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.07277966, 1.91321158, 1.08895338,
+             0.50118381, 0.02916753],
+            [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 3.07277966, 1.91321158,
+             1.08895338, 0.50118381, 0.02916753],
+            [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.19988537, 1.41535246,
+             0.803307, 0.38853383, 0.02916753],
+            [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 4.65472794, 3.46139455, 2.6383388, 1.84880662,
+             1.24153244, 0.72133851, 0.34370604, 0.02916753],
+            [14.61464119, 12.96784878, 10.90732002, 8.75849152, 7.49001646, 5.85520077, 4.65472794, 3.46139455,
+             2.6383388, 1.84880662, 1.24153244, 0.72133851, 0.34370604, 0.02916753],
+            [14.61464119, 12.96784878, 10.90732002, 8.75849152, 7.49001646, 6.14220476, 4.86714602, 3.75677586,
+             2.95596409, 2.19988537, 1.56271636, 1.05362725, 0.64427125, 0.32104823, 0.02916753],
+            [14.61464119, 12.96784878, 10.90732002, 8.75849152, 7.49001646, 6.44769001, 5.58536053, 4.65472794,
+             3.60512662, 2.95596409, 2.19988537, 1.56271636, 1.05362725, 0.64427125, 0.32104823, 0.02916753],
+            [14.61464119, 12.96784878, 11.54541874, 9.24142551, 8.30717278, 7.49001646, 6.44769001, 5.58536053,
+             4.65472794, 3.60512662, 2.95596409, 2.19988537, 1.56271636, 1.05362725, 0.64427125, 0.32104823,
+             0.02916753],
+            [14.61464119, 12.96784878, 11.54541874, 9.24142551, 8.30717278, 7.49001646, 6.44769001, 5.58536053,
+             4.65472794, 3.75677586, 3.07277966, 2.45070267, 1.78698075, 1.24153244, 0.83188516, 0.50118381, 0.22545385,
+             0.02916753],
+            [14.61464119, 12.96784878, 11.54541874, 9.24142551, 8.30717278, 7.49001646, 6.77309084, 5.85520077,
+             5.09240818, 4.45427561, 3.60512662, 2.95596409, 2.36326075, 1.72759056, 1.24153244, 0.83188516, 0.50118381,
+             0.22545385, 0.02916753],
+            [14.61464119, 13.76078796, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.77309084,
+             5.85520077, 5.09240818, 4.45427561, 3.60512662, 2.95596409, 2.36326075, 1.72759056, 1.24153244, 0.83188516,
+             0.50118381, 0.22545385, 0.02916753],
+            [14.61464119, 13.76078796, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.77309084,
+             5.85520077, 5.09240818, 4.45427561, 3.75677586, 3.07277966, 2.45070267, 1.91321158, 1.46270394, 1.05362725,
+             0.72133851, 0.43325692, 0.19894916, 0.02916753],
+        ],
+        1.00: [
+            [14.61464119, 1.56271636, 0.02916753],
+            [14.61464119, 6.77309084, 0.95350921, 0.02916753],
+            [14.61464119, 6.77309084, 2.36326075, 0.803307, 0.02916753],
+            [14.61464119, 7.11996698, 3.07277966, 1.56271636, 0.59516323, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 2.84484982, 1.41535246, 0.57119018, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 2.84484982, 1.61558151, 0.86115354, 0.38853383, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 4.86714602, 2.84484982, 1.61558151, 0.86115354, 0.38853383,
+             0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 4.86714602, 3.07277966, 1.98035145, 1.24153244, 0.72133851,
+             0.34370604, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.07277966, 1.98035145, 1.24153244,
+             0.72133851, 0.34370604, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.27973175, 1.51179266,
+             0.95350921, 0.54755926, 0.25053367, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.36326075, 1.61558151,
+             1.08895338, 0.72133851, 0.41087446, 0.17026083, 0.02916753],
+            [14.61464119, 11.54541874, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.36326075,
+             1.61558151, 1.08895338, 0.72133851, 0.41087446, 0.17026083, 0.02916753],
+            [14.61464119, 11.54541874, 8.75849152, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.84484982,
+             2.12350607, 1.56271636, 1.08895338, 0.72133851, 0.41087446, 0.17026083, 0.02916753],
+            [14.61464119, 11.54541874, 8.75849152, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.84484982,
+             2.19988537, 1.61558151, 1.162866, 0.803307, 0.50118381, 0.27464288, 0.09824532, 0.02916753],
+            [14.61464119, 11.54541874, 8.75849152, 7.49001646, 5.85520077, 4.65472794, 3.75677586, 3.07277966,
+             2.45070267, 1.84880662, 1.36964464, 1.01931262, 0.72133851, 0.45573691, 0.25053367, 0.09824532,
+             0.02916753],
+            [14.61464119, 11.54541874, 8.75849152, 7.49001646, 6.14220476, 5.09240818, 4.26497746, 3.46139455,
+             2.84484982, 2.19988537, 1.67050016, 1.24153244, 0.92192322, 0.64427125, 0.43325692, 0.25053367, 0.09824532,
+             0.02916753],
+            [14.61464119, 11.54541874, 8.75849152, 7.49001646, 6.14220476, 5.09240818, 4.26497746, 3.60512662,
+             2.95596409, 2.45070267, 1.91321158, 1.51179266, 1.12534678, 0.83188516, 0.59516323, 0.38853383, 0.22545385,
+             0.09824532, 0.02916753],
+            [14.61464119, 12.2308979, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 5.09240818, 4.26497746,
+             3.60512662, 2.95596409, 2.45070267, 1.91321158, 1.51179266, 1.12534678, 0.83188516, 0.59516323, 0.38853383,
+             0.22545385, 0.09824532, 0.02916753],
+            [14.61464119, 12.2308979, 9.24142551, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 5.09240818,
+             4.26497746, 3.60512662, 2.95596409, 2.45070267, 1.91321158, 1.51179266, 1.12534678, 0.83188516, 0.59516323,
+             0.38853383, 0.22545385, 0.09824532, 0.02916753],
+        ],
+        1.05: [
+            [14.61464119, 0.95350921, 0.02916753],
+            [14.61464119, 6.77309084, 0.89115214, 0.02916753],
+            [14.61464119, 6.77309084, 2.05039096, 0.72133851, 0.02916753],
+            [14.61464119, 6.77309084, 2.84484982, 1.28281462, 0.52423614, 0.02916753],
+            [14.61464119, 6.77309084, 3.07277966, 1.61558151, 0.803307, 0.34370604, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 2.84484982, 1.56271636, 0.803307, 0.34370604, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 2.84484982, 1.61558151, 0.95350921, 0.52423614, 0.22545385,
+             0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.07277966, 1.98035145, 1.24153244, 0.74807048, 0.41087446,
+             0.17026083, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.27973175, 1.51179266, 0.95350921, 0.59516323, 0.34370604,
+             0.13792117, 0.02916753],
+            [14.61464119, 7.49001646, 5.09240818, 3.46139455, 2.45070267, 1.61558151, 1.08895338, 0.72133851,
+             0.45573691, 0.25053367, 0.09824532, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.09240818, 3.46139455, 2.45070267, 1.61558151, 1.08895338,
+             0.72133851, 0.45573691, 0.25053367, 0.09824532, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.36326075, 1.61558151,
+             1.08895338, 0.72133851, 0.45573691, 0.25053367, 0.09824532, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.45070267, 1.72759056,
+             1.24153244, 0.86115354, 0.59516323, 0.38853383, 0.22545385, 0.09824532, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.84484982, 2.19988537,
+             1.61558151, 1.162866, 0.83188516, 0.59516323, 0.38853383, 0.22545385, 0.09824532, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.84484982, 2.19988537,
+             1.67050016, 1.28281462, 0.95350921, 0.72133851, 0.52423614, 0.34370604, 0.19894916, 0.09824532,
+             0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.95596409, 2.36326075,
+             1.84880662, 1.41535246, 1.08895338, 0.83188516, 0.61951244, 0.45573691, 0.32104823, 0.19894916, 0.09824532,
+             0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.95596409, 2.45070267,
+             1.91321158, 1.51179266, 1.20157266, 0.95350921, 0.74807048, 0.57119018, 0.43325692, 0.29807833, 0.19894916,
+             0.09824532, 0.02916753],
+            [14.61464119, 11.54541874, 8.30717278, 7.11996698, 5.85520077, 4.65472794, 3.60512662, 2.95596409,
+             2.45070267, 1.91321158, 1.51179266, 1.20157266, 0.95350921, 0.74807048, 0.57119018, 0.43325692, 0.29807833,
+             0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 11.54541874, 8.30717278, 7.11996698, 5.85520077, 4.65472794, 3.60512662, 2.95596409,
+             2.45070267, 1.98035145, 1.61558151, 1.32549286, 1.08895338, 0.86115354, 0.69515091, 0.54755926, 0.41087446,
+             0.29807833, 0.19894916, 0.09824532, 0.02916753],
+        ],
+        1.10: [
+            [14.61464119, 0.89115214, 0.02916753],
+            [14.61464119, 2.36326075, 0.72133851, 0.02916753],
+            [14.61464119, 5.85520077, 1.61558151, 0.57119018, 0.02916753],
+            [14.61464119, 6.77309084, 2.45070267, 1.08895338, 0.45573691, 0.02916753],
+            [14.61464119, 6.77309084, 2.95596409, 1.56271636, 0.803307, 0.34370604, 0.02916753],
+            [14.61464119, 6.77309084, 3.07277966, 1.61558151, 0.89115214, 0.4783645, 0.19894916, 0.02916753],
+            [14.61464119, 6.77309084, 3.07277966, 1.84880662, 1.08895338, 0.64427125, 0.34370604, 0.13792117,
+             0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 2.84484982, 1.61558151, 0.95350921, 0.54755926, 0.27464288,
+             0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 2.95596409, 1.91321158, 1.24153244, 0.803307, 0.4783645, 0.25053367,
+             0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.07277966, 2.05039096, 1.41535246, 0.95350921, 0.64427125,
+             0.41087446, 0.22545385, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.27973175, 1.61558151, 1.12534678, 0.803307, 0.54755926,
+             0.36617002, 0.22545385, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.32507086, 2.45070267, 1.72759056, 1.24153244, 0.89115214,
+             0.64427125, 0.45573691, 0.32104823, 0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 5.09240818, 3.60512662, 2.84484982, 2.05039096, 1.51179266, 1.08895338, 0.803307,
+             0.59516323, 0.43325692, 0.29807833, 0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 5.09240818, 3.60512662, 2.84484982, 2.12350607, 1.61558151, 1.24153244,
+             0.95350921, 0.72133851, 0.54755926, 0.41087446, 0.29807833, 0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.45070267, 1.84880662, 1.41535246, 1.08895338,
+             0.83188516, 0.64427125, 0.50118381, 0.36617002, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.45070267, 1.91321158, 1.51179266, 1.20157266,
+             0.95350921, 0.74807048, 0.59516323, 0.45573691, 0.34370604, 0.25053367, 0.17026083, 0.09824532,
+             0.02916753],
+            [14.61464119, 7.49001646, 5.85520077, 4.45427561, 3.46139455, 2.84484982, 2.19988537, 1.72759056,
+             1.36964464, 1.08895338, 0.86115354, 0.69515091, 0.54755926, 0.43325692, 0.34370604, 0.25053367, 0.17026083,
+             0.09824532, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.46139455, 2.84484982, 2.19988537,
+             1.72759056, 1.36964464, 1.08895338, 0.86115354, 0.69515091, 0.54755926, 0.43325692, 0.34370604, 0.25053367,
+             0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.46139455, 2.84484982, 2.19988537,
+             1.72759056, 1.36964464, 1.08895338, 0.89115214, 0.72133851, 0.59516323, 0.4783645, 0.38853383, 0.29807833,
+             0.22545385, 0.17026083, 0.09824532, 0.02916753],
+        ],
+        1.15: [
+            [14.61464119, 0.83188516, 0.02916753],
+            [14.61464119, 1.84880662, 0.59516323, 0.02916753],
+            [14.61464119, 5.85520077, 1.56271636, 0.52423614, 0.02916753],
+            [14.61464119, 5.85520077, 1.91321158, 0.83188516, 0.34370604, 0.02916753],
+            [14.61464119, 5.85520077, 2.45070267, 1.24153244, 0.59516323, 0.25053367, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.51179266, 0.803307, 0.41087446, 0.17026083, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.56271636, 0.89115214, 0.50118381, 0.25053367, 0.09824532,
+             0.02916753],
+            [14.61464119, 6.77309084, 3.07277966, 1.84880662, 1.12534678, 0.72133851, 0.43325692, 0.22545385,
+             0.09824532, 0.02916753],
+            [14.61464119, 6.77309084, 3.07277966, 1.91321158, 1.24153244, 0.803307, 0.52423614, 0.34370604, 0.19894916,
+             0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 2.95596409, 1.91321158, 1.24153244, 0.803307, 0.52423614, 0.34370604,
+             0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.07277966, 2.05039096, 1.36964464, 0.95350921, 0.69515091, 0.4783645,
+             0.32104823, 0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.07277966, 2.12350607, 1.51179266, 1.08895338, 0.803307, 0.59516323,
+             0.43325692, 0.29807833, 0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.07277966, 2.12350607, 1.51179266, 1.08895338, 0.803307, 0.59516323,
+             0.45573691, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.07277966, 2.19988537, 1.61558151, 1.24153244, 0.95350921,
+             0.74807048, 0.59516323, 0.45573691, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.45070267, 1.78698075, 1.32549286, 1.01931262, 0.803307,
+             0.64427125, 0.50118381, 0.38853383, 0.29807833, 0.22545385, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.45070267, 1.78698075, 1.32549286, 1.01931262, 0.803307,
+             0.64427125, 0.52423614, 0.41087446, 0.32104823, 0.25053367, 0.19894916, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.45070267, 1.84880662, 1.41535246, 1.12534678, 0.89115214,
+             0.72133851, 0.59516323, 0.4783645, 0.38853383, 0.32104823, 0.25053367, 0.19894916, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.45070267, 1.84880662, 1.41535246, 1.12534678, 0.89115214,
+             0.72133851, 0.59516323, 0.50118381, 0.41087446, 0.34370604, 0.27464288, 0.22545385, 0.17026083, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.45070267, 1.84880662, 1.41535246, 1.12534678, 0.89115214,
+             0.72133851, 0.59516323, 0.50118381, 0.41087446, 0.34370604, 0.29807833, 0.25053367, 0.19894916, 0.17026083,
+             0.13792117, 0.09824532, 0.02916753],
+        ],
+        1.20: [
+            [14.61464119, 0.803307, 0.02916753],
+            [14.61464119, 1.56271636, 0.52423614, 0.02916753],
+            [14.61464119, 2.36326075, 0.92192322, 0.36617002, 0.02916753],
+            [14.61464119, 2.84484982, 1.24153244, 0.59516323, 0.25053367, 0.02916753],
+            [14.61464119, 5.85520077, 2.05039096, 0.95350921, 0.45573691, 0.17026083, 0.02916753],
+            [14.61464119, 5.85520077, 2.45070267, 1.24153244, 0.64427125, 0.29807833, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.45070267, 1.36964464, 0.803307, 0.45573691, 0.25053367, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.61558151, 0.95350921, 0.59516323, 0.36617002, 0.19894916,
+             0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.67050016, 1.08895338, 0.74807048, 0.50118381, 0.32104823,
+             0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.95596409, 1.84880662, 1.24153244, 0.83188516, 0.59516323, 0.41087446,
+             0.27464288, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 3.07277966, 1.98035145, 1.36964464, 0.95350921, 0.69515091, 0.50118381,
+             0.36617002, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 6.77309084, 3.46139455, 2.36326075, 1.56271636, 1.08895338, 0.803307, 0.59516323, 0.45573691,
+             0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 6.77309084, 3.46139455, 2.45070267, 1.61558151, 1.162866, 0.86115354, 0.64427125, 0.50118381,
+             0.38853383, 0.29807833, 0.22545385, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.12350607, 1.51179266, 1.08895338, 0.83188516,
+             0.64427125, 0.50118381, 0.38853383, 0.29807833, 0.22545385, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.12350607, 1.51179266, 1.08895338, 0.83188516,
+             0.64427125, 0.50118381, 0.41087446, 0.32104823, 0.25053367, 0.19894916, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.12350607, 1.51179266, 1.08895338, 0.83188516,
+             0.64427125, 0.50118381, 0.41087446, 0.34370604, 0.27464288, 0.22545385, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.19988537, 1.61558151, 1.20157266, 0.92192322,
+             0.72133851, 0.57119018, 0.45573691, 0.36617002, 0.29807833, 0.25053367, 0.19894916, 0.17026083, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.19988537, 1.61558151, 1.24153244, 0.95350921,
+             0.74807048, 0.59516323, 0.4783645, 0.38853383, 0.32104823, 0.27464288, 0.22545385, 0.19894916, 0.17026083,
+             0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.19988537, 1.61558151, 1.24153244, 0.95350921,
+             0.74807048, 0.59516323, 0.50118381, 0.41087446, 0.34370604, 0.29807833, 0.25053367, 0.22545385, 0.19894916,
+             0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        ],
+        1.25: [
+            [14.61464119, 0.72133851, 0.02916753],
+            [14.61464119, 1.56271636, 0.50118381, 0.02916753],
+            [14.61464119, 2.05039096, 0.803307, 0.32104823, 0.02916753],
+            [14.61464119, 2.36326075, 0.95350921, 0.43325692, 0.17026083, 0.02916753],
+            [14.61464119, 2.84484982, 1.24153244, 0.59516323, 0.27464288, 0.09824532, 0.02916753],
+            [14.61464119, 3.07277966, 1.51179266, 0.803307, 0.43325692, 0.22545385, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.36326075, 1.24153244, 0.72133851, 0.41087446, 0.22545385, 0.09824532,
+             0.02916753],
+            [14.61464119, 5.85520077, 2.45070267, 1.36964464, 0.83188516, 0.52423614, 0.34370604, 0.19894916,
+             0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.61558151, 0.98595673, 0.64427125, 0.43325692, 0.27464288,
+             0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.67050016, 1.08895338, 0.74807048, 0.52423614, 0.36617002,
+             0.25053367, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.72759056, 1.162866, 0.803307, 0.59516323, 0.45573691, 0.34370604,
+             0.25053367, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.95596409, 1.84880662, 1.24153244, 0.86115354, 0.64427125, 0.4783645, 0.36617002,
+             0.27464288, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.95596409, 1.84880662, 1.28281462, 0.92192322, 0.69515091, 0.52423614,
+             0.41087446, 0.32104823, 0.25053367, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.95596409, 1.91321158, 1.32549286, 0.95350921, 0.72133851, 0.54755926,
+             0.43325692, 0.34370604, 0.27464288, 0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.95596409, 1.91321158, 1.32549286, 0.95350921, 0.72133851, 0.57119018,
+             0.45573691, 0.36617002, 0.29807833, 0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 5.85520077, 2.95596409, 1.91321158, 1.32549286, 0.95350921, 0.74807048, 0.59516323, 0.4783645,
+             0.38853383, 0.32104823, 0.27464288, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 5.85520077, 3.07277966, 2.05039096, 1.41535246, 1.05362725, 0.803307, 0.61951244, 0.50118381,
+             0.41087446, 0.34370604, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 5.85520077, 3.07277966, 2.05039096, 1.41535246, 1.05362725, 0.803307, 0.64427125, 0.52423614,
+             0.43325692, 0.36617002, 0.32104823, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 3.07277966, 2.05039096, 1.46270394, 1.08895338, 0.83188516, 0.66947293,
+             0.54755926, 0.45573691, 0.38853383, 0.34370604, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916,
+             0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        ],
+        1.30: [
+            [14.61464119, 0.72133851, 0.02916753],
+            [14.61464119, 1.24153244, 0.43325692, 0.02916753],
+            [14.61464119, 1.56271636, 0.59516323, 0.22545385, 0.02916753],
+            [14.61464119, 1.84880662, 0.803307, 0.36617002, 0.13792117, 0.02916753],
+            [14.61464119, 2.36326075, 1.01931262, 0.52423614, 0.25053367, 0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.36964464, 0.74807048, 0.41087446, 0.22545385, 0.09824532, 0.02916753],
+            [14.61464119, 3.07277966, 1.56271636, 0.89115214, 0.54755926, 0.34370604, 0.19894916, 0.09824532,
+             0.02916753],
+            [14.61464119, 3.07277966, 1.61558151, 0.95350921, 0.61951244, 0.41087446, 0.27464288, 0.17026083,
+             0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.45070267, 1.36964464, 0.83188516, 0.54755926, 0.36617002, 0.25053367,
+             0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.45070267, 1.41535246, 0.92192322, 0.64427125, 0.45573691, 0.34370604,
+             0.25053367, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.6383388, 1.56271636, 1.01931262, 0.72133851, 0.50118381, 0.36617002, 0.27464288,
+             0.19894916, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.61558151, 1.05362725, 0.74807048, 0.54755926, 0.41087446,
+             0.32104823, 0.25053367, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.61558151, 1.08895338, 0.77538133, 0.57119018, 0.43325692,
+             0.34370604, 0.27464288, 0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.61558151, 1.08895338, 0.803307, 0.59516323, 0.45573691, 0.36617002,
+             0.29807833, 0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.61558151, 1.08895338, 0.803307, 0.59516323, 0.4783645, 0.38853383,
+             0.32104823, 0.27464288, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.72759056, 1.162866, 0.83188516, 0.64427125, 0.50118381, 0.41087446,
+             0.34370604, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.72759056, 1.162866, 0.83188516, 0.64427125, 0.52423614, 0.43325692,
+             0.36617002, 0.32104823, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.78698075, 1.24153244, 0.92192322, 0.72133851, 0.57119018,
+             0.45573691, 0.38853383, 0.34370604, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083,
+             0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.84484982, 1.78698075, 1.24153244, 0.92192322, 0.72133851, 0.57119018, 0.4783645,
+             0.41087446, 0.36617002, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083,
+             0.13792117, 0.09824532, 0.02916753],
+        ],
+        1.35: [
+            [14.61464119, 0.69515091, 0.02916753],
+            [14.61464119, 0.95350921, 0.34370604, 0.02916753],
+            [14.61464119, 1.56271636, 0.57119018, 0.19894916, 0.02916753],
+            [14.61464119, 1.61558151, 0.69515091, 0.29807833, 0.09824532, 0.02916753],
+            [14.61464119, 1.84880662, 0.83188516, 0.43325692, 0.22545385, 0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.162866, 0.64427125, 0.36617002, 0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.36964464, 0.803307, 0.50118381, 0.32104823, 0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.41535246, 0.83188516, 0.54755926, 0.36617002, 0.25053367, 0.17026083,
+             0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.56271636, 0.95350921, 0.64427125, 0.45573691, 0.32104823, 0.22545385,
+             0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.56271636, 0.95350921, 0.64427125, 0.45573691, 0.34370604, 0.25053367,
+             0.19894916, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 3.07277966, 1.61558151, 1.01931262, 0.72133851, 0.52423614, 0.38853383, 0.29807833,
+             0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 3.07277966, 1.61558151, 1.01931262, 0.72133851, 0.52423614, 0.41087446, 0.32104823,
+             0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 3.07277966, 1.61558151, 1.05362725, 0.74807048, 0.54755926, 0.43325692, 0.34370604,
+             0.27464288, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 3.07277966, 1.72759056, 1.12534678, 0.803307, 0.59516323, 0.45573691, 0.36617002, 0.29807833,
+             0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 3.07277966, 1.72759056, 1.12534678, 0.803307, 0.59516323, 0.4783645, 0.38853383, 0.32104823,
+             0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.45070267, 1.51179266, 1.01931262, 0.74807048, 0.57119018, 0.45573691,
+             0.36617002, 0.32104823, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 5.85520077, 2.6383388, 1.61558151, 1.08895338, 0.803307, 0.61951244, 0.50118381, 0.41087446,
+             0.34370604, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 5.85520077, 2.6383388, 1.61558151, 1.08895338, 0.803307, 0.64427125, 0.52423614, 0.43325692,
+             0.36617002, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 5.85520077, 2.6383388, 1.61558151, 1.08895338, 0.803307, 0.64427125, 0.52423614, 0.45573691,
+             0.38853383, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083,
+             0.13792117, 0.09824532, 0.02916753],
+        ],
+        1.40: [
+            [14.61464119, 0.59516323, 0.02916753],
+            [14.61464119, 0.95350921, 0.34370604, 0.02916753],
+            [14.61464119, 1.08895338, 0.43325692, 0.13792117, 0.02916753],
+            [14.61464119, 1.56271636, 0.64427125, 0.27464288, 0.09824532, 0.02916753],
+            [14.61464119, 1.61558151, 0.803307, 0.43325692, 0.22545385, 0.09824532, 0.02916753],
+            [14.61464119, 2.05039096, 0.95350921, 0.54755926, 0.34370604, 0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.24153244, 0.72133851, 0.43325692, 0.27464288, 0.17026083, 0.09824532,
+             0.02916753],
+            [14.61464119, 2.45070267, 1.24153244, 0.74807048, 0.50118381, 0.34370604, 0.25053367, 0.17026083,
+             0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.28281462, 0.803307, 0.52423614, 0.36617002, 0.27464288, 0.19894916, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.28281462, 0.803307, 0.54755926, 0.38853383, 0.29807833, 0.22545385, 0.17026083,
+             0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.41535246, 0.86115354, 0.59516323, 0.43325692, 0.32104823, 0.25053367,
+             0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.51179266, 0.95350921, 0.64427125, 0.45573691, 0.34370604, 0.27464288,
+             0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.51179266, 0.95350921, 0.64427125, 0.4783645, 0.36617002, 0.29807833, 0.25053367,
+             0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.56271636, 0.98595673, 0.69515091, 0.52423614, 0.41087446, 0.34370604,
+             0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.56271636, 1.01931262, 0.72133851, 0.54755926, 0.43325692, 0.36617002,
+             0.32104823, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 2.84484982, 1.61558151, 1.05362725, 0.74807048, 0.57119018, 0.45573691, 0.38853383,
+             0.34370604, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 2.84484982, 1.61558151, 1.08895338, 0.803307, 0.61951244, 0.50118381, 0.41087446, 0.36617002,
+             0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 2.84484982, 1.61558151, 1.08895338, 0.803307, 0.61951244, 0.50118381, 0.43325692, 0.38853383,
+             0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.61558151, 1.08895338, 0.803307, 0.64427125, 0.52423614, 0.45573691, 0.41087446,
+             0.36617002, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083,
+             0.13792117, 0.09824532, 0.02916753],
+        ],
+        1.45: [
+            [14.61464119, 0.59516323, 0.02916753],
+            [14.61464119, 0.803307, 0.25053367, 0.02916753],
+            [14.61464119, 0.95350921, 0.34370604, 0.09824532, 0.02916753],
+            [14.61464119, 1.24153244, 0.54755926, 0.25053367, 0.09824532, 0.02916753],
+            [14.61464119, 1.56271636, 0.72133851, 0.36617002, 0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 1.61558151, 0.803307, 0.45573691, 0.27464288, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 1.91321158, 0.95350921, 0.57119018, 0.36617002, 0.25053367, 0.17026083, 0.09824532,
+             0.02916753],
+            [14.61464119, 2.19988537, 1.08895338, 0.64427125, 0.41087446, 0.27464288, 0.19894916, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.24153244, 0.74807048, 0.50118381, 0.34370604, 0.25053367, 0.19894916,
+             0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.24153244, 0.74807048, 0.50118381, 0.36617002, 0.27464288, 0.22545385,
+             0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.28281462, 0.803307, 0.54755926, 0.41087446, 0.32104823, 0.25053367, 0.19894916,
+             0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.28281462, 0.803307, 0.57119018, 0.43325692, 0.34370604, 0.27464288, 0.22545385,
+             0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.28281462, 0.83188516, 0.59516323, 0.45573691, 0.36617002, 0.29807833,
+             0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.28281462, 0.83188516, 0.59516323, 0.45573691, 0.36617002, 0.32104823,
+             0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.51179266, 0.95350921, 0.69515091, 0.52423614, 0.41087446, 0.34370604,
+             0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 2.84484982, 1.51179266, 0.95350921, 0.69515091, 0.52423614, 0.43325692, 0.36617002,
+             0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 2.84484982, 1.56271636, 0.98595673, 0.72133851, 0.54755926, 0.45573691, 0.38853383,
+             0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.56271636, 1.01931262, 0.74807048, 0.57119018, 0.4783645, 0.41087446, 0.36617002,
+             0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 2.84484982, 1.56271636, 1.01931262, 0.74807048, 0.59516323, 0.50118381, 0.43325692,
+             0.38853383, 0.36617002, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916,
+             0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        ],
+        1.50: [
+            [14.61464119, 0.54755926, 0.02916753],
+            [14.61464119, 0.803307, 0.25053367, 0.02916753],
+            [14.61464119, 0.86115354, 0.32104823, 0.09824532, 0.02916753],
+            [14.61464119, 1.24153244, 0.54755926, 0.25053367, 0.09824532, 0.02916753],
+            [14.61464119, 1.56271636, 0.72133851, 0.36617002, 0.19894916, 0.09824532, 0.02916753],
+            [14.61464119, 1.61558151, 0.803307, 0.45573691, 0.27464288, 0.17026083, 0.09824532, 0.02916753],
+            [14.61464119, 1.61558151, 0.83188516, 0.52423614, 0.34370604, 0.25053367, 0.17026083, 0.09824532,
+             0.02916753],
+            [14.61464119, 1.84880662, 0.95350921, 0.59516323, 0.38853383, 0.27464288, 0.19894916, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 1.84880662, 0.95350921, 0.59516323, 0.41087446, 0.29807833, 0.22545385, 0.17026083,
+             0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 1.84880662, 0.95350921, 0.61951244, 0.43325692, 0.32104823, 0.25053367, 0.19894916,
+             0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.19988537, 1.12534678, 0.72133851, 0.50118381, 0.36617002, 0.27464288, 0.22545385,
+             0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.19988537, 1.12534678, 0.72133851, 0.50118381, 0.36617002, 0.29807833, 0.25053367,
+             0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.36326075, 1.24153244, 0.803307, 0.57119018, 0.43325692, 0.34370604, 0.29807833, 0.25053367,
+             0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.36326075, 1.24153244, 0.803307, 0.57119018, 0.43325692, 0.34370604, 0.29807833, 0.27464288,
+             0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.36326075, 1.24153244, 0.803307, 0.59516323, 0.45573691, 0.36617002, 0.32104823, 0.29807833,
+             0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.36326075, 1.24153244, 0.803307, 0.59516323, 0.45573691, 0.38853383, 0.34370604, 0.32104823,
+             0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532,
+             0.02916753],
+            [14.61464119, 2.45070267, 1.32549286, 0.86115354, 0.64427125, 0.50118381, 0.41087446, 0.36617002,
+             0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117,
+             0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.36964464, 0.92192322, 0.69515091, 0.54755926, 0.45573691, 0.41087446,
+             0.36617002, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083,
+             0.13792117, 0.09824532, 0.02916753],
+            [14.61464119, 2.45070267, 1.41535246, 0.95350921, 0.72133851, 0.57119018, 0.4783645, 0.43325692, 0.38853383,
+             0.36617002, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083,
+             0.13792117, 0.09824532, 0.02916753],
+        ],
+    }
+
+    def get_sigmas(self, coeff, steps, denoise):
+        total_steps = steps
+        if denoise < 1.0:
+            if denoise <= 0.0:
+                return (torch.FloatTensor([]),)
+            total_steps = round(steps * denoise)
+
+        if steps <= 20:
+            sigmas = self.NOISE_LEVELS[round(coeff, 2)][steps-2][:]
+        else:
+            sigmas = self.NOISE_LEVELS[round(coeff, 2)][-1][:]
+            sigmas = loglinear_interp(sigmas, steps + 1)
+
+        sigmas = sigmas[-(total_steps + 1):]
+        sigmas[-1] = 0
+        return (torch.FloatTensor(sigmas), )

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/styleAlign.py

@@ -0,0 +1,148 @@
+import torch
+import torch.nn as nn
+from comfy.model_patcher import ModelPatcher
+from typing import Union
+
+T = torch.Tensor
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d
+
+
+class StyleAlignedArgs:
+    def __init__(self, share_attn: str) -> None:
+        self.adain_keys = "k" in share_attn
+        self.adain_values = "v" in share_attn
+        self.adain_queries = "q" in share_attn
+
+    share_attention: bool = True
+    adain_queries: bool = True
+    adain_keys: bool = True
+    adain_values: bool = True
+
+
+def expand_first(
+    feat: T,
+    scale=1.0,
+) -> T:
+    """
+    Expand the first element so it has the same shape as the rest of the batch.
+    """
+    b = feat.shape[0]
+    feat_style = torch.stack((feat[0], feat[b // 2])).unsqueeze(1)
+    if scale == 1:
+        feat_style = feat_style.expand(2, b // 2, *feat.shape[1:])
+    else:
+        feat_style = feat_style.repeat(1, b // 2, 1, 1, 1)
+        feat_style = torch.cat([feat_style[:, :1], scale * feat_style[:, 1:]], dim=1)
+    return feat_style.reshape(*feat.shape)
+
+
+def concat_first(feat: T, dim=2, scale=1.0) -> T:
+    """
+    concat the the feature and the style feature expanded above
+    """
+    feat_style = expand_first(feat, scale=scale)
+    return torch.cat((feat, feat_style), dim=dim)
+
+
+def calc_mean_std(feat, eps: float = 1e-5) -> "tuple[T, T]":
+    feat_std = (feat.var(dim=-2, keepdims=True) + eps).sqrt()
+    feat_mean = feat.mean(dim=-2, keepdims=True)
+    return feat_mean, feat_std
+
+def adain(feat: T) -> T:
+    feat_mean, feat_std = calc_mean_std(feat)
+    feat_style_mean = expand_first(feat_mean)
+    feat_style_std = expand_first(feat_std)
+    feat = (feat - feat_mean) / feat_std
+    feat = feat * feat_style_std + feat_style_mean
+    return feat
+
+class SharedAttentionProcessor:
+    def __init__(self, args: StyleAlignedArgs, scale: float):
+        self.args = args
+        self.scale = scale
+
+    def __call__(self, q, k, v, extra_options):
+        if self.args.adain_queries:
+            q = adain(q)
+        if self.args.adain_keys:
+            k = adain(k)
+        if self.args.adain_values:
+            v = adain(v)
+        if self.args.share_attention:
+            k = concat_first(k, -2, scale=self.scale)
+            v = concat_first(v, -2)
+
+        return q, k, v
+
+
+def get_norm_layers(
+    layer: nn.Module,
+    norm_layers_: "dict[str, list[Union[nn.GroupNorm, nn.LayerNorm]]]",
+    share_layer_norm: bool,
+    share_group_norm: bool,
+):
+    if isinstance(layer, nn.LayerNorm) and share_layer_norm:
+        norm_layers_["layer"].append(layer)
+    if isinstance(layer, nn.GroupNorm) and share_group_norm:
+        norm_layers_["group"].append(layer)
+    else:
+        for child_layer in layer.children():
+            get_norm_layers(
+                child_layer, norm_layers_, share_layer_norm, share_group_norm
+            )
+
+
+def register_norm_forward(
+    norm_layer: Union[nn.GroupNorm, nn.LayerNorm],
+) -> Union[nn.GroupNorm, nn.LayerNorm]:
+    if not hasattr(norm_layer, "orig_forward"):
+        setattr(norm_layer, "orig_forward", norm_layer.forward)
+    orig_forward = norm_layer.orig_forward
+
+    def forward_(hidden_states: T) -> T:
+        n = hidden_states.shape[-2]
+        hidden_states = concat_first(hidden_states, dim=-2)
+        hidden_states = orig_forward(hidden_states)  # type: ignore
+        return hidden_states[..., :n, :]
+
+    norm_layer.forward = forward_  # type: ignore
+    return norm_layer
+
+
+def register_shared_norm(
+    model: ModelPatcher,
+    share_group_norm: bool = True,
+    share_layer_norm: bool = True,
+):
+    norm_layers = {"group": [], "layer": []}
+    get_norm_layers(model.model, norm_layers, share_layer_norm, share_group_norm)
+    print(
+        f"Patching {len(norm_layers['group'])} group norms, {len(norm_layers['layer'])} layer norms."
+    )
+    return [register_norm_forward(layer) for layer in norm_layers["group"]] + [
+        register_norm_forward(layer) for layer in norm_layers["layer"]
+    ]
+
+
+SHARE_NORM_OPTIONS = ["both", "group", "layer", "disabled"]
+SHARE_ATTN_OPTIONS = ["q+k", "q+k+v", "disabled"]
+
+
+def styleAlignBatch(model, share_norm, share_attn, scale=1.0):
+    m = model.clone()
+    share_group_norm = share_norm in ["group", "both"]
+    share_layer_norm = share_norm in ["layer", "both"]
+    register_shared_norm(model, share_group_norm, share_layer_norm)
+    args = StyleAlignedArgs(share_attn)
+    m.set_model_attn1_patch(SharedAttentionProcessor(args, scale))
+    return m

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/translate.py

@@ -0,0 +1,247 @@
+#credit to shadowcz007 for this module
+#from https://github.com/shadowcz007/comfyui-mixlab-nodes/blob/main/nodes/TextGenerateNode.py
+import re
+import os
+import folder_paths
+
+import comfy.utils
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+from .utils import install_package
+try:
+    from lark import Lark, Transformer, v_args
+except:
+    print('install lark...')
+    install_package('lark')
+    from lark import Lark, Transformer, v_args
+
+model_path = os.path.join(folder_paths.models_dir, 'prompt_generator')
+zh_en_model_path = os.path.join(model_path, 'opus-mt-zh-en')
+zh_en_model, zh_en_tokenizer = None, None
+
+def correct_prompt_syntax(prompt=""):
+    # print("input prompt",prompt)
+    corrected_elements = []
+    # 处理成统一的英文标点
+    prompt = prompt.replace('（', '(').replace('）', ')').replace('，', ',').replace(';', ',').replace('。', '.').replace('：',':').replace('\\',',')
+    # 删除多余的空格
+    prompt = re.sub(r'\s+', ' ', prompt).strip()
+    prompt = prompt.replace("< ","<").replace(" >",">").replace("( ","(").replace(" )",")").replace("[ ","[").replace(' ]',']')
+
+    # 分词
+    prompt_elements = prompt.split(',')
+
+    def balance_brackets(element, open_bracket, close_bracket):
+        open_brackets_count = element.count(open_bracket)
+        close_brackets_count = element.count(close_bracket)
+        return element + close_bracket * (open_brackets_count - close_brackets_count)
+
+    for element in prompt_elements:
+        element = element.strip()
+
+        # 处理空元素
+        if not element:
+            continue
+
+        # 检查并处理圆括号、方括号、尖括号
+        if element[0] in '([':
+            corrected_element = balance_brackets(element, '(', ')') if element[0] == '(' else balance_brackets(element, '[', ']')
+        elif element[0] == '<':
+            corrected_element = balance_brackets(element, '<', '>')
+        else:
+            # 删除开头的右括号或右方括号
+            corrected_element = element.lstrip(')]')
+
+        corrected_elements.append(corrected_element)
+
+    # 重组修正后的prompt
+    return  ','.join(corrected_elements)
+
+def detect_language(input_str):
+    # 统计中文和英文字符的数量
+    count_cn = count_en = 0
+    for char in input_str:
+        if '\u4e00' <= char <= '\u9fff':
+            count_cn += 1
+        elif char.isalpha():
+            count_en += 1
+
+    # 根据统计的字符数量判断主要语言
+    if count_cn > count_en:
+        return "cn"
+    elif count_en > count_cn:
+        return "en"
+    else:
+        return "unknow"
+
+def has_chinese(text):
+    has_cn = False
+    _text = text
+    _text = re.sub(r'<.*?>', '', _text)
+    _text = re.sub(r'__.*?__', '', _text)
+    _text = re.sub(r'embedding:.*?$', '', _text)
+    for char in _text:
+        if '\u4e00' <= char <= '\u9fff':
+            has_cn = True
+            break
+        elif char.isalpha():
+            continue
+    return has_cn
+
+def translate(text):
+    global zh_en_model_path, zh_en_model, zh_en_tokenizer
+
+    if not os.path.exists(zh_en_model_path):
+        zh_en_model_path = 'Helsinki-NLP/opus-mt-zh-en'
+
+    if zh_en_model is None:
+
+        zh_en_model = AutoModelForSeq2SeqLM.from_pretrained(zh_en_model_path).eval()
+        zh_en_tokenizer = AutoTokenizer.from_pretrained(zh_en_model_path, padding=True, truncation=True)
+
+    zh_en_model.to("cuda" if torch.cuda.is_available() else "cpu")
+    with torch.no_grad():
+        encoded = zh_en_tokenizer([text], return_tensors="pt")
+        encoded.to(zh_en_model.device)
+        sequences = zh_en_model.generate(**encoded)
+        return zh_en_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]
+
+@v_args(inline=True)  # Decorator to flatten the tree directly into the function arguments
+class ChinesePromptTranslate(Transformer):
+
+    def sentence(self, *args):
+        return ", ".join(args)
+
+    def phrase(self, *args):
+        return "".join(args)
+
+    def emphasis(self, *args):
+        # Reconstruct the emphasis with translated content
+        return "(" + "".join(args) + ")"
+
+    def weak_emphasis(self, *args):
+        print('weak_emphasis:', args)
+        return "[" + "".join(args) + "]"
+
+    def embedding(self, *args):
+        print('prompt embedding', args[0])
+        if len(args) == 1:
+            embedding_name = str(args[0])
+            return f"embedding:{embedding_name}"
+        elif len(args) > 1:
+            embedding_name, *numbers = args
+
+            if len(numbers) == 2:
+                return f"embedding:{embedding_name}:{numbers[0]}:{numbers[1]}"
+            elif len(numbers) == 1:
+                return f"embedding:{embedding_name}:{numbers[0]}"
+            else:
+                return f"embedding:{embedding_name}"
+
+    def lora(self, *args):
+        if len(args) == 1:
+            return f"<lora:{args[0]}>"
+        elif len(args) > 1:
+            # print('lora', args)
+            _, loar_name, *numbers = args
+            loar_name = str(loar_name).strip()
+            if len(numbers) == 2:
+                return f"<lora:{loar_name}:{numbers[0]}:{numbers[1]}>"
+            elif len(numbers) == 1:
+                return f"<lora:{loar_name}:{numbers[0]}>"
+            else:
+                return f"<lora:{loar_name}>"
+
+    def weight(self, word, number):
+        translated_word = translate(str(word)).rstrip('.')
+        return f"({translated_word}:{str(number).strip()})"
+
+    def schedule(self, *args):
+        print('prompt schedule', args)
+        data = [str(arg).strip() for arg in args]
+
+        return f"[{':'.join(data)}]"
+
+    def word(self, word):
+        # Translate each word using the dictionary
+        word = str(word)
+        match_cn = re.search(r'@.*?@', word)
+        if re.search(r'__.*?__', word):
+            return word.rstrip('.')
+        elif match_cn:
+            chinese = match_cn.group()
+            before = word.split('@', 1)
+            before = before[0] if len(before) > 0 else ''
+            before = translate(str(before)).rstrip('.') if before else ''
+            after = word.rsplit('@', 1)
+            after = after[len(after)-1] if len(after) > 1 else ''
+            after = translate(after).rstrip('.') if after else ''
+            return before + chinese.replace('@', '').rstrip('.') + after
+        elif detect_language(word) == "cn":
+            return translate(word).rstrip('.')
+        else:
+            return word.rstrip('.')
+
+
+#定义Prompt文法
+grammar = r"""
+start: sentence
+sentence: phrase ("," phrase)*
+phrase: emphasis | weight | word | lora | embedding | schedule 
+emphasis: "(" sentence ")" -> emphasis
+        | "[" sentence "]" -> weak_emphasis
+weight: "(" word ":" NUMBER ")"
+schedule: "[" word ":" word ":" NUMBER "]"
+lora: "<" WORD ":" WORD (":" NUMBER)? (":" NUMBER)? ">"
+embedding: "embedding" ":" WORD (":" NUMBER)? (":" NUMBER)?
+word: WORD
+
+NUMBER: /\s*-?\d+(\.\d+)?\s*/
+WORD: /[^,:\(\)\[\]<>]+/
+"""
+def zh_to_en(text):
+    global zh_en_model_path, zh_en_model, zh_en_tokenizer
+    # 进度条
+    pbar = comfy.utils.ProgressBar(len(text) + 1)
+    texts = [correct_prompt_syntax(t) for t in text]
+
+    install_package('sentencepiece', '0.2.0')
+
+    if not os.path.exists(zh_en_model_path):
+        zh_en_model_path = 'Helsinki-NLP/opus-mt-zh-en'
+
+    if zh_en_model is None:
+        zh_en_model = AutoModelForSeq2SeqLM.from_pretrained(zh_en_model_path).eval()
+        zh_en_tokenizer = AutoTokenizer.from_pretrained(zh_en_model_path, padding=True, truncation=True)
+
+    zh_en_model.to("cuda" if torch.cuda.is_available() else "cpu")
+
+    prompt_result = []
+
+    en_texts = []
+
+    for t in texts:
+        if t:
+            # translated_text =  translated_word = translate(zh_en_tokenizer,zh_en_model,str(t))
+            parser = Lark(grammar, start="start", parser="lalr", transformer=ChinesePromptTranslate())
+            # print('t',t)
+            result = parser.parse(t).children
+            # print('en_result',result)
+            # en_text=translate(zh_en_tokenizer,zh_en_model,text_without_syntax)
+            en_texts.append(result[0])
+
+    zh_en_model.to('cpu')
+    # print("test en_text", en_texts)
+    # en_text.to("cuda" if torch.cuda.is_available() else "cpu")
+
+    pbar.update(1)
+    for t in en_texts:
+        prompt_result.append(t)
+        pbar.update(1)
+
+    # print('prompt_result', prompt_result, )
+    if len(prompt_result) == 0:
+        prompt_result = [""]
+
+    return prompt_result

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/utils.py

@@ -0,0 +1,300 @@
+class AlwaysEqualProxy(str):
+    def __eq__(self, _):
+        return True
+
+    def __ne__(self, _):
+        return False
+
+class TautologyStr(str):
+    def __ne__(self, other):
+        return False
+
+class ByPassTypeTuple(tuple):
+    def __getitem__(self, index):
+        if index>0:
+            index=0
+        item = super().__getitem__(index)
+        if isinstance(item, str):
+            return TautologyStr(item)
+        return item
+
+comfy_ui_revision = None
+def get_comfyui_revision():
+    try:
+        import git
+        import os
+        import folder_paths
+        repo = git.Repo(os.path.dirname(folder_paths.__file__))
+        comfy_ui_revision = len(list(repo.iter_commits('HEAD')))
+    except:
+        comfy_ui_revision = "Unknown"
+    return comfy_ui_revision
+
+
+import sys
+import importlib.util
+import importlib.metadata
+import comfy.model_management as mm
+import logging
+import gc
+from packaging import version
+from server import PromptServer
+
+LOG = logging.getLogger(__name__)
+
+def is_package_installed(package):
+    try:
+        module = importlib.util.find_spec(package)
+        return module is not None
+    except ImportError as e:
+        print(e)
+        return False
+
+def install_package(package, v=None, compare=True, compare_version=None):
+    run_install = True
+    if is_package_installed(package):
+        try:
+            installed_version = importlib.metadata.version(package)
+            if v is not None:
+                if compare_version is None:
+                    compare_version = v
+                if not compare or version.parse(installed_version) >= version.parse(compare_version):
+                    run_install = False
+            else:
+                run_install = False
+        except:
+            run_install = False
+
+    if run_install:
+        import subprocess
+        package_command = package + '==' + v if v is not None else package
+        PromptServer.instance.send_sync("easyuse-toast", {'content': f"Installing {package_command}...", 'duration': 5000})
+        result = subprocess.run([sys.executable, '-s', '-m', 'pip', 'install', package_command], capture_output=True, text=True)
+        if result.returncode == 0:
+            PromptServer.instance.send_sync("easyuse-toast", {'content': f"{package} installed successfully", 'type': 'success', 'duration': 5000})
+            print(f"Package {package} installed successfully")
+            return True
+        else:
+            PromptServer.instance.send_sync("easyuse-toast", {'content': f"{package} installed failed", 'type': 'error', 'duration': 5000})
+            print(f"Package {package} installed failed")
+            return False
+    else:
+        return False
+
+def compare_revision(num):
+    global comfy_ui_revision
+    if not comfy_ui_revision:
+        comfy_ui_revision = get_comfyui_revision()
+    return True if comfy_ui_revision == 'Unknown' or int(comfy_ui_revision) >= num else False
+
+def find_tags(string: str, sep="/") -> list[str]:
+    """
+    find tags from string use the sep for split
+    Note: string may contain the \\ or / for path separator
+    """
+    if not string:
+        return []
+    string = string.replace("\\", "/")
+    while "//" in string:
+        string = string.replace("//", "/")
+    if string and sep in string:
+        return string.split(sep)[:-1]
+    return []
+
+
+from comfy.model_base import BaseModel
+import comfy.supported_models
+import comfy.supported_models_base
+def get_sd_version(model):
+    base: BaseModel = model.model
+    model_config: comfy.supported_models.supported_models_base.BASE = base.model_config
+    if isinstance(model_config, comfy.supported_models.SDXL):
+        return 'sdxl'
+    elif isinstance(model_config, comfy.supported_models.SDXLRefiner):
+        return 'sdxl_refiner'
+    elif isinstance(
+            model_config, (comfy.supported_models.SD15, comfy.supported_models.SD20)
+    ):
+        return 'sd1'
+    elif isinstance(
+            model_config, (comfy.supported_models.SVD_img2vid)
+    ):
+        return 'svd'
+    elif isinstance(model_config, comfy.supported_models.SD3):
+        return 'sd3'
+    elif isinstance(model_config, comfy.supported_models.HunyuanDiT):
+        return 'hydit'
+    elif isinstance(model_config, comfy.supported_models.Flux):
+        return 'flux'
+    elif isinstance(model_config, comfy.supported_models.GenmoMochi):
+        return 'mochi'
+    else:
+        return 'unknown'
+
+def find_nearest_steps(clip_id, prompt):
+    """Find the nearest KSampler or preSampling node that references the given id."""
+    def check_link_to_clip(node_id, clip_id, visited=None, node=None):
+        """Check if a given node links directly or indirectly to a loader node."""
+        if visited is None:
+            visited = set()
+
+        if node_id in visited:
+            return False
+        visited.add(node_id)
+        if "pipe" in node["inputs"]:
+            link_ids = node["inputs"]["pipe"]
+            for id in link_ids:
+                if id != 0 and id == str(clip_id):
+                    return True
+        return False
+
+    for id in prompt:
+        node = prompt[id]
+        if "Sampler" in node["class_type"] or "sampler" in node["class_type"] or "Sampling" in node["class_type"]:
+            # Check if this KSampler node directly or indirectly references the given CLIPTextEncode node
+            if check_link_to_clip(id, clip_id, None, node):
+                steps = node["inputs"]["steps"] if "steps" in node["inputs"] else 1
+                return steps
+    return 1
+
+def find_wildcards_seed(clip_id, text, prompt):
+    """ Find easy wildcards seed value"""
+    def find_link_clip_id(id, seed, wildcard_id):
+        node = prompt[id]
+        if "positive" in node['inputs']:
+            link_ids = node["inputs"]["positive"]
+            if type(link_ids) == list:
+                for id in link_ids:
+                    if id != 0:
+                        if id == wildcard_id:
+                            wildcard_node = prompt[wildcard_id]
+                            seed = wildcard_node["inputs"]["seed"] if "seed" in wildcard_node["inputs"] else None
+                            if seed is None:
+                                seed = wildcard_node["inputs"]["seed_num"] if "seed_num" in wildcard_node["inputs"] else None
+                            return seed
+                        else:
+                            return find_link_clip_id(id, seed, wildcard_id)
+            else:
+                return None
+        else:
+            return None
+    if "__" in text:
+        seed = None
+        for id in prompt:
+            node = prompt[id]
+            if "wildcards" in node["class_type"]:
+                wildcard_id = id
+                return find_link_clip_id(str(clip_id), seed, wildcard_id)
+        return seed
+    else:
+        return None
+
+def is_linked_styles_selector(prompt, unique_id, prompt_type='positive'):
+    unique_id = unique_id.split('.')[len(unique_id.split('.')) - 1] if "." in unique_id else unique_id
+    inputs_values = prompt[unique_id]['inputs'][prompt_type] if prompt_type in prompt[unique_id][
+        'inputs'] else None
+    if type(inputs_values) == list and inputs_values != 'undefined' and inputs_values[0]:
+        return True if prompt[inputs_values[0]] and prompt[inputs_values[0]]['class_type'] == 'easy stylesSelector' else False
+    else:
+        return False
+
+use_mirror = False
+def get_local_filepath(url, dirname, local_file_name=None):
+    """Get local file path when is already downloaded or download it"""
+    import os
+    from server import PromptServer
+    from urllib.parse import urlparse
+    from torch.hub import download_url_to_file
+    global use_mirror
+    if not os.path.exists(dirname):
+        os.makedirs(dirname)
+    if not local_file_name:
+        parsed_url = urlparse(url)
+        local_file_name = os.path.basename(parsed_url.path)
+    destination = os.path.join(dirname, local_file_name)
+    if not os.path.exists(destination):
+        try:
+            if use_mirror:
+                url = url.replace('huggingface.co', 'hf-mirror.com')
+            print(f'downloading {url} to {destination}')
+            PromptServer.instance.send_sync("easyuse-toast", {'content': f'Downloading model to {destination}, please wait...', 'duration': 10000})
+            download_url_to_file(url, destination)
+        except Exception as e:
+            use_mirror = True
+            url = url.replace('huggingface.co', 'hf-mirror.com')
+            print(f'Unable to download from huggingface, trying mirror: {url}')
+            PromptServer.instance.send_sync("easyuse-toast", {'content': f'Unable to connect to huggingface, trying mirror: {url}', 'duration': 10000})
+            try:
+                download_url_to_file(url, destination)
+            except Exception as err:
+                error_msg = str(err.args[0]) if err.args else str(err)
+                PromptServer.instance.send_sync("easyuse-toast",
+                                                {'content': f'Unable to download model from {url}', 'type':'error'})
+                raise Exception(f'Download failed. Original URL and mirror both failed.\nError: {error_msg}')
+    return destination
+
+def to_lora_patch_dict(state_dict: dict) -> dict:
+    """ Convert raw lora state_dict to patch_dict that can be applied on
+    modelpatcher."""
+    patch_dict = {}
+    for k, w in state_dict.items():
+        model_key, patch_type, weight_index = k.split('::')
+        if model_key not in patch_dict:
+            patch_dict[model_key] = {}
+        if patch_type not in patch_dict[model_key]:
+            patch_dict[model_key][patch_type] = [None] * 16
+        patch_dict[model_key][patch_type][int(weight_index)] = w
+
+    patch_flat = {}
+    for model_key, v in patch_dict.items():
+        for patch_type, weight_list in v.items():
+            patch_flat[model_key] = (patch_type, weight_list)
+
+    return patch_flat
+
+def easySave(images, filename_prefix, output_type, prompt=None, extra_pnginfo=None):
+    """Save or Preview Image"""
+    from nodes import PreviewImage, SaveImage
+    if output_type in ["Hide", "None"]:
+        return list()
+    elif output_type in ["Preview", "Preview&Choose"]:
+        filename_prefix = 'easyPreview'
+        results = PreviewImage().save_images(images, filename_prefix, prompt, extra_pnginfo)
+        return results['ui']['images']
+    else:
+        results = SaveImage().save_images(images, filename_prefix, prompt, extra_pnginfo)
+        return results['ui']['images']
+
+def getMetadata(filepath):
+    with open(filepath, "rb") as file:
+        # https://github.com/huggingface/safetensors#format
+        # 8 bytes: N, an unsigned little-endian 64-bit integer, containing the size of the header
+        header_size = int.from_bytes(file.read(8), "little", signed=False)
+
+        if header_size <= 0:
+            raise BufferError("Invalid header size")
+
+        header = file.read(header_size)
+        if header_size <= 0:
+            raise BufferError("Invalid header")
+
+        return header
+
+def cleanGPUUsedForce():
+    from .cache import remove_cache
+
+    remove_cache("*")
+    gc.collect()
+    try:
+        import torch
+    except (ImportError, OSError, RuntimeError) as exc:
+        LOG.debug("Skipping CUDA synchronize during cleanGPUUsedForce: torch import failed: %s", exc)
+    else:
+        try:
+            if torch.cuda.is_available():
+                torch.cuda.synchronize()
+        except (AttributeError, OSError, RuntimeError) as exc:
+            LOG.debug("Skipping CUDA synchronize during cleanGPUUsedForce: %s", exc)
+
+    mm.unload_all_models()
+    mm.soft_empty_cache()

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/wildcards.py

@@ -0,0 +1,476 @@
+import json
+import os
+import random
+import re
+from math import prod
+
+import yaml
+
+import folder_paths
+
+from .log import log_node_info
+
+easy_wildcard_dict = {}
+
+def get_wildcard_list():
+    return [f"__{x}__" for x in easy_wildcard_dict.keys()]
+
+def wildcard_normalize(x):
+    return x.replace("\\", "/").lower()
+
+def read_wildcard(k, v):
+    if isinstance(v, list):
+        k = wildcard_normalize(k)
+        easy_wildcard_dict[k] = v
+    elif isinstance(v, dict):
+        for k2, v2 in v.items():
+            new_key = f"{k}/{k2}"
+            new_key = wildcard_normalize(new_key)
+            read_wildcard(new_key, v2)
+
+def read_wildcard_dict(wildcard_path):
+    global easy_wildcard_dict
+    for root, directories, files in os.walk(wildcard_path, followlinks=True):
+        for file in files:
+            if file.endswith('.txt'):
+                file_path = os.path.join(root, file)
+                rel_path = os.path.relpath(file_path, wildcard_path)
+                key = os.path.splitext(rel_path)[0].replace('\\', '/').lower()
+
+                try:
+                    with open(file_path, 'r', encoding="UTF-8", errors="ignore") as f:
+                        lines = f.read().splitlines()
+                        easy_wildcard_dict[key] = lines
+                except UnicodeDecodeError:
+                    with open(file_path, 'r', encoding="ISO-8859-1") as f:
+                        lines = f.read().splitlines()
+                        easy_wildcard_dict[key] = lines
+            elif file.endswith('.yaml'):
+                file_path = os.path.join(root, file)
+                with open(file_path, 'r', encoding="utf-8") as f:
+                    yaml_data = yaml.load(f, Loader=yaml.FullLoader)
+
+                    for k, v in yaml_data.items():
+                        read_wildcard(k, v)
+            elif file.endswith('.json'):
+                file_path = os.path.join(root, file)
+                try:
+                    with open(file_path, 'r', encoding="utf-8") as f:
+                        json_data = json.load(f)
+                        for key, value in json_data.items():
+                            key = wildcard_normalize(key)
+                            easy_wildcard_dict[key] = value
+                except ValueError:
+                    print('json files load error')
+    return easy_wildcard_dict
+
+
+def process(text, seed=None):
+
+    if seed is not None:
+        random.seed(seed)
+
+    def replace_options(string):
+        replacements_found = False
+
+        def replace_option(match):
+            nonlocal replacements_found
+            options = match.group(1).split('|')
+
+            multi_select_pattern = options[0].split('$$')
+            select_range = None
+            select_sep = ' '
+            range_pattern = r'(\d+)(-(\d+))?'
+            range_pattern2 = r'-(\d+)'
+
+            if len(multi_select_pattern) > 1:
+                r = re.match(range_pattern, options[0])
+
+                if r is None:
+                    r = re.match(range_pattern2, options[0])
+                    a = '1'
+                    b = r.group(1).strip()
+                else:
+                    a = r.group(1).strip()
+                    b = r.group(3).strip()
+
+                if r is not None:
+                    if b is not None and is_numeric_string(a) and is_numeric_string(b):
+                        # PATTERN: num1-num2
+                        select_range = int(a), int(b)
+                    elif is_numeric_string(a):
+                        # PATTERN: num
+                        x = int(a)
+                        select_range = (x, x)
+
+                    if select_range is not None and len(multi_select_pattern) == 2:
+                        # PATTERN: count$$
+                        options[0] = multi_select_pattern[1]
+                    elif select_range is not None and len(multi_select_pattern) == 3:
+                        # PATTERN: count$$ sep $$
+                        select_sep = multi_select_pattern[1]
+                        options[0] = multi_select_pattern[2]
+
+            adjusted_probabilities = []
+
+            total_prob = 0
+
+            for option in options:
+                parts = option.split('::', 1)
+                if len(parts) == 2 and is_numeric_string(parts[0].strip()):
+                    config_value = float(parts[0].strip())
+                else:
+                    config_value = 1  # Default value if no configuration is provided
+
+                adjusted_probabilities.append(config_value)
+                total_prob += config_value
+
+            normalized_probabilities = [prob / total_prob for prob in adjusted_probabilities]
+
+            if select_range is None:
+                select_count = 1
+            else:
+                select_count = random.randint(select_range[0], select_range[1])
+
+            if select_count > len(options):
+                selected_items = options
+            else:
+                selected_items = random.choices(options, weights=normalized_probabilities, k=select_count)
+                selected_items = set(selected_items)
+
+                try_count = 0
+                while len(selected_items) < select_count and try_count < 10:
+                    remaining_count = select_count - len(selected_items)
+                    additional_items = random.choices(options, weights=normalized_probabilities, k=remaining_count)
+                    selected_items |= set(additional_items)
+                    try_count += 1
+
+            selected_items2 = [re.sub(r'^\s*[0-9.]+::', '', x, 1) for x in selected_items]
+            replacement = select_sep.join(selected_items2)
+            if '::' in replacement:
+                pass
+
+            replacements_found = True
+            return replacement
+
+        pattern = r'{([^{}]*?)}'
+        replaced_string = re.sub(pattern, replace_option, string)
+
+        return replaced_string, replacements_found
+
+    def replace_wildcard(string):
+        global easy_wildcard_dict
+        pattern = r"__([\w\s.\-+/*\\]+?)__"
+        matches = re.findall(pattern, string)
+        replacements_found = False
+
+        for match in matches:
+            keyword = match.lower()
+            keyword = wildcard_normalize(keyword)
+            if keyword in easy_wildcard_dict:
+                replacement = random.choice(easy_wildcard_dict[keyword])
+                replacements_found = True
+                string = string.replace(f"__{match}__", replacement, 1)
+            elif '*' in keyword:
+                subpattern = keyword.replace('*', '.*').replace('+', r'\+')
+                total_patterns = []
+                found = False
+                for k, v in easy_wildcard_dict.items():
+                    if re.match(subpattern, k) is not None:
+                        total_patterns += v
+                        found = True
+
+                if found:
+                    replacement = random.choice(total_patterns)
+                    replacements_found = True
+                    string = string.replace(f"__{match}__", replacement, 1)
+            elif '/' not in keyword:
+                string_fallback = string.replace(f"__{match}__", f"__*/{match}__", 1)
+                string, replacements_found = replace_wildcard(string_fallback)
+
+        return string, replacements_found
+
+    replace_depth = 100
+    stop_unwrap = False
+    while not stop_unwrap and replace_depth > 1:
+        replace_depth -= 1  # prevent infinite loop
+
+        # pass1: replace options
+        pass1, is_replaced1 = replace_options(text)
+
+        while is_replaced1:
+            pass1, is_replaced1 = replace_options(pass1)
+
+        # pass2: replace wildcards
+        text, is_replaced2 = replace_wildcard(pass1)
+        stop_unwrap = not is_replaced1 and not is_replaced2
+
+    return text
+
+
+def is_numeric_string(input_str):
+    return re.match(r'^-?\d+(\.\d+)?$', input_str) is not None
+
+
+def safe_float(x):
+    if is_numeric_string(x):
+        return float(x)
+    else:
+        return 1.0
+
+
+def extract_lora_values(string):
+    pattern = r'<lora:([^>]+)>'
+    matches = re.findall(pattern, string)
+
+    def touch_lbw(text):
+        return re.sub(r'LBW=[A-Za-z][A-Za-z0-9_-]*:', r'LBW=', text)
+
+    items = [touch_lbw(match.strip(':')) for match in matches]
+
+    added = set()
+    result = []
+    for item in items:
+        item = item.split(':')
+
+        lora = None
+        a = None
+        b = None
+        lbw = None
+        lbw_a = None
+        lbw_b = None
+
+        if len(item) > 0:
+            lora = item[0]
+
+            for sub_item in item[1:]:
+                if is_numeric_string(sub_item):
+                    if a is None:
+                        a = float(sub_item)
+                    elif b is None:
+                        b = float(sub_item)
+                elif sub_item.startswith("LBW="):
+                    for lbw_item in sub_item[4:].split(';'):
+                        if lbw_item.startswith("A="):
+                            lbw_a = safe_float(lbw_item[2:].strip())
+                        elif lbw_item.startswith("B="):
+                            lbw_b = safe_float(lbw_item[2:].strip())
+                        elif lbw_item.strip() != '':
+                            lbw = lbw_item
+
+        if a is None:
+            a = 1.0
+        if b is None:
+            b = 1.0
+
+        if lora is not None and lora not in added:
+            result.append((lora, a, b, lbw, lbw_a, lbw_b))
+            added.add(lora)
+
+    return result
+
+
+def remove_lora_tags(string):
+    pattern = r'<lora:[^>]+>'
+    result = re.sub(pattern, '', string)
+
+    return result
+
+def process_with_loras(wildcard_opt, model, clip, title="Positive", seed=None, can_load_lora=True, pipe_lora_stack=[], easyCache=None):
+    pass1 = process(wildcard_opt, seed)
+    loras = extract_lora_values(pass1)
+    pass2 = remove_lora_tags(pass1)
+
+    has_noodle_key = True if "__" in wildcard_opt else False
+    has_loras = True if loras != [] else False
+    show_wildcard_prompt = True if has_noodle_key or has_loras else False
+
+    if can_load_lora and has_loras:
+        for lora_name, model_weight, clip_weight, lbw, lbw_a, lbw_b in loras:
+            if (lora_name.split('.')[-1]) not in folder_paths.supported_pt_extensions:
+                lora_name = lora_name+".safetensors"
+            lora = {
+                "lora_name": lora_name, "model": model, "clip": clip, "model_strength": model_weight,
+                "clip_strength": clip_weight,
+                "lbw_a": lbw_a,
+                "lbw_b": lbw_b,
+                "lbw": lbw
+            }
+            model, clip = easyCache.load_lora(lora)
+            lora["model"] = model
+            lora["clip"] = clip
+            pipe_lora_stack.append(lora)
+
+    log_node_info("easy wildcards",f"{title}: {pass2}")
+    if pass1 != pass2:
+        log_node_info("easy wildcards",f'{title}_decode: {pass1}')
+
+    return model, clip, pass2, pass1, show_wildcard_prompt, pipe_lora_stack
+
+
+def expand_wildcard(keyword: str) -> tuple[str]:
+    """传入文件通配符的关键词，从 easy_wildcard_dict 中获取通配符的所有选项。"""
+    global easy_wildcard_dict
+    if keyword in easy_wildcard_dict:
+        return tuple(easy_wildcard_dict[keyword])
+    elif '*' in keyword:
+        subpattern = keyword.replace('*', '.*').replace('+', r"\+")
+        total_pattern = []
+        for k, v in easy_wildcard_dict.items():
+            if re.match(subpattern, k) is not None:
+                total_pattern.extend(v)
+        if total_pattern:
+            return tuple(total_pattern)
+    elif '/' not in keyword:
+        return expand_wildcard(f"*/{keyword}")
+
+def expand_options(options: str) -> tuple[str]:
+    """传入去掉 {} 的选项。
+    展开选项通配符，返回该选项中的每一项，这里的每一项都是一个替换项。
+    不会对选项内容进行任何处理，即便存在空格或特殊符号，也会原样返回。"""
+    return tuple(options.split("|"))
+
+
+def decimal_to_irregular(n, bases):
+    """
+    将十进制数转换为不规则进制
+
+    :param n: 十进制数
+    :param bases: 各位置的基数列表，从低位到高位
+    :return: 不规则进制表示的列表，从低位到高位
+    """
+    if n == 0:
+        return [0] * len(bases) if bases else [0]
+
+    digits = []
+    remaining = n
+
+    # 从低位到高位处理
+    for base in bases:
+        digit = remaining % base
+        digits.append(digit)
+        remaining = remaining // base
+
+    return digits
+
+
+class WildcardProcessor:
+    """通配符处理器
+
+    通配符格式：
+    + option  :   {a|b}
+    + wildcard:   __keyword__ 通配符内容将从 Easy-Use 插件提供的 easy_wildcard_dict 中获取
+    """
+
+    RE_OPTIONS = re.compile(r"{([^{}]*?)}")
+    RE_WILDCARD = re.compile(r"__([\w\s.\-+/*\\]+?)__")
+    RE_REPLACER = re.compile(r"{([^{}]*?)}|__([\w\s.\-+/*\\]+?)__")
+
+    # 将输入的提示词转化成符合 python str.format 要求格式的模板，并将 option 和 wildcard 按照顺序在模板中留下 {0}, {1} 等占位符
+    template: str
+    # option、wildcard 的替换项列表，按照在模板中出现的顺序排列，相同的替换项列表只保留第一份
+    replacers: dict[int, tuple[str]]
+    # 占位符的编号和替换项列表的索引的映射，占位符编号按照在模板中出现的顺序排列，方便减少替换项的存储占用
+    placeholder_mapping: dict[str, int]  # placeholder_id => replacer_id
+    # 各替换项列表的项数，按照在模板中出现的顺序排列，提前计算，方便后续使用
+    placeholder_choices: dict[str, int]  # placeholder_id => len(replacer)
+
+    def __init__(self, text: str):
+        self.__make_template(text)
+        self.__total = None
+
+    def random(self, seed=None) -> str:
+        "从所有可能性中随机获取一个"
+        if seed is not None:
+            random.seed(seed)
+        return self.getn(random.randint(0, self.total() - 1))
+
+    def getn(self, n: int) -> str:
+        "从所有可能性中获取第 n 个，以 self.total() 为周期循环"
+        n = n % self.total()
+        indice = decimal_to_irregular(n, self.placeholder_choices.values())
+        replacements = {
+            placeholder_id: self.replacers[self.placeholder_mapping[placeholder_id]][i]
+            for placeholder_id, i in zip(self.placeholder_mapping.keys(), indice)
+        }
+        return self.template.format(**replacements)
+
+    def getmany(self, limit: int, offset: int = 0) -> list[str]:
+        """返回一组可能性组成的列表，为了避免结果太长导致内存占用超限，使用 limit 限制列表的长度，使用 offset 调整偏移。
+        若 limit 和 offset 的设置导致预期的结果长度超过剩下的实际长度，则会回到开头。
+        """
+        return [self.getn(n) for n in range(offset, offset + limit)]
+
+    def total(self) -> int:
+        "计算可能性的数目"
+        if self.__total is None:
+            self.__total = prod(self.placeholder_choices.values())
+        return self.__total
+
+    def __make_template(self, text: str):
+        """将输入的提示词转化成符合 python str.format 要求格式的模板，
+        并将 option 和 wildcard 按照顺序在模板中留下 {r0}, {r1} 等占位符，
+        即使遇到相同的 option 或 wildcard，留下的占位符编号也不同，从而使每项都独立变化。
+        """
+        self.placeholder_mapping = {}
+        placeholder_id = 0
+        replacer_id = 0
+        replacers_rev = {}  # replacers => id
+        blocks = []
+        # 记录所处理过的通配符末尾在文本中的位置，用于拼接完整的模板
+        tail = 0
+        for match in self.RE_REPLACER.finditer(text):
+            # 提取并展开通配符内容
+            m = match.group(0)
+            if m.startswith("{"):
+                choices = expand_options(m[1:-1])
+            elif m.startswith("__"):
+                keyword = m[2:-2].lower()
+                keyword = wildcard_normalize(keyword)
+                choices = expand_wildcard(keyword)
+            else:
+                raise ValueError(f"{m!r} is not a wildcard or option")
+
+            # 记录通配符的替换项列表和ID，相同的通配符只保留第一个
+            if choices not in replacers_rev:
+                replacers_rev[choices] = replacer_id
+                replacer_id += 1
+
+            # 拼接通配符前方文本
+            start, end = match.span()
+            blocks.append(text[tail:start])
+            tail = end
+            # 将通配符替换为占位符，并记录占位符和替换项列表的索引的映射
+            blocks.append(f"{{r{placeholder_id}}}")
+            self.placeholder_mapping[f"r{placeholder_id}"] = replacers_rev[choices]
+            placeholder_id += 1
+
+        if tail < len(text):
+            blocks.append(text[tail:])
+        self.template = "".join(blocks)
+        self.replacers = {v: k for k, v in replacers_rev.items()}
+        self.placeholder_choices = {
+            placeholder_id: len(self.replacers[replacer_id])
+            for placeholder_id, replacer_id in self.placeholder_mapping.items()
+        }
+
+
+def test_option():
+    text = "{|a|b|c}"
+    answer = ["", "a", "b", "c"]
+    p = WildcardProcessor(text)
+    assert p.total() == len(answer)
+    assert p.getn(0) == answer[0]
+    assert p.getmany(4) == answer
+    assert p.getmany(4, 1) == answer[1:]
+
+
+def test_same():
+    text = "{a|b},{a|b}"
+    answer = ["a,a", "b,a", "a,b", "b,b"]
+    p = WildcardProcessor(text)
+    assert p.total() == len(answer)
+    assert p.getn(0) == answer[0]
+    assert p.getmany(4) == answer
+    assert p.getmany(4, 1) == answer[1:]
+

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/libs/xyplot.py

@@ -0,0 +1,719 @@
+import os, torch
+from pathlib import Path
+from PIL import Image, ImageDraw, ImageFont
+from .utils import easySave, get_sd_version
+from .adv_encode import advanced_encode
+from .controlnet import easyControlnet
+from .log import log_node_warn
+from ..modules.layer_diffuse import LayerDiffuse
+from ..config import RESOURCES_DIR
+from nodes import CLIPTextEncode
+import pprint
+try:
+    from comfy_extras.nodes_flux import FluxGuidance
+except:
+    FluxGuidance = None
+
+class easyXYPlot():
+
+    def __init__(self, xyPlotData, save_prefix, image_output, prompt, extra_pnginfo, my_unique_id, sampler, easyCache):
+        self.x_node_type, self.x_type = sampler.safe_split(xyPlotData.get("x_axis"), ': ')
+        self.y_node_type, self.y_type = sampler.safe_split(xyPlotData.get("y_axis"), ': ')
+        self.x_values = xyPlotData.get("x_vals") if self.x_type != "None" else []
+        self.y_values = xyPlotData.get("y_vals") if self.y_type != "None" else []
+        self.custom_font = xyPlotData.get("custom_font")
+
+        self.grid_spacing = xyPlotData.get("grid_spacing")
+        self.latent_id = 0
+        self.output_individuals = xyPlotData.get("output_individuals")
+
+        self.x_label, self.y_label = [], []
+        self.max_width, self.max_height = 0, 0
+        self.latents_plot = []
+        self.image_list = []
+
+        self.num_cols = len(self.x_values) if len(self.x_values) > 0 else 1
+        self.num_rows = len(self.y_values) if len(self.y_values) > 0 else 1
+
+        self.total = self.num_cols * self.num_rows
+        self.num = 0
+
+        self.save_prefix = save_prefix
+        self.image_output = image_output
+        self.prompt = prompt
+        self.extra_pnginfo = extra_pnginfo
+        self.my_unique_id = my_unique_id
+
+        self.sampler = sampler
+        self.easyCache = easyCache
+
+    # Helper Functions
+    @staticmethod
+    def define_variable(plot_image_vars, value_type, value, index):
+
+        plot_image_vars[value_type] = value
+        if value_type in ["seed", "Seeds++ Batch"]:
+            value_label = f"seed: {value}"
+        else:
+            value_label = f"{value_type}: {value}"
+
+        if "ControlNet" in value_type:
+            value_label = f"ControlNet {index + 1}"
+
+        if value_type in ['Lora', 'Checkpoint']:
+            arr = value.split(',')
+            model_name = os.path.basename(os.path.splitext(arr[0])[0])
+            trigger_words = ' ' + arr[3] if value_type == 'Lora' and len(arr) > 3 and len(arr[3]) > 2 else ''
+            lora_weight = float(arr[1]) if value_type == 'Lora' and len(arr) > 1 else 0
+            lora_weight_desc = f" w:{lora_weight:.2f}" if value_type == 'Lora' and lora_weight != 1.0 else ''
+            value_label = f"{model_name[:25]}{lora_weight_desc}{trigger_words}"
+
+        if value_type in ["ModelMergeBlocks"]:
+            if ":" in value:
+                line = value.split(':')
+                value_label = f"{line[0]}"
+            elif len(value) > 16:
+                value_label = f"ModelMergeBlocks {index + 1}"
+            else:
+                value_label = f"MMB: {value}"
+
+        if value_type in ["Pos Condition"]:
+            value_label = f"pos cond {index + 1}" if index>0 else f"pos cond"
+        if value_type in ["Neg Condition"]:
+            value_label = f"neg cond {index + 1}" if index>0 else f"neg cond"
+
+        if value_type in ["Positive Prompt S/R"]:
+            value_label = f"pos prompt {index + 1}" if index>0 else f"pos prompt"
+        if value_type in ["Negative Prompt S/R"]:
+            value_label = f"neg prompt {index + 1}" if index>0 else f"neg prompt"
+
+        if value_type in ["steps", "cfg", "denoise", "clip_skip",
+                          "lora_model_strength", "lora_clip_strength"]:
+            value_label = f"{value_type}: {value}"
+
+        if value_type == "positive":
+            value_label = f"pos prompt {index + 1}"
+        elif value_type == "negative":
+            value_label = f"neg prompt {index + 1}"
+
+        return plot_image_vars, value_label
+
+    @staticmethod
+    def get_font(font_size, font_path=None):
+        if font_path is None:
+            font_path = str(Path(os.path.join(RESOURCES_DIR, 'OpenSans-Medium.ttf')))
+        return ImageFont.truetype(font_path, font_size)
+
+    @staticmethod
+    def update_label(label, value, num_items):
+        if len(label) < num_items:
+            return [*label, value]
+        return label
+
+    @staticmethod
+    def rearrange_tensors(latent, num_cols, num_rows):
+        new_latent = []
+        for i in range(num_rows):
+            for j in range(num_cols):
+                index = j * num_rows + i
+                new_latent.append(latent[index])
+        return new_latent
+
+    def calculate_background_dimensions(self):
+        border_size = int((self.max_width // 8) * 1.5) if self.y_type != "None" or self.x_type != "None" else 0
+
+        bg_width = self.num_cols * (self.max_width + self.grid_spacing) - self.grid_spacing + border_size * (
+                    self.y_type != "None")
+        bg_height = self.num_rows * (self.max_height + self.grid_spacing) - self.grid_spacing + border_size * (
+                    self.x_type != "None")
+        
+        # Add space at the bottom of the image for common informaiton about the image
+        bg_height = bg_height + (border_size*2)
+#        print(f"Grid Size: width = {bg_width} height = {bg_height} border_size = {border_size}")
+
+        x_offset_initial = border_size if self.y_type != "None" else 0
+        y_offset = border_size if self.x_type != "None" else 0
+
+        return bg_width, bg_height, x_offset_initial, y_offset
+
+
+    def adjust_font_size(self, text, initial_font_size, label_width):
+        font = self.get_font(initial_font_size, self.custom_font)
+        text_width = font.getbbox(text)
+#        pprint.pp(f"Initial font size: {initial_font_size}, text: {text}, text_width: {text_width}")
+        if text_width and text_width[2]:
+            text_width = text_width[2]
+
+        scaling_factor = 0.9
+        if text_width > (label_width * scaling_factor):
+#            print(f"Adjusting font size from {initial_font_size} to fit text width {text_width} into label width {label_width} scaling_factor {scaling_factor}")
+            return int(initial_font_size * (label_width / text_width) * scaling_factor)
+        else:
+            return initial_font_size
+
+    def textsize(self, d, text, font):
+        _, _, width, height = d.textbbox((0, 0), text=text, font=font)
+        return width, height
+
+    def create_label(self, img, text, initial_font_size, is_x_label=True, max_font_size=70, min_font_size=10, label_width=0, label_height=0):
+
+        # if the label_width is specified, leave it along.  Otherwise do the old logic.
+        if label_width == 0:          
+            label_width = img.width if is_x_label else img.height
+
+        text_lines = text.split('\n')
+        longest_line = max(text_lines, key=len)
+                
+        # Adjust font size
+        font_size = self.adjust_font_size(longest_line, initial_font_size, label_width)
+        font_size = min(max_font_size, font_size)  # Ensure font isn't too large
+        font_size = max(min_font_size, font_size)  # Ensure font isn't too small
+
+        if label_height == 0:
+            label_height = int(font_size * 1.5) if is_x_label else font_size
+
+        label_bg = Image.new('RGBA', (label_width, label_height), color=(255, 255, 255, 0))
+        d = ImageDraw.Draw(label_bg)
+
+        font = self.get_font(font_size, self.custom_font)
+
+        # Check if text will fit, if not insert ellipsis and reduce text
+        if self.textsize(d, text, font=font)[0] > label_width:
+            while self.textsize(d, text + '...', font=font)[0] > label_width and len(text) > 0:
+                text = text[:-1]
+            text = text + '...'
+
+        # Compute text width and height for multi-line text
+  
+        text_widths, text_heights = zip(*[self.textsize(d, line, font=font) for line in text_lines])
+        max_text_width = max(text_widths)
+        total_text_height = sum(text_heights)
+
+        # Compute position for each line of text
+        lines_positions = []
+        current_y = 0
+        for line, line_width, line_height in zip(text_lines, text_widths, text_heights):
+            text_x = (label_width - line_width) // 2
+            text_y = current_y + (label_height - total_text_height) // 2
+            current_y += line_height
+            lines_positions.append((line, (text_x, text_y)))
+
+        # Draw each line of text
+        for line, (text_x, text_y) in lines_positions:
+            d.text((text_x, text_y), line, fill='black', font=font)
+
+        return label_bg
+
+    def sample_plot_image(self, plot_image_vars, samples, preview_latent, latents_plot, image_list, disable_noise,
+                          start_step, last_step, force_full_denoise, x_value=None, y_value=None):
+        model, clip, vae, positive, negative, seed, steps, cfg = None, None, None, None, None, None, None, None
+        sampler_name, scheduler, denoise = None, None, None
+
+        a1111_prompt_style = plot_image_vars['a1111_prompt_style'] if "a1111_prompt_style" in plot_image_vars else False
+        clip = clip if clip is not None else plot_image_vars["clip"]
+        steps = plot_image_vars['steps'] if "steps" in plot_image_vars else 1
+
+        sd_version = get_sd_version(plot_image_vars['model'])          
+        # 高级用法
+        if plot_image_vars["x_node_type"] == "advanced" or plot_image_vars["y_node_type"] == "advanced":
+            if self.x_type == "Seeds++ Batch" or self.y_type == "Seeds++ Batch":
+                seed = int(x_value) if self.x_type == "Seeds++ Batch" else int(y_value)
+            if self.x_type == "Steps" or self.y_type == "Steps":
+                steps = int(x_value) if self.x_type == "Steps" else int(y_value)
+            if self.x_type == "StartStep" or self.y_type == "StartStep":
+                start_step = int(x_value) if self.x_type == "StartStep" else int(y_value)
+            if self.x_type == "EndStep" or self.y_type == "EndStep":
+                last_step = int(x_value) if self.x_type == "EndStep" else int(y_value)
+            if self.x_type == "CFG Scale" or self.y_type == "CFG Scale":
+                cfg = float(x_value) if self.x_type == "CFG Scale" else float(y_value)
+            if self.x_type == "Sampler" or self.y_type == "Sampler":
+                sampler_name = x_value if self.x_type == "Sampler" else y_value
+            if self.x_type == "Scheduler" or self.y_type == "Scheduler":
+                scheduler = x_value if self.x_type == "Scheduler" else y_value
+            if self.x_type == "Sampler&Scheduler" or self.y_type == "Sampler&Scheduler":
+                arr = x_value.split(',') if self.x_type == "Sampler&Scheduler" else y_value.split(',')
+                if arr[0] and arr[0]!= 'None':
+                    sampler_name = arr[0]
+                if arr[1] and arr[1]!= 'None':
+                    scheduler = arr[1]
+            if self.x_type == "Denoise" or self.y_type == "Denoise":
+                denoise = float(x_value) if self.x_type == "Denoise" else float(y_value)
+            if self.x_type == "Pos Condition" or self.y_type == "Pos Condition":
+                positive = plot_image_vars['positive_cond_stack'][int(x_value)] if self.x_type == "Pos Condition" else plot_image_vars['positive_cond_stack'][int(y_value)]
+            if self.x_type == "Neg Condition" or self.y_type == "Neg Condition":
+                negative = plot_image_vars['negative_cond_stack'][int(x_value)] if self.x_type == "Neg Condition" else plot_image_vars['negative_cond_stack'][int(y_value)]
+            # 模型叠加
+            if self.x_type == "ModelMergeBlocks" or self.y_type == "ModelMergeBlocks":
+                ckpt_name_1, ckpt_name_2 = plot_image_vars['models']
+                model1, clip1, vae1, clip_vision = self.easyCache.load_checkpoint(ckpt_name_1)
+                model2, clip2, vae2, clip_vision = self.easyCache.load_checkpoint(ckpt_name_2)
+                xy_values = x_value if self.x_type == "ModelMergeBlocks" else y_value
+                if ":" in xy_values:
+                    xy_line = xy_values.split(':')
+                    xy_values = xy_line[1]
+
+                xy_arrs = xy_values.split(',')
+                # ModelMergeBlocks
+                if len(xy_arrs) == 3:
+                    input, middle, out = xy_arrs
+                    kwargs = {
+                        "input": input,
+                        "middle": middle,
+                        "out": out
+                    }
+                elif len(xy_arrs) == 30:
+                    kwargs = {}
+                    kwargs["time_embed."] = xy_arrs[0]
+                    kwargs["label_emb."] = xy_arrs[1]
+
+                    for i in range(12):
+                        kwargs["input_blocks.{}.".format(i)] = xy_arrs[2+i]
+
+                    for i in range(3):
+                        kwargs["middle_block.{}.".format(i)] = xy_arrs[14+i]
+
+                    for i in range(12):
+                        kwargs["output_blocks.{}.".format(i)] = xy_arrs[17+i]
+
+                    kwargs["out."] = xy_arrs[29]
+                else:
+                    raise Exception("ModelMergeBlocks weight length error")
+                default_ratio = next(iter(kwargs.values()))
+
+                m = model1.clone()
+                kp = model2.get_key_patches("diffusion_model.")
+
+                for k in kp:
+                    ratio = float(default_ratio)
+                    k_unet = k[len("diffusion_model."):]
+
+                    last_arg_size = 0
+                    for arg in kwargs:
+                        if k_unet.startswith(arg) and last_arg_size < len(arg):
+                            ratio = float(kwargs[arg])
+                            last_arg_size = len(arg)
+
+                    m.add_patches({k: kp[k]}, 1.0 - ratio, ratio)
+
+                vae_use = plot_image_vars['vae_use']
+
+                clip = clip2 if vae_use == 'Use Model 2' else clip1
+                if vae_use == 'Use Model 2':
+                    vae = vae2
+                elif vae_use == 'Use Model 1':
+                    vae = vae1
+                else:
+                    vae = self.easyCache.load_vae(vae_use)
+                model = m
+
+                # 如果存在lora_stack叠加lora
+                optional_lora_stack = plot_image_vars['lora_stack']
+                if optional_lora_stack is not None and optional_lora_stack != []:
+                    for lora in optional_lora_stack:
+                        model, clip = self.easyCache.load_lora(lora)
+
+                # 处理clip
+                clip = clip.clone()
+                if plot_image_vars['clip_skip'] != 0:
+                    clip.clip_layer(plot_image_vars['clip_skip'])
+
+            # CheckPoint
+            if self.x_type == "Checkpoint" or self.y_type == "Checkpoint":
+                xy_values = x_value if self.x_type == "Checkpoint" else y_value
+                ckpt_name, clip_skip, vae_name = xy_values.split(",")
+                ckpt_name = ckpt_name.replace('*', ',')
+                vae_name = vae_name.replace('*', ',')
+                model, clip, vae, clip_vision = self.easyCache.load_checkpoint(ckpt_name)
+                if vae_name != 'None':
+                    vae = self.easyCache.load_vae(vae_name)
+
+                # 如果存在lora_stack叠加lora
+                optional_lora_stack = plot_image_vars['lora_stack']
+                if optional_lora_stack is not None and optional_lora_stack != []:
+                    for lora in optional_lora_stack:
+                        lora['model'] = model
+                        lora['clip'] = clip
+                        model, clip = self.easyCache.load_lora(lora)
+
+                # 处理clip
+                clip = clip.clone()
+                if clip_skip != 'None':
+                    clip.clip_layer(int(clip_skip))
+                    positive = plot_image_vars['positive']
+                    negative = plot_image_vars['negative']
+                    a1111_prompt_style = plot_image_vars['a1111_prompt_style']
+                    steps = plot_image_vars['steps']
+                    clip = clip if clip is not None else plot_image_vars["clip"]
+                    positive = advanced_encode(clip, positive,
+                                               plot_image_vars['positive_token_normalization'],
+                                               plot_image_vars['positive_weight_interpretation'],
+                                               w_max=1.0,
+                                               apply_to_pooled="enable",
+                                               a1111_prompt_style=a1111_prompt_style, steps=steps)
+
+                    negative = advanced_encode(clip, negative,
+                                               plot_image_vars['negative_token_normalization'],
+                                               plot_image_vars['negative_weight_interpretation'],
+                                               w_max=1.0,
+                                               apply_to_pooled="enable",
+                                               a1111_prompt_style=a1111_prompt_style, steps=steps)
+                    if "positive_cond" in plot_image_vars:
+                        positive = positive + plot_image_vars["positive_cond"]
+                    if "negative_cond" in plot_image_vars:
+                        negative = negative + plot_image_vars["negative_cond"]
+
+            # Lora
+            if self.x_type == "Lora" or self.y_type == "Lora":
+#                print(f"Lora: {x_value} {y_value}")
+                model = model if model is not None else plot_image_vars["model"]
+                clip = clip if clip is not None else plot_image_vars["clip"]
+                
+                # Build lora_stack from both X and Y axes if both are LoRA types
+                lora_stack = []
+                
+                # Add X axis LoRA if present
+                if self.x_type == "Lora":
+                    lora_name, lora_model_strength, lora_clip_strength, _ = x_value.split(",")
+                    lora_stack.append({"lora_name": lora_name, "model": model, "clip": clip, "model_strength": float(lora_model_strength), "clip_strength": float(lora_clip_strength)})
+                
+                # Add Y axis LoRA if present
+                if self.y_type == "Lora":
+                    lora_name, lora_model_strength, lora_clip_strength, _ = y_value.split(",")
+                    lora_stack.append({"lora_name": lora_name, "model": model, "clip": clip, "model_strength": float(lora_model_strength), "clip_strength": float(lora_clip_strength)})
+                
+#                print(f"new_lora_stack: {new_lora_stack}")
+
+                
+                if 'lora_stack' in plot_image_vars:
+                    lora_stack = lora_stack + plot_image_vars['lora_stack']
+                
+                if lora_stack is not None and lora_stack != []:
+                    for lora in lora_stack:
+                        # Each generation of the model, must use the reference to previously created model / clip objects.
+                        lora['model'] = model
+                        lora['clip'] = clip
+                        model, clip = self.easyCache.load_lora(lora)
+
+            # 提示词
+            if "Positive" in self.x_type or "Positive" in self.y_type:
+                if self.x_type == 'Positive Prompt S/R' or self.y_type == 'Positive Prompt S/R':
+                    positive = x_value if self.x_type == "Positive Prompt S/R" else y_value
+
+                if sd_version == 'flux':
+                    positive, = CLIPTextEncode().encode(clip, positive)
+                else:
+                    positive = advanced_encode(clip, positive,
+                                                plot_image_vars['positive_token_normalization'],
+                                                plot_image_vars['positive_weight_interpretation'],
+                                                w_max=1.0,
+                                                apply_to_pooled="enable", a1111_prompt_style=a1111_prompt_style, steps=steps)
+
+                # if "positive_cond" in plot_image_vars:
+                #     positive = positive + plot_image_vars["positive_cond"]
+
+            if "Negative" in self.x_type or "Negative" in self.y_type:
+                if self.x_type == 'Negative Prompt S/R' or self.y_type == 'Negative Prompt S/R':
+                    negative = x_value if self.x_type == "Negative Prompt S/R" else y_value
+
+                if sd_version == 'flux':
+                    negative, = CLIPTextEncode().encode(clip, negative)
+                else:
+                    negative = advanced_encode(clip, negative,
+                                                plot_image_vars['negative_token_normalization'],
+                                                plot_image_vars['negative_weight_interpretation'],
+                                                w_max=1.0,
+                                                apply_to_pooled="enable", a1111_prompt_style=a1111_prompt_style, steps=steps)
+                # if "negative_cond" in plot_image_vars:
+                #     negative = negative + plot_image_vars["negative_cond"]
+
+            # ControlNet
+            if "ControlNet" in self.x_type or "ControlNet" in self.y_type:
+                cnet = plot_image_vars["cnet"] if "cnet" in plot_image_vars else None
+                positive = plot_image_vars["positive_cond"] if "positive" in plot_image_vars else None
+                negative = plot_image_vars["negative_cond"] if "negative" in plot_image_vars else None
+                if cnet:
+                    index = x_value if "ControlNet" in self.x_type else y_value
+                    controlnet = cnet[index]
+                    for index, item in enumerate(controlnet):
+                        control_net_name = item[0]
+                        image = item[1]
+                        strength = item[2]
+                        start_percent = item[3]
+                        end_percent = item[4]
+                        provided_control_net = item[5] if len(item) > 5 else None
+                        positive, negative = easyControlnet().apply(control_net_name, image, positive, negative, strength, start_percent, end_percent, provided_control_net, 1)
+            # Flux guidance
+            if self.x_type == "Flux Guidance" or self.y_type == "Flux Guidance":
+                positive = plot_image_vars["positive_cond"] if "positive" in plot_image_vars else None
+                flux_guidance = float(x_value) if self.x_type == "Flux Guidance" else float(y_value)
+                positive, = FluxGuidance().append(positive, flux_guidance)
+
+        # 简单用法
+        if plot_image_vars["x_node_type"] == "loader" or plot_image_vars["y_node_type"] == "loader":
+            if self.x_type == 'ckpt_name' or self.y_type == 'ckpt_name':
+                ckpt_name = x_value if self.x_type == "ckpt_name" else y_value
+                model, clip, vae, clip_vision = self.easyCache.load_checkpoint(ckpt_name)
+
+            if self.x_type == 'lora_name' or self.y_type == 'lora_name':
+                model, clip, vae, clip_vision = self.easyCache.load_checkpoint(plot_image_vars['ckpt_name'])
+                lora_name = x_value if self.x_type == "lora_name" else y_value
+                lora = {"lora_name": lora_name, "model": model, "clip": clip, "model_strength": 1, "clip_strength": 1}
+                model, clip = self.easyCache.load_lora(lora)
+
+            if self.x_type == 'lora_model_strength' or self.y_type == 'lora_model_strength':
+                model, clip, vae, clip_vision = self.easyCache.load_checkpoint(plot_image_vars['ckpt_name'])
+                lora_model_strength = float(x_value) if self.x_type == "lora_model_strength" else float(y_value)
+                lora = {"lora_name": plot_image_vars['lora_name'], "model": model, "clip": clip, "model_strength": lora_model_strength, "clip_strength": plot_image_vars['lora_clip_strength']}
+                model, clip = self.easyCache.load_lora(lora)
+
+            if self.x_type == 'lora_clip_strength' or self.y_type == 'lora_clip_strength':
+                model, clip, vae, clip_vision = self.easyCache.load_checkpoint(plot_image_vars['ckpt_name'])
+                lora_clip_strength = float(x_value) if self.x_type == "lora_clip_strength" else float(y_value)
+                lora = {"lora_name": plot_image_vars['lora_name'], "model": model, "clip": clip, "model_strength": plot_image_vars['lora_model_strength'], "clip_strength": lora_clip_strength}
+                model, clip = self.easyCache.load_lora(lora)
+
+            # Check for custom VAE
+            if self.x_type == 'vae_name' or self.y_type == 'vae_name':
+                vae_name = x_value if self.x_type == "vae_name" else y_value
+                vae = self.easyCache.load_vae(vae_name)
+
+            # CLIP skip
+            if not clip:
+                raise Exception("No CLIP found")
+            clip = clip.clone()
+            clip.clip_layer(plot_image_vars['clip_skip'])
+
+            if sd_version == 'flux':
+                positive, = CLIPTextEncode().encode(clip, positive)
+            else:
+                positive = advanced_encode(clip, plot_image_vars['positive'],
+                                                            plot_image_vars['positive_token_normalization'],
+                                                            plot_image_vars['positive_weight_interpretation'], w_max=1.0,
+                                                            apply_to_pooled="enable",a1111_prompt_style=a1111_prompt_style, steps=steps)
+
+            if sd_version == 'flux':
+                negative, = CLIPTextEncode().encode(clip, negative)
+            else:
+                negative = advanced_encode(clip, plot_image_vars['negative'],
+                                                            plot_image_vars['negative_token_normalization'],
+                                                            plot_image_vars['negative_weight_interpretation'], w_max=1.0,
+                                                            apply_to_pooled="enable", a1111_prompt_style=a1111_prompt_style, steps=steps)
+
+
+        model = model if model is not None else plot_image_vars["model"]
+        vae = vae if vae is not None else plot_image_vars["vae"]
+        positive = positive if positive is not None else plot_image_vars["positive_cond"]
+        negative = negative if negative is not None else plot_image_vars["negative_cond"]
+
+        seed = seed if seed is not None else plot_image_vars["seed"]
+        steps = steps if steps is not None else plot_image_vars["steps"]
+        cfg = cfg if cfg is not None else plot_image_vars["cfg"]
+        sampler_name = sampler_name if sampler_name is not None else plot_image_vars["sampler_name"]
+        scheduler = scheduler if scheduler is not None else plot_image_vars["scheduler"]
+        denoise = denoise if denoise is not None else plot_image_vars["denoise"]
+
+        noise_device = plot_image_vars["noise_device"] if "noise_device" in plot_image_vars else 'cpu'
+
+        # LayerDiffuse
+        layer_diffusion_method = plot_image_vars["layer_diffusion_method"] if "layer_diffusion_method" in plot_image_vars else None
+        empty_samples = plot_image_vars["empty_samples"] if "empty_samples" in plot_image_vars else None
+
+        if layer_diffusion_method:
+            samp_blend_samples = plot_image_vars["blend_samples"] if "blend_samples" in plot_image_vars else None
+            additional_cond = plot_image_vars["layer_diffusion_cond"] if "layer_diffusion_cond" in plot_image_vars else None
+
+            images = plot_image_vars["images"].movedim(-1, 1) if "images" in plot_image_vars else None
+            weight = plot_image_vars['layer_diffusion_weight'] if 'layer_diffusion_weight' in plot_image_vars else 1.0
+            model, positive, negative = LayerDiffuse().apply_layer_diffusion(model, layer_diffusion_method, weight, samples,
+                                                                                  samp_blend_samples, positive,
+                                                                                  negative, images, additional_cond)
+
+        samples = empty_samples if layer_diffusion_method is not None and empty_samples is not None else samples
+        # Sample
+        samples = self.sampler.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, samples,
+                                          denoise=denoise, disable_noise=disable_noise, preview_latent=preview_latent,
+                                          start_step=start_step, last_step=last_step,
+                                          force_full_denoise=force_full_denoise, noise_device=noise_device)
+
+        # Decode images and store
+        latent = samples["samples"]
+
+        # Add the latent tensor to the tensors list
+        latents_plot.append(latent)
+
+        # Decode the image
+        image = vae.decode(latent).cpu()
+
+        if self.output_individuals in [True, "True"]:
+            easySave(image, self.save_prefix, self.image_output)
+
+        # Convert the image from tensor to PIL Image and add it to the list
+        pil_image = self.sampler.tensor2pil(image)
+        image_list.append(pil_image)
+
+        # Update max dimensions
+        self.max_width = max(self.max_width, pil_image.width)
+        self.max_height = max(self.max_height, pil_image.height)
+
+        # Return the touched variables
+        return image_list, self.max_width, self.max_height, latents_plot
+
+    # Process Functions
+    def validate_xy_plot(self):
+        if self.x_type == 'None' and self.y_type == 'None':
+            log_node_warn(f'#{self.my_unique_id}','No Valid Plot Types - Reverting to default sampling...')
+            return False
+        else:
+            return True
+
+    def get_latent(self, samples):
+        # Extract the 'samples' tensor from the dictionary
+        latent_image_tensor = samples["samples"]
+
+        # Split the tensor into individual image tensors
+        image_tensors = torch.split(latent_image_tensor, 1, dim=0)
+
+        # Create a list of dictionaries containing the individual image tensors
+        latent_list = [{'samples': image} for image in image_tensors]
+
+        # Set latent only to the first latent of batch
+        if self.latent_id >= len(latent_list):
+            log_node_warn(f'#{self.my_unique_id}',f'The selected latent_id ({self.latent_id}) is out of range.')
+            log_node_warn(f'#{self.my_unique_id}', f'Automatically setting the latent_id to the last image in the list (index: {len(latent_list) - 1}).')
+
+            self.latent_id = len(latent_list) - 1
+
+        return latent_list[self.latent_id]
+
+    def get_labels_and_sample(self, plot_image_vars, latent_image, preview_latent, start_step, last_step,
+                              force_full_denoise, disable_noise):
+        # Handle X-only variation (Y is "None")
+        if self.y_type == 'None':
+            for x_index, x_value in enumerate(self.x_values):
+                plot_image_vars, x_value_label = self.define_variable(plot_image_vars, self.x_type, x_value, x_index)
+                self.x_label = self.update_label(self.x_label, x_value_label, len(self.x_values))
+                
+                self.image_list, self.max_width, self.max_height, self.latents_plot = self.sample_plot_image(
+                    plot_image_vars, latent_image, preview_latent, self.latents_plot, self.image_list,
+                    disable_noise, start_step, last_step, force_full_denoise, x_value)
+                self.num += 1
+        # Handle Y-only variation (X is "None")
+        elif self.x_type == 'None':
+            for y_index, y_value in enumerate(self.y_values):
+                plot_image_vars, y_value_label = self.define_variable(plot_image_vars, self.y_type, y_value, y_index)
+                self.y_label = self.update_label(self.y_label, y_value_label, len(self.y_values))
+                
+                self.image_list, self.max_width, self.max_height, self.latents_plot = self.sample_plot_image(
+                    plot_image_vars, latent_image, preview_latent, self.latents_plot, self.image_list,
+                    disable_noise, start_step, last_step, force_full_denoise, y_value=y_value)
+                self.num += 1
+        # Handle both X and Y variation
+        else:
+            for x_index, x_value in enumerate(self.x_values):
+                plot_image_vars, x_value_label = self.define_variable(plot_image_vars, self.x_type, x_value, x_index)
+                self.x_label = self.update_label(self.x_label, x_value_label, len(self.x_values))
+                
+                for y_index, y_value in enumerate(self.y_values):
+                    plot_image_vars, y_value_label = self.define_variable(plot_image_vars, self.y_type, y_value, y_index)
+                    self.y_label = self.update_label(self.y_label, y_value_label, len(self.y_values))
+                    
+                    self.image_list, self.max_width, self.max_height, self.latents_plot = self.sample_plot_image(
+                        plot_image_vars, latent_image, preview_latent, self.latents_plot, self.image_list,
+                        disable_noise, start_step, last_step, force_full_denoise, x_value, y_value)
+                    self.num += 1
+
+        # Rearrange latent array to match preview image grid
+        self.latents_plot = self.rearrange_tensors(self.latents_plot, self.num_cols, self.num_rows)
+
+        # Concatenate the tensors along the first dimension (dim=0)
+        self.latents_plot = torch.cat(self.latents_plot, dim=0)
+
+        return self.latents_plot
+
+    def plot_images_and_labels(self, plot_image_vars):
+                    
+        bg_width, bg_height, x_offset_initial, y_offset = self.calculate_background_dimensions()
+
+        background = Image.new('RGBA', (int(bg_width), int(bg_height)), color=(255, 255, 255, 255))
+
+        output_image = []
+        for row_index in range(self.num_rows):
+            x_offset = x_offset_initial
+
+            for col_index in range(self.num_cols):
+                index = col_index * self.num_rows + row_index
+                img = self.image_list[index]
+                output_image.append(self.sampler.pil2tensor(img))
+                background.paste(img, (x_offset, y_offset))
+
+                # Handle X label
+                if row_index == 0 and self.x_type != "None":
+                    label_bg = self.create_label(img, self.x_label[col_index], int(48 * img.width / 512))
+                    label_y = (y_offset - label_bg.height) // 2
+                    background.alpha_composite(label_bg, (x_offset, label_y))
+
+                # Handle Y label
+                if col_index == 0 and self.y_type != "None":
+                    label_bg = self.create_label(img, self.y_label[row_index], int(48 * img.height / 512), False)
+                    label_bg = label_bg.rotate(90, expand=True)
+
+                    label_x = (x_offset - label_bg.width) // 2
+                    label_y = y_offset + (img.height - label_bg.height) // 2
+                    background.alpha_composite(label_bg, (label_x, label_y))
+
+                x_offset += img.width + self.grid_spacing
+
+            y_offset += img.height + self.grid_spacing
+
+        # lookup used models in the image
+        common_label = ""
+        # Update to add a function to do the heavy lifting. Parameters are plot_image_vars name, label to use, names of the axis, 
+
+        # pprint.pp(plot_image_vars)
+
+        # We don't process LORAs here because there can be multiple of them.
+        labels = [
+            {"id": "ckpt_name", "id_desc": "ckpt", "axis_type" : "Checkpoint"},
+            {"id": "vae_name", "id_desc": '', "axis_type" : "vae_name"},
+            {"id": "sampler_name", "id_desc": "sampler", "axis_type" : "Sampler"},
+            {"id": "scheduler", "id_desc": '', "axis_type" : "Scheduler"},
+            {"id": "steps", "id_desc": '', "axis_type" : "Steps"},
+            {"id": "Flux Guidance", "id_desc": 'guidance', "axis_type" : "Flux Guidance"},     
+            {"id": "seed", "id_desc": '', "axis_type" : "Seeds++ Batch"}
+        ]
+        
+        for item in labels:
+            # Only add the label if it's not one of the axis
+            # print(f"Checking item: {item['id']} axis_type {item['axis_type']} x_type: {self.x_type} y_type: {self.y_type}")
+            if self.x_type != item['axis_type'] and self.y_type != item['axis_type']:
+                common_label += self.add_common_label(item['id'], plot_image_vars, item['id_desc'])
+        common_label += f"\n"
+                
+        if plot_image_vars['lora_stack'] is not None and plot_image_vars['lora_stack'] != []:
+#            print(f"lora_stack: {plot_image_vars['lora_stack']}")
+            for lora in plot_image_vars['lora_stack']:
+
+                lora_name = lora['lora_name']
+                lora_weight = lora['model_strength']
+                if lora_name is not None and len(lora_name) > 0 and lora_weight > 0:
+                    common_label += f"LORA: {lora_name} weight: {lora_weight:.2f} \n"
+                    
+        common_label = common_label.strip()
+        
+        if len(common_label) > 0:
+            label_height = background.height - y_offset
+            label_bg = self.create_label(background, common_label, int(48 * background.width / 512), label_width=background.width, label_height=label_height)
+            label_x = (background.width - label_bg.width) // 2
+            label_y = y_offset
+#            print(f"Adding common label: {common_label} x = {label_x} y = {label_y}")
+            background.alpha_composite(label_bg, (label_x, label_y))
+
+        return (self.sampler.pil2tensor(background), output_image)
+
+    def add_common_label(self, tag, plot_image_vars, description = ''):
+        label = ''
+        if description == '': description = tag
+        if tag in plot_image_vars and plot_image_vars[tag] is not None and plot_image_vars[tag] != 'None':
+            label += f"{description}: {plot_image_vars[tag]} "
+#        print(f"add_common_label: {tag} description: {description} label: {label}" )
+        return label

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/ben/model.py

@@ -0,0 +1,1310 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+import torch.utils.checkpoint as checkpoint
+import numpy as np
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from PIL import Image, ImageOps
+from torchvision import transforms
+import numpy as np
+import random
+import cv2
+import os
+import subprocess
+import time
+import tempfile
+
+
+def set_random_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+set_random_seed(9)
+
+torch.set_float32_matmul_precision('highest')
+
+
+class Mlp(nn.Module):
+    """ Multilayer perceptron."""
+
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    """ Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """ Forward function.
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class SwinTransformerBlock(nn.Module):
+    """ Swin Transformer Block.
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, num_heads, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        self.H = None
+        self.W = None
+
+    def forward(self, x, mask_matrix):
+        """ Forward function.
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+            mask_matrix: Attention mask for cyclic shift.
+        """
+        B, L, C = x.shape
+        H, W = self.H, self.W
+        assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # pad feature maps to multiples of window size
+        pad_l = pad_t = 0
+        pad_r = (self.window_size - W % self.window_size) % self.window_size
+        pad_b = (self.window_size - H % self.window_size) % self.window_size
+        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
+        _, Hp, Wp, _ = x.shape
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            attn_mask = mask_matrix
+        else:
+            shifted_x = x
+            attn_mask = None
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, mask=attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+
+        if pad_r > 0 or pad_b > 0:
+            x = x[:, :H, :W, :].contiguous()
+
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+
+class PatchMerging(nn.Module):
+    """ Patch Merging Layer
+    Args:
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x, H, W):
+        """ Forward function.
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+        """
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        x = x.view(B, H, W, C)
+
+        # padding
+        pad_input = (H % 2 == 1) or (W % 2 == 1)
+        if pad_input:
+            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+
+class BasicLayer(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+    Args:
+        dim (int): Number of feature channels
+        depth (int): Depths of this stage.
+        num_heads (int): Number of attention head.
+        window_size (int): Local window size. Default: 7.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self,
+                 dim,
+                 depth,
+                 num_heads,
+                 window_size=7,
+                 mlp_ratio=4.,
+                 qkv_bias=True,
+                 qk_scale=None,
+                 drop=0.,
+                 attn_drop=0.,
+                 drop_path=0.,
+                 norm_layer=nn.LayerNorm,
+                 downsample=None,
+                 use_checkpoint=False):
+        super().__init__()
+        self.window_size = window_size
+        self.shift_size = window_size // 2
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(
+                dim=dim,
+                num_heads=num_heads,
+                window_size=window_size,
+                shift_size=0 if (i % 2 == 0) else window_size // 2,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop,
+                attn_drop=attn_drop,
+                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                norm_layer=norm_layer)
+            for i in range(depth)])
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x, H, W):
+        """ Forward function.
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+        """
+
+        # calculate attention mask for SW-MSA
+        Hp = int(np.ceil(H / self.window_size)) * self.window_size
+        Wp = int(np.ceil(W / self.window_size)) * self.window_size
+        img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # 1 Hp Wp 1
+        h_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        w_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+
+        for blk in self.blocks:
+            blk.H, blk.W = H, W
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x, attn_mask)
+            else:
+                x = blk(x, attn_mask)
+        if self.downsample is not None:
+            x_down = self.downsample(x, H, W)
+            Wh, Ww = (H + 1) // 2, (W + 1) // 2
+            return x, H, W, x_down, Wh, Ww
+        else:
+            return x, H, W, x, H, W
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    Args:
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        patch_size = to_2tuple(patch_size)
+        self.patch_size = patch_size
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        """Forward function."""
+        # padding
+        _, _, H, W = x.size()
+        if W % self.patch_size[1] != 0:
+            x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
+        if H % self.patch_size[0] != 0:
+            x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
+
+        x = self.proj(x)  # B C Wh Ww
+        if self.norm is not None:
+            Wh, Ww = x.size(2), x.size(3)
+            x = x.flatten(2).transpose(1, 2)
+            x = self.norm(x)
+            x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
+
+        return x
+
+
+class SwinTransformer(nn.Module):
+    """ Swin Transformer backbone.
+        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
+          https://arxiv.org/pdf/2103.14030
+    Args:
+        pretrain_img_size (int): Input image size for training the pretrained model,
+            used in absolute postion embedding. Default 224.
+        patch_size (int | tuple(int)): Patch size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        depths (tuple[int]): Depths of each Swin Transformer stage.
+        num_heads (tuple[int]): Number of attention head of each stage.
+        window_size (int): Window size. Default: 7.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
+        drop_rate (float): Dropout rate.
+        attn_drop_rate (float): Attention dropout rate. Default: 0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.2.
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False.
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True.
+        out_indices (Sequence[int]): Output from which stages.
+        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
+            -1 means not freezing any parameters.
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self,
+                 pretrain_img_size=224,
+                 patch_size=4,
+                 in_chans=3,
+                 embed_dim=96,
+                 depths=[2, 2, 6, 2],
+                 num_heads=[3, 6, 12, 24],
+                 window_size=7,
+                 mlp_ratio=4.,
+                 qkv_bias=True,
+                 qk_scale=None,
+                 drop_rate=0.,
+                 attn_drop_rate=0.,
+                 drop_path_rate=0.2,
+                 norm_layer=nn.LayerNorm,
+                 ape=False,
+                 patch_norm=True,
+                 out_indices=(0, 1, 2, 3),
+                 frozen_stages=-1,
+                 use_checkpoint=False):
+        super().__init__()
+
+        self.pretrain_img_size = pretrain_img_size
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.out_indices = out_indices
+        self.frozen_stages = frozen_stages
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+
+        # absolute position embedding
+        if self.ape:
+            pretrain_img_size = to_2tuple(pretrain_img_size)
+            patch_size = to_2tuple(patch_size)
+            patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]]
+
+            self.absolute_pos_embed = nn.Parameter(
+                torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1]))
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayer(
+                dim=int(embed_dim * 2 ** i_layer),
+                depth=depths[i_layer],
+                num_heads=num_heads[i_layer],
+                window_size=window_size,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
+                norm_layer=norm_layer,
+                downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+                use_checkpoint=use_checkpoint)
+            self.layers.append(layer)
+
+        num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
+        self.num_features = num_features
+
+        # add a norm layer for each output
+        for i_layer in out_indices:
+            layer = norm_layer(num_features[i_layer])
+            layer_name = f'norm{i_layer}'
+            self.add_module(layer_name, layer)
+
+        self._freeze_stages()
+
+    def _freeze_stages(self):
+        if self.frozen_stages >= 0:
+            self.patch_embed.eval()
+            for param in self.patch_embed.parameters():
+                param.requires_grad = False
+
+        if self.frozen_stages >= 1 and self.ape:
+            self.absolute_pos_embed.requires_grad = False
+
+        if self.frozen_stages >= 2:
+            self.pos_drop.eval()
+            for i in range(0, self.frozen_stages - 1):
+                m = self.layers[i]
+                m.eval()
+                for param in m.parameters():
+                    param.requires_grad = False
+
+    def forward(self, x):
+
+        x = self.patch_embed(x)
+
+        Wh, Ww = x.size(2), x.size(3)
+        if self.ape:
+            # interpolate the position embedding to the corresponding size
+            absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic')
+            x = (x + absolute_pos_embed)  # B Wh*Ww C
+
+        outs = [x.contiguous()]
+        x = x.flatten(2).transpose(1, 2)
+        x = self.pos_drop(x)
+
+        for i in range(self.num_layers):
+            layer = self.layers[i]
+            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
+
+            if i in self.out_indices:
+                norm_layer = getattr(self, f'norm{i}')
+                x_out = norm_layer(x_out)
+
+                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
+                outs.append(out)
+
+        return tuple(outs)
+
+
+def get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "gelu":
+        return F.gelu
+
+    raise RuntimeError(F"activation should be gelu, not {activation}.")
+
+
+def make_cbr(in_dim, out_dim):
+    return nn.Sequential(nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1), nn.InstanceNorm2d(out_dim), nn.GELU())
+
+
+def make_cbg(in_dim, out_dim):
+    return nn.Sequential(nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1), nn.InstanceNorm2d(out_dim), nn.GELU())
+
+
+def rescale_to(x, scale_factor: float = 2, interpolation='nearest'):
+    return F.interpolate(x, scale_factor=scale_factor, mode=interpolation)
+
+
+def resize_as(x, y, interpolation='bilinear'):
+    return F.interpolate(x, size=y.shape[-2:], mode=interpolation)
+
+
+def image2patches(x):
+    """b c (hg h) (wg w) -> (hg wg b) c h w"""
+    x = rearrange(x, 'b c (hg h) (wg w) -> (hg wg b) c h w', hg=2, wg=2)
+    return x
+
+
+def patches2image(x):
+    """(hg wg b) c h w -> b c (hg h) (wg w)"""
+    x = rearrange(x, '(hg wg b) c h w -> b c (hg h) (wg w)', hg=2, wg=2)
+    return x
+
+
+class PositionEmbeddingSine:
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+        self.dim_t = torch.arange(0, self.num_pos_feats, dtype=torch.float32)
+
+    def __call__(self, b, h, w):
+        device = self.dim_t.device
+        mask = torch.zeros([b, h, w], dtype=torch.bool, device=device)
+        assert mask is not None
+        not_mask = ~mask
+        y_embed = not_mask.cumsum(dim=1, dtype=torch.float32)
+        x_embed = not_mask.cumsum(dim=2, dtype=torch.float32)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = self.temperature ** (2 * (self.dim_t.to(device) // 2) / self.num_pos_feats)
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+
+        return torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+
+
+class PositionEmbeddingSine:
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+        self.dim_t = torch.arange(0, self.num_pos_feats, dtype=torch.float32)
+
+    def __call__(self, b, h, w):
+        device = self.dim_t.device
+        mask = torch.zeros([b, h, w], dtype=torch.bool, device=device)
+        assert mask is not None
+        not_mask = ~mask
+        y_embed = not_mask.cumsum(dim=1, dtype=torch.float32)
+        x_embed = not_mask.cumsum(dim=2, dtype=torch.float32)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = self.temperature ** (2 * (self.dim_t.to(device) // 2) / self.num_pos_feats)
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+
+        return torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+
+
+class MCLM(nn.Module):
+    def __init__(self, d_model, num_heads, pool_ratios=[1, 4, 8]):
+        super(MCLM, self).__init__()
+        self.attention = nn.ModuleList([
+            nn.MultiheadAttention(d_model, num_heads, dropout=0.1),
+            nn.MultiheadAttention(d_model, num_heads, dropout=0.1),
+            nn.MultiheadAttention(d_model, num_heads, dropout=0.1),
+            nn.MultiheadAttention(d_model, num_heads, dropout=0.1),
+            nn.MultiheadAttention(d_model, num_heads, dropout=0.1)
+        ])
+
+        self.linear1 = nn.Linear(d_model, d_model * 2)
+        self.linear2 = nn.Linear(d_model * 2, d_model)
+        self.linear3 = nn.Linear(d_model, d_model * 2)
+        self.linear4 = nn.Linear(d_model * 2, d_model)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(0.1)
+        self.dropout1 = nn.Dropout(0.1)
+        self.dropout2 = nn.Dropout(0.1)
+        self.activation = get_activation_fn('gelu')
+        self.pool_ratios = pool_ratios
+        self.p_poses = []
+        self.g_pos = None
+        self.positional_encoding = PositionEmbeddingSine(num_pos_feats=d_model // 2, normalize=True)
+
+    def forward(self, l, g):
+        """
+        l: 4,c,h,w
+        g: 1,c,h,w
+        """
+        self.p_poses = []
+        self.g_pos = None
+        b, c, h, w = l.size()
+        # 4,c,h,w -> 1,c,2h,2w
+        concated_locs = rearrange(l, '(hg wg b) c h w -> b c (hg h) (wg w)', hg=2, wg=2)
+
+        pools = []
+        for pool_ratio in self.pool_ratios:
+            # b,c,h,w
+            tgt_hw = (round(h / pool_ratio), round(w / pool_ratio))
+            pool = F.adaptive_avg_pool2d(concated_locs, tgt_hw)
+            pools.append(rearrange(pool, 'b c h w -> (h w) b c'))
+            if self.g_pos is None:
+                pos_emb = self.positional_encoding(pool.shape[0], pool.shape[2], pool.shape[3])
+                pos_emb = rearrange(pos_emb, 'b c h w -> (h w) b c')
+                self.p_poses.append(pos_emb)
+        pools = torch.cat(pools, 0)
+        if self.g_pos is None:
+            self.p_poses = torch.cat(self.p_poses, dim=0)
+            pos_emb = self.positional_encoding(g.shape[0], g.shape[2], g.shape[3])
+            self.g_pos = rearrange(pos_emb, 'b c h w -> (h w) b c')
+
+        device = pools.device
+        self.p_poses = self.p_poses.to(device)
+        self.g_pos = self.g_pos.to(device)
+
+        # attention between glb (q) & multisensory concated-locs (k,v)
+        g_hw_b_c = rearrange(g, 'b c h w -> (h w) b c')
+
+        g_hw_b_c = g_hw_b_c + self.dropout1(self.attention[0](g_hw_b_c + self.g_pos, pools + self.p_poses, pools)[0])
+        g_hw_b_c = self.norm1(g_hw_b_c)
+        g_hw_b_c = g_hw_b_c + self.dropout2(self.linear2(self.dropout(self.activation(self.linear1(g_hw_b_c)).clone())))
+        g_hw_b_c = self.norm2(g_hw_b_c)
+
+        # attention between origin locs (q) & freashed glb (k,v)
+        l_hw_b_c = rearrange(l, "b c h w -> (h w) b c")
+        _g_hw_b_c = rearrange(g_hw_b_c, '(h w) b c -> h w b c', h=h, w=w)
+        _g_hw_b_c = rearrange(_g_hw_b_c, "(ng h) (nw w) b c -> (h w) (ng nw b) c", ng=2, nw=2)
+        outputs_re = []
+        for i, (_l, _g) in enumerate(zip(l_hw_b_c.chunk(4, dim=1), _g_hw_b_c.chunk(4, dim=1))):
+            outputs_re.append(self.attention[i + 1](_l, _g, _g)[0])  # (h w) 1 c
+        outputs_re = torch.cat(outputs_re, 1)  # (h w) 4 c
+
+        l_hw_b_c = l_hw_b_c + self.dropout1(outputs_re)
+        l_hw_b_c = self.norm1(l_hw_b_c)
+        l_hw_b_c = l_hw_b_c + self.dropout2(self.linear4(self.dropout(self.activation(self.linear3(l_hw_b_c)).clone())))
+        l_hw_b_c = self.norm2(l_hw_b_c)
+
+        l = torch.cat((l_hw_b_c, g_hw_b_c), 1)  # hw,b(5),c
+        return rearrange(l, "(h w) b c -> b c h w", h=h, w=w)  ## (5,c,h*w)
+
+
+class MCRM(nn.Module):
+    def __init__(self, d_model, num_heads, pool_ratios=[4, 8, 16], h=None):
+        super(MCRM, self).__init__()
+        self.attention = nn.ModuleList([
+            nn.MultiheadAttention(d_model, num_heads, dropout=0.1),
+            nn.MultiheadAttention(d_model, num_heads, dropout=0.1),
+            nn.MultiheadAttention(d_model, num_heads, dropout=0.1),
+            nn.MultiheadAttention(d_model, num_heads, dropout=0.1)
+        ])
+        self.linear3 = nn.Linear(d_model, d_model * 2)
+        self.linear4 = nn.Linear(d_model * 2, d_model)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(0.1)
+        self.dropout1 = nn.Dropout(0.1)
+        self.dropout2 = nn.Dropout(0.1)
+        self.sigmoid = nn.Sigmoid()
+        self.activation = get_activation_fn('gelu')
+        self.sal_conv = nn.Conv2d(d_model, 1, 1)
+        self.pool_ratios = pool_ratios
+
+    def forward(self, x):
+        device = x.device
+        b, c, h, w = x.size()
+        loc, glb = x.split([4, 1], dim=0)  # 4,c,h,w; 1,c,h,w
+
+        patched_glb = rearrange(glb, 'b c (hg h) (wg w) -> (hg wg b) c h w', hg=2, wg=2)
+
+        token_attention_map = self.sigmoid(self.sal_conv(glb))
+        token_attention_map = F.interpolate(token_attention_map, size=patches2image(loc).shape[-2:], mode='nearest')
+        loc = loc * rearrange(token_attention_map, 'b c (hg h) (wg w) -> (hg wg b) c h w', hg=2, wg=2)
+
+        pools = []
+        for pool_ratio in self.pool_ratios:
+            tgt_hw = (round(h / pool_ratio), round(w / pool_ratio))
+            pool = F.adaptive_avg_pool2d(patched_glb, tgt_hw)
+            pools.append(rearrange(pool, 'nl c h w -> nl c (h w)'))  # nl(4),c,hw
+
+        pools = rearrange(torch.cat(pools, 2), "nl c nphw -> nl nphw 1 c")
+        loc_ = rearrange(loc, 'nl c h w -> nl (h w) 1 c')
+
+        outputs = []
+        for i, q in enumerate(loc_.unbind(dim=0)):  # traverse all local patches
+            v = pools[i]
+            k = v
+            outputs.append(self.attention[i](q, k, v)[0])
+
+        outputs = torch.cat(outputs, 1)
+        src = loc.view(4, c, -1).permute(2, 0, 1) + self.dropout1(outputs)
+        src = self.norm1(src)
+        src = src + self.dropout2(self.linear4(self.dropout(self.activation(self.linear3(src)).clone())))
+        src = self.norm2(src)
+        src = src.permute(1, 2, 0).reshape(4, c, h, w)  # freshed loc
+        glb = glb + F.interpolate(patches2image(src), size=glb.shape[-2:], mode='nearest')  # freshed glb
+
+        return torch.cat((src, glb), 0), token_attention_map
+
+
+class BEN_Base(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.backbone = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12)
+        emb_dim = 128
+        self.sideout5 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1))
+        self.sideout4 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1))
+        self.sideout3 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1))
+        self.sideout2 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1))
+        self.sideout1 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1))
+
+        self.output5 = make_cbr(1024, emb_dim)
+        self.output4 = make_cbr(512, emb_dim)
+        self.output3 = make_cbr(256, emb_dim)
+        self.output2 = make_cbr(128, emb_dim)
+        self.output1 = make_cbr(128, emb_dim)
+
+        self.multifieldcrossatt = MCLM(emb_dim, 1, [1, 4, 8])
+        self.conv1 = make_cbr(emb_dim, emb_dim)
+        self.conv2 = make_cbr(emb_dim, emb_dim)
+        self.conv3 = make_cbr(emb_dim, emb_dim)
+        self.conv4 = make_cbr(emb_dim, emb_dim)
+        self.dec_blk1 = MCRM(emb_dim, 1, [2, 4, 8])
+        self.dec_blk2 = MCRM(emb_dim, 1, [2, 4, 8])
+        self.dec_blk3 = MCRM(emb_dim, 1, [2, 4, 8])
+        self.dec_blk4 = MCRM(emb_dim, 1, [2, 4, 8])
+
+        self.insmask_head = nn.Sequential(
+            nn.Conv2d(emb_dim, 384, kernel_size=3, padding=1),
+            nn.InstanceNorm2d(384),
+            nn.GELU(),
+            nn.Conv2d(384, 384, kernel_size=3, padding=1),
+            nn.InstanceNorm2d(384),
+            nn.GELU(),
+            nn.Conv2d(384, emb_dim, kernel_size=3, padding=1)
+        )
+
+        self.shallow = nn.Sequential(nn.Conv2d(3, emb_dim, kernel_size=3, padding=1))
+        self.upsample1 = make_cbg(emb_dim, emb_dim)
+        self.upsample2 = make_cbg(emb_dim, emb_dim)
+        self.output = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1))
+
+        for m in self.modules():
+            if isinstance(m, nn.GELU) or isinstance(m, nn.Dropout):
+                m.inplace = True
+
+    @torch.inference_mode()
+    @torch.autocast(device_type="cuda", dtype=torch.float16)
+    def forward(self, x):
+        real_batch = x.size(0)
+
+        shallow_batch = self.shallow(x)
+        glb_batch = rescale_to(x, scale_factor=0.5, interpolation='bilinear')
+
+        final_input = None
+        for i in range(real_batch):
+            start = i * 4
+            end = (i + 1) * 4
+            loc_batch = image2patches(x[i, :, :, :].unsqueeze(dim=0))
+            input_ = torch.cat((loc_batch, glb_batch[i, :, :, :].unsqueeze(dim=0)), dim=0)
+
+            if final_input == None:
+                final_input = input_
+            else:
+                final_input = torch.cat((final_input, input_), dim=0)
+
+        features = self.backbone(final_input)
+        outputs = []
+
+        for i in range(real_batch):
+            start = i * 5
+            end = (i + 1) * 5
+
+            f4 = features[4][start:end, :, :, :]  # shape: [5, C, H, W]
+            f3 = features[3][start:end, :, :, :]
+            f2 = features[2][start:end, :, :, :]
+            f1 = features[1][start:end, :, :, :]
+            f0 = features[0][start:end, :, :, :]
+            e5 = self.output5(f4)
+            e4 = self.output4(f3)
+            e3 = self.output3(f2)
+            e2 = self.output2(f1)
+            e1 = self.output1(f0)
+            loc_e5, glb_e5 = e5.split([4, 1], dim=0)
+            e5 = self.multifieldcrossatt(loc_e5, glb_e5)  # (4,128,16,16)
+
+            e4, tokenattmap4 = self.dec_blk4(e4 + resize_as(e5, e4))
+            e4 = self.conv4(e4)
+            e3, tokenattmap3 = self.dec_blk3(e3 + resize_as(e4, e3))
+            e3 = self.conv3(e3)
+            e2, tokenattmap2 = self.dec_blk2(e2 + resize_as(e3, e2))
+            e2 = self.conv2(e2)
+            e1, tokenattmap1 = self.dec_blk1(e1 + resize_as(e2, e1))
+            e1 = self.conv1(e1)
+
+            loc_e1, glb_e1 = e1.split([4, 1], dim=0)
+
+            output1_cat = patches2image(loc_e1)  # (1,128,256,256)
+
+            # add glb feat in
+            output1_cat = output1_cat + resize_as(glb_e1, output1_cat)
+            # merge
+            final_output = self.insmask_head(output1_cat)  # (1,128,256,256)
+            # shallow feature merge
+            shallow = shallow_batch[i, :, :, :].unsqueeze(dim=0)
+            final_output = final_output + resize_as(shallow, final_output)
+            final_output = self.upsample1(rescale_to(final_output))
+            final_output = rescale_to(final_output + resize_as(shallow, final_output))
+            final_output = self.upsample2(final_output)
+            final_output = self.output(final_output)
+            mask = final_output.sigmoid()
+            outputs.append(mask)
+
+        return torch.cat(outputs, dim=0)
+
+    def loadcheckpoints(self, model_path):
+        model_dict = torch.load(model_path, map_location="cpu", weights_only=True)
+        self.load_state_dict(model_dict['model_state_dict'], strict=True)
+        del model_path
+
+    def inference(self, image, refine_foreground=False):
+
+        set_random_seed(9)
+        # image = ImageOps.exif_transpose(image)
+        if isinstance(image, Image.Image):
+            image, h, w, original_image = rgb_loader_refiner(image)
+            if torch.cuda.is_available():
+
+                img_tensor = img_transform(image).unsqueeze(0).to(next(self.parameters()).device)
+            else:
+                img_tensor = img_transform32(image).unsqueeze(0).to(next(self.parameters()).device)
+
+            with torch.no_grad():
+                res = self.forward(img_tensor)
+
+            # Show Results
+            if refine_foreground == True:
+
+                pred_pil = transforms.ToPILImage()(res.squeeze())
+                image_masked = refine_foreground_process(original_image, pred_pil)
+                mask = pred_pil.resize(original_image.size)
+                image_masked.putalpha(mask)
+                return mask, image_masked
+
+            else:
+                alpha = postprocess_image(res, im_size=[w, h])
+                pred_pil = transforms.ToPILImage()(alpha)
+                mask = pred_pil.resize(original_image.size)
+                original_image.putalpha(mask)
+                # mask = Image.fromarray(alpha)
+
+                return mask, original_image
+
+    def segment_video(self, video_path, output_path="./", fps=0, refine_foreground=False, batch=1,
+                      print_frames_processed=True, webm=False, rgb_value=(0, 255, 0)):
+
+        """
+        Segments the given video to extract the foreground (with alpha) from each frame
+        and saves the result as either a WebM video (with alpha channel) or MP4 (with a
+        color background).
+        Args:
+            video_path (str):
+                Path to the input video file.
+            output_path (str, optional):
+                Directory (or full path) where the output video and/or files will be saved.
+                Defaults to "./".
+            fps (int, optional):
+                The frames per second (FPS) to use for the output video. If 0 (default), the
+                original FPS of the input video is used. Otherwise, overrides it.
+            refine_foreground (bool, optional):
+                Whether to run an additional “refine foreground” process on each frame.
+                Defaults to False.
+            batch (int, optional):
+                Number of frames to process at once (inference batch size). Large batch sizes
+                may require more GPU memory. Defaults to 1.
+            print_frames_processed (bool, optional):
+                If True (default), prints progress (how many frames have been processed) to
+                the console.
+            webm (bool, optional):
+                If True (default), exports a WebM video with alpha channel (VP9 / yuva420p).
+                If False, exports an MP4 video composited over a solid color background.
+            rgb_value (tuple, optional):
+                The RGB background color (e.g., green screen) used to composite frames when
+                saving to MP4. Defaults to (0, 255, 0).
+        Returns:
+            None. Writes the output video(s) to disk in the specified format.
+        """
+
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            raise IOError(f"Cannot open video: {video_path}")
+
+        original_fps = cap.get(cv2.CAP_PROP_FPS)
+        original_fps = 30 if original_fps == 0 else original_fps
+        fps = original_fps if fps == 0 else fps
+
+        ret, first_frame = cap.read()
+        if not ret:
+            raise ValueError("No frames found in the video.")
+        height, width = first_frame.shape[:2]
+        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+
+        foregrounds = []
+        frame_idx = 0
+        processed_count = 0
+        batch_frames = []
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                if batch_frames:
+                    batch_results = self.inference(batch_frames, refine_foreground)
+                    if isinstance(batch_results, Image.Image):
+                        foregrounds.append(batch_results)
+                    else:
+                        foregrounds.extend(batch_results)
+                    if print_frames_processed:
+                        print(f"Processed frames {frame_idx - len(batch_frames) + 1} to {frame_idx} of {total_frames}")
+                break
+
+            # Process every frame instead of using intervals
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            pil_frame = Image.fromarray(frame_rgb)
+            batch_frames.append(pil_frame)
+
+            if len(batch_frames) == batch:
+                batch_results = self.inference(batch_frames, refine_foreground)
+                if isinstance(batch_results, Image.Image):
+                    foregrounds.append(batch_results)
+                else:
+                    foregrounds.extend(batch_results)
+                if print_frames_processed:
+                    print(f"Processed frames {frame_idx - batch + 1} to {frame_idx} of {total_frames}")
+                batch_frames = []
+                processed_count += batch
+
+            frame_idx += 1
+
+        if webm:
+            alpha_webm_path = os.path.join(output_path, "foreground.webm")
+            pil_images_to_webm_alpha(foregrounds, alpha_webm_path, fps=original_fps)
+
+        else:
+            cap.release()
+            fg_output = os.path.join(output_path, 'foreground.mp4')
+
+            pil_images_to_mp4(foregrounds, fg_output, fps=original_fps, rgb_value=rgb_value)
+            cv2.destroyAllWindows()
+
+            try:
+                fg_audio_output = os.path.join(output_path, 'foreground_output_with_audio.mp4')
+                add_audio_to_video(fg_output, video_path, fg_audio_output)
+            except Exception as e:
+                print("No audio found in the original video")
+                print(e)
+
+
+def rgb_loader_refiner(original_image):
+    h, w = original_image.size
+
+    image = original_image
+    # Convert to RGB if necessary
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+
+    # Resize the image
+    image = image.resize((1024, 1024), resample=Image.LANCZOS)
+
+    return image.convert('RGB'), h, w, original_image
+
+
+# Define the image transformation
+img_transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.ConvertImageDtype(torch.float16),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+img_transform32 = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.ConvertImageDtype(torch.float32),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+
+def pil_images_to_mp4(images, output_path, fps=24, rgb_value=(0, 255, 0)):
+    """
+    Converts an array of PIL images to an MP4 video.
+
+    Args:
+        images: List of PIL images
+        output_path: Path to save the MP4 file
+        fps: Frames per second (default: 24)
+        rgb_value: Background RGB color tuple (default: green (0, 255, 0))
+    """
+    if not images:
+        raise ValueError("No images provided to convert to MP4.")
+
+    width, height = images[0].size
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video_writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    for image in images:
+        # If image has alpha channel, composite onto the specified background color
+        if image.mode == 'RGBA':
+            # Create background image with specified RGB color
+            background = Image.new('RGB', image.size, rgb_value)
+            background = background.convert('RGBA')
+            # Composite the image onto the background
+            image = Image.alpha_composite(background, image)
+            image = image.convert('RGB')
+        else:
+            # Ensure RGB format for non-alpha images
+            image = image.convert('RGB')
+
+        # Convert to OpenCV format and write
+        open_cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        video_writer.write(open_cv_image)
+
+    video_writer.release()
+
+
+def pil_images_to_webm_alpha(images, output_path, fps=30):
+    """
+    Converts a list of PIL RGBA images to a VP9 .webm video with alpha channel.
+    NOTE: Not all players will display alpha in WebM.
+          Browsers like Chrome/Firefox typically do support VP9 alpha.
+    """
+    if not images:
+        raise ValueError("No images provided for WebM with alpha.")
+
+    # Ensure output directory exists
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Save frames as PNG (with alpha)
+        for idx, img in enumerate(images):
+            if img.mode != "RGBA":
+                img = img.convert("RGBA")
+            out_path = os.path.join(tmpdir, f"{idx:06d}.png")
+            img.save(out_path, "PNG")
+
+        # Construct ffmpeg command
+        # -c:v libvpx-vp9 => VP9 encoder
+        # -pix_fmt yuva420p => alpha-enabled pixel format
+        # -auto-alt-ref 0 => helps preserve alpha frames (libvpx quirk)
+        ffmpeg_cmd = [
+            "ffmpeg", "-y",
+            "-framerate", str(fps),
+            "-i", os.path.join(tmpdir, "%06d.png"),
+            "-c:v", "libvpx-vp9",
+            "-pix_fmt", "yuva420p",
+            "-auto-alt-ref", "0",
+            output_path
+        ]
+
+        subprocess.run(ffmpeg_cmd, check=True)
+
+    print(f"WebM with alpha saved to {output_path}")
+
+
+def add_audio_to_video(video_without_audio_path, original_video_path, output_path):
+    """
+    Check if the original video has an audio stream. If yes, add it. If not, skip.
+    """
+    # 1) Probe original video for audio streams
+    probe_command = [
+        'ffprobe', '-v', 'error',
+        '-select_streams', 'a:0',
+        '-show_entries', 'stream=index',
+        '-of', 'csv=p=0',
+        original_video_path
+    ]
+    result = subprocess.run(probe_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+
+    # result.stdout is empty if no audio stream found
+    if not result.stdout.strip():
+        print("No audio track found in original video, skipping audio addition.")
+        return
+
+    print("Audio track detected; proceeding to mux audio.")
+    # 2) If audio found, run ffmpeg to add it
+    command = [
+        'ffmpeg', '-y',
+        '-i', video_without_audio_path,
+        '-i', original_video_path,
+        '-c', 'copy',
+        '-map', '0:v:0',
+        '-map', '1:a:0',  # we know there's an audio track now
+        output_path
+    ]
+    subprocess.run(command, check=True)
+    print(f"Audio added successfully => {output_path}")
+
+
+### Thanks to the source: https://huggingface.co/ZhengPeng7/BiRefNet/blob/main/handler.py
+def refine_foreground_process(image, mask, r=90):
+    if mask.size != image.size:
+        mask = mask.resize(image.size)
+    image = np.array(image) / 255.0
+    mask = np.array(mask) / 255.0
+    estimated_foreground = FB_blur_fusion_foreground_estimator_2(image, mask, r=r)
+    image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
+    return image_masked
+
+
+def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90):
+    # Thanks to the source: https://github.com/Photoroom/fast-foreground-estimation
+    alpha = alpha[:, :, None]
+    F, blur_B = FB_blur_fusion_foreground_estimator(image, image, image, alpha, r)
+    return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]
+
+
+def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
+    if isinstance(image, Image.Image):
+        image = np.array(image) / 255.0
+    blurred_alpha = cv2.blur(alpha, (r, r))[:, :, None]
+
+    blurred_FA = cv2.blur(F * alpha, (r, r))
+    blurred_F = blurred_FA / (blurred_alpha + 1e-5)
+
+    blurred_B1A = cv2.blur(B * (1 - alpha), (r, r))
+    blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5)
+    F = blurred_F + alpha * \
+        (image - alpha * blurred_F - (1 - alpha) * blurred_B)
+    F = np.clip(F, 0, 1)
+    return F, blurred_B
+
+
+def postprocess_image(result: torch.Tensor, im_size: list) -> np.ndarray:
+    result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear'), 0)
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result - mi) / (ma - mi)
+    im_array = (result * 255).permute(1, 2, 0).cpu().data.numpy().astype(np.uint8)
+    im_array = np.squeeze(im_array)
+    return im_array
+
+
+def rgb_loader_refiner(original_image):
+    h, w = original_image.size
+    # # Apply EXIF orientation
+
+    image = ImageOps.exif_transpose(original_image)
+
+    if original_image.mode != 'RGB':
+        original_image = original_image.convert('RGB')
+
+    image = original_image
+    # Convert to RGB if necessary
+
+    # Resize the image
+    image = image.resize((1024, 1024), resample=Image.LANCZOS)
+
+    return image, h, w, original_image
+

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/bitsandbytes_NF4/__init__.py

@@ -0,0 +1,167 @@
+#credit to comfyanonymous for this module
+#from https://github.com/comfyanonymous/ComfyUI_bitsandbytes_NF4
+import comfy.ops
+import torch
+import folder_paths
+from ...libs.utils import install_package
+
+try:
+    from bitsandbytes.nn.modules import Params4bit, QuantState
+except ImportError:
+    Params4bit = torch.nn.Parameter
+    raise ImportError("Please install bitsandbytes>=0.43.3")
+
+def functional_linear_4bits(x, weight, bias):
+    try:
+        install_package("bitsandbytes", "0.43.3", True, "0.43.3")
+        import bitsandbytes as bnb
+    except ImportError:
+        raise ImportError("Please install bitsandbytes>=0.43.3")
+
+    out = bnb.matmul_4bit(x, weight.t(), bias=bias, quant_state=weight.quant_state)
+    out = out.to(x)
+    return out
+
+
+def copy_quant_state(state, device: torch.device = None):
+    if state is None:
+        return None
+
+    device = device or state.absmax.device
+
+    state2 = (
+        QuantState(
+            absmax=state.state2.absmax.to(device),
+            shape=state.state2.shape,
+            code=state.state2.code.to(device),
+            blocksize=state.state2.blocksize,
+            quant_type=state.state2.quant_type,
+            dtype=state.state2.dtype,
+        )
+        if state.nested
+        else None
+    )
+
+    return QuantState(
+        absmax=state.absmax.to(device),
+        shape=state.shape,
+        code=state.code.to(device),
+        blocksize=state.blocksize,
+        quant_type=state.quant_type,
+        dtype=state.dtype,
+        offset=state.offset.to(device) if state.nested else None,
+        state2=state2,
+    )
+
+
+class ForgeParams4bit(Params4bit):
+
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
+        if device is not None and device.type == "cuda" and not self.bnb_quantized:
+            return self._quantize(device)
+        else:
+            n = ForgeParams4bit(
+                torch.nn.Parameter.to(self, device=device, dtype=dtype, non_blocking=non_blocking),
+                requires_grad=self.requires_grad,
+                quant_state=copy_quant_state(self.quant_state, device),
+                blocksize=self.blocksize,
+                compress_statistics=self.compress_statistics,
+                quant_type=self.quant_type,
+                quant_storage=self.quant_storage,
+                bnb_quantized=self.bnb_quantized,
+                module=self.module
+            )
+            self.module.quant_state = n.quant_state
+            self.data = n.data
+            self.quant_state = n.quant_state
+            return n
+
+class ForgeLoader4Bit(torch.nn.Module):
+    def __init__(self, *, device, dtype, quant_type, **kwargs):
+        super().__init__()
+        self.dummy = torch.nn.Parameter(torch.empty(1, device=device, dtype=dtype))
+        self.weight = None
+        self.quant_state = None
+        self.bias = None
+        self.quant_type = quant_type
+
+    def _save_to_state_dict(self, destination, prefix, keep_vars):
+        super()._save_to_state_dict(destination, prefix, keep_vars)
+        quant_state = getattr(self.weight, "quant_state", None)
+        if quant_state is not None:
+            for k, v in quant_state.as_dict(packed=True).items():
+                destination[prefix + "weight." + k] = v if keep_vars else v.detach()
+        return
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        quant_state_keys = {k[len(prefix + "weight."):] for k in state_dict.keys() if k.startswith(prefix + "weight.")}
+
+        if any('bitsandbytes' in k for k in quant_state_keys):
+            quant_state_dict = {k: state_dict[prefix + "weight." + k] for k in quant_state_keys}
+
+            self.weight = ForgeParams4bit().from_prequantized(
+                data=state_dict[prefix + 'weight'],
+                quantized_stats=quant_state_dict,
+                requires_grad=False,
+                device=self.dummy.device,
+                module=self
+            )
+            self.quant_state = self.weight.quant_state
+
+            if prefix + 'bias' in state_dict:
+                self.bias = torch.nn.Parameter(state_dict[prefix + 'bias'].to(self.dummy))
+
+            del self.dummy
+        elif hasattr(self, 'dummy'):
+            if prefix + 'weight' in state_dict:
+                self.weight = ForgeParams4bit(
+                    state_dict[prefix + 'weight'].to(self.dummy),
+                    requires_grad=False,
+                    compress_statistics=True,
+                    quant_type=self.quant_type,
+                    quant_storage=torch.uint8,
+                    module=self,
+                )
+                self.quant_state = self.weight.quant_state
+
+            if prefix + 'bias' in state_dict:
+                self.bias = torch.nn.Parameter(state_dict[prefix + 'bias'].to(self.dummy))
+
+            del self.dummy
+        else:
+            super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+current_device = None
+current_dtype = None
+current_manual_cast_enabled = False
+current_bnb_dtype = None
+
+class OPS(comfy.ops.manual_cast):
+    class Linear(ForgeLoader4Bit):
+        def __init__(self, *args, device=None, dtype=None, **kwargs):
+            super().__init__(device=device, dtype=dtype, quant_type=current_bnb_dtype)
+            self.parameters_manual_cast = current_manual_cast_enabled
+
+        def forward(self, x):
+            self.weight.quant_state = self.quant_state
+
+            if self.bias is not None and self.bias.dtype != x.dtype:
+                # Maybe this can also be set to all non-bnb ops since the cost is very low.
+                # And it only invokes one time, and most linear does not have bias
+                self.bias.data = self.bias.data.to(x.dtype)
+
+            if not self.parameters_manual_cast:
+                return functional_linear_4bits(x, self.weight, self.bias)
+            elif not self.weight.bnb_quantized:
+                assert x.device.type == 'cuda', 'BNB Must Use CUDA as Computation Device!'
+                layer_original_device = self.weight.device
+                self.weight = self.weight._quantize(x.device)
+                bias = self.bias.to(x.device) if self.bias is not None else None
+                out = functional_linear_4bits(x, self.weight, bias)
+                self.weight = self.weight.to(layer_original_device)
+                return out
+            else:
+                weight, bias, signal = weights_manual_cast(self, x, skip_weight_dtype=True, skip_bias_dtype=True)
+                with main_stream_worker(weight, bias, signal):
+                    return functional_linear_4bits(x, weight, bias)

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/briaai/rembg.py

@@ -0,0 +1,475 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+import numpy as np
+
+
+class REBNCONV(nn.Module):
+    def __init__(self,in_ch=3,out_ch=3,dirate=1,stride=1):
+        super(REBNCONV,self).__init__()
+
+        self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate,stride=stride)
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src,tar):
+
+    src = F.interpolate(src,size=tar.shape[2:],mode='bilinear')
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3, img_size=512):
+        super(RSU7,self).__init__()
+
+        self.in_ch = in_ch
+        self.mid_ch = mid_ch
+        self.out_ch = out_ch
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1) ## 1 -> 1/2
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+        b, c, h, w = x.shape
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d =  self.rebnconv6d(torch.cat((hx7,hx6),1))
+        hx6dup = _upsample_like(hx6d,hx5)
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6dup,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+
+### RSU-6 ###
+class RSU6(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-5 ###
+class RSU5(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4 ###
+class RSU4(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4F ###
+class RSU4F(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1))
+
+        return hx1d + hxin
+
+
+class myrebnconv(nn.Module):
+    def __init__(self, in_ch=3,
+                       out_ch=1,
+                       kernel_size=3,
+                       stride=1,
+                       padding=1,
+                       dilation=1,
+                       groups=1):
+        super(myrebnconv,self).__init__()
+
+        self.conv = nn.Conv2d(in_ch,
+                              out_ch,
+                              kernel_size=kernel_size,
+                              stride=stride,
+                              padding=padding,
+                              dilation=dilation,
+                              groups=groups)
+        self.bn = nn.BatchNorm2d(out_ch)
+        self.rl = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+        return self.rl(self.bn(self.conv(x)))
+
+def preprocess_image(im, model_input_size: list) -> torch.Tensor:
+    # im = im.resize(model_input_size, Image.BILINEAR)
+    im_np = np.array(im)
+    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
+    im_tensor = F.interpolate(torch.unsqueeze(im_tensor,0), size=model_input_size, mode='bilinear').type(torch.uint8)
+    image = torch.divide(im_tensor,255.0)
+    image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])
+    return image
+
+def postprocess_image(result: torch.Tensor, im_size: list)-> np.ndarray:
+    result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear') ,0)
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result-mi)/(ma-mi)
+    im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)
+    im_array = np.squeeze(im_array)
+    return im_array
+
+class BriaRMBG(nn.Module):
+
+    def __init__(self, config:dict={"in_ch":3,"out_ch":1}):
+        super(BriaRMBG,self).__init__()
+        in_ch = config["in_ch"]
+        out_ch = config["out_ch"]
+        self.conv_in = nn.Conv2d(in_ch,64,3,stride=2,padding=1)
+        self.pool_in = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage1 = RSU7(64,32,64)
+        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage2 = RSU6(64,32,128)
+        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage3 = RSU5(128,64,256)
+        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage4 = RSU4(256,128,512)
+        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage5 = RSU4F(512,256,512)
+        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage6 = RSU4F(512,256,512)
+
+        # decoder
+        self.stage5d = RSU4F(1024,256,512)
+        self.stage4d = RSU4(1024,128,256)
+        self.stage3d = RSU5(512,64,128)
+        self.stage2d = RSU6(256,32,64)
+        self.stage1d = RSU7(128,16,64)
+
+        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
+        self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
+        self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
+        self.side6 = nn.Conv2d(512,out_ch,3,padding=1)
+
+        # self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.conv_in(hx)
+        #hx = self.pool_in(hxin)
+
+        #stage 1
+        hx1 = self.stage1(hxin)
+        hx = self.pool12(hx1)
+
+        #stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        #stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        #stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        #stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        #stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6,hx5)
+
+        #-------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
+
+
+        #side output
+        d1 = self.side1(hx1d)
+        d1 = _upsample_like(d1,x)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2,x)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3,x)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4,x)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5,x)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6,x)
+
+        return [F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)],[hx1d,hx2d,hx3d,hx4d,hx5d,hx6]

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/brushnet/__init__.py

@@ -0,0 +1,822 @@
+#credit to nullquant for this module
+#from https://github.com/nullquant/ComfyUI-BrushNet
+
+import os
+import types
+
+import torch
+try:
+    from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+except:
+    init_empty_weights, load_checkpoint_and_dispatch = None, None
+
+import comfy
+
+try:
+    from .model import BrushNetModel, PowerPaintModel
+    from .model_patch import add_model_patch_option, patch_model_function_wrapper
+    from .powerpaint_utils import TokenizerWrapper, add_tokens
+except:
+    BrushNetModel, PowerPaintModel = None, None
+    add_model_patch_option, patch_model_function_wrapper = None, None
+    TokenizerWrapper, add_tokens = None, None
+
+cwd_path = os.path.dirname(os.path.realpath(__file__))
+brushnet_config_file = os.path.join(cwd_path, 'config', 'brushnet.json')
+brushnet_xl_config_file = os.path.join(cwd_path, 'config', 'brushnet_xl.json')
+powerpaint_config_file = os.path.join(cwd_path, 'config', 'powerpaint.json')
+
+sd15_scaling_factor = 0.18215
+sdxl_scaling_factor = 0.13025
+
+ModelsToUnload = [comfy.sd1_clip.SD1ClipModel, comfy.ldm.models.autoencoder.AutoencoderKL]
+
+class BrushNet:
+
+    # Check models compatibility
+    def check_compatibilty(self, model, brushnet):
+        is_SDXL = False
+        is_PP = False
+        if isinstance(model.model.model_config, comfy.supported_models.SD15):
+            print('Base model type: SD1.5')
+            is_SDXL = False
+            if brushnet["SDXL"]:
+                raise Exception("Base model is SD15, but BrushNet is SDXL type")
+            if brushnet["PP"]:
+                is_PP = True
+        elif isinstance(model.model.model_config, comfy.supported_models.SDXL):
+            print('Base model type: SDXL')
+            is_SDXL = True
+            if not brushnet["SDXL"]:
+                raise Exception("Base model is SDXL, but BrushNet is SD15 type")
+        else:
+            print('Base model type: ', type(model.model.model_config))
+            raise Exception("Unsupported model type: " + str(type(model.model.model_config)))
+
+        return (is_SDXL, is_PP)
+
+    def check_image_mask(self, image, mask, name):
+        if len(image.shape) < 4:
+            # image tensor shape should be [B, H, W, C], but batch somehow is missing
+            image = image[None, :, :, :]
+
+        if len(mask.shape) > 3:
+            # mask tensor shape should be [B, H, W] but we get [B, H, W, C], image may be?
+            # take first mask, red channel
+            mask = (mask[:, :, :, 0])[:, :, :]
+        elif len(mask.shape) < 3:
+            # mask tensor shape should be [B, H, W] but batch somehow is missing
+            mask = mask[None, :, :]
+
+        if image.shape[0] > mask.shape[0]:
+            print(name, "gets batch of images (%d) but only %d masks" % (image.shape[0], mask.shape[0]))
+            if mask.shape[0] == 1:
+                print(name, "will copy the mask to fill batch")
+                mask = torch.cat([mask] * image.shape[0], dim=0)
+            else:
+                print(name, "will add empty masks to fill batch")
+                empty_mask = torch.zeros([image.shape[0] - mask.shape[0], mask.shape[1], mask.shape[2]])
+                mask = torch.cat([mask, empty_mask], dim=0)
+        elif image.shape[0] < mask.shape[0]:
+            print(name, "gets batch of images (%d) but too many (%d) masks" % (image.shape[0], mask.shape[0]))
+            mask = mask[:image.shape[0], :, :]
+
+        return (image, mask)
+
+    # Prepare image and mask
+    def prepare_image(self, image, mask):
+
+        image, mask = self.check_image_mask(image, mask, 'BrushNet')
+
+        print("BrushNet image.shape =", image.shape, "mask.shape =", mask.shape)
+
+        if mask.shape[2] != image.shape[2] or mask.shape[1] != image.shape[1]:
+            raise Exception("Image and mask should be the same size")
+
+        # As a suggestion of inferno46n2 (https://github.com/nullquant/ComfyUI-BrushNet/issues/64)
+        mask = mask.round()
+
+        masked_image = image * (1.0 - mask[:, :, :, None])
+
+        return (masked_image, mask)
+
+    # Get origin of the mask
+    def cut_with_mask(self, mask, width, height):
+        iy, ix = (mask == 1).nonzero(as_tuple=True)
+
+        h0, w0 = mask.shape
+
+        if iy.numel() == 0:
+            x_c = w0 / 2.0
+            y_c = h0 / 2.0
+        else:
+            x_min = ix.min().item()
+            x_max = ix.max().item()
+            y_min = iy.min().item()
+            y_max = iy.max().item()
+
+            if x_max - x_min > width or y_max - y_min > height:
+                raise Exception("Mask is bigger than provided dimensions")
+
+            x_c = (x_min + x_max) / 2.0
+            y_c = (y_min + y_max) / 2.0
+
+        width2 = width / 2.0
+        height2 = height / 2.0
+
+        if w0 <= width:
+            x0 = 0
+            w = w0
+        else:
+            x0 = max(0, x_c - width2)
+            w = width
+            if x0 + width > w0:
+                x0 = w0 - width
+
+        if h0 <= height:
+            y0 = 0
+            h = h0
+        else:
+            y0 = max(0, y_c - height2)
+            h = height
+            if y0 + height > h0:
+                y0 = h0 - height
+
+        return (int(x0), int(y0), int(w), int(h))
+
+    # Prepare conditioning_latents
+    @torch.inference_mode()
+    def get_image_latents(self, masked_image, mask, vae, scaling_factor):
+        processed_image = masked_image.to(vae.device)
+        image_latents = vae.encode(processed_image[:, :, :, :3]) * scaling_factor
+        processed_mask = 1. - mask[:, None, :, :]
+        interpolated_mask = torch.nn.functional.interpolate(
+            processed_mask,
+            size=(
+                image_latents.shape[-2],
+                image_latents.shape[-1]
+            )
+        )
+        interpolated_mask = interpolated_mask.to(image_latents.device)
+
+        conditioning_latents = [image_latents, interpolated_mask]
+
+        print('BrushNet CL: image_latents shape =', image_latents.shape, 'interpolated_mask shape =',
+              interpolated_mask.shape)
+
+        return conditioning_latents
+
+    def brushnet_blocks(self, sd):
+        brushnet_down_block = 0
+        brushnet_mid_block = 0
+        brushnet_up_block = 0
+        for key in sd:
+            if 'brushnet_down_block' in key:
+                brushnet_down_block += 1
+            if 'brushnet_mid_block' in key:
+                brushnet_mid_block += 1
+            if 'brushnet_up_block' in key:
+                brushnet_up_block += 1
+        return (brushnet_down_block, brushnet_mid_block, brushnet_up_block, len(sd))
+
+    def get_model_type(self, brushnet_file):
+        sd = comfy.utils.load_torch_file(brushnet_file)
+        brushnet_down_block, brushnet_mid_block, brushnet_up_block, keys = self.brushnet_blocks(sd)
+        del sd
+        if brushnet_down_block == 24 and brushnet_mid_block == 2 and brushnet_up_block == 30:
+            is_SDXL = False
+            if keys == 322:
+                is_PP = False
+                print('BrushNet model type: SD1.5')
+            else:
+                is_PP = True
+                print('PowerPaint model type: SD1.5')
+        elif brushnet_down_block == 18 and brushnet_mid_block == 2 and brushnet_up_block == 22:
+            print('BrushNet model type: Loading SDXL')
+            is_SDXL = True
+            is_PP = False
+        else:
+            raise Exception("Unknown BrushNet model")
+        return is_SDXL, is_PP
+
+    def load_brushnet_model(self, brushnet_file, dtype='float16'):
+        is_SDXL, is_PP = self.get_model_type(brushnet_file)
+        with init_empty_weights():
+            if is_SDXL:
+                brushnet_config = BrushNetModel.load_config(brushnet_xl_config_file)
+                brushnet_model = BrushNetModel.from_config(brushnet_config)
+            elif is_PP:
+                brushnet_config = PowerPaintModel.load_config(powerpaint_config_file)
+                brushnet_model = PowerPaintModel.from_config(brushnet_config)
+            else:
+                brushnet_config = BrushNetModel.load_config(brushnet_config_file)
+                brushnet_model = BrushNetModel.from_config(brushnet_config)
+        if is_PP:
+            print("PowerPaint model file:", brushnet_file)
+        else:
+            print("BrushNet model file:", brushnet_file)
+
+        if dtype == 'float16':
+            torch_dtype = torch.float16
+        elif dtype == 'bfloat16':
+            torch_dtype = torch.bfloat16
+        elif dtype == 'float32':
+            torch_dtype = torch.float32
+        else:
+            torch_dtype = torch.float64
+
+        brushnet_model = load_checkpoint_and_dispatch(
+            brushnet_model,
+            brushnet_file,
+            device_map="sequential",
+            max_memory=None,
+            offload_folder=None,
+            offload_state_dict=False,
+            dtype=torch_dtype,
+            force_hooks=False,
+        )
+
+        if is_PP:
+            print("PowerPaint model is loaded")
+        elif is_SDXL:
+            print("BrushNet SDXL model is loaded")
+        else:
+            print("BrushNet SD1.5 model is loaded")
+
+        return ({"brushnet": brushnet_model, "SDXL": is_SDXL, "PP": is_PP, "dtype": torch_dtype},)
+
+    def brushnet_model_update(self, model, vae, image, mask, brushnet, positive, negative, scale, start_at, end_at):
+
+        is_SDXL, is_PP = self.check_compatibilty(model, brushnet)
+
+        if is_PP:
+            raise Exception("PowerPaint model was loaded, please use PowerPaint node")
+
+            # Make a copy of the model so that we're not patching it everywhere in the workflow.
+        model = model.clone()
+
+        # prepare image and mask
+        # no batches for original image and mask
+        masked_image, mask = self.prepare_image(image, mask)
+
+        batch = masked_image.shape[0]
+        width = masked_image.shape[2]
+        height = masked_image.shape[1]
+
+        if hasattr(model.model.model_config, 'latent_format') and hasattr(model.model.model_config.latent_format,
+                                                                          'scale_factor'):
+            scaling_factor = model.model.model_config.latent_format.scale_factor
+        elif is_SDXL:
+            scaling_factor = sdxl_scaling_factor
+        else:
+            scaling_factor = sd15_scaling_factor
+
+        torch_dtype = brushnet['dtype']
+
+        # prepare conditioning latents
+        conditioning_latents = self.get_image_latents(masked_image, mask, vae, scaling_factor)
+        conditioning_latents[0] = conditioning_latents[0].to(dtype=torch_dtype).to(brushnet['brushnet'].device)
+        conditioning_latents[1] = conditioning_latents[1].to(dtype=torch_dtype).to(brushnet['brushnet'].device)
+
+        # unload vae
+        del vae
+        # for loaded_model in comfy.model_management.current_loaded_models:
+        #     if type(loaded_model.model.model) in ModelsToUnload:
+        #         comfy.model_management.current_loaded_models.remove(loaded_model)
+        #         loaded_model.model_unload()
+        #         del loaded_model
+
+        # prepare embeddings
+        prompt_embeds = positive[0][0].to(dtype=torch_dtype).to(brushnet['brushnet'].device)
+        negative_prompt_embeds = negative[0][0].to(dtype=torch_dtype).to(brushnet['brushnet'].device)
+
+        max_tokens = max(prompt_embeds.shape[1], negative_prompt_embeds.shape[1])
+        if prompt_embeds.shape[1] < max_tokens:
+            multiplier = max_tokens // 77 - prompt_embeds.shape[1] // 77
+            prompt_embeds = torch.concat([prompt_embeds] + [prompt_embeds[:, -77:, :]] * multiplier, dim=1)
+            print('BrushNet: negative prompt more than 75 tokens:', negative_prompt_embeds.shape,
+                  'multiplying prompt_embeds')
+        if negative_prompt_embeds.shape[1] < max_tokens:
+            multiplier = max_tokens // 77 - negative_prompt_embeds.shape[1] // 77
+            negative_prompt_embeds = torch.concat(
+                [negative_prompt_embeds] + [negative_prompt_embeds[:, -77:, :]] * multiplier, dim=1)
+            print('BrushNet: positive prompt more than 75 tokens:', prompt_embeds.shape,
+                  'multiplying negative_prompt_embeds')
+
+        if len(positive[0]) > 1 and 'pooled_output' in positive[0][1] and positive[0][1]['pooled_output'] is not None:
+            pooled_prompt_embeds = positive[0][1]['pooled_output'].to(dtype=torch_dtype).to(brushnet['brushnet'].device)
+        else:
+            print('BrushNet: positive conditioning has not pooled_output')
+            if is_SDXL:
+                print('BrushNet will not produce correct results')
+            pooled_prompt_embeds = torch.empty([2, 1280], device=brushnet['brushnet'].device).to(dtype=torch_dtype)
+
+        if len(negative[0]) > 1 and 'pooled_output' in negative[0][1] and negative[0][1]['pooled_output'] is not None:
+            negative_pooled_prompt_embeds = negative[0][1]['pooled_output'].to(dtype=torch_dtype).to(
+                brushnet['brushnet'].device)
+        else:
+            print('BrushNet: negative conditioning has not pooled_output')
+            if is_SDXL:
+                print('BrushNet will not produce correct results')
+            negative_pooled_prompt_embeds = torch.empty([1, pooled_prompt_embeds.shape[1]],
+                                                        device=brushnet['brushnet'].device).to(dtype=torch_dtype)
+
+        time_ids = torch.FloatTensor([[height, width, 0., 0., height, width]]).to(dtype=torch_dtype).to(
+            brushnet['brushnet'].device)
+
+        if not is_SDXL:
+            pooled_prompt_embeds = None
+            negative_pooled_prompt_embeds = None
+            time_ids = None
+
+        # apply patch to model
+        brushnet_conditioning_scale = scale
+        control_guidance_start = start_at
+        control_guidance_end = end_at
+
+        add_brushnet_patch(model,
+                           brushnet['brushnet'],
+                           torch_dtype,
+                           conditioning_latents,
+                           (brushnet_conditioning_scale, control_guidance_start, control_guidance_end),
+                           prompt_embeds, negative_prompt_embeds,
+                           pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids,
+                           False)
+
+        latent = torch.zeros([batch, 4, conditioning_latents[0].shape[2], conditioning_latents[0].shape[3]],
+                             device=brushnet['brushnet'].device)
+
+        return (model, positive, negative, {"samples": latent},)
+
+    #powperpaint
+    def load_powerpaint_clip(self, base_clip_file, pp_clip_file):
+        pp_clip = comfy.sd.load_clip(ckpt_paths=[base_clip_file])
+
+        print('PowerPaint base CLIP file: ', base_clip_file)
+
+        pp_tokenizer = TokenizerWrapper(pp_clip.tokenizer.clip_l.tokenizer)
+        pp_text_encoder = pp_clip.patcher.model.clip_l.transformer
+
+        add_tokens(
+            tokenizer=pp_tokenizer,
+            text_encoder=pp_text_encoder,
+            placeholder_tokens=["P_ctxt", "P_shape", "P_obj"],
+            initialize_tokens=["a", "a", "a"],
+            num_vectors_per_token=10,
+        )
+
+        pp_text_encoder.load_state_dict(comfy.utils.load_torch_file(pp_clip_file), strict=False)
+
+        print('PowerPaint CLIP file: ', pp_clip_file)
+
+        pp_clip.tokenizer.clip_l.tokenizer = pp_tokenizer
+        pp_clip.patcher.model.clip_l.transformer = pp_text_encoder
+
+        return (pp_clip,)
+
+    def powerpaint_model_update(self, model, vae, image, mask, powerpaint, clip, positive, negative, fitting, function, scale, start_at, end_at, save_memory):
+        is_SDXL, is_PP = self.check_compatibilty(model, powerpaint)
+        if not is_PP:
+            raise Exception("BrushNet model was loaded, please use BrushNet node")
+
+            # Make a copy of the model so that we're not patching it everywhere in the workflow.
+        model = model.clone()
+
+        # prepare image and mask
+        # no batches for original image and mask
+        masked_image, mask = self.prepare_image(image, mask)
+
+        batch = masked_image.shape[0]
+        # width = masked_image.shape[2]
+        # height = masked_image.shape[1]
+
+        if hasattr(model.model.model_config, 'latent_format') and hasattr(model.model.model_config.latent_format,
+                                                                          'scale_factor'):
+            scaling_factor = model.model.model_config.latent_format.scale_factor
+        else:
+            scaling_factor = sd15_scaling_factor
+
+        torch_dtype = powerpaint['dtype']
+
+        # prepare conditioning latents
+        conditioning_latents = self.get_image_latents(masked_image, mask, vae, scaling_factor)
+        conditioning_latents[0] = conditioning_latents[0].to(dtype=torch_dtype).to(powerpaint['brushnet'].device)
+        conditioning_latents[1] = conditioning_latents[1].to(dtype=torch_dtype).to(powerpaint['brushnet'].device)
+
+        # prepare embeddings
+
+        if function == "object removal":
+            promptA = "P_ctxt"
+            promptB = "P_ctxt"
+            negative_promptA = "P_obj"
+            negative_promptB = "P_obj"
+            print('You should add to positive prompt: "empty scene blur"')
+            # positive = positive + " empty scene blur"
+        elif function == "context aware":
+            promptA = "P_ctxt"
+            promptB = "P_ctxt"
+            negative_promptA = ""
+            negative_promptB = ""
+            # positive = positive + " empty scene"
+            print('You should add to positive prompt: "empty scene"')
+        elif function == "shape guided":
+            promptA = "P_shape"
+            promptB = "P_ctxt"
+            negative_promptA = "P_shape"
+            negative_promptB = "P_ctxt"
+        elif function == "image outpainting":
+            promptA = "P_ctxt"
+            promptB = "P_ctxt"
+            negative_promptA = "P_obj"
+            negative_promptB = "P_obj"
+            # positive = positive + " empty scene"
+            print('You should add to positive prompt: "empty scene"')
+        else:
+            promptA = "P_obj"
+            promptB = "P_obj"
+            negative_promptA = "P_obj"
+            negative_promptB = "P_obj"
+
+        tokens = clip.tokenize(promptA)
+        prompt_embedsA = clip.encode_from_tokens(tokens, return_pooled=False)
+
+        tokens = clip.tokenize(negative_promptA)
+        negative_prompt_embedsA = clip.encode_from_tokens(tokens, return_pooled=False)
+
+        tokens = clip.tokenize(promptB)
+        prompt_embedsB = clip.encode_from_tokens(tokens, return_pooled=False)
+
+        tokens = clip.tokenize(negative_promptB)
+        negative_prompt_embedsB = clip.encode_from_tokens(tokens, return_pooled=False)
+
+        prompt_embeds_pp = (prompt_embedsA * fitting + (1.0 - fitting) * prompt_embedsB).to(dtype=torch_dtype).to(
+            powerpaint['brushnet'].device)
+        negative_prompt_embeds_pp = (negative_prompt_embedsA * fitting + (1.0 - fitting) * negative_prompt_embedsB).to(
+            dtype=torch_dtype).to(powerpaint['brushnet'].device)
+
+        # unload vae and CLIPs
+        del vae
+        del clip
+        # for loaded_model in comfy.model_management.current_loaded_models:
+        #     if type(loaded_model.model.model) in ModelsToUnload:
+        #         comfy.model_management.current_loaded_models.remove(loaded_model)
+        #         loaded_model.model_unload()
+        #         del loaded_model
+
+        # apply patch to model
+
+        brushnet_conditioning_scale = scale
+        control_guidance_start = start_at
+        control_guidance_end = end_at
+
+        if save_memory != 'none':
+            powerpaint['brushnet'].set_attention_slice(save_memory)
+
+        add_brushnet_patch(model,
+                           powerpaint['brushnet'],
+                           torch_dtype,
+                           conditioning_latents,
+                           (brushnet_conditioning_scale, control_guidance_start, control_guidance_end),
+                           negative_prompt_embeds_pp, prompt_embeds_pp,
+                           None, None, None,
+                           False)
+
+        latent = torch.zeros([batch, 4, conditioning_latents[0].shape[2], conditioning_latents[0].shape[3]],
+                             device=powerpaint['brushnet'].device)
+
+        return (model, positive, negative, {"samples": latent},)
+@torch.inference_mode()
+def brushnet_inference(x, timesteps, transformer_options, debug):
+    if 'model_patch' not in transformer_options:
+        print('BrushNet inference: there is no model_patch key in transformer_options')
+        return ([], 0, [])
+    mp = transformer_options['model_patch']
+    if 'brushnet' not in mp:
+        print('BrushNet inference: there is no brushnet key in mdel_patch')
+        return ([], 0, [])
+    bo = mp['brushnet']
+    if 'model' not in bo:
+        print('BrushNet inference: there is no model key in brushnet')
+        return ([], 0, [])
+    brushnet = bo['model']
+    if not (isinstance(brushnet, BrushNetModel) or isinstance(brushnet, PowerPaintModel)):
+        print('BrushNet model is not a BrushNetModel class')
+        return ([], 0, [])
+
+    torch_dtype = bo['dtype']
+    cl_list = bo['latents']
+    brushnet_conditioning_scale, control_guidance_start, control_guidance_end = bo['controls']
+    pe = bo['prompt_embeds']
+    npe = bo['negative_prompt_embeds']
+    ppe, nppe, time_ids = bo['add_embeds']
+
+    #do_classifier_free_guidance = mp['free_guidance']
+    do_classifier_free_guidance = len(transformer_options['cond_or_uncond']) > 1
+
+    x = x.detach().clone()
+    x = x.to(torch_dtype).to(brushnet.device)
+
+    timesteps = timesteps.detach().clone()
+    timesteps = timesteps.to(torch_dtype).to(brushnet.device)
+
+    total_steps = mp['total_steps']
+    step = mp['step']
+
+    added_cond_kwargs = {}
+
+    if do_classifier_free_guidance and step == 0:
+        print('BrushNet inference: do_classifier_free_guidance is True')
+
+    sub_idx = None
+    if 'ad_params' in transformer_options and 'sub_idxs' in transformer_options['ad_params']:
+        sub_idx = transformer_options['ad_params']['sub_idxs']
+
+    # we have batch input images
+    batch = cl_list[0].shape[0]
+    # we have incoming latents
+    latents_incoming = x.shape[0]
+    # and we already got some
+    latents_got = bo['latent_id']
+    if step == 0 or batch > 1:
+        print('BrushNet inference, step = %d: image batch = %d, got %d latents, starting from %d' \
+                % (step, batch, latents_incoming, latents_got))
+
+    image_latents = []
+    masks = []
+    prompt_embeds = []
+    negative_prompt_embeds = []
+    pooled_prompt_embeds = []
+    negative_pooled_prompt_embeds = []
+    if sub_idx:
+        # AnimateDiff indexes detected
+        if step == 0:
+            print('BrushNet inference: AnimateDiff indexes detected and applied')
+
+        batch = len(sub_idx)
+
+        if do_classifier_free_guidance:
+            for i in sub_idx:
+                image_latents.append(cl_list[0][i][None,:,:,:])
+                masks.append(cl_list[1][i][None,:,:,:])
+                prompt_embeds.append(pe)
+                negative_prompt_embeds.append(npe)
+                pooled_prompt_embeds.append(ppe)
+                negative_pooled_prompt_embeds.append(nppe)
+            for i in sub_idx:
+                image_latents.append(cl_list[0][i][None,:,:,:])
+                masks.append(cl_list[1][i][None,:,:,:])
+        else:
+            for i in sub_idx:
+                image_latents.append(cl_list[0][i][None,:,:,:])
+                masks.append(cl_list[1][i][None,:,:,:])
+                prompt_embeds.append(pe)
+                pooled_prompt_embeds.append(ppe)
+    else:
+        # do_classifier_free_guidance = 2 passes, 1st pass is cond, 2nd is uncond
+        continue_batch = True
+        for i in range(latents_incoming):
+            number = latents_got + i
+            if number < batch:
+                # 1st pass, cond
+                image_latents.append(cl_list[0][number][None,:,:,:])
+                masks.append(cl_list[1][number][None,:,:,:])
+                prompt_embeds.append(pe)
+                pooled_prompt_embeds.append(ppe)
+            elif do_classifier_free_guidance and number < batch * 2:
+                # 2nd pass, uncond
+                image_latents.append(cl_list[0][number-batch][None,:,:,:])
+                masks.append(cl_list[1][number-batch][None,:,:,:])
+                negative_prompt_embeds.append(npe)
+                negative_pooled_prompt_embeds.append(nppe)
+            else:
+                # latent batch
+                image_latents.append(cl_list[0][0][None,:,:,:])
+                masks.append(cl_list[1][0][None,:,:,:])
+                prompt_embeds.append(pe)
+                pooled_prompt_embeds.append(ppe)
+                latents_got = -i
+                continue_batch = False
+
+        if continue_batch:
+            # we don't have full batch yet
+            if do_classifier_free_guidance:
+                if number < batch * 2 - 1:
+                    bo['latent_id'] = number + 1
+                else:
+                    bo['latent_id'] = 0
+            else:
+                if number < batch - 1:
+                    bo['latent_id'] = number + 1
+                else:
+                    bo['latent_id'] = 0
+        else:
+            bo['latent_id'] = 0
+
+    cl = []
+    for il, m in zip(image_latents, masks):
+        cl.append(torch.concat([il, m], dim=1))
+    cl2apply = torch.concat(cl, dim=0)
+
+    conditioning_latents = cl2apply.to(torch_dtype).to(brushnet.device)
+
+    prompt_embeds.extend(negative_prompt_embeds)
+    prompt_embeds = torch.concat(prompt_embeds, dim=0).to(torch_dtype).to(brushnet.device)
+
+    if ppe is not None:
+        added_cond_kwargs = {}
+        added_cond_kwargs['time_ids'] = torch.concat([time_ids] * latents_incoming, dim = 0).to(torch_dtype).to(brushnet.device)
+
+        pooled_prompt_embeds.extend(negative_pooled_prompt_embeds)
+        pooled_prompt_embeds = torch.concat(pooled_prompt_embeds, dim=0).to(torch_dtype).to(brushnet.device)
+        added_cond_kwargs['text_embeds'] = pooled_prompt_embeds
+    else:
+        added_cond_kwargs = None
+
+    if x.shape[2] != conditioning_latents.shape[2] or x.shape[3] != conditioning_latents.shape[3]:
+        if step == 0:
+            print('BrushNet inference: image', conditioning_latents.shape, 'and latent', x.shape, 'have different size, resizing image')
+        conditioning_latents = torch.nn.functional.interpolate(
+            conditioning_latents, size=(
+                x.shape[2],
+                x.shape[3],
+            ), mode='bicubic',
+        ).to(torch_dtype).to(brushnet.device)
+
+    if step == 0:
+        print('BrushNet inference: sample', x.shape, ', CL', conditioning_latents.shape, 'dtype', torch_dtype)
+
+    if debug: print('BrushNet: step =', step)
+
+    if step < control_guidance_start or step > control_guidance_end:
+        cond_scale = 0.0
+    else:
+        cond_scale = brushnet_conditioning_scale
+
+    return brushnet(x,
+        encoder_hidden_states=prompt_embeds,
+        brushnet_cond=conditioning_latents,
+        timestep = timesteps,
+        conditioning_scale=cond_scale,
+        guess_mode=False,
+        added_cond_kwargs=added_cond_kwargs,
+        return_dict=False,
+        debug=debug,
+    )
+
+def add_brushnet_patch(model, brushnet, torch_dtype, conditioning_latents,
+                       controls,
+                       prompt_embeds, negative_prompt_embeds,
+                       pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids,
+                       debug):
+
+    is_SDXL = isinstance(model.model.model_config, comfy.supported_models.SDXL)
+
+    if model.model.model_config.custom_operations is None:
+        fp8 = model.model.model_config.optimizations.get("fp8", model.model.model_config.scaled_fp8 is not None)
+        operations = comfy.ops.pick_operations(model.model.model_config.unet_config.get("dtype", None), model.model.manual_cast_dtype,
+                                               fp8_optimizations=fp8, scaled_fp8=model.model.model_config.scaled_fp8)
+    else:
+        # such as gguf
+        operations = model.model.model_config.custom_operations
+
+    if is_SDXL:
+        input_blocks = [[0, operations.Conv2d],
+                        [1, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
+                        [2, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
+                        [3, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample],
+                        [4, comfy.ldm.modules.attention.SpatialTransformer],
+                        [5, comfy.ldm.modules.attention.SpatialTransformer],
+                        [6, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample],
+                        [7, comfy.ldm.modules.attention.SpatialTransformer],
+                        [8, comfy.ldm.modules.attention.SpatialTransformer]]
+        middle_block = [0, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock]
+        output_blocks = [[0, comfy.ldm.modules.attention.SpatialTransformer],
+                         [1, comfy.ldm.modules.attention.SpatialTransformer],
+                         [2, comfy.ldm.modules.attention.SpatialTransformer],
+                         [2, comfy.ldm.modules.diffusionmodules.openaimodel.Upsample],
+                         [3, comfy.ldm.modules.attention.SpatialTransformer],
+                         [4, comfy.ldm.modules.attention.SpatialTransformer],
+                         [5, comfy.ldm.modules.attention.SpatialTransformer],
+                         [5, comfy.ldm.modules.diffusionmodules.openaimodel.Upsample],
+                         [6, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
+                         [7, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
+                         [8, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock]]
+    else:
+        input_blocks = [[0, operations.Conv2d],
+                        [1, comfy.ldm.modules.attention.SpatialTransformer],
+                        [2, comfy.ldm.modules.attention.SpatialTransformer],
+                        [3, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample],
+                        [4, comfy.ldm.modules.attention.SpatialTransformer],
+                        [5, comfy.ldm.modules.attention.SpatialTransformer],
+                        [6, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample],
+                        [7, comfy.ldm.modules.attention.SpatialTransformer],
+                        [8, comfy.ldm.modules.attention.SpatialTransformer],
+                        [9, comfy.ldm.modules.diffusionmodules.openaimodel.Downsample],
+                        [10, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
+                        [11, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock]]
+        middle_block = [0, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock]
+        output_blocks = [[0, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
+                         [1, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
+                         [2, comfy.ldm.modules.diffusionmodules.openaimodel.ResBlock],
+                         [2, comfy.ldm.modules.diffusionmodules.openaimodel.Upsample],
+                         [3, comfy.ldm.modules.attention.SpatialTransformer],
+                         [4, comfy.ldm.modules.attention.SpatialTransformer],
+                         [5, comfy.ldm.modules.attention.SpatialTransformer],
+                         [5, comfy.ldm.modules.diffusionmodules.openaimodel.Upsample],
+                         [6, comfy.ldm.modules.attention.SpatialTransformer],
+                         [7, comfy.ldm.modules.attention.SpatialTransformer],
+                         [8, comfy.ldm.modules.attention.SpatialTransformer],
+                         [8, comfy.ldm.modules.diffusionmodules.openaimodel.Upsample],
+                         [9, comfy.ldm.modules.attention.SpatialTransformer],
+                         [10, comfy.ldm.modules.attention.SpatialTransformer],
+                         [11, comfy.ldm.modules.attention.SpatialTransformer]]
+
+    def last_layer_index(block, tp):
+        layer_list = []
+        for layer in block:
+            layer_list.append(type(layer))
+        layer_list.reverse()
+        if tp not in layer_list:
+            return -1, layer_list.reverse()
+        return len(layer_list) - 1 - layer_list.index(tp), layer_list
+
+    def brushnet_forward(model, x, timesteps, transformer_options, control):
+        if 'brushnet' not in transformer_options['model_patch']:
+            input_samples = []
+            mid_sample = 0
+            output_samples = []
+        else:
+            # brushnet inference
+            input_samples, mid_sample, output_samples = brushnet_inference(x, timesteps, transformer_options, debug)
+
+        # give additional samples to blocks
+        for i, tp in input_blocks:
+            idx, layer_list = last_layer_index(model.input_blocks[i], tp)
+            if idx < 0:
+                print("BrushNet can't find", tp, "layer in", i, "input block:", layer_list)
+                continue
+            model.input_blocks[i][idx].add_sample_after = input_samples.pop(0) if input_samples else 0
+
+        idx, layer_list = last_layer_index(model.middle_block, middle_block[1])
+        if idx < 0:
+            print("BrushNet can't find", middle_block[1], "layer in middle block", layer_list)
+        model.middle_block[idx].add_sample_after = mid_sample
+
+        for i, tp in output_blocks:
+            idx, layer_list = last_layer_index(model.output_blocks[i], tp)
+            if idx < 0:
+                print("BrushNet can't find", tp, "layer in", i, "outnput block:", layer_list)
+                continue
+            model.output_blocks[i][idx].add_sample_after = output_samples.pop(0) if output_samples else 0
+
+    patch_model_function_wrapper(model, brushnet_forward)
+
+    to = add_model_patch_option(model)
+    mp = to['model_patch']
+    if 'brushnet' not in mp:
+        mp['brushnet'] = {}
+    bo = mp['brushnet']
+
+    bo['model'] = brushnet
+    bo['dtype'] = torch_dtype
+    bo['latents'] = conditioning_latents
+    bo['controls'] = controls
+    bo['prompt_embeds'] = prompt_embeds
+    bo['negative_prompt_embeds'] = negative_prompt_embeds
+    bo['add_embeds'] = (pooled_prompt_embeds, negative_pooled_prompt_embeds, time_ids)
+    bo['latent_id'] = 0
+
+    # patch layers `forward` so we can apply brushnet
+    def forward_patched_by_brushnet(self, x, *args, **kwargs):
+        h = self.original_forward(x, *args, **kwargs)
+        if hasattr(self, 'add_sample_after') and type(self):
+            to_add = self.add_sample_after
+            if torch.is_tensor(to_add):
+                # interpolate due to RAUNet
+                if h.shape[2] != to_add.shape[2] or h.shape[3] != to_add.shape[3]:
+                    to_add = torch.nn.functional.interpolate(to_add, size=(h.shape[2], h.shape[3]), mode='bicubic')
+                h += to_add.to(h.dtype).to(h.device)
+            else:
+                h += self.add_sample_after
+            self.add_sample_after = 0
+        return h
+
+    for i, block in enumerate(model.model.diffusion_model.input_blocks):
+        for j, layer in enumerate(block):
+            if not hasattr(layer, 'original_forward'):
+                layer.original_forward = layer.forward
+            layer.forward = types.MethodType(forward_patched_by_brushnet, layer)
+            layer.add_sample_after = 0
+
+    for j, layer in enumerate(model.model.diffusion_model.middle_block):
+        if not hasattr(layer, 'original_forward'):
+            layer.original_forward = layer.forward
+        layer.forward = types.MethodType(forward_patched_by_brushnet, layer)
+        layer.add_sample_after = 0
+
+    for i, block in enumerate(model.model.diffusion_model.output_blocks):
+        for j, layer in enumerate(block):
+            if not hasattr(layer, 'original_forward'):
+                layer.original_forward = layer.forward
+            layer.forward = types.MethodType(forward_patched_by_brushnet, layer)
+            layer.add_sample_after = 0

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/brushnet/__pycache__/unet_2d_blocks.cpython-312.pyc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1148210ee8f1bc384cd3fbefb1425591976f0e14d66094d494abf175fdf6f255
+size 111447

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/brushnet/config/brushnet.json

@@ -0,0 +1,58 @@
+{
+  "_class_name": "BrushNetModel",
+  "_diffusers_version": "0.27.0.dev0",
+  "_name_or_path": "runs/logs/brushnet_randommask/checkpoint-100000",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "brushnet_conditioning_channel_order": "rgb",
+  "class_embed_type": null,
+  "conditioning_channels": 5,
+  "conditioning_embedding_out_channels": [
+    16,
+    32,
+    96,
+    256
+  ],
+  "cross_attention_dim": 768,
+  "down_block_types": [
+    "DownBlock2D",
+    "DownBlock2D",
+    "DownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "global_pool_conditions": false,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "MidBlock2D",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_time_scale_shift": "default",
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "UpBlock2D",
+    "UpBlock2D",
+    "UpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/brushnet/config/brushnet_xl.json

@@ -0,0 +1,63 @@
+{
+  "_class_name": "BrushNetModel",
+  "_diffusers_version": "0.27.0.dev0",
+  "_name_or_path": "runs/logs/brushnetsdxl_randommask/checkpoint-80000",
+  "act_fn": "silu",
+  "addition_embed_type": "text_time",
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": 256,
+  "attention_head_dim": [
+    5,
+    10,
+    20
+  ],
+  "block_out_channels": [
+    320,
+    640,
+    1280
+  ],
+  "brushnet_conditioning_channel_order": "rgb",
+  "class_embed_type": null,
+  "conditioning_channels": 5,
+  "conditioning_embedding_out_channels": [
+    16,
+    32,
+    96,
+    256
+  ],
+  "cross_attention_dim": 2048,
+  "down_block_types": [
+    "DownBlock2D",
+    "DownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "global_pool_conditions": false,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "MidBlock2D",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "projection_class_embeddings_input_dim": 2816,
+  "resnet_time_scale_shift": "default",
+  "transformer_layers_per_block": [
+    1,
+    2,
+    10
+  ],
+  "up_block_types": [
+    "UpBlock2D",
+    "UpBlock2D",
+    "UpBlock2D"
+  ],
+  "upcast_attention": null,
+  "use_linear_projection": true
+}

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/brushnet/config/powerpaint.json

@@ -0,0 +1,57 @@
+{
+  "_class_name": "BrushNetModel",
+  "_diffusers_version": "0.27.2",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "brushnet_conditioning_channel_order": "rgb",
+  "class_embed_type": null,
+  "conditioning_channels": 5,
+  "conditioning_embedding_out_channels": [
+    16,
+    32,
+    96,
+    256
+  ],
+  "cross_attention_dim": 768,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "global_pool_conditions": false,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_time_scale_shift": "default",
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/brushnet/model.py

@@ -0,0 +1,1688 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from ...libs.utils import install_package
+try:
+    install_package("diffusers", "0.27.2", True, "0.25.0")
+
+    from diffusers.configuration_utils import ConfigMixin, register_to_config
+    from diffusers.utils import BaseOutput, logging
+    from diffusers.models.attention_processor import (
+        ADDED_KV_ATTENTION_PROCESSORS,
+        CROSS_ATTENTION_PROCESSORS,
+        AttentionProcessor,
+        AttnAddedKVProcessor,
+        AttnProcessor,
+    )
+    from diffusers.models.embeddings import TextImageProjection, TextImageTimeEmbedding, TextTimeEmbedding, TimestepEmbedding, Timesteps
+    from diffusers.models.modeling_utils import ModelMixin
+    from diffusers.models.resnet import ResnetBlock2D
+    from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+
+    from diffusers.models.transformers.dual_transformer_2d import DualTransformer2DModel
+    from diffusers.models.transformers.transformer_2d import Transformer2DModel
+
+    from .unet_2d_blocks import (
+        CrossAttnDownBlock2D,
+        DownBlock2D,
+        get_down_block,
+        get_mid_block,
+        get_up_block,
+    )
+
+    from .unet_2d_condition import UNet2DConditionModel
+
+    logger = logging.get_logger(__name__)
+
+    def zero_module(module):
+        for p in module.parameters():
+            nn.init.zeros_(p)
+        return module
+
+    @dataclass
+    class BrushNetOutput(BaseOutput):
+
+        up_block_res_samples: Tuple[torch.Tensor]
+        down_block_res_samples: Tuple[torch.Tensor]
+        mid_block_res_sample: torch.Tensor
+
+    # BrushNetModel
+    class BrushNetModel(ModelMixin, ConfigMixin):
+        """A BrushNet model."""
+        _supports_gradient_checkpointing = True
+
+        @register_to_config
+        def __init__(
+                self,
+                in_channels: int = 4,
+                conditioning_channels: int = 5,
+                flip_sin_to_cos: bool = True,
+                freq_shift: int = 0,
+                down_block_types: Tuple[str, ...] = (
+                        "DownBlock2D",
+                        "DownBlock2D",
+                        "DownBlock2D",
+                        "DownBlock2D",
+                ),
+                mid_block_type: Optional[str] = "UNetMidBlock2D",
+                up_block_types: Tuple[str, ...] = (
+                        "UpBlock2D",
+                        "UpBlock2D",
+                        "UpBlock2D",
+                        "UpBlock2D",
+                ),
+                only_cross_attention: Union[bool, Tuple[bool]] = False,
+                block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
+                layers_per_block: int = 2,
+                downsample_padding: int = 1,
+                mid_block_scale_factor: float = 1,
+                act_fn: str = "silu",
+                norm_num_groups: Optional[int] = 32,
+                norm_eps: float = 1e-5,
+                cross_attention_dim: int = 1280,
+                transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1,
+                encoder_hid_dim: Optional[int] = None,
+                encoder_hid_dim_type: Optional[str] = None,
+                attention_head_dim: Union[int, Tuple[int, ...]] = 8,
+                num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None,
+                use_linear_projection: bool = False,
+                class_embed_type: Optional[str] = None,
+                addition_embed_type: Optional[str] = None,
+                addition_time_embed_dim: Optional[int] = None,
+                num_class_embeds: Optional[int] = None,
+                upcast_attention: bool = False,
+                resnet_time_scale_shift: str = "default",
+                projection_class_embeddings_input_dim: Optional[int] = None,
+                brushnet_conditioning_channel_order: str = "rgb",
+                conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
+                global_pool_conditions: bool = False,
+                addition_embed_type_num_heads: int = 64,
+        ):
+            super().__init__()
+
+            # If `num_attention_heads` is not defined (which is the case for most models)
+            # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
+            # The reason for this behavior is to correct for incorrectly named variables that were introduced
+            # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
+            # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
+            # which is why we correct for the naming here.
+            num_attention_heads = num_attention_heads or attention_head_dim
+
+            # Check inputs
+            if len(down_block_types) != len(up_block_types):
+                raise ValueError(
+                    f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
+                )
+
+            if len(block_out_channels) != len(down_block_types):
+                raise ValueError(
+                    f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+                )
+
+            if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
+                raise ValueError(
+                    f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
+                )
+
+            if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
+                raise ValueError(
+                    f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
+                )
+
+            if isinstance(transformer_layers_per_block, int):
+                transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
+            # input
+            conv_in_kernel = 3
+            conv_in_padding = (conv_in_kernel - 1) // 2
+            self.conv_in_condition = nn.Conv2d(
+                in_channels + conditioning_channels, block_out_channels[0], kernel_size=conv_in_kernel,
+                padding=conv_in_padding
+            )
+
+            # time
+            time_embed_dim = block_out_channels[0] * 4
+            self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+            timestep_input_dim = block_out_channels[0]
+            self.time_embedding = TimestepEmbedding(
+                timestep_input_dim,
+                time_embed_dim,
+                act_fn=act_fn,
+            )
+
+            if encoder_hid_dim_type is None and encoder_hid_dim is not None:
+                encoder_hid_dim_type = "text_proj"
+                self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
+                print("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
+
+            if encoder_hid_dim is None and encoder_hid_dim_type is not None:
+                raise ValueError(
+                    f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
+                )
+
+            if encoder_hid_dim_type == "text_proj":
+                self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
+            elif encoder_hid_dim_type == "text_image_proj":
+                # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+                # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
+                # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`
+                self.encoder_hid_proj = TextImageProjection(
+                    text_embed_dim=encoder_hid_dim,
+                    image_embed_dim=cross_attention_dim,
+                    cross_attention_dim=cross_attention_dim,
+                )
+
+            elif encoder_hid_dim_type is not None:
+                raise ValueError(
+                    f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
+                )
+            else:
+                self.encoder_hid_proj = None
+
+            # class embedding
+            if class_embed_type is None and num_class_embeds is not None:
+                self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
+            elif class_embed_type == "timestep":
+                self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+            elif class_embed_type == "identity":
+                self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
+            elif class_embed_type == "projection":
+                if projection_class_embeddings_input_dim is None:
+                    raise ValueError(
+                        "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
+                    )
+                # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
+                # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
+                # 2. it projects from an arbitrary input dimension.
+                #
+                # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
+                # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
+                # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
+                self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+            else:
+                self.class_embedding = None
+
+            if addition_embed_type == "text":
+                if encoder_hid_dim is not None:
+                    text_time_embedding_from_dim = encoder_hid_dim
+                else:
+                    text_time_embedding_from_dim = cross_attention_dim
+
+                self.add_embedding = TextTimeEmbedding(
+                    text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
+                )
+            elif addition_embed_type == "text_image":
+                # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+                # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
+                # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)`
+                self.add_embedding = TextImageTimeEmbedding(
+                    text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim,
+                    time_embed_dim=time_embed_dim
+                )
+            elif addition_embed_type == "text_time":
+                self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
+                self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+
+            elif addition_embed_type is not None:
+                raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")
+
+            self.down_blocks = nn.ModuleList([])
+            self.brushnet_down_blocks = nn.ModuleList([])
+
+            if isinstance(only_cross_attention, bool):
+                only_cross_attention = [only_cross_attention] * len(down_block_types)
+
+            if isinstance(attention_head_dim, int):
+                attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+            if isinstance(num_attention_heads, int):
+                num_attention_heads = (num_attention_heads,) * len(down_block_types)
+
+            # down
+            output_channel = block_out_channels[0]
+
+            brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+            brushnet_block = zero_module(brushnet_block)
+            self.brushnet_down_blocks.append(brushnet_block)
+
+            for i, down_block_type in enumerate(down_block_types):
+                input_channel = output_channel
+                output_channel = block_out_channels[i]
+                is_final_block = i == len(block_out_channels) - 1
+
+                down_block = get_down_block(
+                    down_block_type,
+                    num_layers=layers_per_block,
+                    transformer_layers_per_block=transformer_layers_per_block[i],
+                    in_channels=input_channel,
+                    out_channels=output_channel,
+                    temb_channels=time_embed_dim,
+                    add_downsample=not is_final_block,
+                    resnet_eps=norm_eps,
+                    resnet_act_fn=act_fn,
+                    resnet_groups=norm_num_groups,
+                    cross_attention_dim=cross_attention_dim,
+                    num_attention_heads=num_attention_heads[i],
+                    attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+                    downsample_padding=downsample_padding,
+                    use_linear_projection=use_linear_projection,
+                    only_cross_attention=only_cross_attention[i],
+                    upcast_attention=upcast_attention,
+                    resnet_time_scale_shift=resnet_time_scale_shift,
+                )
+                self.down_blocks.append(down_block)
+
+                for _ in range(layers_per_block):
+                    brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                    brushnet_block = zero_module(brushnet_block)
+                    self.brushnet_down_blocks.append(brushnet_block)
+
+                if not is_final_block:
+                    brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                    brushnet_block = zero_module(brushnet_block)
+                    self.brushnet_down_blocks.append(brushnet_block)
+
+            # mid
+            mid_block_channel = block_out_channels[-1]
+
+            brushnet_block = nn.Conv2d(mid_block_channel, mid_block_channel, kernel_size=1)
+            brushnet_block = zero_module(brushnet_block)
+            self.brushnet_mid_block = brushnet_block
+
+            self.mid_block = get_mid_block(
+                mid_block_type,
+                transformer_layers_per_block=transformer_layers_per_block[-1],
+                in_channels=mid_block_channel,
+                temb_channels=time_embed_dim,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                output_scale_factor=mid_block_scale_factor,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                cross_attention_dim=cross_attention_dim,
+                num_attention_heads=num_attention_heads[-1],
+                resnet_groups=norm_num_groups,
+                use_linear_projection=use_linear_projection,
+                upcast_attention=upcast_attention,
+            )
+
+            # count how many layers upsample the images
+            self.num_upsamplers = 0
+
+            # up
+            reversed_block_out_channels = list(reversed(block_out_channels))
+            reversed_num_attention_heads = list(reversed(num_attention_heads))
+            reversed_transformer_layers_per_block = (list(reversed(transformer_layers_per_block)))
+            only_cross_attention = list(reversed(only_cross_attention))
+
+            output_channel = reversed_block_out_channels[0]
+
+            self.up_blocks = nn.ModuleList([])
+            self.brushnet_up_blocks = nn.ModuleList([])
+
+            for i, up_block_type in enumerate(up_block_types):
+                is_final_block = i == len(block_out_channels) - 1
+
+                prev_output_channel = output_channel
+                output_channel = reversed_block_out_channels[i]
+                input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+                # add upsample block for all BUT final layer
+                if not is_final_block:
+                    add_upsample = True
+                    self.num_upsamplers += 1
+                else:
+                    add_upsample = False
+
+                up_block = get_up_block(
+                    up_block_type,
+                    num_layers=layers_per_block + 1,
+                    transformer_layers_per_block=reversed_transformer_layers_per_block[i],
+                    in_channels=input_channel,
+                    out_channels=output_channel,
+                    prev_output_channel=prev_output_channel,
+                    temb_channels=time_embed_dim,
+                    add_upsample=add_upsample,
+                    resnet_eps=norm_eps,
+                    resnet_act_fn=act_fn,
+                    resolution_idx=i,
+                    resnet_groups=norm_num_groups,
+                    cross_attention_dim=cross_attention_dim,
+                    num_attention_heads=reversed_num_attention_heads[i],
+                    use_linear_projection=use_linear_projection,
+                    only_cross_attention=only_cross_attention[i],
+                    upcast_attention=upcast_attention,
+                    resnet_time_scale_shift=resnet_time_scale_shift,
+                    attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+                )
+                self.up_blocks.append(up_block)
+                prev_output_channel = output_channel
+
+                for _ in range(layers_per_block + 1):
+                    brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                    brushnet_block = zero_module(brushnet_block)
+                    self.brushnet_up_blocks.append(brushnet_block)
+
+                if not is_final_block:
+                    brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                    brushnet_block = zero_module(brushnet_block)
+                    self.brushnet_up_blocks.append(brushnet_block)
+
+        @classmethod
+        def from_unet(
+                cls,
+                unet: UNet2DConditionModel,
+                brushnet_conditioning_channel_order: str = "rgb",
+                conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
+                load_weights_from_unet: bool = True,
+                conditioning_channels: int = 5,
+        ):
+            r"""
+            Instantiate a [`BrushNetModel`] from [`UNet2DConditionModel`].
+
+            Parameters:
+                unet (`UNet2DConditionModel`):
+                    The UNet model weights to copy to the [`BrushNetModel`]. All configuration options are also copied
+                    where applicable.
+            """
+            transformer_layers_per_block = (
+                unet.config.transformer_layers_per_block if "transformer_layers_per_block" in unet.config else 1
+            )
+            encoder_hid_dim = unet.config.encoder_hid_dim if "encoder_hid_dim" in unet.config else None
+            encoder_hid_dim_type = unet.config.encoder_hid_dim_type if "encoder_hid_dim_type" in unet.config else None
+            addition_embed_type = unet.config.addition_embed_type if "addition_embed_type" in unet.config else None
+            addition_time_embed_dim = (
+                unet.config.addition_time_embed_dim if "addition_time_embed_dim" in unet.config else None
+            )
+
+            down_block_types = ["DownBlock2D" for block_name in unet.config.down_block_types]
+            mid_block_type = "MidBlock2D"
+            up_block_types = ["UpBlock2D" for block_name in unet.config.down_block_types]
+
+            brushnet = cls(
+                in_channels=unet.config.in_channels,
+                conditioning_channels=conditioning_channels,
+                flip_sin_to_cos=unet.config.flip_sin_to_cos,
+                freq_shift=unet.config.freq_shift,
+                down_block_types=down_block_types,
+                mid_block_type=mid_block_type,
+                up_block_types=up_block_types,
+                only_cross_attention=unet.config.only_cross_attention,
+                block_out_channels=unet.config.block_out_channels,
+                layers_per_block=unet.config.layers_per_block,
+                downsample_padding=unet.config.downsample_padding,
+                mid_block_scale_factor=unet.config.mid_block_scale_factor,
+                act_fn=unet.config.act_fn,
+                norm_num_groups=unet.config.norm_num_groups,
+                norm_eps=unet.config.norm_eps,
+                cross_attention_dim=unet.config.cross_attention_dim,
+                transformer_layers_per_block=transformer_layers_per_block,
+                encoder_hid_dim=encoder_hid_dim,
+                encoder_hid_dim_type=encoder_hid_dim_type,
+                attention_head_dim=unet.config.attention_head_dim,
+                num_attention_heads=unet.config.num_attention_heads,
+                use_linear_projection=unet.config.use_linear_projection,
+                class_embed_type=unet.config.class_embed_type,
+                addition_embed_type=addition_embed_type,
+                addition_time_embed_dim=addition_time_embed_dim,
+                num_class_embeds=unet.config.num_class_embeds,
+                upcast_attention=unet.config.upcast_attention,
+                resnet_time_scale_shift=unet.config.resnet_time_scale_shift,
+                projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim,
+                brushnet_conditioning_channel_order=brushnet_conditioning_channel_order,
+                conditioning_embedding_out_channels=conditioning_embedding_out_channels,
+            )
+
+            if load_weights_from_unet:
+                conv_in_condition_weight = torch.zeros_like(brushnet.conv_in_condition.weight)
+                conv_in_condition_weight[:, :4, ...] = unet.conv_in.weight
+                conv_in_condition_weight[:, 4:8, ...] = unet.conv_in.weight
+                brushnet.conv_in_condition.weight = torch.nn.Parameter(conv_in_condition_weight)
+                brushnet.conv_in_condition.bias = unet.conv_in.bias
+
+                brushnet.time_proj.load_state_dict(unet.time_proj.state_dict())
+                brushnet.time_embedding.load_state_dict(unet.time_embedding.state_dict())
+
+                if brushnet.class_embedding:
+                    brushnet.class_embedding.load_state_dict(unet.class_embedding.state_dict())
+
+                brushnet.down_blocks.load_state_dict(unet.down_blocks.state_dict(), strict=False)
+                brushnet.mid_block.load_state_dict(unet.mid_block.state_dict(), strict=False)
+                brushnet.up_blocks.load_state_dict(unet.up_blocks.state_dict(), strict=False)
+
+            return brushnet
+
+        @property
+        # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+        def attn_processors(self) -> Dict[str, AttentionProcessor]:
+            r"""
+            Returns:
+                `dict` of attention processors: A dictionary containing all attention processors used in the model with
+                indexed by its weight name.
+            """
+            # set recursively
+            processors = {}
+
+            def fn_recursive_add_processors(name: str, module: torch.nn.Module,
+                                            processors: Dict[str, AttentionProcessor]):
+                if hasattr(module, "get_processor"):
+                    processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
+
+                for sub_name, child in module.named_children():
+                    fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+                return processors
+
+            for name, module in self.named_children():
+                fn_recursive_add_processors(name, module, processors)
+
+            return processors
+
+        # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+        def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+            r"""
+            Sets the attention processor to use to compute attention.
+
+            Parameters:
+                processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                    The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                    for **all** `Attention` layers.
+
+                    If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                    processor. This is strongly recommended when setting trainable attention processors.
+
+            """
+            count = len(self.attn_processors.keys())
+
+            if isinstance(processor, dict) and len(processor) != count:
+                raise ValueError(
+                    f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                    f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+                )
+
+            def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+                if hasattr(module, "set_processor"):
+                    if not isinstance(processor, dict):
+                        module.set_processor(processor)
+                    else:
+                        module.set_processor(processor.pop(f"{name}.processor"))
+
+                for sub_name, child in module.named_children():
+                    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+            for name, module in self.named_children():
+                fn_recursive_attn_processor(name, module, processor)
+
+        # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+        def set_default_attn_processor(self):
+            """
+            Disables custom attention processors and sets the default attention implementation.
+            """
+            if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+                processor = AttnAddedKVProcessor()
+            elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+                processor = AttnProcessor()
+            else:
+                raise ValueError(
+                    f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+                )
+
+            self.set_attn_processor(processor)
+
+        # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice
+        def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None:
+            r"""
+            Enable sliced attention computation.
+
+            When this option is enabled, the attention module splits the input tensor in slices to compute attention in
+            several steps. This is useful for saving some memory in exchange for a small decrease in speed.
+
+            Args:
+                slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
+                    When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
+                    `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
+                    provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
+                    must be a multiple of `slice_size`.
+            """
+            sliceable_head_dims = []
+
+            def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
+                if hasattr(module, "set_attention_slice"):
+                    sliceable_head_dims.append(module.sliceable_head_dim)
+
+                for child in module.children():
+                    fn_recursive_retrieve_sliceable_dims(child)
+
+            # retrieve number of attention layers
+            for module in self.children():
+                fn_recursive_retrieve_sliceable_dims(module)
+
+            num_sliceable_layers = len(sliceable_head_dims)
+
+            if slice_size == "auto":
+                # half the attention head size is usually a good trade-off between
+                # speed and memory
+                slice_size = [dim // 2 for dim in sliceable_head_dims]
+            elif slice_size == "max":
+                # make smallest slice possible
+                slice_size = num_sliceable_layers * [1]
+
+            slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
+
+            if len(slice_size) != len(sliceable_head_dims):
+                raise ValueError(
+                    f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
+                    f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
+                )
+
+            for i in range(len(slice_size)):
+                size = slice_size[i]
+                dim = sliceable_head_dims[i]
+                if size is not None and size > dim:
+                    raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+
+            # Recursively walk through all the children.
+            # Any children which exposes the set_attention_slice method
+            # gets the message
+            def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+                if hasattr(module, "set_attention_slice"):
+                    module.set_attention_slice(slice_size.pop())
+
+                for child in module.children():
+                    fn_recursive_set_attention_slice(child, slice_size)
+
+            reversed_slice_size = list(reversed(slice_size))
+            for module in self.children():
+                fn_recursive_set_attention_slice(module, reversed_slice_size)
+
+        def _set_gradient_checkpointing(self, module, value: bool = False) -> None:
+            if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)):
+                module.gradient_checkpointing = value
+
+        def forward(
+                self,
+                sample: torch.FloatTensor,
+                encoder_hidden_states: torch.Tensor,
+                brushnet_cond: torch.FloatTensor,
+                timestep=None,
+                time_emb=None,
+                conditioning_scale: float = 1.0,
+                class_labels: Optional[torch.Tensor] = None,
+                timestep_cond: Optional[torch.Tensor] = None,
+                attention_mask: Optional[torch.Tensor] = None,
+                added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+                cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+                guess_mode: bool = False,
+                return_dict: bool = True,
+                debug=False,
+        ) -> Union[BrushNetOutput, Tuple[Tuple[torch.FloatTensor, ...], torch.FloatTensor]]:
+
+            # check channel order
+            channel_order = self.config.brushnet_conditioning_channel_order
+
+            if channel_order == "rgb":
+                # in rgb order by default
+                ...
+            elif channel_order == "bgr":
+                brushnet_cond = torch.flip(brushnet_cond, dims=[1])
+            else:
+                raise ValueError(f"unknown `brushnet_conditioning_channel_order`: {channel_order}")
+
+            # prepare attention_mask
+            if attention_mask is not None:
+                attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+                attention_mask = attention_mask.unsqueeze(1)
+
+            if timestep is None and time_emb is None:
+                raise ValueError(f"`timestep` and `emb` are both None")
+
+            # print("BN: sample.device", sample.device)
+            # print("BN: TE.device", self.time_embedding.linear_1.weight.device)
+
+            if timestep is not None:
+                # 1. time
+                timesteps = timestep
+                if not torch.is_tensor(timesteps):
+                    # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+                    # This would be a good case for the `match` statement (Python 3.10+)
+                    is_mps = sample.device.type == "mps"
+                    if isinstance(timestep, float):
+                        dtype = torch.float32 if is_mps else torch.float64
+                    else:
+                        dtype = torch.int32 if is_mps else torch.int64
+                    timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+                elif len(timesteps.shape) == 0:
+                    timesteps = timesteps[None].to(sample.device)
+
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timesteps = timesteps.expand(sample.shape[0])
+
+                t_emb = self.time_proj(timesteps)
+
+                # timesteps does not contain any weights and will always return f32 tensors
+                # but time_embedding might actually be running in fp16. so we need to cast here.
+                # there might be better ways to encapsulate this.
+                t_emb = t_emb.to(dtype=sample.dtype)
+
+                # print("t_emb.device =",t_emb.device)
+
+                emb = self.time_embedding(t_emb, timestep_cond)
+                aug_emb = None
+
+                # print('emb.shape', emb.shape)
+
+                if self.class_embedding is not None:
+                    if class_labels is None:
+                        raise ValueError("class_labels should be provided when num_class_embeds > 0")
+
+                    if self.config.class_embed_type == "timestep":
+                        class_labels = self.time_proj(class_labels)
+
+                    class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
+                    emb = emb + class_emb
+
+                if self.config.addition_embed_type is not None:
+                    if self.config.addition_embed_type == "text":
+                        aug_emb = self.add_embedding(encoder_hidden_states)
+
+                    elif self.config.addition_embed_type == "text_time":
+                        if "text_embeds" not in added_cond_kwargs:
+                            raise ValueError(
+                                f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
+                            )
+                        text_embeds = added_cond_kwargs.get("text_embeds")
+                        if "time_ids" not in added_cond_kwargs:
+                            raise ValueError(
+                                f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
+                            )
+                        time_ids = added_cond_kwargs.get("time_ids")
+                        time_embeds = self.add_time_proj(time_ids.flatten())
+                        time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+
+                        add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+                        add_embeds = add_embeds.to(emb.dtype)
+                        aug_emb = self.add_embedding(add_embeds)
+
+                        # print('text_embeds', text_embeds.shape, 'time_ids', time_ids.shape, 'time_embeds', time_embeds.shape, 'add__embeds', add_embeds.shape, 'aug_emb', aug_emb.shape)
+
+                emb = emb + aug_emb if aug_emb is not None else emb
+            else:
+                emb = time_emb
+
+            # 2. pre-process
+
+            brushnet_cond = torch.concat([sample, brushnet_cond], 1)
+            sample = self.conv_in_condition(brushnet_cond)
+
+            # 3. down
+            down_block_res_samples = (sample,)
+            for downsample_block in self.down_blocks:
+                if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                    sample, res_samples = downsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        encoder_hidden_states=encoder_hidden_states,
+                        attention_mask=attention_mask,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                    )
+                else:
+                    sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+                down_block_res_samples += res_samples
+
+            # 4. PaintingNet down blocks
+            brushnet_down_block_res_samples = ()
+            for down_block_res_sample, brushnet_down_block in zip(down_block_res_samples, self.brushnet_down_blocks):
+                down_block_res_sample = brushnet_down_block(down_block_res_sample)
+                brushnet_down_block_res_samples = brushnet_down_block_res_samples + (down_block_res_sample,)
+
+            # 5. mid
+            if self.mid_block is not None:
+                if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
+                    sample = self.mid_block(
+                        sample,
+                        emb,
+                        encoder_hidden_states=encoder_hidden_states,
+                        attention_mask=attention_mask,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                    )
+                else:
+                    sample = self.mid_block(sample, emb)
+
+            # 6. BrushNet mid blocks
+            brushnet_mid_block_res_sample = self.brushnet_mid_block(sample)
+
+            # 7. up
+            up_block_res_samples = ()
+            for i, upsample_block in enumerate(self.up_blocks):
+                is_final_block = i == len(self.up_blocks) - 1
+
+                res_samples = down_block_res_samples[-len(upsample_block.resnets):]
+                down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+                # if we have not reached the final block and need to forward the
+                # upsample size, we do it here
+                if not is_final_block:
+                    upsample_size = down_block_res_samples[-1].shape[2:]
+
+                if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                    sample, up_res_samples = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        encoder_hidden_states=encoder_hidden_states,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                        upsample_size=upsample_size,
+                        attention_mask=attention_mask,
+                        return_res_samples=True
+                    )
+                else:
+                    sample, up_res_samples = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        upsample_size=upsample_size,
+                        return_res_samples=True
+                    )
+
+                up_block_res_samples += up_res_samples
+
+            # 8. BrushNet up blocks
+            brushnet_up_block_res_samples = ()
+            for up_block_res_sample, brushnet_up_block in zip(up_block_res_samples, self.brushnet_up_blocks):
+                up_block_res_sample = brushnet_up_block(up_block_res_sample)
+                brushnet_up_block_res_samples = brushnet_up_block_res_samples + (up_block_res_sample,)
+
+            # 6. scaling
+            if guess_mode and not self.config.global_pool_conditions:
+                scales = torch.logspace(-1, 0,
+                                        len(brushnet_down_block_res_samples) + 1 + len(brushnet_up_block_res_samples),
+                                        device=sample.device)  # 0.1 to 1.0
+                scales = scales * conditioning_scale
+
+                brushnet_down_block_res_samples = [sample * scale for sample, scale in
+                                                   zip(brushnet_down_block_res_samples,
+                                                       scales[:len(brushnet_down_block_res_samples)])]
+                brushnet_mid_block_res_sample = brushnet_mid_block_res_sample * scales[
+                    len(brushnet_down_block_res_samples)]
+                brushnet_up_block_res_samples = [sample * scale for sample, scale in zip(brushnet_up_block_res_samples,
+                                                                                         scales[
+                                                                                         len(brushnet_down_block_res_samples) + 1:])]
+            else:
+                brushnet_down_block_res_samples = [sample * conditioning_scale for sample in
+                                                   brushnet_down_block_res_samples]
+                brushnet_mid_block_res_sample = brushnet_mid_block_res_sample * conditioning_scale
+                brushnet_up_block_res_samples = [sample * conditioning_scale for sample in
+                                                 brushnet_up_block_res_samples]
+
+            if self.config.global_pool_conditions:
+                brushnet_down_block_res_samples = [
+                    torch.mean(sample, dim=(2, 3), keepdim=True) for sample in brushnet_down_block_res_samples
+                ]
+                brushnet_mid_block_res_sample = torch.mean(brushnet_mid_block_res_sample, dim=(2, 3), keepdim=True)
+                brushnet_up_block_res_samples = [
+                    torch.mean(sample, dim=(2, 3), keepdim=True) for sample in brushnet_up_block_res_samples
+                ]
+
+            if not return_dict:
+                return (brushnet_down_block_res_samples, brushnet_mid_block_res_sample, brushnet_up_block_res_samples)
+
+            return BrushNetOutput(
+                down_block_res_samples=brushnet_down_block_res_samples,
+                mid_block_res_sample=brushnet_mid_block_res_sample,
+                up_block_res_samples=brushnet_up_block_res_samples
+            )
+
+    # PowerPaintModel
+    class PowerPaintModel(ModelMixin, ConfigMixin):
+        _supports_gradient_checkpointing = True
+
+        @register_to_config
+        def __init__(
+            self,
+            in_channels: int = 4,
+            conditioning_channels: int = 5,
+            flip_sin_to_cos: bool = True,
+            freq_shift: int = 0,
+            down_block_types: Tuple[str, ...] = (
+                "CrossAttnDownBlock2D",
+                "CrossAttnDownBlock2D",
+                "CrossAttnDownBlock2D",
+                "DownBlock2D",
+            ),
+            mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
+            up_block_types: Tuple[str, ...] = (
+                "UpBlock2D",
+                "CrossAttnUpBlock2D",
+                "CrossAttnUpBlock2D",
+                "CrossAttnUpBlock2D",
+            ),
+            only_cross_attention: Union[bool, Tuple[bool]] = False,
+            block_out_channels: Tuple[int, ...] = (320, 640, 1280, 1280),
+            layers_per_block: int = 2,
+            downsample_padding: int = 1,
+            mid_block_scale_factor: float = 1,
+            act_fn: str = "silu",
+            norm_num_groups: Optional[int] = 32,
+            norm_eps: float = 1e-5,
+            cross_attention_dim: int = 1280,
+            transformer_layers_per_block: Union[int, Tuple[int, ...]] = 1,
+            encoder_hid_dim: Optional[int] = None,
+            encoder_hid_dim_type: Optional[str] = None,
+            attention_head_dim: Union[int, Tuple[int, ...]] = 8,
+            num_attention_heads: Optional[Union[int, Tuple[int, ...]]] = None,
+            use_linear_projection: bool = False,
+            class_embed_type: Optional[str] = None,
+            addition_embed_type: Optional[str] = None,
+            addition_time_embed_dim: Optional[int] = None,
+            num_class_embeds: Optional[int] = None,
+            upcast_attention: bool = False,
+            resnet_time_scale_shift: str = "default",
+            projection_class_embeddings_input_dim: Optional[int] = None,
+            brushnet_conditioning_channel_order: str = "rgb",
+            conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
+            global_pool_conditions: bool = False,
+            addition_embed_type_num_heads: int = 64,
+        ):
+            super().__init__()
+
+            # If `num_attention_heads` is not defined (which is the case for most models)
+            # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
+            # The reason for this behavior is to correct for incorrectly named variables that were introduced
+            # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
+            # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
+            # which is why we correct for the naming here.
+            num_attention_heads = num_attention_heads or attention_head_dim
+
+            # Check inputs
+            if len(down_block_types) != len(up_block_types):
+                raise ValueError(
+                    f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
+                )
+
+            if len(block_out_channels) != len(down_block_types):
+                raise ValueError(
+                    f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+                )
+
+            if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
+                raise ValueError(
+                    f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
+                )
+
+            if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
+                raise ValueError(
+                    f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
+                )
+
+            if isinstance(transformer_layers_per_block, int):
+                transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
+            # input
+            conv_in_kernel = 3
+            conv_in_padding = (conv_in_kernel - 1) // 2
+            self.conv_in_condition = nn.Conv2d(
+                in_channels + conditioning_channels,
+                block_out_channels[0],
+                kernel_size=conv_in_kernel,
+                padding=conv_in_padding,
+            )
+
+            # time
+            time_embed_dim = block_out_channels[0] * 4
+            self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+            timestep_input_dim = block_out_channels[0]
+            self.time_embedding = TimestepEmbedding(
+                timestep_input_dim,
+                time_embed_dim,
+                act_fn=act_fn,
+            )
+
+            if encoder_hid_dim_type is None and encoder_hid_dim is not None:
+                encoder_hid_dim_type = "text_proj"
+                self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
+                logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
+
+            if encoder_hid_dim is None and encoder_hid_dim_type is not None:
+                raise ValueError(
+                    f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
+                )
+
+            if encoder_hid_dim_type == "text_proj":
+                self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
+            elif encoder_hid_dim_type == "text_image_proj":
+                # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+                # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
+                # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`
+                self.encoder_hid_proj = TextImageProjection(
+                    text_embed_dim=encoder_hid_dim,
+                    image_embed_dim=cross_attention_dim,
+                    cross_attention_dim=cross_attention_dim,
+                )
+
+            elif encoder_hid_dim_type is not None:
+                raise ValueError(
+                    f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
+                )
+            else:
+                self.encoder_hid_proj = None
+
+            # class embedding
+            if class_embed_type is None and num_class_embeds is not None:
+                self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
+            elif class_embed_type == "timestep":
+                self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+            elif class_embed_type == "identity":
+                self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
+            elif class_embed_type == "projection":
+                if projection_class_embeddings_input_dim is None:
+                    raise ValueError(
+                        "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
+                    )
+                # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
+                # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
+                # 2. it projects from an arbitrary input dimension.
+                #
+                # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
+                # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
+                # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
+                self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+            else:
+                self.class_embedding = None
+
+            if addition_embed_type == "text":
+                if encoder_hid_dim is not None:
+                    text_time_embedding_from_dim = encoder_hid_dim
+                else:
+                    text_time_embedding_from_dim = cross_attention_dim
+
+                self.add_embedding = TextTimeEmbedding(
+                    text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
+                )
+            elif addition_embed_type == "text_image":
+                # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+                # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
+                # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)`
+                self.add_embedding = TextImageTimeEmbedding(
+                    text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim
+                )
+            elif addition_embed_type == "text_time":
+                self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
+                self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+
+            elif addition_embed_type is not None:
+                raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")
+
+            self.down_blocks = nn.ModuleList([])
+            self.brushnet_down_blocks = nn.ModuleList([])
+
+            if isinstance(only_cross_attention, bool):
+                only_cross_attention = [only_cross_attention] * len(down_block_types)
+
+            if isinstance(attention_head_dim, int):
+                attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+            if isinstance(num_attention_heads, int):
+                num_attention_heads = (num_attention_heads,) * len(down_block_types)
+
+            # down
+            output_channel = block_out_channels[0]
+
+            brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+            brushnet_block = zero_module(brushnet_block)
+            self.brushnet_down_blocks.append(brushnet_block)
+
+            for i, down_block_type in enumerate(down_block_types):
+                input_channel = output_channel
+                output_channel = block_out_channels[i]
+                is_final_block = i == len(block_out_channels) - 1
+
+                down_block = get_down_block(
+                    down_block_type,
+                    num_layers=layers_per_block,
+                    transformer_layers_per_block=transformer_layers_per_block[i],
+                    in_channels=input_channel,
+                    out_channels=output_channel,
+                    temb_channels=time_embed_dim,
+                    add_downsample=not is_final_block,
+                    resnet_eps=norm_eps,
+                    resnet_act_fn=act_fn,
+                    resnet_groups=norm_num_groups,
+                    cross_attention_dim=cross_attention_dim,
+                    num_attention_heads=num_attention_heads[i],
+                    attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+                    downsample_padding=downsample_padding,
+                    use_linear_projection=use_linear_projection,
+                    only_cross_attention=only_cross_attention[i],
+                    upcast_attention=upcast_attention,
+                    resnet_time_scale_shift=resnet_time_scale_shift,
+                )
+                self.down_blocks.append(down_block)
+
+                for _ in range(layers_per_block):
+                    brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                    brushnet_block = zero_module(brushnet_block)
+                    self.brushnet_down_blocks.append(brushnet_block)
+
+                if not is_final_block:
+                    brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                    brushnet_block = zero_module(brushnet_block)
+                    self.brushnet_down_blocks.append(brushnet_block)
+
+            # mid
+            mid_block_channel = block_out_channels[-1]
+
+            brushnet_block = nn.Conv2d(mid_block_channel, mid_block_channel, kernel_size=1)
+            brushnet_block = zero_module(brushnet_block)
+            self.brushnet_mid_block = brushnet_block
+
+            self.mid_block = get_mid_block(
+                mid_block_type,
+                transformer_layers_per_block=transformer_layers_per_block[-1],
+                in_channels=mid_block_channel,
+                temb_channels=time_embed_dim,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                output_scale_factor=mid_block_scale_factor,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                cross_attention_dim=cross_attention_dim,
+                num_attention_heads=num_attention_heads[-1],
+                resnet_groups=norm_num_groups,
+                use_linear_projection=use_linear_projection,
+                upcast_attention=upcast_attention,
+            )
+
+            # count how many layers upsample the images
+            self.num_upsamplers = 0
+
+            # up
+            reversed_block_out_channels = list(reversed(block_out_channels))
+            reversed_num_attention_heads = list(reversed(num_attention_heads))
+            reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
+            only_cross_attention = list(reversed(only_cross_attention))
+
+            output_channel = reversed_block_out_channels[0]
+
+            self.up_blocks = nn.ModuleList([])
+            self.brushnet_up_blocks = nn.ModuleList([])
+
+            for i, up_block_type in enumerate(up_block_types):
+                is_final_block = i == len(block_out_channels) - 1
+
+                prev_output_channel = output_channel
+                output_channel = reversed_block_out_channels[i]
+                input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+                # add upsample block for all BUT final layer
+                if not is_final_block:
+                    add_upsample = True
+                    self.num_upsamplers += 1
+                else:
+                    add_upsample = False
+
+                up_block = get_up_block(
+                    up_block_type,
+                    num_layers=layers_per_block + 1,
+                    transformer_layers_per_block=reversed_transformer_layers_per_block[i],
+                    in_channels=input_channel,
+                    out_channels=output_channel,
+                    prev_output_channel=prev_output_channel,
+                    temb_channels=time_embed_dim,
+                    add_upsample=add_upsample,
+                    resnet_eps=norm_eps,
+                    resnet_act_fn=act_fn,
+                    resolution_idx=i,
+                    resnet_groups=norm_num_groups,
+                    cross_attention_dim=cross_attention_dim,
+                    num_attention_heads=reversed_num_attention_heads[i],
+                    use_linear_projection=use_linear_projection,
+                    only_cross_attention=only_cross_attention[i],
+                    upcast_attention=upcast_attention,
+                    resnet_time_scale_shift=resnet_time_scale_shift,
+                    attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+                )
+                self.up_blocks.append(up_block)
+                prev_output_channel = output_channel
+
+                for _ in range(layers_per_block + 1):
+                    brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                    brushnet_block = zero_module(brushnet_block)
+                    self.brushnet_up_blocks.append(brushnet_block)
+
+                if not is_final_block:
+                    brushnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
+                    brushnet_block = zero_module(brushnet_block)
+                    self.brushnet_up_blocks.append(brushnet_block)
+
+        @classmethod
+        def from_unet(
+            cls,
+            unet: UNet2DConditionModel,
+            brushnet_conditioning_channel_order: str = "rgb",
+            conditioning_embedding_out_channels: Optional[Tuple[int, ...]] = (16, 32, 96, 256),
+            load_weights_from_unet: bool = True,
+            conditioning_channels: int = 5,
+        ):
+            r"""
+            Instantiate a [`BrushNetModel`] from [`UNet2DConditionModel`].
+
+            Parameters:
+                unet (`UNet2DConditionModel`):
+                    The UNet model weights to copy to the [`BrushNetModel`]. All configuration options are also copied
+                    where applicable.
+            """
+            transformer_layers_per_block = (
+                unet.config.transformer_layers_per_block if "transformer_layers_per_block" in unet.config else 1
+            )
+            encoder_hid_dim = unet.config.encoder_hid_dim if "encoder_hid_dim" in unet.config else None
+            encoder_hid_dim_type = unet.config.encoder_hid_dim_type if "encoder_hid_dim_type" in unet.config else None
+            addition_embed_type = unet.config.addition_embed_type if "addition_embed_type" in unet.config else None
+            addition_time_embed_dim = (
+                unet.config.addition_time_embed_dim if "addition_time_embed_dim" in unet.config else None
+            )
+
+            brushnet = cls(
+                in_channels=unet.config.in_channels,
+                conditioning_channels=conditioning_channels,
+                flip_sin_to_cos=unet.config.flip_sin_to_cos,
+                freq_shift=unet.config.freq_shift,
+                # down_block_types=['DownBlock2D','DownBlock2D','DownBlock2D','DownBlock2D'],
+                down_block_types=[
+                    "CrossAttnDownBlock2D",
+                    "CrossAttnDownBlock2D",
+                    "CrossAttnDownBlock2D",
+                    "DownBlock2D",
+                ],
+                # mid_block_type='MidBlock2D',
+                mid_block_type="UNetMidBlock2DCrossAttn",
+                # up_block_types=['UpBlock2D','UpBlock2D','UpBlock2D','UpBlock2D'],
+                up_block_types=["UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"],
+                only_cross_attention=unet.config.only_cross_attention,
+                block_out_channels=unet.config.block_out_channels,
+                layers_per_block=unet.config.layers_per_block,
+                downsample_padding=unet.config.downsample_padding,
+                mid_block_scale_factor=unet.config.mid_block_scale_factor,
+                act_fn=unet.config.act_fn,
+                norm_num_groups=unet.config.norm_num_groups,
+                norm_eps=unet.config.norm_eps,
+                cross_attention_dim=unet.config.cross_attention_dim,
+                transformer_layers_per_block=transformer_layers_per_block,
+                encoder_hid_dim=encoder_hid_dim,
+                encoder_hid_dim_type=encoder_hid_dim_type,
+                attention_head_dim=unet.config.attention_head_dim,
+                num_attention_heads=unet.config.num_attention_heads,
+                use_linear_projection=unet.config.use_linear_projection,
+                class_embed_type=unet.config.class_embed_type,
+                addition_embed_type=addition_embed_type,
+                addition_time_embed_dim=addition_time_embed_dim,
+                num_class_embeds=unet.config.num_class_embeds,
+                upcast_attention=unet.config.upcast_attention,
+                resnet_time_scale_shift=unet.config.resnet_time_scale_shift,
+                projection_class_embeddings_input_dim=unet.config.projection_class_embeddings_input_dim,
+                brushnet_conditioning_channel_order=brushnet_conditioning_channel_order,
+                conditioning_embedding_out_channels=conditioning_embedding_out_channels,
+            )
+
+            if load_weights_from_unet:
+                conv_in_condition_weight = torch.zeros_like(brushnet.conv_in_condition.weight)
+                conv_in_condition_weight[:, :4, ...] = unet.conv_in.weight
+                conv_in_condition_weight[:, 4:8, ...] = unet.conv_in.weight
+                brushnet.conv_in_condition.weight = torch.nn.Parameter(conv_in_condition_weight)
+                brushnet.conv_in_condition.bias = unet.conv_in.bias
+
+                brushnet.time_proj.load_state_dict(unet.time_proj.state_dict())
+                brushnet.time_embedding.load_state_dict(unet.time_embedding.state_dict())
+
+                if brushnet.class_embedding:
+                    brushnet.class_embedding.load_state_dict(unet.class_embedding.state_dict())
+
+                brushnet.down_blocks.load_state_dict(unet.down_blocks.state_dict(), strict=False)
+                brushnet.mid_block.load_state_dict(unet.mid_block.state_dict(), strict=False)
+                brushnet.up_blocks.load_state_dict(unet.up_blocks.state_dict(), strict=False)
+
+            return brushnet.to(unet.dtype)
+
+        @property
+        # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+        def attn_processors(self) -> Dict[str, AttentionProcessor]:
+            r"""
+            Returns:
+                `dict` of attention processors: A dictionary containing all attention processors used in the model with
+                indexed by its weight name.
+            """
+            # set recursively
+            processors = {}
+
+            def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+                if hasattr(module, "get_processor"):
+                    processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
+
+                for sub_name, child in module.named_children():
+                    fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+                return processors
+
+            for name, module in self.named_children():
+                fn_recursive_add_processors(name, module, processors)
+
+            return processors
+
+        # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+        def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+            r"""
+            Sets the attention processor to use to compute attention.
+
+            Parameters:
+                processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                    The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                    for **all** `Attention` layers.
+
+                    If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                    processor. This is strongly recommended when setting trainable attention processors.
+
+            """
+            count = len(self.attn_processors.keys())
+
+            if isinstance(processor, dict) and len(processor) != count:
+                raise ValueError(
+                    f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                    f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+                )
+
+            def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+                if hasattr(module, "set_processor"):
+                    if not isinstance(processor, dict):
+                        module.set_processor(processor)
+                    else:
+                        module.set_processor(processor.pop(f"{name}.processor"))
+
+                for sub_name, child in module.named_children():
+                    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+            for name, module in self.named_children():
+                fn_recursive_attn_processor(name, module, processor)
+
+        # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+        def set_default_attn_processor(self):
+            """
+            Disables custom attention processors and sets the default attention implementation.
+            """
+            if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+                processor = AttnAddedKVProcessor()
+            elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+                processor = AttnProcessor()
+            else:
+                raise ValueError(
+                    f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+                )
+
+            self.set_attn_processor(processor)
+
+        # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attention_slice
+        def set_attention_slice(self, slice_size: Union[str, int, List[int]]) -> None:
+            r"""
+            Enable sliced attention computation.
+
+            When this option is enabled, the attention module splits the input tensor in slices to compute attention in
+            several steps. This is useful for saving some memory in exchange for a small decrease in speed.
+
+            Args:
+                slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
+                    When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
+                    `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
+                    provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
+                    must be a multiple of `slice_size`.
+            """
+            sliceable_head_dims = []
+
+            def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
+                if hasattr(module, "set_attention_slice"):
+                    sliceable_head_dims.append(module.sliceable_head_dim)
+
+                for child in module.children():
+                    fn_recursive_retrieve_sliceable_dims(child)
+
+            # retrieve number of attention layers
+            for module in self.children():
+                fn_recursive_retrieve_sliceable_dims(module)
+
+            num_sliceable_layers = len(sliceable_head_dims)
+
+            if slice_size == "auto":
+                # half the attention head size is usually a good trade-off between
+                # speed and memory
+                slice_size = [dim // 2 for dim in sliceable_head_dims]
+            elif slice_size == "max":
+                # make smallest slice possible
+                slice_size = num_sliceable_layers * [1]
+
+            slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
+
+            if len(slice_size) != len(sliceable_head_dims):
+                raise ValueError(
+                    f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
+                    f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
+                )
+
+            for i in range(len(slice_size)):
+                size = slice_size[i]
+                dim = sliceable_head_dims[i]
+                if size is not None and size > dim:
+                    raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+
+            # Recursively walk through all the children.
+            # Any children which exposes the set_attention_slice method
+            # gets the message
+            def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+                if hasattr(module, "set_attention_slice"):
+                    module.set_attention_slice(slice_size.pop())
+
+                for child in module.children():
+                    fn_recursive_set_attention_slice(child, slice_size)
+
+            reversed_slice_size = list(reversed(slice_size))
+            for module in self.children():
+                fn_recursive_set_attention_slice(module, reversed_slice_size)
+
+        def _set_gradient_checkpointing(self, module, value: bool = False) -> None:
+            if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)):
+                module.gradient_checkpointing = value
+
+        def forward(
+            self,
+            sample: torch.FloatTensor,
+            timestep: Union[torch.Tensor, float, int],
+            encoder_hidden_states: torch.Tensor,
+            brushnet_cond: torch.FloatTensor,
+            conditioning_scale: float = 1.0,
+            class_labels: Optional[torch.Tensor] = None,
+            timestep_cond: Optional[torch.Tensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            guess_mode: bool = False,
+            return_dict: bool = True,
+            debug=False,
+        ) -> Union[BrushNetOutput, Tuple[Tuple[torch.FloatTensor, ...], torch.FloatTensor]]:
+            """
+            The [`BrushNetModel`] forward method.
+
+            Args:
+                sample (`torch.FloatTensor`):
+                    The noisy input tensor.
+                timestep (`Union[torch.Tensor, float, int]`):
+                    The number of timesteps to denoise an input.
+                encoder_hidden_states (`torch.Tensor`):
+                    The encoder hidden states.
+                brushnet_cond (`torch.FloatTensor`):
+                    The conditional input tensor of shape `(batch_size, sequence_length, hidden_size)`.
+                conditioning_scale (`float`, defaults to `1.0`):
+                    The scale factor for BrushNet outputs.
+                class_labels (`torch.Tensor`, *optional*, defaults to `None`):
+                    Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
+                timestep_cond (`torch.Tensor`, *optional*, defaults to `None`):
+                    Additional conditional embeddings for timestep. If provided, the embeddings will be summed with the
+                    timestep_embedding passed through the `self.time_embedding` layer to obtain the final timestep
+                    embeddings.
+                attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
+                    An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+                    is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+                    negative values to the attention scores corresponding to "discard" tokens.
+                added_cond_kwargs (`dict`):
+                    Additional conditions for the Stable Diffusion XL UNet.
+                cross_attention_kwargs (`dict[str]`, *optional*, defaults to `None`):
+                    A kwargs dictionary that if specified is passed along to the `AttnProcessor`.
+                guess_mode (`bool`, defaults to `False`):
+                    In this mode, the BrushNet encoder tries its best to recognize the input content of the input even if
+                    you remove all prompts. A `guidance_scale` between 3.0 and 5.0 is recommended.
+                return_dict (`bool`, defaults to `True`):
+                    Whether or not to return a [`~models.brushnet.BrushNetOutput`] instead of a plain tuple.
+
+            Returns:
+                [`~models.brushnet.BrushNetOutput`] **or** `tuple`:
+                    If `return_dict` is `True`, a [`~models.brushnet.BrushNetOutput`] is returned, otherwise a tuple is
+                    returned where the first element is the sample tensor.
+            """
+            # check channel order
+            channel_order = self.config.brushnet_conditioning_channel_order
+
+            if channel_order == "rgb":
+                # in rgb order by default
+                ...
+            elif channel_order == "bgr":
+                brushnet_cond = torch.flip(brushnet_cond, dims=[1])
+            else:
+                raise ValueError(f"unknown `brushnet_conditioning_channel_order`: {channel_order}")
+
+            if debug: print('BrushNet CA: attn mask')
+
+            # prepare attention_mask
+            if attention_mask is not None:
+                attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+                attention_mask = attention_mask.unsqueeze(1)
+
+            if debug: print('BrushNet CA: time')
+
+            # 1. time
+            timesteps = timestep
+            if not torch.is_tensor(timesteps):
+                # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+                # This would be a good case for the `match` statement (Python 3.10+)
+                is_mps = sample.device.type == "mps"
+                if isinstance(timestep, float):
+                    dtype = torch.float32 if is_mps else torch.float64
+                else:
+                    dtype = torch.int32 if is_mps else torch.int64
+                timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+            elif len(timesteps.shape) == 0:
+                timesteps = timesteps[None].to(sample.device)
+
+            # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+            timesteps = timesteps.expand(sample.shape[0])
+
+            t_emb = self.time_proj(timesteps)
+
+            # timesteps does not contain any weights and will always return f32 tensors
+            # but time_embedding might actually be running in fp16. so we need to cast here.
+            # there might be better ways to encapsulate this.
+            t_emb = t_emb.to(dtype=sample.dtype)
+
+            emb = self.time_embedding(t_emb, timestep_cond)
+            aug_emb = None
+
+            if self.class_embedding is not None:
+                if class_labels is None:
+                    raise ValueError("class_labels should be provided when num_class_embeds > 0")
+
+                if self.config.class_embed_type == "timestep":
+                    class_labels = self.time_proj(class_labels)
+
+                class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
+                emb = emb + class_emb
+
+            if self.config.addition_embed_type is not None:
+                if self.config.addition_embed_type == "text":
+                    aug_emb = self.add_embedding(encoder_hidden_states)
+
+                elif self.config.addition_embed_type == "text_time":
+                    if "text_embeds" not in added_cond_kwargs:
+                        raise ValueError(
+                            f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
+                        )
+                    text_embeds = added_cond_kwargs.get("text_embeds")
+                    if "time_ids" not in added_cond_kwargs:
+                        raise ValueError(
+                            f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
+                        )
+                    time_ids = added_cond_kwargs.get("time_ids")
+                    time_embeds = self.add_time_proj(time_ids.flatten())
+                    time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+
+                    add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+                    add_embeds = add_embeds.to(emb.dtype)
+                    aug_emb = self.add_embedding(add_embeds)
+
+            emb = emb + aug_emb if aug_emb is not None else emb
+
+            if debug: print('BrushNet CA: pre-process')
+
+
+            # 2. pre-process
+            brushnet_cond = torch.concat([sample, brushnet_cond], 1)
+            sample = self.conv_in_condition(brushnet_cond)
+
+            if debug: print('BrushNet CA: down')
+
+            # 3. down
+            down_block_res_samples = (sample,)
+            for downsample_block in self.down_blocks:
+                if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                    if debug: print('BrushNet CA (down block with XA): ', type(downsample_block))
+                    sample, res_samples = downsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        encoder_hidden_states=encoder_hidden_states,
+                        attention_mask=attention_mask,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                        debug=debug,
+                    )
+                else:
+                    if debug: print('BrushNet CA (down block): ', type(downsample_block))
+                    sample, res_samples = downsample_block(hidden_states=sample, temb=emb, debug=debug)
+
+                down_block_res_samples += res_samples
+
+            if debug: print('BrushNet CA: PP down')
+
+            # 4. PaintingNet down blocks
+            brushnet_down_block_res_samples = ()
+            for down_block_res_sample, brushnet_down_block in zip(down_block_res_samples, self.brushnet_down_blocks):
+                down_block_res_sample = brushnet_down_block(down_block_res_sample)
+                brushnet_down_block_res_samples = brushnet_down_block_res_samples + (down_block_res_sample,)
+
+            if debug: print('BrushNet CA: PP mid')
+
+            # 5. mid
+            if self.mid_block is not None:
+                if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
+                    sample = self.mid_block(
+                        sample,
+                        emb,
+                        encoder_hidden_states=encoder_hidden_states,
+                        attention_mask=attention_mask,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                    )
+                else:
+                    sample = self.mid_block(sample, emb)
+
+            if debug: print('BrushNet CA: mid')
+
+            # 6. BrushNet mid blocks
+            brushnet_mid_block_res_sample = self.brushnet_mid_block(sample)
+
+            if debug: print('BrushNet CA: PP up')
+
+            # 7. up
+            up_block_res_samples = ()
+            for i, upsample_block in enumerate(self.up_blocks):
+                is_final_block = i == len(self.up_blocks) - 1
+
+                res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+                down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+                # if we have not reached the final block and need to forward the
+                # upsample size, we do it here
+                if not is_final_block:
+                    upsample_size = down_block_res_samples[-1].shape[2:]
+
+                if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                    sample, up_res_samples = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        encoder_hidden_states=encoder_hidden_states,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                        upsample_size=upsample_size,
+                        attention_mask=attention_mask,
+                        return_res_samples=True,
+                    )
+                else:
+                    sample, up_res_samples = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        upsample_size=upsample_size,
+                        return_res_samples=True,
+                    )
+
+                up_block_res_samples += up_res_samples
+
+            if debug: print('BrushNet CA: up')
+
+            # 8. BrushNet up blocks
+            brushnet_up_block_res_samples = ()
+            for up_block_res_sample, brushnet_up_block in zip(up_block_res_samples, self.brushnet_up_blocks):
+                up_block_res_sample = brushnet_up_block(up_block_res_sample)
+                brushnet_up_block_res_samples = brushnet_up_block_res_samples + (up_block_res_sample,)
+
+            if debug: print('BrushNet CA: scaling')
+
+            # 6. scaling
+            if guess_mode and not self.config.global_pool_conditions:
+                scales = torch.logspace(
+                    -1,
+                    0,
+                    len(brushnet_down_block_res_samples) + 1 + len(brushnet_up_block_res_samples),
+                    device=sample.device,
+                )  # 0.1 to 1.0
+                scales = scales * conditioning_scale
+
+                brushnet_down_block_res_samples = [
+                    sample * scale
+                    for sample, scale in zip(
+                        brushnet_down_block_res_samples, scales[: len(brushnet_down_block_res_samples)]
+                    )
+                ]
+                brushnet_mid_block_res_sample = (
+                    brushnet_mid_block_res_sample * scales[len(brushnet_down_block_res_samples)]
+                )
+                brushnet_up_block_res_samples = [
+                    sample * scale
+                    for sample, scale in zip(
+                        brushnet_up_block_res_samples, scales[len(brushnet_down_block_res_samples) + 1 :]
+                    )
+                ]
+            else:
+                brushnet_down_block_res_samples = [
+                    sample * conditioning_scale for sample in brushnet_down_block_res_samples
+                ]
+                brushnet_mid_block_res_sample = brushnet_mid_block_res_sample * conditioning_scale
+                brushnet_up_block_res_samples = [sample * conditioning_scale for sample in brushnet_up_block_res_samples]
+
+            if self.config.global_pool_conditions:
+                brushnet_down_block_res_samples = [
+                    torch.mean(sample, dim=(2, 3), keepdim=True) for sample in brushnet_down_block_res_samples
+                ]
+                brushnet_mid_block_res_sample = torch.mean(brushnet_mid_block_res_sample, dim=(2, 3), keepdim=True)
+                brushnet_up_block_res_samples = [
+                    torch.mean(sample, dim=(2, 3), keepdim=True) for sample in brushnet_up_block_res_samples
+                ]
+
+            if debug: print('BrushNet CA: finish')
+
+            if not return_dict:
+                return (brushnet_down_block_res_samples, brushnet_mid_block_res_sample, brushnet_up_block_res_samples)
+
+            return BrushNetOutput(
+                down_block_res_samples=brushnet_down_block_res_samples,
+                mid_block_res_sample=brushnet_mid_block_res_sample,
+                up_block_res_samples=brushnet_up_block_res_samples,
+            )
+
+except ImportError:
+    BrushNetModel = None
+    PowerPaintModel = None
+    # print("\33[33mModule 'diffusers' load failed. If you don't have it installed, do it:\033[0m")
+    # print("\33[33mpip install diffusers\033[0m")

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/brushnet/model_patch.py

@@ -0,0 +1,137 @@
+import torch
+import comfy
+
+# Check and add 'model_patch' to model.model_options['transformer_options']
+def add_model_patch_option(model):
+    if 'transformer_options' not in model.model_options:
+        model.model_options['transformer_options'] = {}
+    to = model.model_options['transformer_options']
+    if "model_patch" not in to:
+        to["model_patch"] = {}
+    return to
+
+
+# Patch model with model_function_wrapper
+def patch_model_function_wrapper(model, forward_patch, remove=False):
+    def brushnet_model_function_wrapper(apply_model_method, options_dict):
+        to = options_dict['c']['transformer_options']
+
+        control = None
+        if 'control' in options_dict['c']:
+            control = options_dict['c']['control']
+
+        x = options_dict['input']
+        timestep = options_dict['timestep']
+
+        # check if there are patches to execute
+        if 'model_patch' not in to or 'forward' not in to['model_patch']:
+            return apply_model_method(x, timestep, **options_dict['c'])
+
+        mp = to['model_patch']
+        unet = mp['unet']
+
+        all_sigmas = mp['all_sigmas']
+        sigma = to['sigmas'][0].item()
+        total_steps = all_sigmas.shape[0] - 1
+        step = torch.argmin((all_sigmas - sigma).abs()).item()
+
+        mp['step'] = step
+        mp['total_steps'] = total_steps
+
+        # comfy.model_base.apply_model
+        xc = model.model.model_sampling.calculate_input(timestep, x)
+        if 'c_concat' in options_dict['c'] and options_dict['c']['c_concat'] is not None:
+            xc = torch.cat([xc] + [options_dict['c']['c_concat']], dim=1)
+        t = model.model.model_sampling.timestep(timestep).float()
+        # execute all patches
+        for method in mp['forward']:
+            method(unet, xc, t, to, control)
+
+        return apply_model_method(x, timestep, **options_dict['c'])
+
+    if "model_function_wrapper" in model.model_options and model.model_options["model_function_wrapper"]:
+        print('BrushNet is going to replace existing model_function_wrapper:',
+              model.model_options["model_function_wrapper"])
+    model.set_model_unet_function_wrapper(brushnet_model_function_wrapper)
+
+    to = add_model_patch_option(model)
+    mp = to['model_patch']
+
+    if isinstance(model.model.model_config, comfy.supported_models.SD15):
+        mp['SDXL'] = False
+    elif isinstance(model.model.model_config, comfy.supported_models.SDXL):
+        mp['SDXL'] = True
+    else:
+        print('Base model type: ', type(model.model.model_config))
+        raise Exception("Unsupported model type: ", type(model.model.model_config))
+
+    if 'forward' not in mp:
+        mp['forward'] = []
+
+    if remove:
+        if forward_patch in mp['forward']:
+            mp['forward'].remove(forward_patch)
+    else:
+        mp['forward'].append(forward_patch)
+
+    mp['unet'] = model.model.diffusion_model
+    mp['step'] = 0
+    mp['total_steps'] = 1
+
+    # apply patches to code
+    if comfy.samplers.sample.__doc__ is None or 'BrushNet' not in comfy.samplers.sample.__doc__:
+        comfy.samplers.original_sample = comfy.samplers.sample
+        comfy.samplers.sample = modified_sample
+
+    if comfy.ldm.modules.diffusionmodules.openaimodel.apply_control.__doc__ is None or \
+            'BrushNet' not in comfy.ldm.modules.diffusionmodules.openaimodel.apply_control.__doc__:
+        comfy.ldm.modules.diffusionmodules.openaimodel.original_apply_control = comfy.ldm.modules.diffusionmodules.openaimodel.apply_control
+        comfy.ldm.modules.diffusionmodules.openaimodel.apply_control = modified_apply_control
+
+
+# Model needs current step number and cfg at inference step. It is possible to write a custom KSampler but I'd like to use ComfyUI's one.
+# The first versions had modified_common_ksampler, but it broke custom KSampler nodes
+def modified_sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={},
+                    latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
+    ''' Modified by BrushNet nodes'''
+    cfg_guider = comfy.samplers.CFGGuider(model)
+    cfg_guider.set_conds(positive, negative)
+    cfg_guider.set_cfg(cfg)
+
+    ### Modified part ######################################################################
+    to = add_model_patch_option(model)
+    to['model_patch']['all_sigmas'] = sigmas
+    #######################################################################################
+
+    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
+
+# To use Controlnet with RAUNet it is much easier to modify apply_control a little
+def modified_apply_control(h, control, name):
+    '''Modified by BrushNet nodes'''
+    if control is not None and name in control and len(control[name]) > 0:
+        ctrl = control[name].pop()
+        if ctrl is not None:
+            if h.shape[2] != ctrl.shape[2] or h.shape[3] != ctrl.shape[3]:
+                ctrl = torch.nn.functional.interpolate(ctrl, size=(h.shape[2], h.shape[3]), mode='bicubic').to(
+                    h.dtype).to(h.device)
+            try:
+                h += ctrl
+            except:
+                print.warning("warning control could not be applied {} {}".format(h.shape, ctrl.shape))
+    return h
+
+def add_model_patch(model):
+    to = add_model_patch_option(model)
+    mp = to['model_patch']
+    if "brushnet" in mp:
+        if isinstance(model.model.model_config, comfy.supported_models.SD15):
+            mp['SDXL'] = False
+        elif isinstance(model.model.model_config, comfy.supported_models.SDXL):
+            mp['SDXL'] = True
+        else:
+            print('Base model type: ', type(model.model.model_config))
+            raise Exception("Unsupported model type: ", type(model.model.model_config))
+
+        mp['unet'] = model.model.diffusion_model
+        mp['step'] = 0
+        mp['total_steps'] = 1

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/brushnet/powerpaint_utils.py

@@ -0,0 +1,467 @@
+import copy
+import random
+
+import torch
+import torch.nn as nn
+from transformers import CLIPTokenizer
+from typing import Any, List, Optional, Union
+
+
+class TokenizerWrapper:
+    """Tokenizer wrapper for CLIPTokenizer. Only support CLIPTokenizer
+    currently. This wrapper is modified from https://github.com/huggingface/dif
+    fusers/blob/e51f19aee82c8dd874b715a09dbc521d88835d68/src/diffusers/loaders.
+    py#L358  # noqa.
+
+    Args:
+        from_pretrained (Union[str, os.PathLike], optional): The *model id*
+            of a pretrained model or a path to a *directory* containing
+            model weights and config. Defaults to None.
+        from_config (Union[str, os.PathLike], optional): The *model id*
+            of a pretrained model or a path to a *directory* containing
+            model weights and config. Defaults to None.
+
+        *args, **kwargs: If `from_pretrained` is passed, *args and **kwargs
+            will be passed to `from_pretrained` function. Otherwise, *args
+            and **kwargs will be used to initialize the model by
+            `self._module_cls(*args, **kwargs)`.
+    """
+
+    def __init__(self, tokenizer: CLIPTokenizer):
+        self.wrapped = tokenizer
+        self.token_map = {}
+
+    def __getattr__(self, name: str) -> Any:
+        if name in self.__dict__:
+            return getattr(self, name)
+        # if name == "wrapped":
+        #    return getattr(self, 'wrapped')#super().__getattr__("wrapped")
+
+        try:
+            return getattr(self.wrapped, name)
+        except AttributeError:
+            raise AttributeError(
+                "'name' cannot be found in both "
+                f"'{self.__class__.__name__}' and "
+                f"'{self.__class__.__name__}.tokenizer'."
+            )
+
+    def try_adding_tokens(self, tokens: Union[str, List[str]], *args, **kwargs):
+        """Attempt to add tokens to the tokenizer.
+
+        Args:
+            tokens (Union[str, List[str]]): The tokens to be added.
+        """
+        num_added_tokens = self.wrapped.add_tokens(tokens, *args, **kwargs)
+        assert num_added_tokens != 0, (
+            f"The tokenizer already contains the token {tokens}. Please pass "
+            "a different `placeholder_token` that is not already in the "
+            "tokenizer."
+        )
+
+    def get_token_info(self, token: str) -> dict:
+        """Get the information of a token, including its start and end index in
+        the current tokenizer.
+
+        Args:
+            token (str): The token to be queried.
+
+        Returns:
+            dict: The information of the token, including its start and end
+                index in current tokenizer.
+        """
+        token_ids = self.__call__(token).input_ids
+        start, end = token_ids[1], token_ids[-2] + 1
+        return {"name": token, "start": start, "end": end}
+
+    def add_placeholder_token(self, placeholder_token: str, *args, num_vec_per_token: int = 1, **kwargs):
+        """Add placeholder tokens to the tokenizer.
+
+        Args:
+            placeholder_token (str): The placeholder token to be added.
+            num_vec_per_token (int, optional): The number of vectors of
+                the added placeholder token.
+            *args, **kwargs: The arguments for `self.wrapped.add_tokens`.
+        """
+        output = []
+        if num_vec_per_token == 1:
+            self.try_adding_tokens(placeholder_token, *args, **kwargs)
+            output.append(placeholder_token)
+        else:
+            output = []
+            for i in range(num_vec_per_token):
+                ith_token = placeholder_token + f"_{i}"
+                self.try_adding_tokens(ith_token, *args, **kwargs)
+                output.append(ith_token)
+
+        for token in self.token_map:
+            if token in placeholder_token:
+                raise ValueError(
+                    f"The tokenizer already has placeholder token {token} "
+                    f"that can get confused with {placeholder_token} "
+                    "keep placeholder tokens independent"
+                )
+        self.token_map[placeholder_token] = output
+
+    def replace_placeholder_tokens_in_text(
+            self, text: Union[str, List[str]], vector_shuffle: bool = False, prop_tokens_to_load: float = 1.0
+    ) -> Union[str, List[str]]:
+        """Replace the keywords in text with placeholder tokens. This function
+        will be called in `self.__call__` and `self.encode`.
+
+        Args:
+            text (Union[str, List[str]]): The text to be processed.
+            vector_shuffle (bool, optional): Whether to shuffle the vectors.
+                Defaults to False.
+            prop_tokens_to_load (float, optional): The proportion of tokens to
+                be loaded. If 1.0, all tokens will be loaded. Defaults to 1.0.
+
+        Returns:
+            Union[str, List[str]]: The processed text.
+        """
+        if isinstance(text, list):
+            output = []
+            for i in range(len(text)):
+                output.append(self.replace_placeholder_tokens_in_text(text[i], vector_shuffle=vector_shuffle))
+            return output
+
+        for placeholder_token in self.token_map:
+            if placeholder_token in text:
+                tokens = self.token_map[placeholder_token]
+                tokens = tokens[: 1 + int(len(tokens) * prop_tokens_to_load)]
+                if vector_shuffle:
+                    tokens = copy.copy(tokens)
+                    random.shuffle(tokens)
+                text = text.replace(placeholder_token, " ".join(tokens))
+        return text
+
+    def replace_text_with_placeholder_tokens(self, text: Union[str, List[str]]) -> Union[str, List[str]]:
+        """Replace the placeholder tokens in text with the original keywords.
+        This function will be called in `self.decode`.
+
+        Args:
+            text (Union[str, List[str]]): The text to be processed.
+
+        Returns:
+            Union[str, List[str]]: The processed text.
+        """
+        if isinstance(text, list):
+            output = []
+            for i in range(len(text)):
+                output.append(self.replace_text_with_placeholder_tokens(text[i]))
+            return output
+
+        for placeholder_token, tokens in self.token_map.items():
+            merged_tokens = " ".join(tokens)
+            if merged_tokens in text:
+                text = text.replace(merged_tokens, placeholder_token)
+        return text
+
+    def __call__(
+            self,
+            text: Union[str, List[str]],
+            *args,
+            vector_shuffle: bool = False,
+            prop_tokens_to_load: float = 1.0,
+            **kwargs,
+    ):
+        """The call function of the wrapper.
+
+        Args:
+            text (Union[str, List[str]]): The text to be tokenized.
+            vector_shuffle (bool, optional): Whether to shuffle the vectors.
+                Defaults to False.
+            prop_tokens_to_load (float, optional): The proportion of tokens to
+                be loaded. If 1.0, all tokens will be loaded. Defaults to 1.0
+            *args, **kwargs: The arguments for `self.wrapped.__call__`.
+        """
+        replaced_text = self.replace_placeholder_tokens_in_text(
+            text, vector_shuffle=vector_shuffle, prop_tokens_to_load=prop_tokens_to_load
+        )
+
+        return self.wrapped.__call__(replaced_text, *args, **kwargs)
+
+    def encode(self, text: Union[str, List[str]], *args, **kwargs):
+        """Encode the passed text to token index.
+
+        Args:
+            text (Union[str, List[str]]): The text to be encode.
+            *args, **kwargs: The arguments for `self.wrapped.__call__`.
+        """
+        replaced_text = self.replace_placeholder_tokens_in_text(text)
+        return self.wrapped(replaced_text, *args, **kwargs)
+
+    def decode(self, token_ids, return_raw: bool = False, *args, **kwargs) -> Union[str, List[str]]:
+        """Decode the token index to text.
+
+        Args:
+            token_ids: The token index to be decoded.
+            return_raw: Whether keep the placeholder token in the text.
+                Defaults to False.
+            *args, **kwargs: The arguments for `self.wrapped.decode`.
+
+        Returns:
+            Union[str, List[str]]: The decoded text.
+        """
+        text = self.wrapped.decode(token_ids, *args, **kwargs)
+        if return_raw:
+            return text
+        replaced_text = self.replace_text_with_placeholder_tokens(text)
+        return replaced_text
+
+    def __repr__(self):
+        """The representation of the wrapper."""
+        s = super().__repr__()
+        prefix = f"Wrapped Module Class: {self._module_cls}\n"
+        prefix += f"Wrapped Module Name: {self._module_name}\n"
+        if self._from_pretrained:
+            prefix += f"From Pretrained: {self._from_pretrained}\n"
+        s = prefix + s
+        return s
+
+
+class EmbeddingLayerWithFixes(nn.Module):
+    """The revised embedding layer to support external embeddings. This design
+    of this class is inspired by https://github.com/AUTOMATIC1111/stable-
+    diffusion-webui/blob/22bcc7be428c94e9408f589966c2040187245d81/modules/sd_hi
+    jack.py#L224  # noqa.
+
+    Args:
+        wrapped (nn.Emebdding): The embedding layer to be wrapped.
+        external_embeddings (Union[dict, List[dict]], optional): The external
+            embeddings added to this layer. Defaults to None.
+    """
+
+    def __init__(self, wrapped: nn.Embedding, external_embeddings: Optional[Union[dict, List[dict]]] = None):
+        super().__init__()
+        self.wrapped = wrapped
+        self.num_embeddings = wrapped.weight.shape[0]
+
+        self.external_embeddings = []
+        if external_embeddings:
+            self.add_embeddings(external_embeddings)
+
+        self.trainable_embeddings = nn.ParameterDict()
+
+    @property
+    def weight(self):
+        """Get the weight of wrapped embedding layer."""
+        return self.wrapped.weight
+
+    def check_duplicate_names(self, embeddings: List[dict]):
+        """Check whether duplicate names exist in list of 'external
+        embeddings'.
+
+        Args:
+            embeddings (List[dict]): A list of embedding to be check.
+        """
+        names = [emb["name"] for emb in embeddings]
+        assert len(names) == len(set(names)), (
+            "Found duplicated names in 'external_embeddings'. Name list: " f"'{names}'"
+        )
+
+    def check_ids_overlap(self, embeddings):
+        """Check whether overlap exist in token ids of 'external_embeddings'.
+
+        Args:
+            embeddings (List[dict]): A list of embedding to be check.
+        """
+        ids_range = [[emb["start"], emb["end"], emb["name"]] for emb in embeddings]
+        ids_range.sort()  # sort by 'start'
+        # check if 'end' has overlapping
+        for idx in range(len(ids_range) - 1):
+            name1, name2 = ids_range[idx][-1], ids_range[idx + 1][-1]
+            assert ids_range[idx][1] <= ids_range[idx + 1][0], (
+                f"Found ids overlapping between embeddings '{name1}' " f"and '{name2}'."
+            )
+
+    def add_embeddings(self, embeddings: Optional[Union[dict, List[dict]]]):
+        """Add external embeddings to this layer.
+        Use case:
+        Args:
+            embeddings (Union[dict, list[dict]]): The external embeddings to
+                be added. Each dict must contain the following 4 fields: 'name'
+                (the name of this embedding), 'embedding' (the embedding
+                tensor), 'start' (the start token id of this embedding), 'end'
+                (the end token id of this embedding). For example:
+                `{name: NAME, start: START, end: END, embedding: torch.Tensor}`
+        """
+        if isinstance(embeddings, dict):
+            embeddings = [embeddings]
+
+        self.external_embeddings += embeddings
+        self.check_duplicate_names(self.external_embeddings)
+        self.check_ids_overlap(self.external_embeddings)
+
+        # set for trainable
+        added_trainable_emb_info = []
+        for embedding in embeddings:
+            trainable = embedding.get("trainable", False)
+            if trainable:
+                name = embedding["name"]
+                embedding["embedding"] = torch.nn.Parameter(embedding["embedding"])
+                self.trainable_embeddings[name] = embedding["embedding"]
+                added_trainable_emb_info.append(name)
+
+        added_emb_info = [emb["name"] for emb in embeddings]
+        added_emb_info = ", ".join(added_emb_info)
+        print(f"Successfully add external embeddings: {added_emb_info}.", "current")
+
+        if added_trainable_emb_info:
+            added_trainable_emb_info = ", ".join(added_trainable_emb_info)
+            print("Successfully add trainable external embeddings: " f"{added_trainable_emb_info}", "current")
+
+    def replace_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """Replace external input ids to 0.
+
+        Args:
+            input_ids (torch.Tensor): The input ids to be replaced.
+
+        Returns:
+            torch.Tensor: The replaced input ids.
+        """
+        input_ids_fwd = input_ids.clone()
+        input_ids_fwd[input_ids_fwd >= self.num_embeddings] = 0
+        return input_ids_fwd
+
+    def replace_embeddings(
+            self, input_ids: torch.Tensor, embedding: torch.Tensor, external_embedding: dict
+    ) -> torch.Tensor:
+        """Replace external embedding to the embedding layer. Noted that, in
+        this function we use `torch.cat` to avoid inplace modification.
+
+        Args:
+            input_ids (torch.Tensor): The original token ids. Shape like
+                [LENGTH, ].
+            embedding (torch.Tensor): The embedding of token ids after
+                `replace_input_ids` function.
+            external_embedding (dict): The external embedding to be replaced.
+
+        Returns:
+            torch.Tensor: The replaced embedding.
+        """
+        new_embedding = []
+
+        name = external_embedding["name"]
+        start = external_embedding["start"]
+        end = external_embedding["end"]
+        target_ids_to_replace = [i for i in range(start, end)]
+        ext_emb = external_embedding["embedding"].to(embedding.device)
+
+        # do not need to replace
+        if not (input_ids == start).any():
+            return embedding
+
+        # start replace
+        s_idx, e_idx = 0, 0
+        while e_idx < len(input_ids):
+            if input_ids[e_idx] == start:
+                if e_idx != 0:
+                    # add embedding do not need to replace
+                    new_embedding.append(embedding[s_idx:e_idx])
+
+                # check if the next embedding need to replace is valid
+                actually_ids_to_replace = [int(i) for i in input_ids[e_idx: e_idx + end - start]]
+                assert actually_ids_to_replace == target_ids_to_replace, (
+                    f"Invalid 'input_ids' in position: {s_idx} to {e_idx}. "
+                    f"Expect '{target_ids_to_replace}' for embedding "
+                    f"'{name}' but found '{actually_ids_to_replace}'."
+                )
+
+                new_embedding.append(ext_emb)
+
+                s_idx = e_idx + end - start
+                e_idx = s_idx + 1
+            else:
+                e_idx += 1
+
+        if e_idx == len(input_ids):
+            new_embedding.append(embedding[s_idx:e_idx])
+
+        return torch.cat(new_embedding, dim=0)
+
+    def forward(self, input_ids: torch.Tensor, external_embeddings: Optional[List[dict]] = None, out_dtype = None):
+        """The forward function.
+
+        Args:
+            input_ids (torch.Tensor): The token ids shape like [bz, LENGTH] or
+                [LENGTH, ].
+            external_embeddings (Optional[List[dict]]): The external
+                embeddings. If not passed, only `self.external_embeddings`
+                will be used.  Defaults to None.
+
+        input_ids: shape like [bz, LENGTH] or [LENGTH].
+        """
+        assert input_ids.ndim in [1, 2]
+        if input_ids.ndim == 1:
+            input_ids = input_ids.unsqueeze(0)
+
+        if external_embeddings is None and not self.external_embeddings:
+            return self.wrapped(input_ids, out_dtype=out_dtype)
+
+        input_ids_fwd = self.replace_input_ids(input_ids)
+        inputs_embeds = self.wrapped(input_ids_fwd)
+
+        vecs = []
+
+        if external_embeddings is None:
+            external_embeddings = []
+        elif isinstance(external_embeddings, dict):
+            external_embeddings = [external_embeddings]
+        embeddings = self.external_embeddings + external_embeddings
+
+        for input_id, embedding in zip(input_ids, inputs_embeds):
+            new_embedding = embedding
+            for external_embedding in embeddings:
+                new_embedding = self.replace_embeddings(input_id, new_embedding, external_embedding)
+            vecs.append(new_embedding)
+
+        return torch.stack(vecs).to(out_dtype)
+
+
+def add_tokens(
+        tokenizer, text_encoder, placeholder_tokens: list, initialize_tokens: list = None,
+        num_vectors_per_token: int = 1
+):
+    """Add token for training.
+
+    # TODO: support add tokens as dict, then we can load pretrained tokens.
+    """
+    if initialize_tokens is not None:
+        assert len(initialize_tokens) == len(
+            placeholder_tokens
+        ), "placeholder_token should be the same length as initialize_token"
+    for ii in range(len(placeholder_tokens)):
+        tokenizer.add_placeholder_token(placeholder_tokens[ii], num_vec_per_token=num_vectors_per_token)
+
+    # text_encoder.set_embedding_layer()
+    embedding_layer = text_encoder.text_model.embeddings.token_embedding
+    text_encoder.text_model.embeddings.token_embedding = EmbeddingLayerWithFixes(embedding_layer)
+    embedding_layer = text_encoder.text_model.embeddings.token_embedding
+
+    assert embedding_layer is not None, (
+        "Do not support get embedding layer for current text encoder. " "Please check your configuration."
+    )
+    initialize_embedding = []
+    if initialize_tokens is not None:
+        for ii in range(len(placeholder_tokens)):
+            init_id = tokenizer(initialize_tokens[ii]).input_ids[1]
+            temp_embedding = embedding_layer.weight[init_id]
+            initialize_embedding.append(temp_embedding[None, ...].repeat(num_vectors_per_token, 1))
+    else:
+        for ii in range(len(placeholder_tokens)):
+            init_id = tokenizer("a").input_ids[1]
+            temp_embedding = embedding_layer.weight[init_id]
+            len_emb = temp_embedding.shape[0]
+            init_weight = (torch.rand(num_vectors_per_token, len_emb) - 0.5) / 2.0
+            initialize_embedding.append(init_weight)
+
+    # initialize_embedding  = torch.cat(initialize_embedding,dim=0)
+
+    token_info_all = []
+    for ii in range(len(placeholder_tokens)):
+        token_info = tokenizer.get_token_info(placeholder_tokens[ii])
+        token_info["embedding"] = initialize_embedding[ii]
+        token_info["trainable"] = True
+        token_info_all.append(token_info)
+    embedding_layer.add_embeddings(token_info_all)

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/brushnet/unet_2d_condition.py

@@ -0,0 +1,1359 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin
+from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.models.activations import get_activation
+from diffusers.models.attention_processor import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    Attention,
+    AttentionProcessor,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+)
+from diffusers.models.embeddings import (
+    GaussianFourierProjection,
+    GLIGENTextBoundingboxProjection,
+    ImageHintTimeEmbedding,
+    ImageProjection,
+    ImageTimeEmbedding,
+    TextImageProjection,
+    TextImageTimeEmbedding,
+    TextTimeEmbedding,
+    TimestepEmbedding,
+    Timesteps,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from .unet_2d_blocks import (
+    get_down_block,
+    get_mid_block,
+    get_up_block,
+)
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class UNet2DConditionOutput(BaseOutput):
+    """
+    The output of [`UNet2DConditionModel`].
+
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
+    """
+
+    sample: torch.FloatTensor = None
+
+
+class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin):
+    r"""
+    A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample
+    shaped output.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+            Height and width of input/output sample.
+        in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.
+        out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
+        center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
+        flip_sin_to_cos (`bool`, *optional*, defaults to `True`):
+            Whether to flip the sin to cos in the time embedding.
+        freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
+            The tuple of downsample blocks to use.
+        mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
+            Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`, `UNetMidBlock2D`, or
+            `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
+            The tuple of upsample blocks to use.
+        only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
+            Whether to include self-attention in the basic transformer blocks, see
+            [`~models.attention.BasicTransformerBlock`].
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+            The tuple of output channels for each block.
+        layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
+        downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
+        mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+        norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
+            If `None`, normalization and activation layers is skipped in post-processing.
+        norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
+        cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
+            The dimension of the cross attention features.
+        transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
+            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
+            [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
+            [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
+        reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None):
+            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling
+            blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for
+            [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
+            [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
+        encoder_hid_dim (`int`, *optional*, defaults to None):
+            If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
+            dimension to `cross_attention_dim`.
+        encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
+            If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
+            embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
+        attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
+        num_attention_heads (`int`, *optional*):
+            The number of attention heads. If not defined, defaults to `attention_head_dim`
+        resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
+            for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.
+        class_embed_type (`str`, *optional*, defaults to `None`):
+            The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
+            `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
+        addition_embed_type (`str`, *optional*, defaults to `None`):
+            Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or
+            "text". "text" will use the `TextTimeEmbedding` layer.
+        addition_time_embed_dim: (`int`, *optional*, defaults to `None`):
+            Dimension for the timestep embeddings.
+        num_class_embeds (`int`, *optional*, defaults to `None`):
+            Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
+            class conditioning with `class_embed_type` equal to `None`.
+        time_embedding_type (`str`, *optional*, defaults to `positional`):
+            The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.
+        time_embedding_dim (`int`, *optional*, defaults to `None`):
+            An optional override for the dimension of the projected time embedding.
+        time_embedding_act_fn (`str`, *optional*, defaults to `None`):
+            Optional activation function to use only once on the time embeddings before they are passed to the rest of
+            the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`.
+        timestep_post_act (`str`, *optional*, defaults to `None`):
+            The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.
+        time_cond_proj_dim (`int`, *optional*, defaults to `None`):
+            The dimension of `cond_proj` layer in the timestep embedding.
+        conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer.
+        conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer.
+        projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when
+            `class_embed_type="projection"`. Required when `class_embed_type="projection"`.
+        class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time
+            embeddings with the class embeddings.
+        mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`):
+            Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If
+            `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the
+            `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False`
+            otherwise.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+            self,
+            sample_size: Optional[int] = None,
+            in_channels: int = 4,
+            out_channels: int = 4,
+            center_input_sample: bool = False,
+            flip_sin_to_cos: bool = True,
+            freq_shift: int = 0,
+            down_block_types: Tuple[str] = (
+                    "CrossAttnDownBlock2D",
+                    "CrossAttnDownBlock2D",
+                    "CrossAttnDownBlock2D",
+                    "DownBlock2D",
+            ),
+            mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
+            up_block_types: Tuple[str] = (
+            "UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
+            only_cross_attention: Union[bool, Tuple[bool]] = False,
+            block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+            layers_per_block: Union[int, Tuple[int]] = 2,
+            downsample_padding: int = 1,
+            mid_block_scale_factor: float = 1,
+            dropout: float = 0.0,
+            act_fn: str = "silu",
+            norm_num_groups: Optional[int] = 32,
+            norm_eps: float = 1e-5,
+            cross_attention_dim: Union[int, Tuple[int]] = 1280,
+            transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
+            reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
+            encoder_hid_dim: Optional[int] = None,
+            encoder_hid_dim_type: Optional[str] = None,
+            attention_head_dim: Union[int, Tuple[int]] = 8,
+            num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
+            dual_cross_attention: bool = False,
+            use_linear_projection: bool = False,
+            class_embed_type: Optional[str] = None,
+            addition_embed_type: Optional[str] = None,
+            addition_time_embed_dim: Optional[int] = None,
+            num_class_embeds: Optional[int] = None,
+            upcast_attention: bool = False,
+            resnet_time_scale_shift: str = "default",
+            resnet_skip_time_act: bool = False,
+            resnet_out_scale_factor: float = 1.0,
+            time_embedding_type: str = "positional",
+            time_embedding_dim: Optional[int] = None,
+            time_embedding_act_fn: Optional[str] = None,
+            timestep_post_act: Optional[str] = None,
+            time_cond_proj_dim: Optional[int] = None,
+            conv_in_kernel: int = 3,
+            conv_out_kernel: int = 3,
+            projection_class_embeddings_input_dim: Optional[int] = None,
+            attention_type: str = "default",
+            class_embeddings_concat: bool = False,
+            mid_block_only_cross_attention: Optional[bool] = None,
+            cross_attention_norm: Optional[str] = None,
+            addition_embed_type_num_heads: int = 64,
+    ):
+        super().__init__()
+
+        self.sample_size = sample_size
+
+        if num_attention_heads is not None:
+            raise ValueError(
+                "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."
+            )
+
+        # If `num_attention_heads` is not defined (which is the case for most models)
+        # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
+        # The reason for this behavior is to correct for incorrectly named variables that were introduced
+        # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
+        # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
+        # which is why we correct for the naming here.
+        num_attention_heads = num_attention_heads or attention_head_dim
+
+        # Check inputs
+        self._check_config(
+            down_block_types=down_block_types,
+            up_block_types=up_block_types,
+            only_cross_attention=only_cross_attention,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            cross_attention_dim=cross_attention_dim,
+            transformer_layers_per_block=transformer_layers_per_block,
+            reverse_transformer_layers_per_block=reverse_transformer_layers_per_block,
+            attention_head_dim=attention_head_dim,
+            num_attention_heads=num_attention_heads,
+        )
+
+        # input
+        conv_in_padding = (conv_in_kernel - 1) // 2
+        self.conv_in = nn.Conv2d(
+            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
+        )
+
+        # time
+        time_embed_dim, timestep_input_dim = self._set_time_proj(
+            time_embedding_type,
+            block_out_channels=block_out_channels,
+            flip_sin_to_cos=flip_sin_to_cos,
+            freq_shift=freq_shift,
+            time_embedding_dim=time_embedding_dim,
+        )
+
+        self.time_embedding = TimestepEmbedding(
+            timestep_input_dim,
+            time_embed_dim,
+            act_fn=act_fn,
+            post_act_fn=timestep_post_act,
+            cond_proj_dim=time_cond_proj_dim,
+        )
+
+        self._set_encoder_hid_proj(
+            encoder_hid_dim_type,
+            cross_attention_dim=cross_attention_dim,
+            encoder_hid_dim=encoder_hid_dim,
+        )
+
+        # class embedding
+        self._set_class_embedding(
+            class_embed_type,
+            act_fn=act_fn,
+            num_class_embeds=num_class_embeds,
+            projection_class_embeddings_input_dim=projection_class_embeddings_input_dim,
+            time_embed_dim=time_embed_dim,
+            timestep_input_dim=timestep_input_dim,
+        )
+
+        self._set_add_embedding(
+            addition_embed_type,
+            addition_embed_type_num_heads=addition_embed_type_num_heads,
+            addition_time_embed_dim=addition_time_embed_dim,
+            cross_attention_dim=cross_attention_dim,
+            encoder_hid_dim=encoder_hid_dim,
+            flip_sin_to_cos=flip_sin_to_cos,
+            freq_shift=freq_shift,
+            projection_class_embeddings_input_dim=projection_class_embeddings_input_dim,
+            time_embed_dim=time_embed_dim,
+        )
+
+        if time_embedding_act_fn is None:
+            self.time_embed_act = None
+        else:
+            self.time_embed_act = get_activation(time_embedding_act_fn)
+
+        self.down_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        if isinstance(only_cross_attention, bool):
+            if mid_block_only_cross_attention is None:
+                mid_block_only_cross_attention = only_cross_attention
+
+            only_cross_attention = [only_cross_attention] * len(down_block_types)
+
+        if mid_block_only_cross_attention is None:
+            mid_block_only_cross_attention = False
+
+        if isinstance(num_attention_heads, int):
+            num_attention_heads = (num_attention_heads,) * len(down_block_types)
+
+        if isinstance(attention_head_dim, int):
+            attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+        if isinstance(cross_attention_dim, int):
+            cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
+
+        if isinstance(layers_per_block, int):
+            layers_per_block = [layers_per_block] * len(down_block_types)
+
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
+        if class_embeddings_concat:
+            # The time embeddings are concatenated with the class embeddings. The dimension of the
+            # time embeddings passed to the down, middle, and up blocks is twice the dimension of the
+            # regular time embeddings
+            blocks_time_embed_dim = time_embed_dim * 2
+        else:
+            blocks_time_embed_dim = time_embed_dim
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block[i],
+                transformer_layers_per_block=transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim[i],
+                num_attention_heads=num_attention_heads[i],
+                downsample_padding=downsample_padding,
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                only_cross_attention=only_cross_attention[i],
+                upcast_attention=upcast_attention,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                attention_type=attention_type,
+                resnet_skip_time_act=resnet_skip_time_act,
+                resnet_out_scale_factor=resnet_out_scale_factor,
+                cross_attention_norm=cross_attention_norm,
+                attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+                dropout=dropout,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = get_mid_block(
+            mid_block_type,
+            temb_channels=blocks_time_embed_dim,
+            in_channels=block_out_channels[-1],
+            resnet_eps=norm_eps,
+            resnet_act_fn=act_fn,
+            resnet_groups=norm_num_groups,
+            output_scale_factor=mid_block_scale_factor,
+            transformer_layers_per_block=transformer_layers_per_block[-1],
+            num_attention_heads=num_attention_heads[-1],
+            cross_attention_dim=cross_attention_dim[-1],
+            dual_cross_attention=dual_cross_attention,
+            use_linear_projection=use_linear_projection,
+            mid_block_only_cross_attention=mid_block_only_cross_attention,
+            upcast_attention=upcast_attention,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            attention_type=attention_type,
+            resnet_skip_time_act=resnet_skip_time_act,
+            cross_attention_norm=cross_attention_norm,
+            attention_head_dim=attention_head_dim[-1],
+            dropout=dropout,
+        )
+
+        # count how many layers upsample the images
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        reversed_num_attention_heads = list(reversed(num_attention_heads))
+        reversed_layers_per_block = list(reversed(layers_per_block))
+        reversed_cross_attention_dim = list(reversed(cross_attention_dim))
+        reversed_transformer_layers_per_block = (
+            list(reversed(transformer_layers_per_block))
+            if reverse_transformer_layers_per_block is None
+            else reverse_transformer_layers_per_block
+        )
+        only_cross_attention = list(reversed(only_cross_attention))
+
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            is_final_block = i == len(block_out_channels) - 1
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+            # add upsample block for all BUT final layer
+            if not is_final_block:
+                add_upsample = True
+                self.num_upsamplers += 1
+            else:
+                add_upsample = False
+
+            up_block = get_up_block(
+                up_block_type,
+                num_layers=reversed_layers_per_block[i] + 1,
+                transformer_layers_per_block=reversed_transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_upsample=add_upsample,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resolution_idx=i,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=reversed_cross_attention_dim[i],
+                num_attention_heads=reversed_num_attention_heads[i],
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                only_cross_attention=only_cross_attention[i],
+                upcast_attention=upcast_attention,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                attention_type=attention_type,
+                resnet_skip_time_act=resnet_skip_time_act,
+                resnet_out_scale_factor=resnet_out_scale_factor,
+                cross_attention_norm=cross_attention_norm,
+                attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+                dropout=dropout,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        if norm_num_groups is not None:
+            self.conv_norm_out = nn.GroupNorm(
+                num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
+            )
+
+            self.conv_act = get_activation(act_fn)
+
+        else:
+            self.conv_norm_out = None
+            self.conv_act = None
+
+        conv_out_padding = (conv_out_kernel - 1) // 2
+        self.conv_out = nn.Conv2d(
+            block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding
+        )
+
+        self._set_pos_net_if_use_gligen(attention_type=attention_type, cross_attention_dim=cross_attention_dim)
+
+    def _check_config(
+            self,
+            down_block_types: Tuple[str],
+            up_block_types: Tuple[str],
+            only_cross_attention: Union[bool, Tuple[bool]],
+            block_out_channels: Tuple[int],
+            layers_per_block: Union[int, Tuple[int]],
+            cross_attention_dim: Union[int, Tuple[int]],
+            transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple[int]]],
+            reverse_transformer_layers_per_block: bool,
+            attention_head_dim: int,
+            num_attention_heads: Optional[Union[int, Tuple[int]]],
+    ):
+        if len(down_block_types) != len(up_block_types):
+            raise ValueError(
+                f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
+            )
+
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
+            )
+        if isinstance(transformer_layers_per_block, list) and reverse_transformer_layers_per_block is None:
+            for layer_number_per_block in transformer_layers_per_block:
+                if isinstance(layer_number_per_block, list):
+                    raise ValueError("Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet.")
+
+    def _set_time_proj(
+            self,
+            time_embedding_type: str,
+            block_out_channels: int,
+            flip_sin_to_cos: bool,
+            freq_shift: float,
+            time_embedding_dim: int,
+    ) -> Tuple[int, int]:
+        if time_embedding_type == "fourier":
+            time_embed_dim = time_embedding_dim or block_out_channels[0] * 2
+            if time_embed_dim % 2 != 0:
+                raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.")
+            self.time_proj = GaussianFourierProjection(
+                time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos
+            )
+            timestep_input_dim = time_embed_dim
+        elif time_embedding_type == "positional":
+            time_embed_dim = time_embedding_dim or block_out_channels[0] * 4
+
+            self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+            timestep_input_dim = block_out_channels[0]
+        else:
+            raise ValueError(
+                f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."
+            )
+
+        return time_embed_dim, timestep_input_dim
+
+    def _set_encoder_hid_proj(
+            self,
+            encoder_hid_dim_type: Optional[str],
+            cross_attention_dim: Union[int, Tuple[int]],
+            encoder_hid_dim: Optional[int],
+    ):
+        if encoder_hid_dim_type is None and encoder_hid_dim is not None:
+            encoder_hid_dim_type = "text_proj"
+            self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
+            logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
+
+        if encoder_hid_dim is None and encoder_hid_dim_type is not None:
+            raise ValueError(
+                f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
+            )
+
+        if encoder_hid_dim_type == "text_proj":
+            self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
+        elif encoder_hid_dim_type == "text_image_proj":
+            # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+            # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
+            # case when `addition_embed_type == "text_image_proj"` (Kandinsky 2.1)`
+            self.encoder_hid_proj = TextImageProjection(
+                text_embed_dim=encoder_hid_dim,
+                image_embed_dim=cross_attention_dim,
+                cross_attention_dim=cross_attention_dim,
+            )
+        elif encoder_hid_dim_type == "image_proj":
+            # Kandinsky 2.2
+            self.encoder_hid_proj = ImageProjection(
+                image_embed_dim=encoder_hid_dim,
+                cross_attention_dim=cross_attention_dim,
+            )
+        elif encoder_hid_dim_type is not None:
+            raise ValueError(
+                f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
+            )
+        else:
+            self.encoder_hid_proj = None
+
+    def _set_class_embedding(
+            self,
+            class_embed_type: Optional[str],
+            act_fn: str,
+            num_class_embeds: Optional[int],
+            projection_class_embeddings_input_dim: Optional[int],
+            time_embed_dim: int,
+            timestep_input_dim: int,
+    ):
+        if class_embed_type is None and num_class_embeds is not None:
+            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
+        elif class_embed_type == "timestep":
+            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn)
+        elif class_embed_type == "identity":
+            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
+        elif class_embed_type == "projection":
+            if projection_class_embeddings_input_dim is None:
+                raise ValueError(
+                    "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
+                )
+            # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
+            # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
+            # 2. it projects from an arbitrary input dimension.
+            #
+            # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
+            # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
+            # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
+            self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+        elif class_embed_type == "simple_projection":
+            if projection_class_embeddings_input_dim is None:
+                raise ValueError(
+                    "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"
+                )
+            self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim)
+        else:
+            self.class_embedding = None
+
+    def _set_add_embedding(
+            self,
+            addition_embed_type: str,
+            addition_embed_type_num_heads: int,
+            addition_time_embed_dim: Optional[int],
+            flip_sin_to_cos: bool,
+            freq_shift: float,
+            cross_attention_dim: Optional[int],
+            encoder_hid_dim: Optional[int],
+            projection_class_embeddings_input_dim: Optional[int],
+            time_embed_dim: int,
+    ):
+        if addition_embed_type == "text":
+            if encoder_hid_dim is not None:
+                text_time_embedding_from_dim = encoder_hid_dim
+            else:
+                text_time_embedding_from_dim = cross_attention_dim
+
+            self.add_embedding = TextTimeEmbedding(
+                text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
+            )
+        elif addition_embed_type == "text_image":
+            # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+            # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
+            # case when `addition_embed_type == "text_image"` (Kandinsky 2.1)`
+            self.add_embedding = TextImageTimeEmbedding(
+                text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim
+            )
+        elif addition_embed_type == "text_time":
+            self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
+            self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+        elif addition_embed_type == "image":
+            # Kandinsky 2.2
+            self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
+        elif addition_embed_type == "image_hint":
+            # Kandinsky 2.2 ControlNet
+            self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
+        elif addition_embed_type is not None:
+            raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")
+
+    def _set_pos_net_if_use_gligen(self, attention_type: str, cross_attention_dim: int):
+        if attention_type in ["gated", "gated-text-image"]:
+            positive_len = 768
+            if isinstance(cross_attention_dim, int):
+                positive_len = cross_attention_dim
+            elif isinstance(cross_attention_dim, tuple) or isinstance(cross_attention_dim, list):
+                positive_len = cross_attention_dim[0]
+
+            feature_type = "text-only" if attention_type == "gated" else "text-image"
+            self.position_net = GLIGENTextBoundingboxProjection(
+                positive_len=positive_len, out_dim=cross_attention_dim, feature_type=feature_type
+            )
+
+    @property
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnAddedKVProcessor()
+        elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    def set_attention_slice(self, slice_size: Union[str, int, List[int]] = "auto"):
+        r"""
+        Enable sliced attention computation.
+
+        When this option is enabled, the attention module splits the input tensor in slices to compute attention in
+        several steps. This is useful for saving some memory in exchange for a small decrease in speed.
+
+        Args:
+            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
+                When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
+                `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
+                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
+                must be a multiple of `slice_size`.
+        """
+        sliceable_head_dims = []
+
+        def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
+            if hasattr(module, "set_attention_slice"):
+                sliceable_head_dims.append(module.sliceable_head_dim)
+
+            for child in module.children():
+                fn_recursive_retrieve_sliceable_dims(child)
+
+        # retrieve number of attention layers
+        for module in self.children():
+            fn_recursive_retrieve_sliceable_dims(module)
+
+        num_sliceable_layers = len(sliceable_head_dims)
+
+        if slice_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            slice_size = [dim // 2 for dim in sliceable_head_dims]
+        elif slice_size == "max":
+            # make smallest slice possible
+            slice_size = num_sliceable_layers * [1]
+
+        slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
+
+        if len(slice_size) != len(sliceable_head_dims):
+            raise ValueError(
+                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
+                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
+            )
+
+        for i in range(len(slice_size)):
+            size = slice_size[i]
+            dim = sliceable_head_dims[i]
+            if size is not None and size > dim:
+                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+
+        # Recursively walk through all the children.
+        # Any children which exposes the set_attention_slice method
+        # gets the message
+        def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+            if hasattr(module, "set_attention_slice"):
+                module.set_attention_slice(slice_size.pop())
+
+            for child in module.children():
+                fn_recursive_set_attention_slice(child, slice_size)
+
+        reversed_slice_size = list(reversed(slice_size))
+        for module in self.children():
+            fn_recursive_set_attention_slice(module, reversed_slice_size)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def enable_freeu(self, s1: float, s2: float, b1: float, b2: float):
+        r"""Enables the FreeU mechanism from https://arxiv.org/abs/2309.11497.
+
+        The suffixes after the scaling factors represent the stage blocks where they are being applied.
+
+        Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that
+        are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL.
+
+        Args:
+            s1 (`float`):
+                Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to
+                mitigate the "oversmoothing effect" in the enhanced denoising process.
+            s2 (`float`):
+                Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to
+                mitigate the "oversmoothing effect" in the enhanced denoising process.
+            b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.
+            b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.
+        """
+        for i, upsample_block in enumerate(self.up_blocks):
+            setattr(upsample_block, "s1", s1)
+            setattr(upsample_block, "s2", s2)
+            setattr(upsample_block, "b1", b1)
+            setattr(upsample_block, "b2", b2)
+
+    def disable_freeu(self):
+        """Disables the FreeU mechanism."""
+        freeu_keys = {"s1", "s2", "b1", "b2"}
+        for i, upsample_block in enumerate(self.up_blocks):
+            for k in freeu_keys:
+                if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None:
+                    setattr(upsample_block, k, None)
+
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)
+
+    def unload_lora(self):
+        """Unloads LoRA weights."""
+        deprecate(
+            "unload_lora",
+            "0.28.0",
+            "Calling `unload_lora()` is deprecated and will be removed in a future version. Please install `peft` and then call `disable_adapters().",
+        )
+        for module in self.modules():
+            if hasattr(module, "set_lora_layer"):
+                module.set_lora_layer(None)
+
+    def get_time_embed(
+            self, sample: torch.Tensor, timestep: Union[torch.Tensor, float, int]
+    ) -> Optional[torch.Tensor]:
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+        # `Timesteps` does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=sample.dtype)
+        return t_emb
+
+    def get_class_embed(self, sample: torch.Tensor, class_labels: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+        class_emb = None
+        if self.class_embedding is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when num_class_embeds > 0")
+
+            if self.config.class_embed_type == "timestep":
+                class_labels = self.time_proj(class_labels)
+
+                # `Timesteps` does not contain any weights and will always return f32 tensors
+                # there might be better ways to encapsulate this.
+                class_labels = class_labels.to(dtype=sample.dtype)
+
+            class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
+        return class_emb
+
+    def get_aug_embed(
+            self, emb: torch.Tensor, encoder_hidden_states: torch.Tensor, added_cond_kwargs: Dict[str, Any]
+    ) -> Optional[torch.Tensor]:
+        aug_emb = None
+        if self.config.addition_embed_type == "text":
+            aug_emb = self.add_embedding(encoder_hidden_states)
+        elif self.config.addition_embed_type == "text_image":
+            # Kandinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+
+            image_embs = added_cond_kwargs.get("image_embeds")
+            text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)
+            aug_emb = self.add_embedding(text_embs, image_embs)
+        elif self.config.addition_embed_type == "text_time":
+            # SDXL - style
+            if "text_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
+                )
+            text_embeds = added_cond_kwargs.get("text_embeds")
+            if "time_ids" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
+                )
+            time_ids = added_cond_kwargs.get("time_ids")
+            time_embeds = self.add_time_proj(time_ids.flatten())
+            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+            add_embeds = add_embeds.to(emb.dtype)
+            aug_emb = self.add_embedding(add_embeds)
+        elif self.config.addition_embed_type == "image":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            aug_emb = self.add_embedding(image_embs)
+        elif self.config.addition_embed_type == "image_hint":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            hint = added_cond_kwargs.get("hint")
+            aug_emb = self.add_embedding(image_embs, hint)
+        return aug_emb
+
+    def process_encoder_hidden_states(
+            self, encoder_hidden_states: torch.Tensor, added_cond_kwargs: Dict[str, Any]
+    ) -> torch.Tensor:
+        if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":
+            # Kandinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(image_embeds)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "ip_image_proj":
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            image_embeds = self.encoder_hid_proj(image_embeds)
+            encoder_hidden_states = (encoder_hidden_states, image_embeds)
+        return encoder_hidden_states
+
+    def forward(
+            self,
+            sample: torch.FloatTensor,
+            timestep: Union[torch.Tensor, float, int],
+            encoder_hidden_states: torch.Tensor,
+            class_labels: Optional[torch.Tensor] = None,
+            timestep_cond: Optional[torch.Tensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+            down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+            mid_block_additional_residual: Optional[torch.Tensor] = None,
+            down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+            encoder_attention_mask: Optional[torch.Tensor] = None,
+            return_dict: bool = True,
+            down_block_add_samples: Optional[Tuple[torch.Tensor]] = None,
+            mid_block_add_sample: Optional[Tuple[torch.Tensor]] = None,
+            up_block_add_samples: Optional[Tuple[torch.Tensor]] = None,
+    ) -> Union[UNet2DConditionOutput, Tuple]:
+        r"""
+        The [`UNet2DConditionModel`] forward method.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The noisy input tensor with the following shape `(batch, channel, height, width)`.
+            timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
+            encoder_hidden_states (`torch.FloatTensor`):
+                The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
+            class_labels (`torch.Tensor`, *optional*, defaults to `None`):
+                Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
+            timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`):
+                Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed
+                through the `self.time_embedding` layer to obtain the timestep embeddings.
+            attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
+                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+                negative values to the attention scores corresponding to "discard" tokens.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            added_cond_kwargs: (`dict`, *optional*):
+                A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
+                are passed along to the UNet blocks.
+            down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
+                A tuple of tensors that if specified are added to the residuals of down unet blocks.
+            mid_block_additional_residual: (`torch.Tensor`, *optional*):
+                A tensor that if specified is added to the residual of the middle unet block.
+            down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*):
+                additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s)
+            encoder_attention_mask (`torch.Tensor`):
+                A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
+                `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
+                which adds large negative values to the attention scores corresponding to "discard" tokens.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
+                If `return_dict` is True, an [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] is returned,
+                otherwise a `tuple` is returned where the first element is the sample tensor.
+        """
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2 ** self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        for dim in sample.shape[-2:]:
+            if dim % default_overall_up_factor != 0:
+                # Forward upsample size to force interpolation output size.
+                forward_upsample_size = True
+                break
+
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        if attention_mask is not None:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 0. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # 1. time
+        t_emb = self.get_time_embed(sample=sample, timestep=timestep)
+        emb = self.time_embedding(t_emb, timestep_cond)
+        aug_emb = None
+
+        class_emb = self.get_class_embed(sample=sample, class_labels=class_labels)
+        if class_emb is not None:
+            if self.config.class_embeddings_concat:
+                emb = torch.cat([emb, class_emb], dim=-1)
+            else:
+                emb = emb + class_emb
+
+        aug_emb = self.get_aug_embed(
+            emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
+        )
+        if self.config.addition_embed_type == "image_hint":
+            aug_emb, hint = aug_emb
+            sample = torch.cat([sample, hint], dim=1)
+
+        emb = emb + aug_emb if aug_emb is not None else emb
+
+        if self.time_embed_act is not None:
+            emb = self.time_embed_act(emb)
+
+        encoder_hidden_states = self.process_encoder_hidden_states(
+            encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
+        )
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        # 2.5 GLIGEN position net
+        if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            gligen_args = cross_attention_kwargs.pop("gligen")
+            cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
+
+        # 3. down
+        # we're popping the `scale` instead of getting it because otherwise `scale` will be propagated
+        # to the internal blocks and will raise deprecation warnings. this will be confusing for our users.
+        if cross_attention_kwargs is not None:
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            lora_scale = cross_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+
+        is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+        # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets
+        is_adapter = down_intrablock_additional_residuals is not None
+        # maintain backward compatibility for legacy usage, where
+        #       T2I-Adapter and ControlNet both use down_block_additional_residuals arg
+        #       but can only use one or the other
+        is_brushnet = down_block_add_samples is not None and mid_block_add_sample is not None and up_block_add_samples is not None
+        if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None:
+            deprecate(
+                "T2I should not use down_block_additional_residuals",
+                "1.3.0",
+                "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \
+                       and will be removed in diffusers 1.3.0.  `down_block_additional_residuals` should only be used \
+                       for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ",
+                standard_warn=False,
+            )
+            down_intrablock_additional_residuals = down_block_additional_residuals
+            is_adapter = True
+
+        down_block_res_samples = (sample,)
+
+        if is_brushnet:
+            sample = sample + down_block_add_samples.pop(0)
+
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                # For t2i-adapter CrossAttnDownBlock2D
+                additional_residuals = {}
+                if is_adapter and len(down_intrablock_additional_residuals) > 0:
+                    additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0)
+
+                i = len(down_block_add_samples)
+
+                if is_brushnet and len(down_block_add_samples) > 0:
+                    additional_residuals["down_block_add_samples"] = [down_block_add_samples.pop(0)
+                                                                      for _ in range(
+                            len(downsample_block.resnets) + (downsample_block.downsamplers != None))]
+
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                    **additional_residuals,
+                )
+            else:
+                additional_residuals = {}
+
+                i = len(down_block_add_samples)
+
+                if is_brushnet and len(down_block_add_samples) > 0:
+                    additional_residuals["down_block_add_samples"] = [down_block_add_samples.pop(0)
+                                                                      for _ in range(
+                            len(downsample_block.resnets) + (downsample_block.downsamplers != None))]
+
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb, **additional_residuals)
+                if is_adapter and len(down_intrablock_additional_residuals) > 0:
+                    sample += down_intrablock_additional_residuals.pop(0)
+
+            down_block_res_samples += res_samples
+
+        if is_controlnet:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                    down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
+                sample = self.mid_block(
+                    sample,
+                    emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
+                sample = self.mid_block(sample, emb)
+
+            # To support T2I-Adapter-XL
+            if (
+                    is_adapter
+                    and len(down_intrablock_additional_residuals) > 0
+                    and sample.shape == down_intrablock_additional_residuals[0].shape
+            ):
+                sample += down_intrablock_additional_residuals.pop(0)
+
+        if is_controlnet:
+            sample = sample + mid_block_additional_residual
+
+        if is_brushnet:
+            sample = sample + mid_block_add_sample
+
+        # 5. up
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets):]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                additional_residuals = {}
+
+                i = len(up_block_add_samples)
+
+                if is_brushnet and len(up_block_add_samples) > 0:
+                    additional_residuals["up_block_add_samples"] = [up_block_add_samples.pop(0)
+                                                                    for _ in range(
+                            len(upsample_block.resnets) + (upsample_block.upsamplers != None))]
+
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                    encoder_attention_mask=encoder_attention_mask,
+                    **additional_residuals,
+                )
+            else:
+                additional_residuals = {}
+
+                i = len(up_block_add_samples)
+
+                if is_brushnet and len(up_block_add_samples) > 0:
+                    additional_residuals["up_block_add_samples"] = [up_block_add_samples.pop(0)
+                                                                    for _ in range(
+                            len(upsample_block.resnets) + (upsample_block.upsamplers != None))]
+
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                    **additional_residuals,
+                )
+
+        # 6. post-process
+        if self.conv_norm_out:
+            sample = self.conv_norm_out(sample)
+            sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet2DConditionOutput(sample=sample)

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/dit/__init__.py

@@ -0,0 +1,2 @@
+#credit to city96 for this module
+#from https://github.com/city96/ComfyUI_ExtraModels/

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/dit/config.py

@@ -0,0 +1,120 @@
+"""
+List of all DiT model types / settings
+"""
+sampling_settings = {
+	"beta_schedule" : "sqrt_linear",
+	"linear_start"  : 0.0001,
+	"linear_end"    : 0.02,
+	"timesteps"     : 1000,
+}
+
+dit_conf = {
+	"XL/2": { # DiT_XL_2
+		"unet_config": {
+			"depth"       :   28,
+			"num_heads"   :   16,
+			"patch_size"  :    2,
+			"hidden_size" : 1152,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"XL/4": { # DiT_XL_4
+		"unet_config": {
+			"depth"       :   28,
+			"num_heads"   :   16,
+			"patch_size"  :    4,
+			"hidden_size" : 1152,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"XL/8": { # DiT_XL_8
+		"unet_config": {
+			"depth"       :   28,
+			"num_heads"   :   16,
+			"patch_size"  :    8,
+			"hidden_size" : 1152,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"L/2": { # DiT_L_2
+		"unet_config": {
+			"depth"       :   24,
+			"num_heads"   :   16,
+			"patch_size"  :    2,
+			"hidden_size" : 1024,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"L/4": { # DiT_L_4
+		"unet_config": {
+			"depth"       :   24,
+			"num_heads"   :   16,
+			"patch_size"  :    4,
+			"hidden_size" : 1024,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"L/8": { # DiT_L_8
+		"unet_config": {
+			"depth"       :   24,
+			"num_heads"   :   16,
+			"patch_size"  :    8,
+			"hidden_size" : 1024,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"B/2": { # DiT_B_2
+		"unet_config": {
+			"depth"       :   12,
+			"num_heads"   :   12,
+			"patch_size"  :    2,
+			"hidden_size" :  768,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"B/4": { # DiT_B_4
+		"unet_config": {
+			"depth"       :   12,
+			"num_heads"   :   12,
+			"patch_size"  :    4,
+			"hidden_size" :  768,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"B/8": { # DiT_B_8
+		"unet_config": {
+			"depth"       :   12,
+			"num_heads"   :   12,
+			"patch_size"  :    8,
+			"hidden_size" :  768,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"S/2": { # DiT_S_2
+		"unet_config": {
+			"depth"       :   12,
+			"num_heads"   :    6,
+			"patch_size"  :    2,
+			"hidden_size" :  384,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"S/4": { # DiT_S_4
+		"unet_config": {
+			"depth"       :   12,
+			"num_heads"   :    6,
+			"patch_size"  :    4,
+			"hidden_size" :  384,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+	"S/8": { # DiT_S_8
+		"unet_config": {
+			"depth"       :   12,
+			"num_heads"   :    6,
+			"patch_size"  :    8,
+			"hidden_size" :  384,
+		},
+		"sampling_settings" : sampling_settings,
+	},
+}

ComfyUI/custom_nodes/ComfyUI-Easy-Use/py/modules/dit/pixArt/LICENSE-Pixart

@@ -0,0 +1,661 @@
+                    GNU AFFERO GENERAL PUBLIC LICENSE
+                       Version 3, 19 November 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+                            Preamble
+
+  The GNU Affero General Public License is a free, copyleft license for
+software and other kinds of works, specifically designed to ensure
+cooperation with the community in the case of network server software.
+
+  The licenses for most software and other practical works are designed
+to take away your freedom to share and change the works.  By contrast,
+our General Public Licenses are intended to guarantee your freedom to
+share and change all versions of a program--to make sure it remains free
+software for all its users.
+
+  When we speak of free software, we are referring to freedom, not
+price.  Our General Public Licenses are designed to make sure that you
+have the freedom to distribute copies of free software (and charge for
+them if you wish), that you receive source code or can get it if you
+want it, that you can change the software or use pieces of it in new
+free programs, and that you know you can do these things.
+
+  Developers that use our General Public Licenses protect your rights
+with two steps: (1) assert copyright on the software, and (2) offer
+you this License which gives you legal permission to copy, distribute
+and/or modify the software.
+
+  A secondary benefit of defending all users' freedom is that
+improvements made in alternate versions of the program, if they
+receive widespread use, become available for other developers to
+incorporate.  Many developers of free software are heartened and
+encouraged by the resulting cooperation.  However, in the case of
+software used on network servers, this result may fail to come about.
+The GNU General Public License permits making a modified version and
+letting the public access it on a server without ever releasing its
+source code to the public.
+
+  The GNU Affero General Public License is designed specifically to
+ensure that, in such cases, the modified source code becomes available
+to the community.  It requires the operator of a network server to
+provide the source code of the modified version running there to the
+users of that server.  Therefore, public use of a modified version, on
+a publicly accessible server, gives the public access to the source
+code of the modified version.
+
+  An older license, called the Affero General Public License and
+published by Affero, was designed to accomplish similar goals.  This is
+a different license, not a version of the Affero GPL, but Affero has
+released a new version of the Affero GPL which permits relicensing under
+this license.
+
+  The precise terms and conditions for copying, distribution and
+modification follow.
+
+                       TERMS AND CONDITIONS
+
+  0. Definitions.
+
+  "This License" refers to version 3 of the GNU Affero General Public License.
+
+  "Copyright" also means copyright-like laws that apply to other kinds of
+works, such as semiconductor masks.
+
+  "The Program" refers to any copyrightable work licensed under this
+License.  Each licensee is addressed as "you".  "Licensees" and
+"recipients" may be individuals or organizations.
+
+  To "modify" a work means to copy from or adapt all or part of the work
+in a fashion requiring copyright permission, other than the making of an
+exact copy.  The resulting work is called a "modified version" of the
+earlier work or a work "based on" the earlier work.
+
+  A "covered work" means either the unmodified Program or a work based
+on the Program.
+
+  To "propagate" a work means to do anything with it that, without
+permission, would make you directly or secondarily liable for
+infringement under applicable copyright law, except executing it on a
+computer or modifying a private copy.  Propagation includes copying,
+distribution (with or without modification), making available to the
+public, and in some countries other activities as well.
+
+  To "convey" a work means any kind of propagation that enables other
+parties to make or receive copies.  Mere interaction with a user through
+a computer network, with no transfer of a copy, is not conveying.
+
+  An interactive user interface displays "Appropriate Legal Notices"
+to the extent that it includes a convenient and prominently visible
+feature that (1) displays an appropriate copyright notice, and (2)
+tells the user that there is no warranty for the work (except to the
+extent that warranties are provided), that licensees may convey the
+work under this License, and how to view a copy of this License.  If
+the interface presents a list of user commands or options, such as a
+menu, a prominent item in the list meets this criterion.
+
+  1. Source Code.
+
+  The "source code" for a work means the preferred form of the work
+for making modifications to it.  "Object code" means any non-source
+form of a work.
+
+  A "Standard Interface" means an interface that either is an official
+standard defined by a recognized standards body, or, in the case of
+interfaces specified for a particular programming language, one that
+is widely used among developers working in that language.
+
+  The "System Libraries" of an executable work include anything, other
+than the work as a whole, that (a) is included in the normal form of
+packaging a Major Component, but which is not part of that Major
+Component, and (b) serves only to enable use of the work with that
+Major Component, or to implement a Standard Interface for which an
+implementation is available to the public in source code form.  A
+"Major Component", in this context, means a major essential component
+(kernel, window system, and so on) of the specific operating system
+(if any) on which the executable work runs, or a compiler used to
+produce the work, or an object code interpreter used to run it.
+
+  The "Corresponding Source" for a work in object code form means all
+the source code needed to generate, install, and (for an executable
+work) run the object code and to modify the work, including scripts to
+control those activities.  However, it does not include the work's
+System Libraries, or general-purpose tools or generally available free
+programs which are used unmodified in performing those activities but
+which are not part of the work.  For example, Corresponding Source
+includes interface definition files associated with source files for
+the work, and the source code for shared libraries and dynamically
+linked subprograms that the work is specifically designed to require,
+such as by intimate data communication or control flow between those
+subprograms and other parts of the work.
+
+  The Corresponding Source need not include anything that users
+can regenerate automatically from other parts of the Corresponding
+Source.
+
+  The Corresponding Source for a work in source code form is that
+same work.
+
+  2. Basic Permissions.
+
+  All rights granted under this License are granted for the term of
+copyright on the Program, and are irrevocable provided the stated
+conditions are met.  This License explicitly affirms your unlimited
+permission to run the unmodified Program.  The output from running a
+covered work is covered by this License only if the output, given its
+content, constitutes a covered work.  This License acknowledges your
+rights of fair use or other equivalent, as provided by copyright law.
+
+  You may make, run and propagate covered works that you do not
+convey, without conditions so long as your license otherwise remains
+in force.  You may convey covered works to others for the sole purpose
+of having them make modifications exclusively for you, or provide you
+with facilities for running those works, provided that you comply with
+the terms of this License in conveying all material for which you do
+not control copyright.  Those thus making or running the covered works
+for you must do so exclusively on your behalf, under your direction
+and control, on terms that prohibit them from making any copies of
+your copyrighted material outside their relationship with you.
+
+  Conveying under any other circumstances is permitted solely under
+the conditions stated below.  Sublicensing is not allowed; section 10
+makes it unnecessary.
+
+  3. Protecting Users' Legal Rights From Anti-Circumvention Law.
+
+  No covered work shall be deemed part of an effective technological
+measure under any applicable law fulfilling obligations under article
+11 of the WIPO copyright treaty adopted on 20 December 1996, or
+similar laws prohibiting or restricting circumvention of such
+measures.
+
+  When you convey a covered work, you waive any legal power to forbid
+circumvention of technological measures to the extent such circumvention
+is effected by exercising rights under this License with respect to
+the covered work, and you disclaim any intention to limit operation or
+modification of the work as a means of enforcing, against the work's
+users, your or third parties' legal rights to forbid circumvention of
+technological measures.
+
+  4. Conveying Verbatim Copies.
+
+  You may convey verbatim copies of the Program's source code as you
+receive it, in any medium, provided that you conspicuously and
+appropriately publish on each copy an appropriate copyright notice;
+keep intact all notices stating that this License and any
+non-permissive terms added in accord with section 7 apply to the code;
+keep intact all notices of the absence of any warranty; and give all
+recipients a copy of this License along with the Program.
+
+  You may charge any price or no price for each copy that you convey,
+and you may offer support or warranty protection for a fee.
+
+  5. Conveying Modified Source Versions.
+
+  You may convey a work based on the Program, or the modifications to
+produce it from the Program, in the form of source code under the
+terms of section 4, provided that you also meet all of these conditions:
+
+    a) The work must carry prominent notices stating that you modified
+    it, and giving a relevant date.
+
+    b) The work must carry prominent notices stating that it is
+    released under this License and any conditions added under section
+    7.  This requirement modifies the requirement in section 4 to
+    "keep intact all notices".
+
+    c) You must license the entire work, as a whole, under this
+    License to anyone who comes into possession of a copy.  This
+    License will therefore apply, along with any applicable section 7
+    additional terms, to the whole of the work, and all its parts,
+    regardless of how they are packaged.  This License gives no
+    permission to license the work in any other way, but it does not
+    invalidate such permission if you have separately received it.
+
+    d) If the work has interactive user interfaces, each must display
+    Appropriate Legal Notices; however, if the Program has interactive
+    interfaces that do not display Appropriate Legal Notices, your
+    work need not make them do so.
+
+  A compilation of a covered work with other separate and independent
+works, which are not by their nature extensions of the covered work,
+and which are not combined with it such as to form a larger program,
+in or on a volume of a storage or distribution medium, is called an
+"aggregate" if the compilation and its resulting copyright are not
+used to limit the access or legal rights of the compilation's users
+beyond what the individual works permit.  Inclusion of a covered work
+in an aggregate does not cause this License to apply to the other
+parts of the aggregate.
+
+  6. Conveying Non-Source Forms.
+
+  You may convey a covered work in object code form under the terms
+of sections 4 and 5, provided that you also convey the
+machine-readable Corresponding Source under the terms of this License,
+in one of these ways:
+
+    a) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by the
+    Corresponding Source fixed on a durable physical medium
+    customarily used for software interchange.
+
+    b) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by a
+    written offer, valid for at least three years and valid for as
+    long as you offer spare parts or customer support for that product
+    model, to give anyone who possesses the object code either (1) a
+    copy of the Corresponding Source for all the software in the
+    product that is covered by this License, on a durable physical
+    medium customarily used for software interchange, for a price no
+    more than your reasonable cost of physically performing this
+    conveying of source, or (2) access to copy the
+    Corresponding Source from a network server at no charge.
+
+    c) Convey individual copies of the object code with a copy of the
+    written offer to provide the Corresponding Source.  This
+    alternative is allowed only occasionally and noncommercially, and
+    only if you received the object code with such an offer, in accord
+    with subsection 6b.
+
+    d) Convey the object code by offering access from a designated
+    place (gratis or for a charge), and offer equivalent access to the
+    Corresponding Source in the same way through the same place at no
+    further charge.  You need not require recipients to copy the
+    Corresponding Source along with the object code.  If the place to
+    copy the object code is a network server, the Corresponding Source
+    may be on a different server (operated by you or a third party)
+    that supports equivalent copying facilities, provided you maintain
+    clear directions next to the object code saying where to find the
+    Corresponding Source.  Regardless of what server hosts the
+    Corresponding Source, you remain obligated to ensure that it is
+    available for as long as needed to satisfy these requirements.
+
+    e) Convey the object code using peer-to-peer transmission, provided
+    you inform other peers where the object code and Corresponding
+    Source of the work are being offered to the general public at no
+    charge under subsection 6d.
+
+  A separable portion of the object code, whose source code is excluded
+from the Corresponding Source as a System Library, need not be
+included in conveying the object code work.
+
+  A "User Product" is either (1) a "consumer product", which means any
+tangible personal property which is normally used for personal, family,
+or household purposes, or (2) anything designed or sold for incorporation
+into a dwelling.  In determining whether a product is a consumer product,
+doubtful cases shall be resolved in favor of coverage.  For a particular
+product received by a particular user, "normally used" refers to a
+typical or common use of that class of product, regardless of the status
+of the particular user or of the way in which the particular user
+actually uses, or expects or is expected to use, the product.  A product
+is a consumer product regardless of whether the product has substantial
+commercial, industrial or non-consumer uses, unless such uses represent
+the only significant mode of use of the product.
+
+  "Installation Information" for a User Product means any methods,
+procedures, authorization keys, or other information required to install
+and execute modified versions of a covered work in that User Product from
+a modified version of its Corresponding Source.  The information must
+suffice to ensure that the continued functioning of the modified object
+code is in no case prevented or interfered with solely because
+modification has been made.
+
+  If you convey an object code work under this section in, or with, or
+specifically for use in, a User Product, and the conveying occurs as
+part of a transaction in which the right of possession and use of the
+User Product is transferred to the recipient in perpetuity or for a
+fixed term (regardless of how the transaction is characterized), the
+Corresponding Source conveyed under this section must be accompanied
+by the Installation Information.  But this requirement does not apply
+if neither you nor any third party retains the ability to install
+modified object code on the User Product (for example, the work has
+been installed in ROM).
+
+  The requirement to provide Installation Information does not include a
+requirement to continue to provide support service, warranty, or updates
+for a work that has been modified or installed by the recipient, or for
+the User Product in which it has been modified or installed.  Access to a
+network may be denied when the modification itself materially and
+adversely affects the operation of the network or violates the rules and
+protocols for communication across the network.
+
+  Corresponding Source conveyed, and Installation Information provided,
+in accord with this section must be in a format that is publicly
+documented (and with an implementation available to the public in
+source code form), and must require no special password or key for
+unpacking, reading or copying.
+
+  7. Additional Terms.
+
+  "Additional permissions" are terms that supplement the terms of this
+License by making exceptions from one or more of its conditions.
+Additional permissions that are applicable to the entire Program shall
+be treated as though they were included in this License, to the extent
+that they are valid under applicable law.  If additional permissions
+apply only to part of the Program, that part may be used separately
+under those permissions, but the entire Program remains governed by
+this License without regard to the additional permissions.
+
+  When you convey a copy of a covered work, you may at your option
+remove any additional permissions from that copy, or from any part of
+it.  (Additional permissions may be written to require their own
+removal in certain cases when you modify the work.)  You may place
+additional permissions on material, added by you to a covered work,
+for which you have or can give appropriate copyright permission.
+
+  Notwithstanding any other provision of this License, for material you
+add to a covered work, you may (if authorized by the copyright holders of
+that material) supplement the terms of this License with terms:
+
+    a) Disclaiming warranty or limiting liability differently from the
+    terms of sections 15 and 16 of this License; or
+
+    b) Requiring preservation of specified reasonable legal notices or
+    author attributions in that material or in the Appropriate Legal
+    Notices displayed by works containing it; or
+
+    c) Prohibiting misrepresentation of the origin of that material, or
+    requiring that modified versions of such material be marked in
+    reasonable ways as different from the original version; or
+
+    d) Limiting the use for publicity purposes of names of licensors or
+    authors of the material; or
+
+    e) Declining to grant rights under trademark law for use of some
+    trade names, trademarks, or service marks; or
+
+    f) Requiring indemnification of licensors and authors of that
+    material by anyone who conveys the material (or modified versions of
+    it) with contractual assumptions of liability to the recipient, for
+    any liability that these contractual assumptions directly impose on
+    those licensors and authors.
+
+  All other non-permissive additional terms are considered "further
+restrictions" within the meaning of section 10.  If the Program as you
+received it, or any part of it, contains a notice stating that it is
+governed by this License along with a term that is a further
+restriction, you may remove that term.  If a license document contains
+a further restriction but permits relicensing or conveying under this
+License, you may add to a covered work material governed by the terms
+of that license document, provided that the further restriction does
+not survive such relicensing or conveying.
+
+  If you add terms to a covered work in accord with this section, you
+must place, in the relevant source files, a statement of the
+additional terms that apply to those files, or a notice indicating
+where to find the applicable terms.
+
+  Additional terms, permissive or non-permissive, may be stated in the
+form of a separately written license, or stated as exceptions;
+the above requirements apply either way.
+
+  8. Termination.
+
+  You may not propagate or modify a covered work except as expressly
+provided under this License.  Any attempt otherwise to propagate or
+modify it is void, and will automatically terminate your rights under
+this License (including any patent licenses granted under the third
+paragraph of section 11).
+
+  However, if you cease all violation of this License, then your
+license from a particular copyright holder is reinstated (a)
+provisionally, unless and until the copyright holder explicitly and
+finally terminates your license, and (b) permanently, if the copyright
+holder fails to notify you of the violation by some reasonable means
+prior to 60 days after the cessation.
+
+  Moreover, your license from a particular copyright holder is
+reinstated permanently if the copyright holder notifies you of the
+violation by some reasonable means, this is the first time you have
+received notice of violation of this License (for any work) from that
+copyright holder, and you cure the violation prior to 30 days after
+your receipt of the notice.
+
+  Termination of your rights under this section does not terminate the
+licenses of parties who have received copies or rights from you under
+this License.  If your rights have been terminated and not permanently
+reinstated, you do not qualify to receive new licenses for the same
+material under section 10.
+
+  9. Acceptance Not Required for Having Copies.
+
+  You are not required to accept this License in order to receive or
+run a copy of the Program.  Ancillary propagation of a covered work
+occurring solely as a consequence of using peer-to-peer transmission
+to receive a copy likewise does not require acceptance.  However,
+nothing other than this License grants you permission to propagate or
+modify any covered work.  These actions infringe copyright if you do
+not accept this License.  Therefore, by modifying or propagating a
+covered work, you indicate your acceptance of this License to do so.
+
+  10. Automatic Licensing of Downstream Recipients.
+
+  Each time you convey a covered work, the recipient automatically
+receives a license from the original licensors, to run, modify and
+propagate that work, subject to this License.  You are not responsible
+for enforcing compliance by third parties with this License.
+
+  An "entity transaction" is a transaction transferring control of an
+organization, or substantially all assets of one, or subdividing an
+organization, or merging organizations.  If propagation of a covered
+work results from an entity transaction, each party to that
+transaction who receives a copy of the work also receives whatever
+licenses to the work the party's predecessor in interest had or could
+give under the previous paragraph, plus a right to possession of the
+Corresponding Source of the work from the predecessor in interest, if
+the predecessor has it or can get it with reasonable efforts.
+
+  You may not impose any further restrictions on the exercise of the
+rights granted or affirmed under this License.  For example, you may
+not impose a license fee, royalty, or other charge for exercise of
+rights granted under this License, and you may not initiate litigation
+(including a cross-claim or counterclaim in a lawsuit) alleging that
+any patent claim is infringed by making, using, selling, offering for
+sale, or importing the Program or any portion of it.
+
+  11. Patents.
+
+  A "contributor" is a copyright holder who authorizes use under this
+License of the Program or a work on which the Program is based.  The
+work thus licensed is called the contributor's "contributor version".
+
+  A contributor's "essential patent claims" are all patent claims
+owned or controlled by the contributor, whether already acquired or
+hereafter acquired, that would be infringed by some manner, permitted
+by this License, of making, using, or selling its contributor version,
+but do not include claims that would be infringed only as a
+consequence of further modification of the contributor version.  For
+purposes of this definition, "control" includes the right to grant
+patent sublicenses in a manner consistent with the requirements of
+this License.
+
+  Each contributor grants you a non-exclusive, worldwide, royalty-free
+patent license under the contributor's essential patent claims, to
+make, use, sell, offer for sale, import and otherwise run, modify and
+propagate the contents of its contributor version.
+
+  In the following three paragraphs, a "patent license" is any express
+agreement or commitment, however denominated, not to enforce a patent
+(such as an express permission to practice a patent or covenant not to
+sue for patent infringement).  To "grant" such a patent license to a
+party means to make such an agreement or commitment not to enforce a
+patent against the party.
+
+  If you convey a covered work, knowingly relying on a patent license,
+and the Corresponding Source of the work is not available for anyone
+to copy, free of charge and under the terms of this License, through a
+publicly available network server or other readily accessible means,
+then you must either (1) cause the Corresponding Source to be so
+available, or (2) arrange to deprive yourself of the benefit of the
+patent license for this particular work, or (3) arrange, in a manner
+consistent with the requirements of this License, to extend the patent
+license to downstream recipients.  "Knowingly relying" means you have
+actual knowledge that, but for the patent license, your conveying the
+covered work in a country, or your recipient's use of the covered work
+in a country, would infringe one or more identifiable patents in that
+country that you have reason to believe are valid.
+
+  If, pursuant to or in connection with a single transaction or
+arrangement, you convey, or propagate by procuring conveyance of, a
+covered work, and grant a patent license to some of the parties
+receiving the covered work authorizing them to use, propagate, modify
+or convey a specific copy of the covered work, then the patent license
+you grant is automatically extended to all recipients of the covered
+work and works based on it.
+
+  A patent license is "discriminatory" if it does not include within
+the scope of its coverage, prohibits the exercise of, or is
+conditioned on the non-exercise of one or more of the rights that are
+specifically granted under this License.  You may not convey a covered
+work if you are a party to an arrangement with a third party that is
+in the business of distributing software, under which you make payment
+to the third party based on the extent of your activity of conveying
+the work, and under which the third party grants, to any of the
+parties who would receive the covered work from you, a discriminatory
+patent license (a) in connection with copies of the covered work
+conveyed by you (or copies made from those copies), or (b) primarily
+for and in connection with specific products or compilations that
+contain the covered work, unless you entered into that arrangement,
+or that patent license was granted, prior to 28 March 2007.
+
+  Nothing in this License shall be construed as excluding or limiting
+any implied license or other defenses to infringement that may
+otherwise be available to you under applicable patent law.
+
+  12. No Surrender of Others' Freedom.
+
+  If conditions are imposed on you (whether by court order, agreement or
+otherwise) that contradict the conditions of this License, they do not
+excuse you from the conditions of this License.  If you cannot convey a
+covered work so as to satisfy simultaneously your obligations under this
+License and any other pertinent obligations, then as a consequence you may
+not convey it at all.  For example, if you agree to terms that obligate you
+to collect a royalty for further conveying from those to whom you convey
+the Program, the only way you could satisfy both those terms and this
+License would be to refrain entirely from conveying the Program.
+
+  13. Remote Network Interaction; Use with the GNU General Public License.
+
+  Notwithstanding any other provision of this License, if you modify the
+Program, your modified version must prominently offer all users
+interacting with it remotely through a computer network (if your version
+supports such interaction) an opportunity to receive the Corresponding
+Source of your version by providing access to the Corresponding Source
+from a network server at no charge, through some standard or customary
+means of facilitating copying of software.  This Corresponding Source
+shall include the Corresponding Source for any work covered by version 3
+of the GNU General Public License that is incorporated pursuant to the
+following paragraph.
+
+  Notwithstanding any other provision of this License, you have
+permission to link or combine any covered work with a work licensed
+under version 3 of the GNU General Public License into a single
+combined work, and to convey the resulting work.  The terms of this
+License will continue to apply to the part which is the covered work,
+but the work with which it is combined will remain governed by version
+3 of the GNU General Public License.
+
+  14. Revised Versions of this License.
+
+  The Free Software Foundation may publish revised and/or new versions of
+the GNU Affero General Public License from time to time.  Such new versions
+will be similar in spirit to the present version, but may differ in detail to
+address new problems or concerns.
+
+  Each version is given a distinguishing version number.  If the
+Program specifies that a certain numbered version of the GNU Affero General
+Public License "or any later version" applies to it, you have the
+option of following the terms and conditions either of that numbered
+version or of any later version published by the Free Software
+Foundation.  If the Program does not specify a version number of the
+GNU Affero General Public License, you may choose any version ever published
+by the Free Software Foundation.
+
+  If the Program specifies that a proxy can decide which future
+versions of the GNU Affero General Public License can be used, that proxy's
+public statement of acceptance of a version permanently authorizes you
+to choose that version for the Program.
+
+  Later license versions may give you additional or different
+permissions.  However, no additional obligations are imposed on any
+author or copyright holder as a result of your choosing to follow a
+later version.
+
+  15. Disclaimer of Warranty.
+
+  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
+APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
+HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
+OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
+THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
+IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
+ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
+
+  16. Limitation of Liability.
+
+  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
+WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
+THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
+GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
+USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
+DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
+PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
+EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGES.
+
+  17. Interpretation of Sections 15 and 16.
+
+  If the disclaimer of warranty and limitation of liability provided
+above cannot be given local legal effect according to their terms,
+reviewing courts shall apply local law that most closely approximates
+an absolute waiver of all civil liability in connection with the
+Program, unless a warranty or assumption of liability accompanies a
+copy of the Program in return for a fee.
+
+                     END OF TERMS AND CONDITIONS
+
+            How to Apply These Terms to Your New Programs
+
+  If you develop a new program, and you want it to be of the greatest
+possible use to the public, the best way to achieve this is to make it
+free software which everyone can redistribute and change under these terms.
+
+  To do so, attach the following notices to the program.  It is safest
+to attach them to the start of each source file to most effectively
+state the exclusion of warranty; and each file should have at least
+the "copyright" line and a pointer to where the full notice is found.
+
+    <one line to give the program's name and a brief idea of what it does.>
+    Copyright (C) <year>  <name of author>
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU Affero General Public License as published
+    by the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU Affero General Public License for more details.
+
+    You should have received a copy of the GNU Affero General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+Also add information on how to contact you by electronic and paper mail.
+
+  If your software can interact with users remotely through a computer
+network, you should also make sure that it provides a way for users to
+get its source.  For example, if your program is a web application, its
+interface could display a "Source" link that leads users to an archive
+of the code.  There are many ways you could offer source, and different
+solutions will be better for different programs; see section 13 for the
+specific requirements.
+
+  You should also get your employer (if you work as a programmer) or school,
+if any, to sign a "copyright disclaimer" for the program, if necessary.
+For more information on this, and how to apply and follow the GNU AGPL, see
+<https://www.gnu.org/licenses/>.