118 lines
3.5 KiB
Python
118 lines
3.5 KiB
Python
import folder_paths
|
|
from PIL import Image, ImageOps
|
|
import numpy as np
|
|
import torch
|
|
from comfy_execution.graph import ExecutionBlocker
|
|
import threading
|
|
from server import PromptServer
|
|
from aiohttp import web
|
|
|
|
thread = None
|
|
g_node_id = None
|
|
g_filename = None
|
|
threads = []
|
|
|
|
routes = PromptServer.instance.routes
|
|
@routes.post('/compositor/done')
|
|
async def receivedDone(request):
|
|
return web.json_response({})
|
|
|
|
class Compositor3:
|
|
file = "new.png"
|
|
result = None
|
|
configCache = None
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
fabricData = kwargs.get("fabricData")
|
|
# print(fabricData)
|
|
return fabricData
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"config": ("COMPOSITOR_CONFIG", {"forceInput": True}),
|
|
"fabricData": ("STRING", {"default": "{}"}),
|
|
"imageName": ("STRING", {"default": "new.png"}),
|
|
},
|
|
"optional": {
|
|
"tools": ("BOOLEAN", {"forceInput": True, "default": True}),
|
|
},
|
|
"hidden": {
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
"node_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "IMAGE")
|
|
RETURN_NAMES = ("transforms", "image")
|
|
FUNCTION = "composite"
|
|
CATEGORY = "image"
|
|
|
|
def composite(self, **kwargs):
|
|
# https://blog.miguelgrinberg.com/post/how-to-make-python-wait
|
|
node_id = kwargs.pop('node_id', None)
|
|
|
|
|
|
imageName = kwargs.get('imageName', "new.png")
|
|
|
|
config = kwargs.get('config', "default")
|
|
padding = config["padding"]
|
|
invertMask = config["invertMask"]
|
|
width = config["width"]
|
|
height = config["height"]
|
|
config_node_id = config["node_id"]
|
|
onConfigChanged = config["onConfigChanged"]
|
|
names = config["names"]
|
|
preset = config["preset"]
|
|
fabricData = kwargs.get("fabricData")
|
|
|
|
configChanged = self.configCache != config
|
|
# print(configChanged)
|
|
# print(config)
|
|
# print(self.configCache)
|
|
|
|
|
|
self.configCache = config
|
|
ui = {
|
|
"test": ("value",),
|
|
"padding": [padding],
|
|
"width": [width],
|
|
"height": [height],
|
|
"config_node_id": [config_node_id],
|
|
"node_id": [node_id],
|
|
"names": names,
|
|
"fabricData": [fabricData],
|
|
"preset": [preset],
|
|
"awaited": [self.result],
|
|
"configChanged": [configChanged],
|
|
"onConfigChanged": [onConfigChanged],
|
|
}
|
|
|
|
# break and send a message to the gui as if it was "executed" below
|
|
detail = {"output": ui, "node": node_id}
|
|
PromptServer.instance.send_sync("compositor_init", detail)
|
|
|
|
imageExists = folder_paths.exists_annotated_filepath(imageName)
|
|
# block when config changed
|
|
if imageName == "new.png" or not imageExists or configChanged:
|
|
return {
|
|
"ui": ui,
|
|
"result": (ExecutionBlocker(None), ExecutionBlocker(None))
|
|
}
|
|
|
|
image_path = folder_paths.get_annotated_filepath(imageName)
|
|
i = Image.open(image_path)
|
|
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, ]
|
|
|
|
return {
|
|
"ui": ui,
|
|
"result": (fabricData, image)
|
|
}
|