processing_dendro_omni.py

"""Image processor for Phillnet Mini Text-Vision.

The processor produces Qwen-compatible visual patch tensors for the retained
transplanted vision encoder and deliberately exposes no image or video synthesis
functionality.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Sequence

import torch
from PIL import Image

from transformers import AutoTokenizer

try:
    from transformers.feature_extraction_utils import BatchFeature
except Exception:  # pragma: no cover
    BatchFeature = dict  # type: ignore[misc,assignment]


IMAGE_MARKER = "\ue000"


class DendroVisionProcessor:
    """Prepare text and still-image inputs for Phillnet Mini Text-Vision.

    This class supports text generation and image understanding only.  It accepts
    one conversation at a time, creates static two-frame visual patches required
    by the retained vision tower, and expands one image placeholder per merged
    visual token.
    """

    model_input_names = [
        "input_ids",
        "attention_mask",
        "pixel_values",
        "image_grid_thw",
        "mm_token_type_ids",
    ]

    @classmethod
    def register_for_auto_class(cls, auto_class: str = "AutoProcessor") -> None:
        """Compatibility hook used by Transformers dynamic-module loading."""
        cls._auto_class = str(auto_class)

    def __init__(
        self,
        tokenizer: Any,
        *,
        image_token_id: int,
        vision_start_token_id: int,
        vision_end_token_id: int,
        patch_size: int = 16,
        temporal_patch_size: int = 2,
        spatial_merge_size: int = 2,
        image_mean: Sequence[float] = (0.5, 0.5, 0.5),
        image_std: Sequence[float] = (0.5, 0.5, 0.5),
        max_side: int = 448,
    ) -> None:
        self.tokenizer = tokenizer
        self.image_token_id = int(image_token_id)
        self.vision_start_token_id = int(vision_start_token_id)
        self.vision_end_token_id = int(vision_end_token_id)
        self.patch_size = int(patch_size)
        self.temporal_patch_size = int(temporal_patch_size)
        self.spatial_merge_size = int(spatial_merge_size)
        self.image_mean = tuple(float(x) for x in image_mean)
        self.image_std = tuple(float(x) for x in image_std)
        self.max_side = int(max_side)
        if self.patch_size < 1 or self.temporal_patch_size < 1 or self.spatial_merge_size < 1:
            raise ValueError("Visual patch and merge dimensions must be positive")

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str | Path, **kwargs: Any) -> "DendroVisionProcessor":
        root = Path(pretrained_model_name_or_path)
        config = json.loads((root / "config.json").read_text(encoding="utf-8"))
        preprocessing = json.loads((root / "preprocessor_config.json").read_text(encoding="utf-8"))
        tokenizer_kwargs = dict(kwargs)
        max_side = int(tokenizer_kwargs.pop("max_side", preprocessing.get("max_side", 448)))
        tokenizer_kwargs.pop("trust_remote_code", None)
        tokenizer_kwargs.pop("_from_auto", None)
        tokenizer = AutoTokenizer.from_pretrained(str(root), trust_remote_code=False, **tokenizer_kwargs)
        return cls(
            tokenizer,
            image_token_id=int(config["image_token_id"]),
            vision_start_token_id=int(config["vision_start_token_id"]),
            vision_end_token_id=int(config["vision_end_token_id"]),
            patch_size=int(preprocessing.get("patch_size", 16)),
            temporal_patch_size=int(preprocessing.get("temporal_patch_size", 2)),
            spatial_merge_size=int(preprocessing.get("merge_size", 2)),
            image_mean=preprocessing.get("image_mean", (0.5, 0.5, 0.5)),
            image_std=preprocessing.get("image_std", (0.5, 0.5, 0.5)),
            max_side=max_side,
        )

    def _encode_text(self, value: str) -> list[int]:
        return list(self.tokenizer.encode(str(value), add_special_tokens=False))

    @staticmethod
    def _to_pil(image: Any) -> Image.Image:
        if isinstance(image, Image.Image):
            return image.convert("RGB")
        if isinstance(image, (str, Path)):
            with Image.open(image) as opened:
                return opened.convert("RGB")
        if torch.is_tensor(image):
            value = image.detach().cpu().float()
            if value.ndim == 4 and value.shape[0] == 1:
                value = value[0]
            if value.ndim != 3:
                raise TypeError("Image tensor must be [C,H,W] or [1,C,H,W]")
            if value.shape[0] in {1, 3, 4}:
                value = value[:3].permute(1, 2, 0)
            value = value.clamp(0, 1).mul(255).byte().numpy()
            return Image.fromarray(value).convert("RGB")
        try:
            import numpy as np
            value = np.asarray(image)
            if value.ndim == 3:
                if value.dtype != np.uint8:
                    value = (value.clip(0, 1) * 255 if float(value.max()) <= 1 else value.clip(0, 255)).astype(np.uint8)
                return Image.fromarray(value).convert("RGB")
        except Exception as error:  # pragma: no cover
            raise TypeError("Unsupported image input") from error
        raise TypeError("Unsupported image input")

    def _resize(self, image: Image.Image) -> Image.Image:
        unit = self.patch_size * self.spatial_merge_size
        width, height = image.size
        if max(width, height) > self.max_side:
            scale = self.max_side / max(width, height)
            width, height = round(width * scale), round(height * scale)
        width = max(unit, round(width / unit) * unit)
        height = max(unit, round(height / unit) * unit)
        return image.resize((width, height), Image.Resampling.BICUBIC)

    def _patchify(self, image: Any) -> tuple[torch.Tensor, torch.Tensor, int]:
        prepared = self._resize(self._to_pil(image))
        try:
            import numpy as np
            values = torch.from_numpy(np.asarray(prepared).copy()).permute(2, 0, 1).float().div_(255.0)
        except Exception as error:  # pragma: no cover
            raise RuntimeError("NumPy is required for image preprocessing") from error
        mean = torch.tensor(self.image_mean).view(3, 1, 1)
        std = torch.tensor(self.image_std).view(3, 1, 1)
        values = (values - mean) / std
        channels, height, width = values.shape
        patch = self.patch_size
        temporal = self.temporal_patch_size
        grid_h, grid_w = height // patch, width // patch
        frames = values.unsqueeze(0).repeat(temporal, 1, 1, 1)
        blocks = frames.reshape(1, temporal, channels, grid_h, patch, grid_w, patch)
        blocks = blocks.permute(0, 3, 5, 2, 1, 4, 6).reshape(-1, channels * temporal * patch * patch)
        grid = torch.tensor([[1, grid_h, grid_w]], dtype=torch.long)
        placeholder_count = grid_h // self.spatial_merge_size * (grid_w // self.spatial_merge_size)
        return blocks, grid, int(placeholder_count)

    def _build_single(self, text: str, images: Sequence[Any] | None) -> dict[str, torch.Tensor]:
        image_list = list(images or [])
        if IMAGE_MARKER not in text and image_list:
            text = text + IMAGE_MARKER * len(image_list)
        chunks = text.split(IMAGE_MARKER)
        if len(chunks) != len(image_list) + 1:
            raise ValueError("Image markers must match the number of supplied images")
        ids: list[int] = []
        types: list[int] = []
        pixel_blocks: list[torch.Tensor] = []
        grids: list[torch.Tensor] = []
        for index, chunk in enumerate(chunks):
            text_ids = self._encode_text(chunk)
            ids.extend(text_ids)
            types.extend([0] * len(text_ids))
            if index == len(image_list):
                continue
            blocks, grid, count = self._patchify(image_list[index])
            ids.append(self.vision_start_token_id)
            types.append(0)
            ids.extend([self.image_token_id] * count)
            types.extend([1] * count)
            ids.append(self.vision_end_token_id)
            types.append(0)
            pixel_blocks.append(blocks)
            grids.append(grid)
        output: dict[str, torch.Tensor] = {
            "input_ids": torch.tensor([ids], dtype=torch.long),
            "attention_mask": torch.ones((1, len(ids)), dtype=torch.long),
        }
        if pixel_blocks:
            output["pixel_values"] = torch.cat(pixel_blocks, dim=0)
            output["image_grid_thw"] = torch.cat(grids, dim=0)
            output["mm_token_type_ids"] = torch.tensor([types], dtype=torch.long)
        return output

    def __call__(
        self,
        text: str | Sequence[str],
        *,
        images: Any | Sequence[Any] | None = None,
        return_tensors: str | None = "pt",
        **_: Any,
    ) -> Any:
        if isinstance(text, Sequence) and not isinstance(text, str):
            if len(text) != 1:
                raise ValueError("DendroVisionProcessor currently accepts a single conversation per call")
            text = text[0]
        if images is None:
            image_list: list[Any] = []
        elif isinstance(images, (str, Path, Image.Image)) or torch.is_tensor(images):
            image_list = [images]
        else:
            image_list = list(images)
        encoded = self._build_single(str(text), image_list)
        if return_tensors not in {None, "pt"}:
            raise ValueError("Only return_tensors='pt' is supported")
        return BatchFeature(data=encoded, tensor_type="pt") if BatchFeature is not dict else encoded

    def apply_chat_template(
        self,
        messages: Sequence[dict[str, Any]],
        *,
        tokenize: bool = True,
        add_generation_prompt: bool = True,
        enable_thinking: bool = False,
        return_dict: bool = True,
        return_tensors: str | None = "pt",
        **_: Any,
    ) -> Any:
        parts: list[str] = []
        images: list[Any] = []
        for message in messages:
            role = str(message.get("role", "user"))
            parts.append(f"<|im_start|>{role}\n")
            content = message.get("content", "")
            if isinstance(content, str):
                parts.append(content)
            else:
                for item in content:
                    item_type = item.get("type") if isinstance(item, dict) else None
                    if item_type == "text":
                        parts.append(str(item.get("text", "")))
                    elif item_type == "image":
                        image = item.get("image", item.get("image_url"))
                        if image is None:
                            raise ValueError("Image content must provide an image object")
                        images.append(image)
                        parts.append(IMAGE_MARKER)
                    else:
                        raise ValueError(f"Unsupported chat content type: {item_type!r}")
            parts.append("<|im_end|>\n")
        if add_generation_prompt:
            parts.append("<|im_start|>assistant\n")
            parts.append("<think>\n" if enable_thinking else "<think>\n\n</think>\n\n")
        prompt = "".join(parts)
        if not tokenize:
            return prompt
        encoded = self(prompt, images=images, return_tensors=return_tensors)
        return encoded if return_dict else encoded["input_ids"]

    def save_pretrained(self, save_directory: str | Path, **_: Any) -> tuple[str]:
        root = Path(save_directory)
        root.mkdir(parents=True, exist_ok=True)
        path = root / "preprocessor_config.json"
        path.write_text(json.dumps({"processor_class": "DendroVisionProcessor"}, indent=2) + "\n", encoding="utf-8")
        return (str(path),)