ComfyUI/comfy_extras/nodes_qwen.py

import node_helpers
import comfy.utils
import math
import torch
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io


class TextEncodeQwenImageEdit(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="TextEncodeQwenImageEdit",
            category="advanced/conditioning",
            inputs=[
                io.Clip.Input("clip"),
                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
                io.Vae.Input("vae", optional=True),
                io.Image.Input("image", optional=True),
            ],
            outputs=[
                io.Conditioning.Output(),
            ],
        )

    @classmethod
    def execute(cls, clip, prompt, vae=None, image=None) -> io.NodeOutput:
        ref_latent = None
        if image is None:
            images = []
        else:
            samples = image.movedim(-1, 1)
            total = int(1024 * 1024)

            scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
            width = round(samples.shape[3] * scale_by)
            height = round(samples.shape[2] * scale_by)

            s = comfy.utils.common_upscale(samples, width, height, "area", "disabled")
            image = s.movedim(1, -1)
            images = [image[:, :, :, :3]]
            if vae is not None:
                ref_latent = vae.encode(image[:, :, :, :3])

        tokens = clip.tokenize(prompt, images=images)
        conditioning = clip.encode_from_tokens_scheduled(tokens)
        if ref_latent is not None:
            conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": [ref_latent]}, append=True)
        return io.NodeOutput(conditioning)

class QwenImageInpaintConditioning(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="QwenImageInpaintConditioning",
            category="advanced/conditioning",
            description=(
                "Prepares conditioning and latents for Qwen Image Edit inpainting."
            ),
            inputs=[
                io.Conditioning.Input("positive"),
                io.Conditioning.Input("negative"),
                io.Vae.Input("vae"),
                io.Image.Input("image"),
                io.Mask.Input("mask"),
                io.Boolean.Input(
                    "use_noise_mask",
                    default=True,
                    tooltip="When enabled, provide the resized mask as noise mask so sampling only affects the painted region.",
                ),
            ],
            outputs=[
                io.Conditioning.Output(display_name="positive"),
                io.Conditioning.Output(display_name="negative"),
                io.Latent.Output(display_name="latent"),
            ],
        )

    @classmethod
    def execute(
        cls,
        positive,
        negative,
        vae,
        image,
        mask,
        use_noise_mask=True,
    ) -> io.NodeOutput:
        if image.ndim != 4:
            raise ValueError("Expected image tensor with shape [B, H, W, C].")

        image = image[:, :, :, :3]
        batch, height, width, _ = image.shape

        requested_width = float(width)
        requested_height = float(height)

        spacial_scale = vae.spacial_compression_encode()
        if isinstance(spacial_scale, tuple):
            spacial_scale = spacial_scale[-1]
        spacial_scale = int(spacial_scale)
        align_multiple = max(1, spacial_scale * 2)

        target_width = max(
            align_multiple,
            round(requested_width / align_multiple) * align_multiple,
        )
        target_height = max(
            align_multiple,
            round(requested_height / align_multiple) * align_multiple,
        )

        samples = image.movedim(-1, 1)
        resized = comfy.utils.common_upscale(samples, target_width, target_height, "area", "disabled")
        resized = resized.movedim(1, -1)

        mask_tensor = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
        mask_tensor = torch.nn.functional.interpolate(mask_tensor, size=(target_height, target_width), mode="bilinear")
        mask_tensor = mask_tensor.clamp(0.0, 1.0)
        mask_tensor = mask_tensor.to(resized.dtype)
        mask_tensor = comfy.utils.resize_to_batch_size(mask_tensor, batch)

        masked_pixels = resized.clone()
        keep_region = (1.0 - mask_tensor.round()).squeeze(1)
        masked_pixels[:, :, :, :3] = (masked_pixels[:, :, :, :3] - 0.5) * keep_region.unsqueeze(-1) + 0.5

        concat_latent = vae.encode(masked_pixels)
        orig_latent = vae.encode(resized)

        out_latent: dict[str, torch.Tensor] = {"samples": orig_latent}
        if use_noise_mask:
            out_latent["noise_mask"] = mask_tensor

        positive = node_helpers.conditioning_set_values(
            positive, {"concat_latent_image": concat_latent, "concat_mask": mask_tensor}
        )
        negative = node_helpers.conditioning_set_values(
            negative, {"concat_latent_image": concat_latent, "concat_mask": mask_tensor}
        )

        return io.NodeOutput(positive, negative, out_latent)

class TextEncodeQwenImageEditPlus(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="TextEncodeQwenImageEditPlus",
            category="advanced/conditioning",
            inputs=[
                io.Clip.Input("clip"),
                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
                io.Vae.Input("vae", optional=True),
                io.Image.Input("image1", optional=True),
                io.Image.Input("image2", optional=True),
                io.Image.Input("image3", optional=True),
            ],
            outputs=[
                io.Conditioning.Output(),
            ],
        )

    @classmethod
    def execute(cls, clip, prompt, vae=None, image1=None, image2=None, image3=None) -> io.NodeOutput:
        ref_latents = []
        images = [image1, image2, image3]
        images_vl = []
        llama_template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
        image_prompt = ""

        for i, image in enumerate(images):
            if image is not None:
                samples = image.movedim(-1, 1)
                total = int(384 * 384)

                scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
                width = round(samples.shape[3] * scale_by)
                height = round(samples.shape[2] * scale_by)

                s = comfy.utils.common_upscale(samples, width, height, "area", "disabled")
                images_vl.append(s.movedim(1, -1))
                if vae is not None:
                    total = int(1024 * 1024)
                    scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
                    width = round(samples.shape[3] * scale_by / 8.0) * 8
                    height = round(samples.shape[2] * scale_by / 8.0) * 8

                    s = comfy.utils.common_upscale(samples, width, height, "area", "disabled")
                    ref_latents.append(vae.encode(s.movedim(1, -1)[:, :, :, :3]))

                image_prompt += "Picture {}: <|vision_start|><|image_pad|><|vision_end|>".format(i + 1)

        tokens = clip.tokenize(image_prompt + prompt, images=images_vl, llama_template=llama_template)
        conditioning = clip.encode_from_tokens_scheduled(tokens)
        if len(ref_latents) > 0:
            conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": ref_latents}, append=True)
        return io.NodeOutput(conditioning)


class QwenExtension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [
            TextEncodeQwenImageEdit,
            QwenImageInpaintConditioning,
            TextEncodeQwenImageEditPlus,
        ]


async def comfy_entrypoint() -> QwenExtension:
    return QwenExtension()