# Copyright 2026 SARMAE Authors and The HuggingFace Inc. team.
"""SARMAE image feature extraction pipeline (self-contained for trust_remote_code)."""

from typing import Any, Union

from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
from transformers.utils import add_end_docstrings, is_vision_available


if is_vision_available():
    from transformers.image_utils import load_image


@add_end_docstrings(
    build_pipeline_init_args(has_image_processor=True),
    """
        pool (`bool`, *optional*, defaults to `False`):
            Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
    """,
)
class SarmaeImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
    def _sanitize_parameters(
        self,
        image_processor_kwargs=None,
        return_tensors=None,
        pool=None,
        **kwargs,
    ):
        preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
        if "timeout" in kwargs:
            preprocess_params["timeout"] = kwargs["timeout"]
        postprocess_params = {}
        if pool is not None:
            postprocess_params["pool"] = pool
        if return_tensors is not None:
            postprocess_params["return_tensors"] = return_tensors
        return preprocess_params, {}, postprocess_params

    def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
        if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
            image = load_image(image, timeout=timeout)
        model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
        model_inputs = model_inputs.to(self.dtype)
        return model_inputs

    def __call__(self, *args: Union[str, Any, list[Any]], **kwargs: Any) -> list[Any]:
        return super().__call__(*args, **kwargs)


__all__ = ["SarmaeImageFeatureExtractionPipeline"]