ComfyUI-Image-Promptor/nodes.py

"""Detailed image-to-prompt node for ComfyUI.

The provider protocol follows the original MiniMax H3 Promptor package, but
this node is intentionally image-only and does not expose video modes.
"""

from __future__ import annotations

import base64
import io
import json
import os
import re
from typing import Any

import requests


PROVIDERS = ["ollama", "openai_compatible", "gemini", "claude"]

SYSTEM_PROMPT = r"""You are an expert visual analyst, photographer, art director, and prompt engineer for image-generation models.

Your task is to inspect the supplied image(s) and write one exceptionally detailed, self-contained prompt for generating a still image that reproduces the visible content. The result will be pasted directly into an image-generation model. Analyze only what is visible. Never invent an identity, brand, logo, text, object, body part, or event that cannot be supported by the image. When a detail is uncertain, use a visually grounded description such as "appears to be" or omit it.

Return ONLY the final image-generation prompt as continuous prose in English. Do not include a title, preface, analysis, bullet list, markdown, quotation marks around the whole answer, negative prompt, camera settings outside the prompt, or commentary about your process. Do not mention that you saw an image. Do not describe time, duration, movement over time, transitions, animation, video, or sound. Describe a single frozen moment.

Build the prompt from broad composition to fine detail, using natural language and clear spatial relationships. Cover every applicable category below:

1. Scene summary: the complete subject and situation in one precise sentence; identify the image type (photograph, cinematic still, editorial fashion image, illustration, 3D render, painting, etc.) only when visually supported.
2. People, separately and in order from foreground to background: approximate age group, gender presentation only when visually apparent, body build, skin tone, face shape, hair style/color/length, facial hair, visible facial features, gaze direction, facial expression, emotion conveyed, head orientation, exact body orientation, pose, weight distribution, arm and hand positions, leg and foot positions, and the specific action or interaction. For multiple people, label them with clear spatial identifiers such as "the woman on the left" or "the man in the background" and explain their relationships, eye-lines, touch, gestures, and relative positions. Describe only visible anatomy and do not infer private attributes.
3. Clothing and styling for each person: garment type, cut, fit, layers, sleeve length, neckline, fabric, texture, pattern, colors, closures, folds, wear, accessories, jewelry, eyewear, headwear, shoes, makeup, and hairstyle. Preserve exact visible color relationships and distinguish clothing from the background.
4. Objects and interactions: every important prop, furniture item, vehicle, animal, food, tool, or architectural element; material, shape, size, color, condition, placement, orientation, and what each person is holding, touching, using, sitting on, leaning against, or looking at.
5. Environment and setting: indoor or outdoor setting, architecture, room layout, landscape, street or location cues, foreground/midground/background layers, surfaces, walls, floor, ceiling, windows, foliage, weather, haze, reflections, and small contextual details that affect the scene.
6. Composition and spatial layout: aspect-ratio impression, orientation, subject placement, foreground framing, symmetry or asymmetry, negative space, leading lines, visual hierarchy, overlap, depth, scale, horizon or vanishing point, and what is cropped by the frame. Use relative positions rather than invented measurements.
7. Camera and optics, only when inferable: viewpoint height, camera angle, shot scale (extreme close-up, close-up, medium, full-body, wide shot), perspective, focal-length impression (wide-angle, normal, telephoto, macro), depth of field, focus plane, bokeh, motion freeze, and lens distortion. Do not fabricate exact camera model or settings.
8. Lighting: apparent light sources and direction, hard or soft quality, key/fill/rim light, natural or artificial character, shadows, highlights, reflections, contrast, exposure, color temperature, and how light falls across faces, clothing, objects, and background.
9. Color and visual finish: dominant and secondary colors, palette, saturation, contrast, tonal range, skin-tone rendering, color grading, texture, grain, sharpness, realism level, and the visual style supported by the image. Mention cinematic, documentary, studio, fashion, vintage, painterly, anime, or other style labels only if they match the visible result.
10. Text and symbols: transcribe visible text exactly only when legible; otherwise state its approximate placement and that it is illegible. Never invent readable wording. Include logos or symbols only when clearly visible, and do not guess brand names.

Use specific nouns and descriptive adjectives, not vague phrases such as "nice image" or "beautiful scene." Preserve the original subject count, identity-neutral appearance, pose, action, clothing, object relationships, camera framing, lighting direction, and background structure. Make the prompt richly detailed but coherent, with no contradictions and no unsupported additions. End with a concise quality/style clause grounded in the image, such as realistic high-detail photography or detailed digital illustration, only when appropriate."""


class ProviderError(RuntimeError):
    pass


def _tensor_to_jpeg_b64(image: Any) -> list[str]:
    """Convert a ComfyUI IMAGE batch (BHWC float tensor) to JPEG base64."""
    try:
        import torch
        from PIL import Image
    except Exception as exc:
        raise ProviderError(f"Image conversion dependencies unavailable: {exc}") from exc

    if not isinstance(image, torch.Tensor):
        raise ProviderError("The image input is not a ComfyUI tensor.")
    tensor = image.detach().cpu().float().clamp(0.0, 1.0)
    if tensor.ndim == 3:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim != 4 or tensor.shape[-1] not in (1, 3, 4):
        raise ProviderError(f"Expected IMAGE tensor in BHWC format, got shape {tuple(tensor.shape)}.")

    result = []
    for frame in tensor:
        pixels = (frame * 255.0).round().byte().numpy()
        if pixels.shape[-1] == 1:
            pixels = pixels[:, :, 0]
        elif pixels.shape[-1] == 4:
            pixels = pixels[:, :, :3]
        with io.BytesIO() as buffer:
            Image.fromarray(pixels).save(buffer, format="JPEG", quality=92, optimize=True)
            result.append(base64.b64encode(buffer.getvalue()).decode("ascii"))
    return result


def _first_content(value: Any) -> str:
    if isinstance(value, str):
        return value
    if isinstance(value, list):
        return "\n".join(_first_content(item) for item in value)
    if isinstance(value, dict):
        for key in ("text", "content", "response", "output"):
            if key in value:
                return _first_content(value[key])
    return ""


def _clean_response(text: str) -> str:
    text = (text or "").strip()
    text = re.sub(r"^```(?:text|markdown)?\s*|\s*```$", "", text, flags=re.IGNORECASE | re.DOTALL).strip()
    text = re.sub(r"^(?:final prompt|image prompt|prompt)\s*:\s*", "", text, flags=re.IGNORECASE)
    return text.strip().strip('"').strip()


def _load_config() -> dict[str, Any]:
    path = os.path.join(os.path.dirname(__file__), "config.json")
    defaults = {
        "ollama": {"api_base": "http://127.0.0.1:11434", "model": "SetneufPT/Gemma4-12B-IT-QAT_Q4_64K_16GB-GPU:latest"},
        "openai_compatible": {"api_base": "https://api.openai.com/v1", "api_key": "", "model": "gpt-4o"},
        "gemini": {"api_base": "https://generativelanguage.googleapis.com/v1beta", "api_key": "", "model": "gemini-2.5-flash"},
        "claude": {"api_base": "https://api.anthropic.com/v1", "api_key": "", "model": "claude-3-5-sonnet-20241022"},
    }
    try:
        with open(path, "r", encoding="utf-8") as handle:
            user = json.load(handle)
        for name, values in user.get("providers", user).items():
            if isinstance(values, dict) and name in defaults:
                defaults[name].update(values)
    except FileNotFoundError:
        pass
    except Exception as exc:
        print(f"[Image Promptor] config warning: {exc}")
    return defaults


def _post_json(url: str, payload: dict, headers: dict | None = None, timeout: int = 300) -> dict:
    try:
        response = requests.post(url, json=payload, headers=headers or {}, timeout=timeout)
    except requests.RequestException as exc:
        raise ProviderError(f"Cannot connect to {url}: {exc}") from exc
    if not response.ok:
        raise ProviderError(f"HTTP {response.status_code}: {response.text[:400]}")
    try:
        return response.json()
    except ValueError as exc:
        raise ProviderError(f"Invalid JSON response from {url}: {response.text[:400]}") from exc


def _call_ollama(cfg: dict, user_message: str, images: list[str], temperature: float, max_tokens: int) -> str:
    payload = {"model": cfg["model"], "stream": False, "messages": [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_message, "images": images},
    ], "options": {"temperature": temperature, "num_predict": max_tokens}}
    data = _post_json(cfg["api_base"].rstrip("/") + "/api/chat", payload)
    return _first_content(data.get("message", {}).get("content", ""))


def _call_openai(cfg: dict, user_message: str, images: list[str], temperature: float, max_tokens: int) -> str:
    content: list[dict] = [{"type": "text", "text": user_message}]
    content.extend({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img}"}} for img in images)
    headers = {"Content-Type": "application/json"}
    if cfg.get("api_key"):
        headers["Authorization"] = f"Bearer {cfg['api_key']}"
    payload = {"model": cfg["model"], "messages": [
        {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": content}
    ], "temperature": temperature, "max_tokens": max_tokens}
    data = _post_json(cfg["api_base"].rstrip("/") + "/chat/completions", payload, headers)
    return _first_content(data.get("choices", [{}])[0].get("message", {}).get("content", ""))


def _call_gemini(cfg: dict, user_message: str, images: list[str], temperature: float, max_tokens: int) -> str:
    parts: list[dict] = [{"text": SYSTEM_PROMPT + "\n\n" + user_message}]
    parts.extend({"inline_data": {"mime_type": "image/jpeg", "data": img}} for img in images)
    url = cfg["api_base"].rstrip("/") + f"/models/{cfg['model']}:generateContent"
    data = _post_json(url, {"contents": [{"role": "user", "parts": parts}], "generationConfig": {"temperature": temperature, "maxOutputTokens": max_tokens}}, {"x-goog-api-key": cfg.get("api_key", "")})
    return _first_content(data.get("candidates", [{}])[0].get("content", {}).get("parts", []))


def _call_claude(cfg: dict, user_message: str, images: list[str], temperature: float, max_tokens: int) -> str:
    content: list[dict] = [{"type": "text", "text": user_message}]
    content.extend({"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": img}} for img in images)
    headers = {"content-type": "application/json", "x-api-key": cfg.get("api_key", ""), "anthropic-version": "2023-06-01"}
    data = _post_json(cfg["api_base"].rstrip("/") + "/messages", {"model": cfg["model"], "system": SYSTEM_PROMPT, "messages": [{"role": "user", "content": content}], "temperature": temperature, "max_tokens": max_tokens}, headers)
    return _first_content(data.get("content", []))


class ImagePromptor:
    CATEGORY = "image/Prompting"
    FUNCTION = "generate_prompt"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("prompt",)
    OUTPUT_NODE = True
    DESCRIPTION = "Generate a detailed still-image generation prompt from one or more images. No video analysis."

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "image": ("IMAGE", {"tooltip": "One or more images to analyze as a single visual reference."}),
            "provider": (PROVIDERS, {"default": "ollama"}),
            "model": ("STRING", {"default": "SetneufPT/Gemma4-12B-IT-QAT_Q4_64K_16GB-GPU:latest", "tooltip": "Provider model name. Ollama is recommended for local use."}),
            "api_base": ("STRING", {"default": "http://127.0.0.1:11434", "tooltip": "Provider API base URL."}),
            "api_key": ("STRING", {"default": "", "multiline": False, "tooltip": "Optional API key; prefer config.json or environment variables."}),
            "temperature": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.05}),
            "max_tokens": ("INT", {"default": 2048, "min": 256, "max": 8192, "step": 128}),
        }, "optional": {
            "additional_instruction": ("STRING", {"default": "", "multiline": True, "tooltip": "Optional image-only focus, for example: emphasize clothing and hand placement."}),
        }}

    def generate_prompt(self, image, provider, model, api_base, api_key, temperature, max_tokens, additional_instruction=""):
        try:
            images = _tensor_to_jpeg_b64(image)
            if not images:
                return ("[Image Promptor Error] The IMAGE input is empty.",)
            config = _load_config().get(provider, {})
            config.update({"api_base": api_base.strip() or config.get("api_base", ""), "model": model.strip() or config.get("model", ""), "api_key": api_key.strip() or config.get("api_key", "")})
            focus = additional_instruction.strip() or "Describe the complete visible composition with maximum useful detail."
            user_message = f"Analyze the supplied still image reference(s) as one composition. {focus}\nThere are {len(images)} image frame(s); do not describe them as a video."
            if provider == "ollama":
                result = _call_ollama(config, user_message, images, temperature, max_tokens)
            elif provider == "openai_compatible":
                result = _call_openai(config, user_message, images, temperature, max_tokens)
            elif provider == "gemini":
                result = _call_gemini(config, user_message, images, temperature, max_tokens)
            else:
                result = _call_claude(config, user_message, images, temperature, max_tokens)
            result = _clean_response(result)
            if not result:
                return ("[Image Promptor Error] Provider returned an empty prompt.",)
            print(f"[Image Promptor] generated {len(result)} characters with {provider}/{config['model']}")
            return (result,)
        except Exception as exc:
            print(f"[Image Promptor] error: {exc}")
            return (f"[Image Promptor Error] {exc}",)


NODE_CLASS_MAPPINGS = {"ComfyUIImagePromptor": ImagePromptor}
NODE_DISPLAY_NAME_MAPPINGS = {"ComfyUIImagePromptor": "Detailed Image Promptor"}