Mister-Link-ComfyUI-Qwen-TeleStyle/nodes_appmatch.py

"""Qwen TeleStyle Apply node that reproduces /home/linkray/TeleStyleV2/app.py's `infer()`
function exactly - same prompt-enhancer generation calls, same resize math, same
prompt-merge logic, same pipe() call. It only knows how to *generate*; all model
loading/downloading lives in nodes_loader.py's Qwen TeleStyle Loader, which hands this node
a ready TELESTYLE_MODELS pipe.
"""

import numpy as np
import torch
from PIL import Image


def _process_vision_info(messages):
    # imported lazily so a missing qwen_vl_utils only breaks the prompt-enhancer feature,
    # not registration of this node (and every other node in the same package).
    try:
        from qwen_vl_utils import process_vision_info
    except ImportError as exc:
        raise ImportError(
            "qwen_vl_utils is required for infer_content_prompt/infer_style_prompt. "
            "Install it in ComfyUI's own venv: pip install qwen-vl-utils"
        ) from exc
    return process_vision_info(messages)


MAX_SEED = np.iinfo(np.int32).max

CONTENT_TEXT_INPUT = (
    "describe main objects (fewer than 3) with separated words, each word is separated by comma,  "
    "the total number of words is strictly fewer than 3"
)
STYLE_TEXT_INPUT = "describe only the artistic style, material, stroke in 5 words, not objects."


def _comfy_image_to_numpy_uint8(image):
    # comfy IMAGE tensors are [B, H, W, C] float in [0, 1]; app.py's gr.Image(type="numpy")
    # hands infer() an (H, W, C) uint8 array - reproduce that exact array before doing
    # anything else, so every downstream line of app.py's infer() can be copied verbatim.
    array = image[0].clamp(0.0, 1.0).mul(255.0).round().to(torch.uint8).cpu().numpy()
    return array


def _tensor_from_pil(image: Image.Image) -> torch.Tensor:
    array = np.array(image.convert("RGB")).astype(np.float32) / 255.0
    return torch.from_numpy(array).unsqueeze(0)


def _resolve_runtime(models):
    if isinstance(models, dict):
        pipe = models.get("pipe")
        if pipe is None:
            raise ValueError("TELESTYLE_MODELS handle is missing 'pipe'.")
        return models, pipe
    return {"mode": "legacy", "pipe": models}, models


class QwenTeleStyleApply:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "models": ("TELESTYLE_MODELS",),
                "prompt": (
                    "STRING",
                    {
                        "default": "Style Transfer the style of Figure 2 to Figure 1, and keep the content and characteristics of Figure 1.",
                        "multiline": True,
                    },
                ),
                "seed": ("INT", {"default": 123, "min": 0, "max": MAX_SEED}),
                "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}),
                "steps": ("INT", {"default": 4, "min": 1, "max": 50}),
                "resolution": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 8}),
                "infer_content_prompt": ("BOOLEAN", {"default": False}),
                "infer_style_prompt": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "content_ref": ("IMAGE",),
                "style_ref": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING", "STRING", "STRING")
    RETURN_NAMES = ("image", "content_prompt", "style_prompt", "prompt")
    FUNCTION = "infer"
    CATEGORY = "image/TeleStyleV2"

    def infer(
        self,
        models,
        prompt,
        seed,
        cfg,
        steps,
        resolution,
        infer_content_prompt,
        infer_style_prompt,
        content_ref=None,
        style_ref=None,
    ):
        runtime, pipe = _resolve_runtime(models)
        if content_ref is None and style_ref is None:
            raise ValueError("Connect at least one of content_ref or style_ref.")
        if (infer_content_prompt or infer_style_prompt) and not pipe.has_lm_head:
            raise RuntimeError(
                "The text encoder this pipe was built from has no lm_head weight, so it "
                "can't run infer_content_prompt/infer_style_prompt. Write the prompt "
                "manually instead, or disable these toggles."
            )
        device = runtime.get("resolved_device", getattr(pipe, "resolved_device", "cpu"))

        # from here down this is app.py's infer(), unmodified line-for-line except the
        # gradio-numpy inputs are swapped for comfy IMAGE tensors converted just above the
        # point app.py itself first calls Image.fromarray(), and `minedge`/`true_guidance_scale`/
        # `num_inference_steps`/`randomize_seed` are renamed/dropped to match comfy conventions
        # (seed randomization is comfy's own "control_after_generate" widget, not this node's job).
        checkbox = []
        if infer_content_prompt:
            checkbox.append("infer with content prompt")
        if infer_style_prompt:
            checkbox.append("infer with style prompt")

        content_prompt = ""
        style_prompt = ""

        if content_ref is not None:
            content_ref = Image.fromarray(_comfy_image_to_numpy_uint8(content_ref))
            content_messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": content_ref},
                        {"type": "text", "text": CONTENT_TEXT_INPUT},
                    ],
                }
            ]
            content_text = pipe.processor.apply_chat_template(
                content_messages, tokenize=False, add_generation_prompt=True
            )
            image_inputs, video_inputs = _process_vision_info(content_messages)
            inputs = pipe.processor(
                text=[content_text],
                images=image_inputs,
                videos=video_inputs,
                padding=True,
                return_tensors="pt",
            )
            inputs = inputs.to(device)

            generated_ids = pipe.text_encoder.generate(**inputs, max_new_tokens=1024)
            generated_ids_trimmed = [
                out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
            ]
            content_prompt = pipe.processor.batch_decode(
                generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
            )[0]

        if style_ref is not None:
            style_ref = Image.fromarray(_comfy_image_to_numpy_uint8(style_ref))
            style_messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": style_ref},
                        {"type": "text", "text": STYLE_TEXT_INPUT},
                    ],
                }
            ]
            style_text = pipe.processor.apply_chat_template(
                style_messages, tokenize=False, add_generation_prompt=True
            )
            image_inputs, video_inputs = _process_vision_info(style_messages)
            inputs = pipe.processor(
                text=[style_text],
                images=image_inputs,
                videos=video_inputs,
                padding=True,
                return_tensors="pt",
            )
            inputs = inputs.to(device)

            generated_ids = pipe.text_encoder.generate(**inputs, max_new_tokens=1024)
            generated_ids_trimmed = [
                out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
            ]
            style_prompt = pipe.processor.batch_decode(
                generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
            )[0]

        sw, sh, w, h = 0, 0, 0, 0
        if content_ref:
            w, h = content_ref.size
            if w > h:
                r = w / h
                h = resolution
                w = int(h * r) - int(h * r) % 16
            else:
                r = h / w
                w = resolution
                h = int(w * r) - int(w * r) % 16
        if style_ref:
            sw, sh = style_ref.size
            if sw > sh:
                r = sw / sh
                sh = resolution
                sw = int(sh * r) - int(sh * r) % 16
            else:
                r = sh / sw
                sw = resolution
                sh = int(sw * r) - int(sw * r) % 16

        if content_ref and style_ref:
            images = [content_ref.resize((w, h)), style_ref.resize((sw, sh))]
        elif content_ref:
            images = [content_ref.resize((w, h))]
        elif style_ref:
            images = [style_ref.resize((sw, sh))]

        if "infer with content prompt" in checkbox and content_prompt not in prompt:
            prompt = ",".join([prompt, content_prompt])
        if "infer with style prompt" in checkbox and style_prompt not in prompt:
            prompt = ",".join([prompt, style_prompt])
        if "infer with content prompt" not in checkbox and content_prompt in prompt:
            prompt = prompt.replace(content_prompt.strip(","), "")
        if "infer with style prompt" not in checkbox and style_prompt in prompt:
            prompt = prompt.replace(style_prompt.strip(), "")
        prompt = prompt.strip(",")

        pipe_inputs = {
            "image": images,
            "prompt": prompt,
            "generator": torch.manual_seed(seed),
            "true_cfg_scale": cfg,
            "negative_prompt": " ",
            "num_inference_steps": steps,
            "guidance_scale": cfg,
            "num_images_per_prompt": 1,
            "width": w or sw,
            "height": h or sh,
        }
        with torch.inference_mode():
            result = pipe(**pipe_inputs)
        output_image = result.images[0]

        return (_tensor_from_pil(output_image), content_prompt, style_prompt, prompt)


NODE_CLASS_MAPPINGS = {
    "Qwen TeleStyle Apply": QwenTeleStyleApply,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Qwen TeleStyle Apply": "Qwen TeleStyle Apply",
}