ComfyUI/comfy_extras/nodes_lingbot_video.py

import math
import nodes
import numpy as np
import torch
from PIL import Image

import comfy.latent_formats
import comfy.model_management
import comfy.utils
from comfy_api.latest import ComfyExtension, io
from typing_extensions import override


IMAGE_MIN_TOKEN_NUM = 4
IMAGE_MAX_TOKEN_NUM = 16384
MAX_RATIO = 200
SPATIAL_MERGE_SIZE = 2
VISION_PATCH_SIZE = 16


def _crop_image(image, width, height):
    image = image[:1].movedim(-1, 1)
    image = comfy.utils.common_upscale(image, width, height, "bilinear", "center")
    return image.movedim(1, -1)[:, :, :, :3]


def _round_by_factor(number, factor):
    return round(number / factor) * factor


def _ceil_by_factor(number, factor):
    return math.ceil(number / factor) * factor


def _floor_by_factor(number, factor):
    return math.floor(number / factor) * factor


def _smart_resize(height, width, factor, min_pixels=None, max_pixels=None):
    max_pixels = max_pixels if max_pixels is not None else IMAGE_MAX_TOKEN_NUM * factor ** 2
    min_pixels = min_pixels if min_pixels is not None else IMAGE_MIN_TOKEN_NUM * factor ** 2
    if max_pixels < min_pixels:
        raise ValueError("max_pixels must be greater than or equal to min_pixels.")
    if max(height, width) / min(height, width) > MAX_RATIO:
        raise ValueError(f"LingBotVideo image aspect ratio must be smaller than {MAX_RATIO}.")

    resized_height = max(factor, _round_by_factor(height, factor))
    resized_width = max(factor, _round_by_factor(width, factor))
    if resized_height * resized_width > max_pixels:
        beta = math.sqrt((height * width) / max_pixels)
        resized_height = _floor_by_factor(height / beta, factor)
        resized_width = _floor_by_factor(width / beta, factor)
    elif resized_height * resized_width < min_pixels:
        beta = math.sqrt(min_pixels / (height * width))
        resized_height = _ceil_by_factor(height * beta, factor)
        resized_width = _ceil_by_factor(width * beta, factor)
    return resized_height, resized_width


def _vlm_image(image):
    factor = VISION_PATCH_SIZE * SPATIAL_MERGE_SIZE
    height, width = image.shape[1:3]
    resized_height, resized_width = _smart_resize(height, width, factor)
    array = image[0].detach().cpu().clamp(0, 1).mul(255).byte().numpy()
    pil_image = Image.fromarray(array, mode="RGB").resize((resized_width, resized_height))
    array = np.asarray(pil_image).astype(np.float32) / 255.0
    return torch.from_numpy(array).unsqueeze(0)


class TextEncodeLingBotVideoI2V(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="TextEncodeLingBotVideoI2V",
            category="model/conditioning/lingbot_video",
            inputs=[
                io.Clip.Input("clip"),
                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
                io.String.Input("negative_prompt", multiline=True, dynamic_prompts=True, default=""),
                io.Int.Input("width", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
                io.Int.Input("height", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
                io.Image.Input("image", optional=True),
            ],
            outputs=[
                io.Conditioning.Output(display_name="positive"),
                io.Conditioning.Output(display_name="negative"),
            ],
        )

    @classmethod
    def execute(cls, clip, prompt, negative_prompt, width, height, image=None) -> io.NodeOutput:
        if height % 16 != 0 or width % 16 != 0:
            raise ValueError(f"LingBotVideo width and height must be multiples of 16, got {width}x{height}.")
        if image is None:
            images = []
        else:
            image = _crop_image(image, width, height)
            images = [_vlm_image(image)]
        positive = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=images))
        negative = clip.encode_from_tokens_scheduled(clip.tokenize(negative_prompt, images=images))
        return io.NodeOutput(positive, negative)


class LingBotImageToVideo(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="LingBotImageToVideo",
            category="model/conditioning/lingbot_video",
            inputs=[
                io.Conditioning.Input("positive"),
                io.Conditioning.Input("negative"),
                io.Vae.Input("vae"),
                io.Int.Input("width", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
                io.Int.Input("height", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
                io.Int.Input("length", default=81, min=1, max=nodes.MAX_RESOLUTION, step=4),
                io.Int.Input("batch_size", default=1, min=1, max=4096),
                io.Image.Input("start_image", optional=True),
            ],
            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, width, height, length, batch_size, start_image=None) -> io.NodeOutput:
        if length != 1 and (length - 1) % 4 != 0:
            raise ValueError(f"LingBotVideo length must be 1 or 4n+1, got {length}.")
        if height % 16 != 0 or width % 16 != 0:
            raise ValueError(f"LingBotVideo width and height must be multiples of 16, got {width}x{height}.")
        latent_frames = ((length - 1) // 4) + 1
        latent = torch.zeros(
            [batch_size, 16, latent_frames, height // 8, width // 8],
            device=comfy.model_management.intermediate_device(),
        )
        out_latent = {"samples": latent}

        if start_image is not None:
            start_image = _crop_image(start_image, width, height)
            cond_latent = comfy.latent_formats.LingBotVideo().process_in(vae.encode(start_image))
            cond_latent = comfy.utils.resize_to_batch_size(cond_latent, batch_size)
            cond_t = min(cond_latent.shape[2], latent.shape[2])
            latent[:, :, :cond_t] = cond_latent[:, :, :cond_t].to(device=latent.device, dtype=latent.dtype)
            noise_mask = torch.ones(
                (batch_size, 1, latent.shape[2], latent.shape[3], latent.shape[4]),
                device=latent.device,
                dtype=latent.dtype,
            )
            noise_mask[:, :, :cond_t] = 0.0
            out_latent["noise_mask"] = noise_mask

        return io.NodeOutput(positive, negative, out_latent)


class LingBotVideoExtension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [
            TextEncodeLingBotVideoI2V,
            LingBotImageToVideo,
        ]


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