GenorTG-ComfyUI-PromptForge/load_image.py

"""PromptForge Load Image: pick any image from the input folder (subfolders included) in a thumbnail browser,
and keep a text description per image that the node outputs next to the picture.

The descriptions live in the user directory (`promptforge/image_descriptions.json`, keyed by the path inside
the input folder), so switching to another character also switches to its description. A description typed into
the node wins over the stored one, which makes the node usable with one-off text as well.
"""
import asyncio
import hashlib
from pathlib import Path

import numpy as np
import torch
from aiohttp import web
from PIL import Image, ImageOps
from server import PromptServer

import node_helpers

from . import image_library

EXTENSIONS = image_library.EXTENSIONS
_resolve = image_library.resolve
_descriptions = image_library.descriptions


@PromptServer.instance.routes.get("/promptforge/images")
async def _images_route(request):
    described, known = _descriptions(), image_library.tags()
    images = await asyncio.get_running_loop().run_in_executor(None, image_library.scan)
    for image in images:
        text = described.get(image["name"], "").strip()
        entry = known.get(image["name"])
        image.update({"described": bool(text), "description": text[:300], "tagged": entry is not None,
                      "theme": entry["theme"] if entry else "", "nsfw": bool(entry and entry.get("nsfw")),
                      "tags": entry["tags"] if entry else [], "caption": entry["caption"] if entry else ""})
    return web.json_response({"images": images})


def _thumbnail(path, size):
    stat = path.stat()
    cached = image_library.store_dir() / "thumbs" / (hashlib.sha1(f"{path}|{stat.st_mtime_ns}|{size}".encode()).hexdigest() + ".jpg")
    if not cached.exists():
        cached.parent.mkdir(parents=True, exist_ok=True)
        with node_helpers.pillow(Image.open, str(path)) as image:
            image = node_helpers.pillow(ImageOps.exif_transpose, image)
            if image.mode in ("RGBA", "LA", "P"):
                image = image.convert("RGBA")
                backdrop = Image.new("RGBA", image.size, (40, 40, 40, 255))
                image = Image.alpha_composite(backdrop, image)
            image = image.convert("RGB")
            image.thumbnail((size, size))
            image.save(cached, "JPEG", quality=82)
    return cached


@PromptServer.instance.routes.get("/promptforge/image_thumb")
async def _thumb_route(request):
    try:
        path = _resolve(request.query.get("name", ""))
        size = min(1024, max(32, int(request.query.get("size", 200))))
        cached = await asyncio.get_running_loop().run_in_executor(None, _thumbnail, path, size)
    except (ValueError, OSError):
        return web.Response(status=404)
    return web.FileResponse(cached, headers={"Cache-Control": "max-age=86400"})


@PromptServer.instance.routes.get("/promptforge/image_description")
async def _description_route(request):
    return web.json_response({"description": _descriptions().get(request.query.get("name", ""), "")})


@PromptServer.instance.routes.post("/promptforge/image_description")
async def _save_description_route(request):
    body = await request.json()
    try:
        name = _resolve(str(body.get("name", ""))).relative_to(image_library.input_root()).as_posix()
    except ValueError as error:
        return web.json_response({"error": str(error)}, status=400)
    text = str(body.get("description", "")).strip()
    with image_library._store_lock:
        store = _descriptions()
        if text:
            store[name] = text
        else:
            store.pop(name, None)
        image_library.save_descriptions(store)
    return web.json_response({"saved": bool(text)})


class PROMPTFORGE_LoadImage:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("STRING", {"default": "", "tooltip": "Path inside the input folder. Use Choose image to pick it from thumbnails."}),
                "description": ("STRING", {"default": "", "multiline": True, "placeholder": "Description of this image / character. Empty = the stored description.", "tooltip": "Output as DESCRIPTION. The stored description of the image is used when this is empty; Save description stores this text for the image."}),
                "enabled": ("BOOLEAN", {"default": True, "tooltip": "Off = the node outputs nothing, as if it were bypassed, so it can stay in a workflow and be switched on when needed."}),
            },
        }

    RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING")
    RETURN_NAMES = ("IMAGE", "MASK", "DESCRIPTION", "NAME")
    FUNCTION = "load"
    CATEGORY = "🔮 PromptForge/Load"
    DESCRIPTION = "Load an image from a thumbnail browser of the input folder and output a per-image description (stored, or typed in the node) as text. With no image chosen the image output is empty, so the node can stay in a workflow unused."

    @classmethod
    def IS_CHANGED(cls, image, description, enabled=True):
        try:
            if not image or not enabled:
                return description
            path = _resolve(image)
            return f"{path.stat().st_mtime_ns}|{description}|{_descriptions().get(image, '')}"
        except (ValueError, OSError):
            return float("nan")

    @classmethod
    def VALIDATE_INPUTS(cls, image, enabled=True):
        if not image.strip() or not enabled:
            return True
        try:
            return True if _resolve(image).is_file() else f"{image} does not exist"
        except ValueError as error:
            return str(error)

    def load(self, image, description, enabled=True):
        if not image or not enabled:
            return (None, None, "", "")
        path = _resolve(image)
        with node_helpers.pillow(Image.open, str(path)) as opened:
            frame = node_helpers.pillow(ImageOps.exif_transpose, opened)
            if frame.mode == "I":
                frame = frame.point(lambda value: value * (1 / 255))
            rgb = np.asarray(frame.convert("RGB"), dtype=np.float32) / 255.0
            alpha = np.asarray(frame.getchannel("A"), dtype=np.float32) / 255.0 if "A" in frame.getbands() else None
        mask = 1.0 - torch.from_numpy(alpha) if alpha is not None else torch.zeros(rgb.shape[:2], dtype=torch.float32)
        text = description.strip() or _descriptions().get(image, "").strip()
        return (torch.from_numpy(rgb)[None], mask[None], text, path.stem)


NODE_CLASS_MAPPINGS = {"PROMPTFORGE_LoadImage": PROMPTFORGE_LoadImage}
NODE_DISPLAY_NAME_MAPPINGS = {"PROMPTFORGE_LoadImage": "PromptForge Load Image"}