ipadapter.py

"""
This module provides nodes load and apply IP-Adapter models
to FLUX pipelines, enabling image-based conditioning for generative models.
"""

import logging
import os
from typing import Any, List, Optional

import torch
from diffusers import FluxPipeline
from torchvision import transforms

from nunchaku.models.ip_adapter.diffusers_adapters import apply_IPA_on_pipe
from nunchaku.models.ip_adapter.utils import undo_all_mods_on_transformer

from .utils import set_extra_config_model_path

log_level = os.getenv("LOG_LEVEL", "INFO").upper()

logging.basicConfig(level=getattr(logging, log_level, logging.INFO), format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)


class IPAFluxPipelineWrapper(FluxPipeline):
    """
    FluxPipeline wrapper with IP-Adapter support.
    """

    @torch.no_grad()
    def get_image_embeds(
        self,
        num_images_per_prompt: int = 1,
        ip_adapter_image: Optional[Any] = None,
        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
        negative_ip_adapter_image: Optional[Any] = None,
        negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
    ) -> (Optional[torch.Tensor], Optional[torch.Tensor]):
        """
        Compute image and negative image embeddings for IP-Adapter.

        Parameters
        ----------
        num_images_per_prompt : int, optional
            Number of images per prompt (default is 1).
        ip_adapter_image : Any, optional
            Input image for positive conditioning.
        ip_adapter_image_embeds : list of torch.Tensor, optional
            Precomputed positive image embeddings.
        negative_ip_adapter_image : Any, optional
            Input image for negative conditioning.
        negative_ip_adapter_image_embeds : list of torch.Tensor, optional
            Precomputed negative image embeddings.

        Returns
        -------
        image_embeds : torch.Tensor or None
            Positive image embeddings.
        negative_image_embeds : torch.Tensor or None
            Negative image embeddings.
        """
        batch_size = 1

        device = self.transformer.device

        image_embeds = None
        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
            image_embeds = self.prepare_ip_adapter_image_embeds(
                ip_adapter_image=ip_adapter_image,
                ip_adapter_image_embeds=ip_adapter_image_embeds,
                device=device,
                num_images_per_prompt=batch_size * num_images_per_prompt,
            )
            image_embeds = self.transformer.encoder_hid_proj(image_embeds)

        negative_image_embeds = None
        if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None:
            negative_image_embeds = self.prepare_ip_adapter_image_embeds(
                ip_adapter_image=negative_ip_adapter_image,
                ip_adapter_image_embeds=negative_ip_adapter_image_embeds,
                device=device,
                num_images_per_prompt=batch_size * num_images_per_prompt,
            )
            negative_image_embeds = self.transformer.encoder_hid_proj(negative_image_embeds)

        return image_embeds, negative_image_embeds


set_extra_config_model_path("ipadapter", "ipadapter")
set_extra_config_model_path("clip", "clip")


class NunchakuIPAdapterLoader:
    """
    Node for loading Nunchaku IP-Adapter pipelines.

    .. warning::
        This node will automatically download the IP-Adapter and associated CLIP models from Hugging Face.
        Custom model paths are not supported for now.
    """

    @classmethod
    def INPUT_TYPES(s):
        """
        Defines the input types and tooltips for the node.

        Returns
        -------
        dict
            A dictionary specifying the required inputs and their descriptions for the node interface.
        """
        return {
            "required": {
                "model": ("MODEL", {"tooltip": "The nunchaku model."}),
            }
        }

    RETURN_TYPES = ("MODEL", "IPADAPTER_PIPELINE")
    FUNCTION = "load"
    CATEGORY = "Nunchaku"
    TITLE = "Nunchaku IP-Adapter Loader"

    def load(self, model):
        """
        Load the IP-Adapter pipeline and attach it to the given model.

        Parameters
        ----------
        model : object
            The Nunchaku model to which the IP-Adapter will be attached.
            It should be loaded with :class:`~comfyui_nunchaku.nodes.models.flux.NunchakuFluxDiTLoader`.

        Returns
        -------
        tuple
            The original model and the loaded IP-Adapter pipeline.
        """
        device = model.model.diffusion_model.model.device
        pipeline = IPAFluxPipelineWrapper.from_pretrained(
            "black-forest-labs/FLUX.1-dev", transformer=model.model.diffusion_model.model, torch_dtype=torch.bfloat16
        ).to(device)

        pipeline.load_ip_adapter(
            pretrained_model_name_or_path_or_dict="XLabs-AI/flux-ip-adapter-v2",
            weight_name="ip_adapter.safetensors",
            image_encoder_pretrained_model_name_or_path="openai/clip-vit-large-patch14",
        )
        return (model, pipeline)


class NunchakuFluxIPAdapterApply:
    """
    Node for applying IP-Adapter to a Nunchaku model using a given image and weight.
    """

    @classmethod
    def INPUT_TYPES(s):
        """
        Defines the input types and tooltips for the node.

        Returns
        -------
        dict
            A dictionary specifying the required inputs and their descriptions for the node interface.
        """
        return {
            "required": {
                "model": ("MODEL",),
                "ipadapter_pipeline": ("IPADAPTER_PIPELINE",),
                "image": ("IMAGE",),
                "weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.05}),
            },
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "apply_ipa"
    CATEGORY = "Nunchaku"
    TITLE = "Nunchaku FLUX IP-Adapter Apply"

    def apply_ipa(
        self,
        model,
        ipadapter_pipeline: IPAFluxPipelineWrapper,
        image,
        weight: float,
    ):
        """
        Apply the IP-Adapter to the given model using the provided image and weight.

        Parameters
        ----------
        model : object
            The Nunchaku model to modify.
        ipadapter_pipeline : IPAFluxPipelineWrapper
            The IP-Adapter pipeline.
        image : torch.Tensor
            The input image tensor.
        weight : float
            The scale/weight for the IP-Adapter.

        Returns
        -------
        tuple
            The modified model.
        """
        to_pil_transformer = transforms.ToPILImage()
        image_tensor_chw = image[0].permute(2, 0, 1)
        pil_image = to_pil_transformer(image_tensor_chw)

        image_embeds, _ = ipadapter_pipeline.get_image_embeds(
            ip_adapter_image=pil_image,
        )

        undo_all_mods_on_transformer(ipadapter_pipeline.transformer)
        apply_IPA_on_pipe(ipadapter_pipeline, ip_adapter_scale=weight, repo_id="XLabs-AI/flux-ip-adapter-v2")

        ipadapter_pipeline.transformer.transformer_blocks[0].set_ip_hidden_states(image_embeds=image_embeds)

        model.model.diffusion_model.model = ipadapter_pipeline.transformer

        return (model,)