experiment/knowledge_editing/llava15_compat.py

"""
LLaVA-1.5 compatibility layer for EasyEdit.

EasyEdit's multimodal editing methods (WISE, GRACE, LoRA) internally call
``processor.apply_chat_template()`` which exists for LLaVA-OneVision but
NOT for LLaVA-1.5.  This module provides a thin wrapper that adds the
missing method, allowing EasyEdit algorithms to run on LLaVA-1.5 unchanged.

LLaVA-1.5 prompt format:
    USER: <image>\\n{prompt}\\nASSISTANT:

LLaVA-OneVision (Qwen-style chat template):
    <|im_start|>user\\n<image>\\n{prompt}<|im_end|>\\n<|im_start|>assistant\\n
"""

from __future__ import annotations

from PIL import Image
from typing import Union, List


class LLaVA15ProcessorWrapper:
    """Wraps LLaVA-1.5 AutoProcessor to add apply_chat_template().

    EasyEdit methods call three things on the processor object:
      1. processor.apply_chat_template(messages, ...) → str
      2. processor(images=..., text=..., return_tensors="pt") → BatchEncoding
      3. processor.tokenizer → underlying tokenizer

    (1) doesn't exist for LLaVA-1.5.  This wrapper intercepts it and
    produces the correct LLaVA-1.5 prompt format.  (2) and (3) are
    delegated to the underlying processor.
    """

    def __init__(self, processor):
        self._processor = processor

    # ------ Delegate everything to underlying processor ------
    def __getattr__(self, name):
        return getattr(self._processor, name)

    def __call__(self, *args, **kwargs):
        return self._processor(*args, **kwargs)

    # ------ The missing method ------
    def apply_chat_template(
        self,
        messages: list[dict],
        add_generation_prompt: bool = True,
        tokenize: bool = False,
        **kwargs,
    ) -> str:
        """Convert EasyEdit's message format to LLaVA-1.5 prompt string.

        EasyEdit passes messages like:
            [{"role": "user", "content": [
                {"type": "image"},
                {"type": "text", "text": "Describe this image."}
            ]}]

        We convert to LLaVA-1.5 format:
            USER: <image>\\nDescribe this image.\\nASSISTANT:
        """
        parts = []

        for msg in messages:
            role = msg.get("role", "user").upper()
            content = msg.get("content", "")

            if isinstance(content, str):
                # Simple string content (e.g. locality prompts with no image)
                parts.append(f"{role}: {content}")
            elif isinstance(content, list):
                # Structured content with image/text entries
                text_parts = []
                n_images = 0
                for item in content:
                    if item.get("type") == "image":
                        n_images += 1
                    elif item.get("type") == "video":
                        # LLaVA-1.5 doesn't support video; treat as image
                        n_images += 1
                    elif item.get("type") == "text":
                        text_parts.append(item["text"])

                # Build prompt: <image> tokens first, then text
                image_tokens = "\n".join(["<image>"] * n_images)
                text = " ".join(text_parts)

                if n_images > 0:
                    parts.append(f"{role}: {image_tokens}\n{text}")
                else:
                    parts.append(f"{role}: {text}")
            else:
                parts.append(f"{role}: {content}")

        prompt = "\n".join(parts)
        if add_generation_prompt:
            prompt += "\nASSISTANT:"

        return prompt


class LLaVA15ImageProcessor:
    """Minimal image processor matching EasyEdit's vis_tok interface.

    EasyEdit's _prepare_requests() calls vis_tok(file_path, file_type)
    to load and pre-process images.  For LLaVA-1.5 we just load PIL.
    """

    def __call__(self, file: Union[List[str], str], file_type=None) -> Image.Image:
        if isinstance(file, list):
            return [Image.open(f).convert("RGB") for f in file]
        return Image.open(file).convert("RGB")