| import os
|
| import time
|
| import torch
|
| import folder_paths
|
| from server import PromptServer
|
| from aiohttp import web
|
| from PIL import Image
|
| import numpy as np
|
|
|
| CACHE_DIR = os.path.join(folder_paths.get_temp_directory(), "image_chooser_cache")
|
| os.makedirs(CACHE_DIR, exist_ok=True)
|
|
|
|
|
| CHOICES = {}
|
|
|
| class ImageChooser:
|
| def __init__(self):
|
| pass
|
|
|
| @classmethod
|
| def INPUT_TYPES(cls):
|
| return {
|
| "required": {
|
| "image_A": ("IMAGE",),
|
| "image_B": ("IMAGE",),
|
| },
|
| "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", "my_id": "UNIQUE_ID"},
|
| }
|
|
|
| RETURN_TYPES = ("IMAGE",)
|
| FUNCTION = "execute"
|
| CATEGORY = "utils"
|
|
|
|
|
|
|
|
|
| @classmethod
|
| def IS_CHANGED(cls, **kwargs):
|
| return time.time()
|
|
|
| def execute(self, image_A, image_B, prompt=None, extra_pnginfo=None, my_id=None):
|
| my_id_str = str(my_id)
|
|
|
|
|
| self.save_temp_image(image_A, f"{my_id_str}_A.png")
|
| self.save_temp_image(image_B, f"{my_id_str}_B.png")
|
|
|
|
|
| CHOICES[my_id_str] = None
|
|
|
|
|
| PromptServer.instance.send_sync("image_chooser_wait", {"node_id": my_id_str})
|
|
|
| print(f"[Image Chooser] Workflow paused. Waiting for user selection on node {my_id_str}...")
|
|
|
|
|
| while CHOICES.get(my_id_str) is None:
|
| time.sleep(0.1)
|
|
|
| choice = CHOICES.pop(my_id_str)
|
|
|
| print(f"[Image Chooser] Pause lifted. Result received: {choice}")
|
|
|
| if choice == "CANCEL":
|
| raise Exception("Execution cancelled by user via Image Chooser.")
|
|
|
| if choice == "A":
|
| return (image_A,)
|
| else:
|
| return (image_B,)
|
|
|
| def save_temp_image(self, tensor, filename):
|
| if len(tensor.shape) == 4:
|
| img_np = tensor[0].cpu().numpy()
|
| else:
|
| img_np = tensor.cpu().numpy()
|
|
|
| img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8)
|
| img = Image.fromarray(img_np)
|
| img.save(os.path.join(CACHE_DIR, filename))
|
|
|
|
|
| @PromptServer.instance.routes.post("/image_chooser/select")
|
| async def handle_select(request):
|
| json_data = await request.json()
|
| node_id = str(json_data.get("node_id"))
|
| choice = json_data.get("choice")
|
|
|
| CHOICES[node_id] = choice
|
| return web.json_response({"status": "success"})
|
|
|
|
|
| NODE_CLASS_MAPPINGS = {
|
| "ImageChooser": ImageChooser
|
| }
|
|
|
| NODE_DISPLAY_NAME_MAPPINGS = {
|
| "ImageChooser": "Image Chooser"
|
| }
|
|
|
| WEB_DIRECTORY = "." |
