optimum/intel/openvino/modeling_sam.py

import logging
import os
import warnings
from pathlib import Path
from typing import Dict, Optional, Union

import numpy as np
import openvino as ov
import torch
from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
from openvino._offline_transformations import apply_moc_transformations, compress_model_transformation
from transformers import AutoConfig, PretrainedConfig, SamModel
from transformers.modeling_outputs import ModelOutput
from transformers.models.sam.modeling_sam import SamImageSegmentationOutput, SamPositionalEmbedding

from ...exporters.openvino.utils import save_config
from .modeling_base import OVBaseModel, OVModelPart
from .utils import (
    ONNX_PROMPT_ENCODER_MASK_DECODER_MODEL_NAME,
    ONNX_VISION_ENCODER_MODEL_NAME,
    OV_PROMPT_ENCODER_MASK_DECODER_MODEL_NAME,
    OV_VISION_ENCODER_MODEL_NAME,
    TemporaryDirectory,
    model_has_dynamic_inputs,
)


logger = logging.getLogger(__name__)

core = ov.Core()


class OVSamVisionEncoder(OVModelPart):
    _model_name = "vision_encoder"

    def __init__(self, model: ov.Model, parent_model: OVBaseModel) -> None:
        super().__init__(model, parent_model, model_name=self._model_name)

    def forward(self, pixel_values):
        self._compile()
        inputs = {"pixel_values": pixel_values}
        result = self.request(inputs)
        image_embeddings = result["image_embeddings"]
        image_positional_embeddings = result["image_positional_embeddings"]
        return ModelOutput(image_embeddings=image_embeddings, image_positional_embeddings=image_positional_embeddings)


class OVSamPromptEncoder(OVModelPart):
    _model_name = "prompt_encoder_mask_decoder"

    def __init__(self, model: ov.Model, parent_model: OVBaseModel) -> None:
        super().__init__(model, parent_model, model_name=self._model_name)

    def forward(self, image_embeddings, image_positional_embeddings, input_points, input_labels=None):
        self._compile()
        inputs = {
            "image_embeddings": image_embeddings,
            "image_positional_embeddings": image_positional_embeddings,
            "input_points": input_points,
        }

        if input_labels is None:
            input_labels = np.ones(input_points[:, :, :, 0].shape, dtype=int)
        inputs["input_labels"] = input_labels

        result = self.request(inputs)
        iou_scores = torch.from_numpy(result["iou_scores"])
        pred_masks = torch.from_numpy(result["pred_masks"])

        return SamImageSegmentationOutput(iou_scores=iou_scores, pred_masks=pred_masks)


class OVSamModel(OVBaseModel):
    export_feature = "feature-extraction"
    auto_model_class = SamModel

    def __init__(
        self,
        vision_encoder_model: ov.Model,
        prompt_encoder_mask_decoder_model: ov.Model,
        config: PretrainedConfig = None,
        device: str = "CPU",
        dynamic_shapes: bool = True,
        ov_config: Optional[Dict[str, str]] = None,
        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
        **kwargs,
    ):
        self.config = config
        self._model_save_dir = model_save_dir
        self._device = device.upper()
        self.ov_config = {} if ov_config is None else {**ov_config}
        self.preprocessors = kwargs.get("preprocessors", [])
        self.vision_encoder_model = vision_encoder_model
        self.prompt_encoder_mask_decoder_model = prompt_encoder_mask_decoder_model
        self._compile_only = kwargs.get("compile_only", False)
        enable_compilation = kwargs.get("compile", True)
        self.vision_encoder = OVSamVisionEncoder(self.vision_encoder_model, self)
        self.prompt_encoder_mask_decoder = OVSamPromptEncoder(self.prompt_encoder_mask_decoder_model, self)

        if dynamic_shapes and not self.is_dynamic and not self._compile_only:
            self.reshape()

        if enable_compilation and not self._compile_only:
            self.compile()

        # Avoid warnings when creating a transformers pipeline
        AutoConfig.register(self.base_model_prefix, AutoConfig)
        try:
            self.auto_model_class.register(AutoConfig, self.__class__)
        except AttributeError:
            pass
        self.shared_image_embedding = SamPositionalEmbedding(config.vision_config)
        self.positional_embeddings = self.get_image_wide_positional_embeddings()

    def clear_requests(self):
        if self._compile_only:
            raise ValueError(
                "`clear_requests()` is not supported with `compile_only` mode, please initialize model without this option"
            )

        for _, component in self.components.items():
            component.clear_requests()

    def compile(self):
        self.vision_encoder._compile()
        self.prompt_encoder_mask_decoder._compile()

    def _save_config(self, save_directory):
        """
        Saves a model configuration into a directory, so that it can be re-loaded using the
        [`from_pretrained`] class method.
        """
        save_config(self.config, save_directory)

    def _save_pretrained(self, save_directory: Union[str, Path]):
        """
        Saves the model to the OpenVINO IR format so that it can be re-loaded using the
        [`~optimum.intel.openvino.modeling.OVModel.from_pretrained`] class method.

        Arguments:
            save_directory (`str` or `Path`):
                The directory where to save the model files.
        """
        src_models = self.ov_submodels
        dst_file_names = {
            "vision_encoder_model": OV_VISION_ENCODER_MODEL_NAME,
            "prompt_encoder_mask_decoder_model": OV_PROMPT_ENCODER_MASK_DECODER_MODEL_NAME,
        }

        for name in self._ov_submodel_names:
            model = src_models[name]
            dst_file_name = dst_file_names[name]
            dst_path = os.path.join(save_directory, dst_file_name)
            ov.save_model(model, dst_path, compress_to_fp16=False)

    @classmethod
    def _from_pretrained(
        cls,
        model_id: Union[str, Path],
        config: PretrainedConfig,
        use_auth_token: Optional[Union[bool, str]] = None,
        token: Optional[Union[bool, str]] = None,
        revision: Optional[str] = None,
        force_download: bool = False,
        cache_dir: str = HUGGINGFACE_HUB_CACHE,
        vision_encoder_file_name: Optional[str] = None,
        prompt_encoder_mask_decoder_file_name: Optional[str] = None,
        local_files_only: bool = False,
        **kwargs,
    ):
        """
        Loads a model and its configuration file from a directory or the HF Hub.

        Arguments:
            model_id (`str` or `Path`):
                The directory from which to load the model.
                Can be either:
                    - The model id of a pretrained model hosted inside a model repo on huggingface.co.
                    - The path to a directory containing the model weights.
            use_auth_token (Optional[Union[bool, str]], defaults to `None`):
                Deprecated. Please use `token` instead.
            token (Optional[Union[bool, str]], defaults to `None`):
                The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
                when running `huggingface-cli login` (stored in `~/.huggingface`).
            revision (`str`):
                The specific model version to use. It can be a branch name, a tag name, or a commit id.
            force_download (`bool`, *optional*, defaults to `False`):
                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
                cached versions if they exist.
            cache_dir (`Union[str, Path]`, *optional*):
                The path to a directory in which a downloaded pretrained model configuration should be cached if the
                standard cache should not be used.
            vision_encoder_file_name(`str`, *optional*):
                The vision encoder model file name. Overwrites the default file name openvino_vision_encoder.xml and allows one to
                load the vision encoder model with a different name.
            prompt_encoder_mask_decoder_file_name(`str`, *optional*):
                The prompt encdoer and mask decoder model file name overwriting the default file name
                openvino_prompt_encoder_mask_decoder.xml, allowing to load the decoder model with a different name.
            local_files_only(`bool`, *optional*, defaults to `False`):
                Whether or not to only look at local files (i.e., do not try to download the model).
        """
        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed soon. Please use the `token` argument instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError("You cannot use both `use_auth_token` and `token` arguments at the same time.")
            token = use_auth_token

        from_onnx = kwargs.get("from_onnx", False)

        default_vision_encoder_file_name = (
            ONNX_VISION_ENCODER_MODEL_NAME if from_onnx else OV_VISION_ENCODER_MODEL_NAME
        )
        default_prompt_encoder_mask_decoder_file_name = (
            ONNX_PROMPT_ENCODER_MASK_DECODER_MODEL_NAME if from_onnx else OV_PROMPT_ENCODER_MASK_DECODER_MODEL_NAME
        )
        vision_encoder_file_name = vision_encoder_file_name or default_vision_encoder_file_name
        prompt_encoder_mask_decoder_file_name = (
            prompt_encoder_mask_decoder_file_name or default_prompt_encoder_mask_decoder_file_name
        )

        model_file_names = {
            "vision_encoder_model": vision_encoder_file_name,
            "prompt_encoder_mask_decoder_model": prompt_encoder_mask_decoder_file_name,
        }

        if not from_onnx:
            model_file_names.update(
                {
                    "vision_encoder_model_bin": vision_encoder_file_name.replace(".xml", ".bin"),
                    "prompt_encoder_mask_decoder_model_bin": prompt_encoder_mask_decoder_file_name.replace(
                        ".xml", ".bin"
                    ),
                }
            )

        compile_only = kwargs.get("compile_only", False)
        if os.path.isdir(model_id):
            # Load model from a local directory
            model_save_dir = Path(model_id)
            file_names = {k: os.path.join(model_id, model_file_names[k]) for k in model_file_names}
        else:
            file_names = {}
            for name, file_name in model_file_names.items():
                model_cache_path = cls._cached_file(
                    model_path=model_id,
                    file_name=file_name,
                    token=token,
                    revision=revision,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    local_files_only=local_files_only,
                )
                file_names[name] = model_cache_path
            model_save_dir = Path(model_cache_path).parent
        if not compile_only:
            vision_encoder_model = cls.load_model(file_names["vision_encoder_model"])
            prompt_encoder_model = cls.load_model(file_names["prompt_encoder_mask_decoder_model"])
        else:
            vision_encoder_model = cls._compile_model(
                file_names["vision_encoder_model"],
                kwargs.get("device", "CPU"),
                kwargs.get("ov_config"),
                model_save_dir,
            )
            prompt_encoder_model = cls._compile_model(
                file_names["prompt_encoder_mask_decoder_model"],
                kwargs.get("device", "CPU"),
                kwargs.get("ov_config"),
                model_save_dir,
            )

        model = cls(
            vision_encoder_model=vision_encoder_model,
            prompt_encoder_mask_decoder_model=prompt_encoder_model,
            config=config,
            model_save_dir=model_save_dir,
            **kwargs,
        )

        return model

    @property
    def _ov_submodel_names(self):
        model_names = ["vision_encoder_model", "prompt_encoder_mask_decoder_model"]
        return model_names

    def reshape(self, batch_size: int = -1, point_batch_size: int = -1, num_points_per_image: int = -1):
        """
        Propagates the given input shapes on the model's layers, fixing the inputs shapes of the model.

        Arguments:
            batch_size (`int`):
                The batch size.
            point_batch_size (`int`):
                The point batch size.
            num_points_per_image (`int`, *optional*):
                The number of points per image
        """
        if self._compile_only:
            raise ValueError(
                "`reshape()` is not supported with `compile_only` mode, please initialize model without this option"
            )
        vision_encoder_shapes = {}
        for inputs in self.vision_encoder_model.inputs:
            vision_encoder_shapes[inputs] = inputs.get_partial_shape()
            vision_encoder_shapes[inputs][0] = batch_size
        self.vision_encoder_model.reshape(vision_encoder_shapes)
        self.vision_encoder.request = None
        mask_decoder_shapes = {}
        for inputs in self.prompt_encoder_mask_decoder_model.inputs:
            mask_decoder_shapes[inputs] = inputs.get_partial_shape()
            mask_decoder_shapes[inputs][0] = batch_size
            if inputs.get_any_name() in ["input_points", "input_labels"]:
                mask_decoder_shapes[inputs][1] = point_batch_size
                mask_decoder_shapes[inputs][2] = num_points_per_image
        self.prompt_encoder_mask_decoder_model.reshape(mask_decoder_shapes)
        self.prompt_encoder_mask_decoder.request = None
        return self

    def half(self):
        """
        Converts all the model weights to FP16 for more efficient inference on GPU.
        """
        for submodel in self.ov_submodels.values():
            apply_moc_transformations(submodel, cf=False)
            compress_model_transformation(submodel)
        return self

    def forward(
        self,
        pixel_values: Optional[torch.LongTensor] = None,
        input_points: Optional[torch.LongTensor] = None,
        input_labels: Optional[torch.LongTensor] = None,
        image_embeddings: Optional[torch.FloatTensor] = None,
        **kwargs,
    ):
        if pixel_values is None and image_embeddings is None:
            raise ValueError("Either `pixel_values` or `image_embeddings` must be provided.")

        if pixel_values is not None and image_embeddings is not None:
            raise ValueError("Only one of `pixel_values` and `image_embeddings` can be provided.")

        if input_points is None:
            raise ValueError("`input_points` must be provided.")

        if len(input_points.shape) != 4:
            raise ValueError(
                "The input_points must be a 4D tensor. Of shape `batch_size`, `point_batch_size`, `nb_points_per_image`, `2`.",
                " got {}.".format(input_points.shape),
            )

        if pixel_values is not None:
            vision_out = self.vision_encoder(pixel_values)
            image_embeddings = vision_out.image_embeddings
            image_positional_embeddings = vision_out.image_positional_embeddings
        else:
            image_positional_embeddings = self.positional_embeddings
            # repeat with batch size
            batch_size = image_embeddings.shape[0]
            image_positional_embeddings = image_positional_embeddings.repeat(batch_size, 1, 1, 1)

        if input_points is not None and input_labels is None:
            input_labels = torch.ones_like(input_points[:, :, :, 0], dtype=torch.int, device=input_points.device)

        if input_points is not None and image_embeddings.shape[0] != input_points.shape[0]:
            raise ValueError(
                "The batch size of the image embeddings and the input points must be the same. ",
                "Got {} and {} respectively.".format(image_embeddings.shape[0], input_points.shape[0]),
                " if you want to pass multiple points for the same image, make sure that you passed ",
                " input_points of shape (batch_size, point_batch_size, num_points_per_image, 3) and ",
                " input_labels of shape (batch_size, point_batch_size, num_points_per_image)",
            )
        return self.prompt_encoder_mask_decoder(
            image_embeddings=image_embeddings,
            image_positional_embeddings=image_positional_embeddings,
            input_points=input_points,
            input_labels=input_labels,
        )

    def get_image_wide_positional_embeddings(self):
        size = self.config.prompt_encoder_config.image_embedding_size
        target_device = self.shared_image_embedding.positional_embedding.device
        target_dtype = self.shared_image_embedding.positional_embedding.dtype
        grid = torch.ones((size, size), device=target_device, dtype=target_dtype)
        y_embed = grid.cumsum(dim=0) - 0.5
        x_embed = grid.cumsum(dim=1) - 0.5
        y_embed = y_embed / size
        x_embed = x_embed / size

        positional_embedding = self.shared_image_embedding(torch.stack([x_embed, y_embed], dim=-1))
        return positional_embedding.permute(2, 0, 1).unsqueeze(0)

    def get_image_features(self, pixel_values, *args, **kwargs):
        return torch.from_numpy(self.vision_encoder(pixel_values).image_embeddings)

    @property
    def is_dynamic(self):
        return model_has_dynamic_inputs(self.vision_encoder_model) or model_has_dynamic_inputs(
            self.prompt_encoder_mask_decoder_model
        )