import hashlib import os import json from server import PromptServer from aiohttp import web import base64 from io import BytesIO import asyncio from PIL import Image, ImageOps import torch import numpy as np import folder_paths # Directory node save settings CHUNK_SIZE = 1024 dir_painter_node = os.path.dirname(__file__) extension_path = os.path.join(os.path.abspath(dir_painter_node)) nodes_settings_path = os.path.join(extension_path, "settings_nodes") # Create directory settings_nodes if not exists if not os.path.exists(nodes_settings_path): os.mkdir(nodes_settings_path) tipsfile = os.path.join(nodes_settings_path, "Stores painter nodes settings.txt") with open(tipsfile, "w+", encoding="utf-8") as tipsfile: tipsfile.write("Painter node saved settings!") # Function create file json file PREFIX = "_setting.json" def isFileName(filename): if ( not filename and filename is not None and (type(filename) == str and filename.strip() == "") ): print("Filename is incorrect") return False return True def create_settings_json(filename): try: json_file = os.path.join(nodes_settings_path, filename) if not os.path.isfile(json_file): print(f"File settings for '{filename}' is not found! Create file!") with open(json_file, "w") as f: json.dump({}, f) except Exception as e: print(f"Error: ${e}") def get_settings_json(filename, notExistCreate=True): if not isFileName(filename): return {} json_file = os.path.join(nodes_settings_path, filename) if os.path.isfile(json_file): f = open(json_file, "rb") try: load_data = json.load(f) return load_data except Exception as e: print("Error load json file: ", e) if notExistCreate: f.close() os.remove(json_file) create_settings_json(filename) finally: f.close() else: create_settings_json(filename) return {} # Load json file @PromptServer.instance.routes.get("/alekpet/loading_node_settings/{nodeName}") async def loadingSettings(request): filename = request.match_info.get("nodeName", None) if not isFileName(filename): load_data = {} else: load_data = get_settings_json(filename + PREFIX) return web.json_response({"settings_nodes": load_data}) # Save data to json file @PromptServer.instance.routes.post("/alekpet/save_node_settings") async def saveSettings(request): try: if not request.content_type.startswith("multipart/"): return web.json_response( {"error": "multipart/* content type expected"}, status=400 ) reader = await request.multipart() filename_reader = await reader.next() filename = await filename_reader.text() data_reader = await reader.next() if isFileName(filename): filename = filename + PREFIX json_file = os.path.join(nodes_settings_path, filename) if os.path.isfile(json_file): with open(json_file, "wb") as f: while True: chunk = await data_reader.read_chunk(size=CHUNK_SIZE) if not chunk: break f.write(chunk) return web.json_response( {"message": "Painter data saved successfully"}, status=200 ) else: create_settings_json(filename) return web.json_response( {"message": "Painter file settings created!"}, status=200 ) else: raise Exception("Filename is not found or incorrect!") except Exception as e: print("Error save json file: ", e) return web.json_response({"error": str(e)}, status=500) # Remove file settings painter node data @PromptServer.instance.routes.post("/alekpet/remove_node_settings") async def saveSettings(request): try: json_data = await request.json() filename = json_data.get("name") if isFileName(filename): filename = filename + PREFIX json_file = os.path.join(nodes_settings_path, filename) os.remove(json_file) return web.json_response( {"message": "Painter data removed successfully"}, status=200 ) except OSError as e: return web.json_response( {"error": "Error: %s - %s." % (e.filename, e.strerror)}, status=500 ) # Piping image PAINTER_DICT = {} # Painter nodes dict instances def toBase64ImgUrl(img): bytesIO = BytesIO() img.save(bytesIO, format="PNG") img_types = bytesIO.getvalue() img_base64 = base64.b64encode(img_types) return f"data:image/png;base64,{img_base64.decode('utf-8')}" @PromptServer.instance.routes.post("/alekpet/check_canvas_changed") async def check_canvas_changed(request): json_data = await request.json() unique_id = json_data.get("unique_id", None) is_ok = json_data.get("is_ok", False) if unique_id is not None and unique_id in PAINTER_DICT and is_ok == True: PAINTER_DICT[unique_id].canvas_set = True return web.json_response({"status": "Ok"}) return web.json_response({"status": "Error"}) async def wait_canvas_change(unique_id, time_out=40): for _ in range(time_out): if ( hasattr(PAINTER_DICT[unique_id], "canvas_set") and PAINTER_DICT[unique_id].canvas_set == True ): PAINTER_DICT[unique_id].canvas_set = False return True await asyncio.sleep(0.1) return False # end - Piping image class PainterNode(object): @classmethod def INPUT_TYPES(self): self.canvas_set = False work_dir = folder_paths.get_input_directory() imgs = [ img for img in os.listdir(work_dir) if os.path.isfile(os.path.join(work_dir, img)) ] return { "required": {"image": (sorted(imgs),)}, "hidden": {"unique_id": "UNIQUE_ID"}, "optional": {"images": ("IMAGE",), "update_node": (([True, False],))}, } RETURN_TYPES = ("IMAGE", "MASK") FUNCTION = "painter_execute" CATEGORY = "AlekPet Nodes/image" def painter_execute(self, image, unique_id, update_node=True, images=None): # Piping image input if unique_id not in PAINTER_DICT: PAINTER_DICT[unique_id] = self if update_node == True and images is not None: input_images = [] for imgs in images: i = 255.0 * imgs.cpu().numpy() i = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) input_images.append(toBase64ImgUrl(i)) PAINTER_DICT[unique_id].canvas_set = False PromptServer.instance.send_sync( "alekpet_get_image", {"unique_id": unique_id, "images": input_images} ) if not asyncio.run(wait_canvas_change(unique_id)): print(f"Painter_{unique_id}: Failed to get image!") else: print(f"Painter_{unique_id}: Image received, canvas changed!") # end - Piping image input image_path = folder_paths.get_annotated_filepath(image) i = Image.open(image_path) i = ImageOps.exif_transpose(i) 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.0 - torch.from_numpy(mask) else: mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") return (image, mask.unsqueeze(0)) @classmethod def IS_CHANGED(self, image, unique_id, update_node=True, images=None): 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() @classmethod def VALIDATE_INPUTS(self, image, unique_id, update_node=True, images=None): if not folder_paths.exists_annotated_filepath(image): return "Invalid image file: {}".format(image) return True