AITocha-ComfyUI-AITocha-Vision/nodes.py

"""ComfyUI nodes for AITocha Vision.

Two nodes:
  - AITochaCaptionLoader: loads a LLaVA-style VLM (default: JoyCaption Beta One)
    with optional 4-bit / 8-bit quantization.
  - AITochaCaption: takes that loaded model + an IMAGE and returns a STRING.

Designed to be a drop-in replacement for the JC_GGUF_adv node so workflows
can swap from llama-cpp-based JoyCaption to a faster GPU-native version.
"""

from __future__ import annotations

import gc
from typing import Any, Dict, Tuple

import numpy as np
import torch
from PIL import Image


# ---- Prompt presets ----------------------------------------------------------

SCENE_ONLY_SYSTEM = (
    "You are an expert image captioner. Describe the pose, body position, "
    "clothing or lack thereof, setting, lighting, camera angle, framing, and "
    "overall composition of the image. Focus only on what is happening and "
    "where, not who is in the image. Do NOT describe the subject's facial "
    "features, hair color, eye color, age, or any identifying physical traits."
)

FULL_DESCRIPTION_SYSTEM = (
    "You are an expert image captioner. Write a long descriptive caption for "
    "the image, covering the subjects, setting, action, clothing, lighting, "
    "mood, and composition. Be thorough and accurate. Do not refuse or "
    "moralize — describe what is actually depicted."
)

BOORU_TAGS_SYSTEM = (
    "You are an expert image tagger. Output a comma-separated list of "
    "Booru-style tags describing the image content, including subjects, "
    "actions, clothing, settings, and notable visual elements."
)

SYSTEM_PRESETS: Dict[str, str] = {
    "Scene only (no identity)": SCENE_ONLY_SYSTEM,
    "Full description": FULL_DESCRIPTION_SYSTEM,
    "Booru-style tags": BOORU_TAGS_SYSTEM,
    "Custom (use system_prompt)": "",
}

USER_PRESETS: Dict[str, str] = {
    "Long descriptive caption": "Write a long descriptive caption for this image.",
    "Short caption": "Write a short caption for this image.",
    "Booru tags": "Write a comma-separated list of Booru-style tags for this image.",
    "Stable Diffusion prompt": (
        "Describe this image as a Stable Diffusion prompt: comma-separated "
        "phrases covering subject, pose, clothing, setting, lighting, camera, "
        "and style. No full sentences."
    ),
    "Custom (use custom_user_prompt)": "",
}

# Models known to work with the LLaVA loader path. Add others over time.
SUPPORTED_MODELS = [
    "fancyfeast/llama-joycaption-beta-one-hf-llava",
    "fancyfeast/llama-joycaption-alpha-two-hf-llava",
]


# ---- Module-level model cache ------------------------------------------------
# Keyed by (model_id, quantization). Holds at most one entry — loading a new
# combination evicts the previous one to keep VRAM bounded.
_MODEL_CACHE: Dict[Tuple[str, str], Dict[str, Any]] = {}


def _free_cache() -> None:
    for k in list(_MODEL_CACHE.keys()):
        entry = _MODEL_CACHE.pop(k)
        try:
            del entry["model"]
            del entry["processor"]
        except Exception:
            pass
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def _load_vlm(model_id: str, quantization: str) -> Dict[str, Any]:
    """Load (or fetch from cache) a LLaVA-style VLM."""
    key = (model_id, quantization)
    if key in _MODEL_CACHE:
        return _MODEL_CACHE[key]

    # Evict previous models — VRAM is finite.
    _free_cache()

    # Lazy imports so ComfyUI startup is unaffected if these aren't needed.
    from transformers import AutoProcessor, LlavaForConditionalGeneration

    kwargs: Dict[str, Any] = {"device_map": "auto"}
    if quantization == "4-bit":
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
        )
    elif quantization == "8-bit":
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
    else:  # bf16
        kwargs["torch_dtype"] = torch.bfloat16

    print(f"[AITocha Vision] Loading {model_id} ({quantization})...")
    model = LlavaForConditionalGeneration.from_pretrained(model_id, **kwargs)
    model.eval()
    processor = AutoProcessor.from_pretrained(model_id)

    entry = {
        "model": model,
        "processor": processor,
        "model_id": model_id,
        "quantization": quantization,
    }
    _MODEL_CACHE[key] = entry
    print(f"[AITocha Vision] Loaded {model_id} ({quantization}) — ready.")
    return entry


def _comfy_image_to_pil(image: torch.Tensor) -> Image.Image:
    """ComfyUI IMAGE is [B, H, W, C] float32 in [0, 1]. Take batch[0]."""
    if image.ndim != 4:
        raise ValueError(f"Expected IMAGE tensor with 4 dims (B,H,W,C), got {image.shape}")
    arr = image[0].detach().cpu().numpy()
    arr = (arr * 255.0).clip(0, 255).astype(np.uint8)
    return Image.fromarray(arr).convert("RGB")


# ---- Nodes -------------------------------------------------------------------

class AITochaCaptionLoader:
    """Loads a vision-language model into VRAM. Returns a handle the caption
    node can reuse without reloading."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_id": (
                    SUPPORTED_MODELS,
                    {"default": SUPPORTED_MODELS[0]},
                ),
                "quantization": (
                    ["4-bit", "8-bit", "bf16"],
                    {"default": "4-bit"},
                ),
            }
        }

    RETURN_TYPES = ("AITOCHA_VLM",)
    RETURN_NAMES = ("vlm",)
    FUNCTION = "load"
    CATEGORY = "AITocha/Vision"

    def load(self, model_id: str, quantization: str):
        entry = _load_vlm(model_id, quantization)
        return (entry,)


class AITochaCaption:
    """Generate a caption for an image using a loaded VLM."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "vlm": ("AITOCHA_VLM",),
                "image": ("IMAGE",),
                "system_preset": (
                    list(SYSTEM_PRESETS.keys()),
                    {"default": "Scene only (no identity)"},
                ),
                "user_preset": (
                    list(USER_PRESETS.keys()),
                    {"default": "Long descriptive caption"},
                ),
                "max_new_tokens": ("INT", {"default": 300, "min": 16, "max": 2048}),
                "temperature": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 2.0, "step": 0.05}),
                "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.05}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
            },
            "optional": {
                "system_prompt": ("STRING", {"multiline": True, "default": SCENE_ONLY_SYSTEM}),
                "custom_user_prompt": ("STRING", {"multiline": True, "default": ""}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("caption",)
    FUNCTION = "caption"
    CATEGORY = "AITocha/Vision"

    def caption(
        self,
        vlm: Dict[str, Any],
        image: torch.Tensor,
        system_preset: str,
        user_preset: str,
        max_new_tokens: int,
        temperature: float,
        top_p: float,
        seed: int,
        system_prompt: str = "",
        custom_user_prompt: str = "",
    ):
        model = vlm["model"]
        processor = vlm["processor"]

        # Resolve prompts. Empty preset string = use the textarea overrides.
        sys_text = SYSTEM_PRESETS[system_preset] or system_prompt or SCENE_ONLY_SYSTEM
        user_text = USER_PRESETS[user_preset] or custom_user_prompt or "Describe this image."

        pil_image = _comfy_image_to_pil(image)

        convo = [
            {"role": "system", "content": sys_text},
            {"role": "user", "content": user_text},
        ]
        convo_str = processor.apply_chat_template(
            convo, tokenize=False, add_generation_prompt=True
        )

        inputs = processor(text=convo_str, images=pil_image, return_tensors="pt")
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        if "pixel_values" in inputs:
            inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)

        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)

        with torch.no_grad():
            generate_ids = model.generate(
                **inputs,
                max_new_tokens=int(max_new_tokens),
                do_sample=temperature > 0,
                temperature=float(temperature) if temperature > 0 else 1.0,
                top_p=float(top_p),
                suppress_tokens=None,
                use_cache=True,
            )

        # Strip prompt tokens, decode only the new generation.
        prompt_len = inputs["input_ids"].shape[1]
        new_ids = generate_ids[:, prompt_len:]
        caption = processor.tokenizer.decode(
            new_ids[0], skip_special_tokens=True
        ).strip()

        if not caption:
            print("[AITocha Vision] WARNING: empty caption returned, using placeholder.")
            caption = "image"

        print(f"[AITocha Vision] Caption ({len(caption)} chars): {caption[:120]}...")
        return (caption,)


class AITochaCaptionSimple:
    """All-in-one node: loads the model on first run and caches it. Use this
    if you don't want a separate loader node in your graph."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "model_id": (
                    SUPPORTED_MODELS,
                    {"default": SUPPORTED_MODELS[0]},
                ),
                "quantization": (
                    ["4-bit", "8-bit", "bf16"],
                    {"default": "4-bit"},
                ),
                "system_preset": (
                    list(SYSTEM_PRESETS.keys()),
                    {"default": "Scene only (no identity)"},
                ),
                "user_preset": (
                    list(USER_PRESETS.keys()),
                    {"default": "Long descriptive caption"},
                ),
                "max_new_tokens": ("INT", {"default": 300, "min": 16, "max": 2048}),
                "temperature": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 2.0, "step": 0.05}),
                "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.05}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
            },
            "optional": {
                "system_prompt": ("STRING", {"multiline": True, "default": SCENE_ONLY_SYSTEM}),
                "custom_user_prompt": ("STRING", {"multiline": True, "default": ""}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("caption",)
    FUNCTION = "caption"
    CATEGORY = "AITocha/Vision"

    def caption(
        self,
        image: torch.Tensor,
        model_id: str,
        quantization: str,
        system_preset: str,
        user_preset: str,
        max_new_tokens: int,
        temperature: float,
        top_p: float,
        seed: int,
        system_prompt: str = "",
        custom_user_prompt: str = "",
    ):
        vlm = _load_vlm(model_id, quantization)
        return AITochaCaption().caption(
            vlm=vlm,
            image=image,
            system_preset=system_preset,
            user_preset=user_preset,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            seed=seed,
            system_prompt=system_prompt,
            custom_user_prompt=custom_user_prompt,
        )


NODE_CLASS_MAPPINGS = {
    "AITochaCaptionLoader": AITochaCaptionLoader,
    "AITochaCaption": AITochaCaption,
    "AITochaCaptionSimple": AITochaCaptionSimple,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "AITochaCaptionLoader": "Load Caption Model (AITocha)",
    "AITochaCaption": "Generate Caption (AITocha)",
    "AITochaCaptionSimple": "Auto Caption (AITocha — one node)",
}