ComfyUI-BlenderAI-node/SDNode/custom_nodes/ComfyUI-CUP/cup.py

import os
import json
import sys
import traceback
import numpy as np
import builtins
import torch
import shutil
import hashlib
import atexit
import server
import random
import gc
import execution
import folder_paths
import nodes
import time
import asyncio
import logging as logger
from comfy.cli_args import args
try:
    from comfy.comfy_types import IO
except (ImportError, ModuleNotFoundError):
    class IO:
        STRING = "STRING"
try:
    from comfy.model_management import unload_all_models, soft_empty_cache
except ImportError:
    unload_all_models = lambda: None
    soft_empty_cache = lambda: None
from threading import Thread
from aiohttp import web
from pathlib import Path
from PIL import Image, ImageOps, ImageSequence
from PIL.PngImagePlugin import PngInfo

FORCE_LOG = False
CATEGORY_ = "Blender"
TEMPDIR = Path(__file__).parent.parent.parent / "SDNodeTemp"
HOST_PATH = Path("XXXHOST-PATHXXX")
if not HOST_PATH.parent.exists():
    HOST_PATH = Path(__file__).parent


def remove_old_version():
    cup = Path(__file__).parent.parent.joinpath("cup.py")
    if cup.exists():
        cup.unlink()


remove_old_version()


def removetemp():
    try:
        if TEMPDIR.exists():
            shutil.rmtree(TEMPDIR, ignore_errors=True)
        TEMPDIR.mkdir(parents=True)
    except Exception as e:
        sys.stdout.write(e)


removetemp()


def execute_wrap():
    def exec_wrap(func):
        def wrap(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except Exception:
                traceback.print_exc()
                sys.stdout.flush()
        return wrap

    execution.PromptExecutor.execute = exec_wrap(execution.PromptExecutor.execute)


execute_wrap()

atexit.register(removetemp)


def hk(func):
    def __print_wrap__(*args, **kwargs):
        try:
            func(*args, **kwargs)
        except BaseException:
            ...
        sys.stdout.flush()
        # sys.stderr.flush()
    return __print_wrap__


if FORCE_LOG:
    __print_wrap__ = hk(print)
    globals()["__print_wrap__"] = __print_wrap__
    builtins.print = __print_wrap__

    sys.stdout.write = hk(sys.stdout.write)
    sys.stderr.write = hk(sys.stderr.write)
    # sys.stderr.write = builtins.print


def try_write_config():
    config_path = HOST_PATH.joinpath("PATH_CFG.json")
    from folder_paths import folder_names_and_paths
    config = {}
    for k in folder_names_and_paths:
        config[k] = list(folder_names_and_paths[k])
        if isinstance(config[k][1], set):
            config[k][1] = list(config[k][1])
    Path(config_path).write_text(json.dumps(config, indent=4))


try:
    try_write_config()
except Exception as e:
    sys.stdout.write("Config Export Error")
    sys.stdout.write(str(e))
    sys.stdout.flush()

CACHED_EXECUTOR = []


@server.PromptServer.instance.routes.post("/cup/clear_cache")
async def clear_cache(request):
    inst = server.PromptServer.instance
    if inst.prompt_queue:
        inst.prompt_queue.history.clear()
    if not CACHED_EXECUTOR:
        CACHED_EXECUTOR.extend([ob for ob in gc.get_objects() if isinstance(ob, execution.PromptExecutor)])
    for ob in CACHED_EXECUTOR:
        print("Clear Node Tree Cache", ob)
        ob.outputs.clear()
        ob.outputs_ui.clear()
        ob.old_prompt.clear()
    return web.Response(status=200)


@server.PromptServer.instance.routes.post("/cup/clear_vram")
async def clearn_vram(request):
    try:
        gc.collect()
        unload_all_models()
        soft_empty_cache()
    except Exception as e:
        sys.stdout.write("Clear VRAM Error")
        sys.stdout.write(str(e))
        sys.stdout.flush()
    return web.Response(status=200)

def node_info(node_class):
    """
    ref: ComfyUI/server.py PromptServer
    """
    obj_class = nodes.NODE_CLASS_MAPPINGS[node_class]
    info = {}
    try:
        info['input'] = obj_class.INPUT_TYPES()
    except Exception:
        ...
    info['output'] = obj_class.RETURN_TYPES
    info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES)
    info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output']
    info['name'] = node_class
    info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[node_class] if node_class in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else node_class
    info['description'] = ''
    info['category'] = 'sd'
    if hasattr(obj_class, 'OUTPUT_NODE') and obj_class.OUTPUT_NODE is True:
        info['output_node'] = True
    else:
        info['output_node'] = False
    if hasattr(obj_class, 'CATEGORY'):
        info['category'] = obj_class.CATEGORY
    return info


class NodeCacheManager:
    def __init__(self):
        self.cached_nodes = {}
        self.diff = {}
        self.filter_node = {"Note", "PrimitiveNode", "Cache Node"}

    def calc_diff(self):
        if self.diff:
            return self.diff

        for x in list(nodes.NODE_CLASS_MAPPINGS):
            if x in self.filter_node:
                continue
            ni = node_info(x)
            if x not in self.cached_nodes or self.cached_nodes[x] != ni:
                self.diff[x] = ni
            self.cached_nodes[x] = ni
        return self.diff

    def update_cached_nodes(self, remote=False):
        # 本地部署 但差异路径不存在
        if not remote and not HOST_PATH.parent.exists():
            return {}
        from copy import deepcopy
        self.calc_diff()
        diff = self.diff
        if remote:
            diff = deepcopy(self.diff)
            self.diff.clear()
        else:
            try:
                # 本地部署时写入差异后清理
                diff_path = HOST_PATH.joinpath("diff_object_info.json")
                with diff_path.open("w", encoding="utf-8") as f:
                    f.write(json.dumps(self.diff))
                self.diff.clear()
            except Exception as e:
                # 写入失败后
                # print(f"Failed to write diff: {e}")
                ...
        return diff


cache_manager = NodeCacheManager()


def diff_listen_loop():
    while True:
        time.sleep(5)
        cache_manager.update_cached_nodes()


Thread(target=diff_listen_loop, daemon=True).start()


@server.PromptServer.instance.routes.post("/cup/get_temp_directory")
async def get_temp_directory(request):
    return web.Response(status=200, body=folder_paths.get_temp_directory())


async def queue_msg():
    queue = {}
    current_queue = server.PromptServer.instance.prompt_queue.get_current_queue()
    queue['queue_running'] = current_queue[0]
    queue['queue_pending'] = current_queue[1]
    server.PromptServer.instance.send_sync("cup.queue", queue)


async def diff_msg():
    diff = cache_manager.update_cached_nodes(remote=True)
    server.PromptServer.instance.send_sync("cup.diff", diff)


async def msg_loop():
    while True:
        await queue_msg()
        await diff_msg()
        await asyncio.sleep(1)

Thread(target=asyncio.run, args=(msg_loop(),), daemon=True).start()


class ToBlender:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "批次": ("INT", {
                    "default": 0,
                    "min": 0,  # Minimum value
                    "max": 4096,  # Maximum value
                    "step": 64  # Slider's step
                }),
                "降噪": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
                "日志": (["enable", "disable"],)
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "saveto"
    CATEGORY = CATEGORY_

    def saveto(self, image, 批次, 降噪, 日志):
        image = 1.0 - image
        return (image,)


class SaveImage:
    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()
        self.type = "output"
        self.prefix_append = ""
        self.compress_level = 4

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                {"images": ("IMAGE", ),
                 "filename_prefix": ("STRING", {"default": "ComfyUI"}),
                 "output_dir": ("STRING", {"default": ""}),
                 },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }

    RETURN_TYPES = ()
    FUNCTION = "save_images"

    OUTPUT_NODE = True

    CATEGORY = CATEGORY_

    def save_images(self, images, filename_prefix="ComfyUI", output_dir="", prompt=None, extra_pnginfo=None):
        filename_prefix += self.prefix_append
        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
        results = list()
        for (batch_number, image) in enumerate(images):
            i = 255. * image.cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
            metadata = None
            if not args.disable_metadata:
                metadata = PngInfo()
                if prompt is not None:
                    metadata.add_text("prompt", json.dumps(prompt))
                if extra_pnginfo is not None:
                    for x in extra_pnginfo:
                        metadata.add_text(x, json.dumps(extra_pnginfo[x]))

            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
            file = f"{filename_with_batch_num}_{counter:05}_.png"
            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level)
            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type
            })
            counter += 1

        return {"ui": {"images": results}}


class PreviewImage(SaveImage):
    def __init__(self):
        self.output_dir = folder_paths.get_temp_directory()
        self.type = "temp"
        self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for _ in range(5))
        self.compress_level = 1

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                {"images": ("IMAGE", ), },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }


class CupException(Exception):
    pass


class LoadImage:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("STRING", {"default": ""}),
                "mode": (["输入", "渲染", "序列图", "视口"], ),
            },
        }

    CATEGORY = CATEGORY_

    RETURN_TYPES = ("IMAGE", "MASK")
    FUNCTION = "load_image"

    def load_image(self, image, mode=None):
        if not image:
            raise CupException({
                "type": "sdn_no_image_provided",
                "message": "No Image Provided",
                "details": "",
                "extra_info": {}
            })
        if image:
            image = image.replace("\\\\", "/").replace("\\", "/")
        image = Path(image).name
        image = f"SDN/{image}"
        image_path = folder_paths.get_annotated_filepath(image)
        if not Path(image_path).exists():
            raise CupException({
                "type": "sdn_image_not_found",
                "message": "Image Not Found",
                "details": Path(image_path).as_posix(),
                "extra_info": {}
            })
        # image_path = image
        img = Image.open(image_path)
        output_images = []
        output_masks = []
        for i in ImageSequence.Iterator(img):
            i = ImageOps.exif_transpose(i)
            if i.mode == 'I':
                i = i.point(lambda i: i * (1 / 255))
            image = i.convert("RGB")
            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[None,]
            if 'A' in i.getbands():
                mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
                mask = 1. - torch.from_numpy(mask)
            else:
                mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
            output_images.append(image)
            output_masks.append(mask.unsqueeze(0))

        if len(output_images) > 1:
            output_image = torch.cat(output_images, dim=0)
            output_mask = torch.cat(output_masks, dim=0)
        else:
            output_image = output_images[0]
            output_mask = output_masks[0]

        return (output_image, output_mask)
        try:
            image_path = image
            i = Image.open(image_path)
            if 'P' in i.getbands():
                i = i.convert("RGBA")
                image = i.convert("RGB")
            image = i.convert("RGB")
            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[None,]
        except Exception as e:
            sys.stderr.write(f"|已忽略| Load Image Error -> {e}")
            sys.stderr.flush()
            image = np.zeros(shape=(64, 64, 3)).astype(np.float32)
            image = torch.from_numpy(image)[None,]
            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
            return (image, mask)
        if 'A' in i.getbands():
            mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
            mask = 1. - torch.from_numpy(mask)
        else:
            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
        return (image, mask)

    @classmethod
    def IS_CHANGED(s, image, mode):
        if image:
            image = image.replace("\\\\", "/").replace("\\", "/")
        image = Path(image).name
        image = f"SDN/{image}"
        image_path = folder_paths.get_annotated_filepath(image)
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()


class SaveAudio:
    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()
        self.type = "output"
        self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for _ in range(5))
        self.compress_level = 4

    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"audio": ("AUDIO", ), },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }

    RETURN_TYPES = ()
    FUNCTION = "save_audio"

    OUTPUT_NODE = True

    CATEGORY = CATEGORY_

    def save_audio(self, audio, prompt=None, extra_pnginfo=None):
        import torchaudio
        filename_prefix = "audio/ComfyUI"
        filename_prefix += self.prefix_append
        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
        results = list()
        for (batch_number, waveform) in enumerate(audio["waveform"]):
            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
            file = f"{filename_with_batch_num}_{counter:05}_.flac"
            torchaudio.save(os.path.join(full_output_folder, file), waveform, audio["sample_rate"], format="FLAC")
            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type
            })
            counter += 1

        return {"ui": {"audio": results}}


class SaveModel:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": (
                    IO.STRING,
                    {
                        "default": None,
                        "tooltip": "模型.",
                    },
                ),
            },
            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
        }

    RETURN_TYPES = ()
    FUNCTION = "save"

    OUTPUT_NODE = True

    CATEGORY = CATEGORY_

    def save(self, model, prompt=None, extra_pnginfo=None):
        return {"ui": {"result": [model]}}


class PreviewAudio:
    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()
        self.type = "output"
        self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for _ in range(5))
        self.compress_level = 4

    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"audio": ("AUDIO", ), },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }

    RETURN_TYPES = ()
    FUNCTION = "save_audio"

    OUTPUT_NODE = True

    CATEGORY = CATEGORY_

    def save_audio(self, audio, prompt=None, extra_pnginfo=None):
        import torchaudio
        filename_prefix = "audio/ComfyUI"
        filename_prefix += self.prefix_append
        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
        results = list()
        for (batch_number, waveform) in enumerate(audio["waveform"]):
            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
            file = f"{filename_with_batch_num}_{counter:05}_.flac"
            torchaudio.save(os.path.join(full_output_folder, file), waveform, audio["sample_rate"], format="FLAC")
            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type
            })
            counter += 1

        return {"ui": {"audio": results}}


class MatImage(LoadImage):
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("STRING", {"default": ""}),
            },
        }


class Screenshot(LoadImage):
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("STRING", {"default": ""}),
            },
        }


class Mask:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("STRING", {"default": ""}),
                "channel": (["alpha", "red", "green", "blue"], ),
            }
        }

    CATEGORY = CATEGORY_

    RETURN_TYPES = ("MASK",)
    FUNCTION = "load_image"

    def load_image(self, image, channel):
        image_path = image
        try:
            i = Image.open(image_path)
        except FileNotFoundError:
            print(f"FileNotFound -> {image_path}")
            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
            return (mask,)
        mask = None
        c = channel[0].upper()
        if c in i.getbands():
            mask = np.array(i.getchannel(c)).astype(np.float32) / 255.0
            mask = torch.from_numpy(mask)
            if c == 'A':
                mask = 1. - mask
        else:
            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
        return (mask,)

    @classmethod
    def IS_CHANGED(s, image, channel):
        image_path = image
        image_path = Path(image_path)
        if not image or not image_path.exists():
            return ""
        return Path(image_path).stat().st_mtime_ns


class OpenPoseBase:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("STRING", {"default": ""}),
                "frame": ("INT", {
                    "default": 0,
                    "min": -2**31,  # Minimum value
                    "max": 2**31,  # Maximum value
                    "step": 1  # Slider's step
                }),
            },
        }

    CATEGORY = "OpenPose"

    RETURN_TYPES = ("IMAGE", "MASK")
    FUNCTION = "load_image"

    posetype = ""

    def load_image(self, image, frame):
        try:
            img_dir = Path(image) / self.posetype
            find_img = ""
            for file in img_dir.iterdir():
                if not file.name.startswith("Image"):
                    continue

                f = int(file.name[len("Image"): -len(file.suffix)])
                if f == frame:
                    find_img = file.as_posix()
            if not find_img:
                sys.stderr.write(f"|错误| Frame Not Found -> Image{frame:04}.png")
                sys.stderr.flush()
            image_path = find_img

            i = Image.open(image_path)
            image = i.convert("RGB")
            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[None,]
        except Exception as e:
            sys.stderr.write(f"|已忽略| Load Image Error -> {e}")
            sys.stderr.flush()
            image = np.zeros(shape=(64, 64, 3)).astype(np.float32)
            image = torch.from_numpy(image)[None,]
            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
            return (image, mask)
        if 'A' in i.getbands():
            mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
            mask = 1. - torch.from_numpy(mask)
        else:
            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
        return (image, mask)

    @classmethod
    def IS_CHANGED(s, image, frame):
        image_path = image
        if not os.path.exists(image_path):
            return ""
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()


class OpenPoseFull(OpenPoseBase):
    posetype = "openpose_full"


class OpenPoseHand(OpenPoseBase):
    posetype = "openpose_hand"


class OpenPoseMediaPipeFace(OpenPoseBase):
    posetype = "MediaPipe_face"


class OpenPoseDepth(OpenPoseBase):
    posetype = "depth"


class OpenPose(OpenPoseBase):
    posetype = "openpose"


class OpenPoseFace(OpenPoseBase):
    posetype = "openpose_face"


class OpenPoseLineart(OpenPoseBase):
    posetype = "Lineart"


class OpenPoseFullExtraLimb(OpenPoseBase):
    posetype = "openpose_full_Extra_Limb"


class OpenPoseKeyPose(OpenPoseBase):
    posetype = "keypose"


class OpenPoseCanny(OpenPoseBase):
    posetype = "canny"


[
    'openpose_full',
    'openpose_hand',
    'MediaPipe_face',
    'depth',
    'openpose',
    'openpose_face',
    'Lineart',
    'openpose_full_Extra_Limb',
    'keypose',
    'canny'
]

# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
    "输入图像": LoadImage,
    "材质图": MatImage,
    "截图": Screenshot,
    "Mask": Mask,
    "存储": SaveImage,
    # "导入": ToBlender,
    "预览": PreviewImage,
    "SaveAudioBL": SaveAudio,
    "SaveModel": SaveModel,
    "PreviewAudio": PreviewAudio,
    'OpenPoseFull': OpenPoseFull,
    'OpenPoseHand': OpenPoseHand,
    'OpenPoseMediaPipeFace': OpenPoseMediaPipeFace,
    'OpenPoseDepth': OpenPoseDepth,
    'OpenPose': OpenPose,
    'OpenPoseFace': OpenPoseFace,
    'OpenPoseLineart': OpenPoseLineart,
    'OpenPoseFullExtraLimb': OpenPoseFullExtraLimb,
    'OpenPoseKeyPose': OpenPoseKeyPose,
    'OpenPoseCanny': OpenPoseCanny,
}


NODE_DISPLAY_NAME_MAPPINGS = {
    # Sampling
    "输入图像": "Input Image",
    "材质图": "Mat Image",
    "截图": "Screenshot",
    "Mask": "Mask",
    "存储": "Save",
    "预览": "Preview",
}
WEB_DIRECTORY = "./"