optimum/amd/ryzenai/pipelines/__init__.py

# Copyright 2023 The HuggingFace Team. All rights reserved.
# Licensed under the MIT License.


from typing import TYPE_CHECKING, Any, Optional, Union

from optimum.exporters import TasksManager
from transformers import ImageClassificationPipeline, Pipeline, PretrainedConfig
from transformers import pipeline as transformers_pipeline
from transformers.image_processing_utils import BaseImageProcessor
from transformers.onnx.utils import get_preprocessor

from ..modeling import (
    RyzenAIModel,
    RyzenAIModelForImageClassification,
    RyzenAIModelForObjectDetection,
    RyzenAIModelForSemanticSegmentation,
)
from ..models import (
    HRNetImageProcessor,
    SemanticFPNImageProcessor,
    YoloV3ImageProcessor,
    YoloV5ImageProcessor,
    YoloV8ImageProcessor,
    YoloXImageProcessor,
)
from .image_classification import TimmImageClassificationPipeline
from .image_segmentation import ImageSegmentationPipeline
from .object_detection import YoloObjectDetectionPipeline


if TYPE_CHECKING:
    from transformers.feature_extraction_utils import PreTrainedFeatureExtractor

pipeline_map = {
    "yolox": {"preprocessor": YoloXImageProcessor, "impl": YoloObjectDetectionPipeline},
    "yolov5": {"preprocessor": YoloV5ImageProcessor, "impl": YoloObjectDetectionPipeline},
    "yolov3": {"preprocessor": YoloV3ImageProcessor, "impl": YoloObjectDetectionPipeline},
    "yolov8": {"preprocessor": YoloV8ImageProcessor, "impl": YoloObjectDetectionPipeline},
    "semantic_fpn": {"preprocessor": SemanticFPNImageProcessor, "impl": ImageSegmentationPipeline},
    "hrnet": {"preprocessor": HRNetImageProcessor, "impl": ImageSegmentationPipeline},
}

RYZENAI_SUPPORTED_TASKS = {
    "image-classification": {
        "impl": ImageClassificationPipeline,
        "class": (RyzenAIModelForImageClassification,),
        "default": "amd/resnet50",
        "type": "image",
    },
    "object-detection": {
        "impl": YoloObjectDetectionPipeline,
        "class": (RyzenAIModelForObjectDetection,),
        "default": "amd/yolox-s",
        "type": "image",
        "model_type": "yolox",
    },
    "image-segmentation": {
        "impl": ImageSegmentationPipeline,
        "class": (RyzenAIModelForSemanticSegmentation,),
        "default": "amd/SemanticFPN",
        "type": "image",
        "model_type": "semantic_fpn",
    },
}


def load_model(
    model,
    task,
    model_type,
    vaip_config,
    SUPPORTED_TASKS,
    token: Optional[Union[bool, str]] = None,
    revision: str = "main",
):
    if model is None:
        if task != "object-detection":
            raise ValueError("Creating pipeline without model for the task is not supported!")

        model_id = SUPPORTED_TASKS[task]["default"]
        if model_type is None:
            model_type = SUPPORTED_TASKS[task].get("model_type", None)
        model = SUPPORTED_TASKS[task]["class"][0].from_pretrained(
            model_id, vaip_config=vaip_config, use_auth_token=token, revision=revision
        )
    elif isinstance(model, str):
        model_id = model
        ort_model_class = SUPPORTED_TASKS[task]["class"][0]

        model = ort_model_class.from_pretrained(
            model_id, vaip_config=vaip_config, use_auth_token=token, revision=revision
        )
    elif isinstance(model, RyzenAIModel):
        model_id = None
    else:
        raise ValueError(
            f"Model {model} is not supported. Please provide a valid model either as string or RyzenAIModel."
            "You can also provide non model then a default one will be used"
        )
    return model, model_id, model_type


def pipeline(
    task,
    model: Optional[Any] = None,
    vaip_config: Optional[str] = None,
    model_type: Optional[str] = None,
    feature_extractor: Optional[Union[str, "PreTrainedFeatureExtractor"]] = None,
    image_processor: Optional[Union[str, BaseImageProcessor]] = None,
    use_fast: bool = True,
    token: Optional[Union[str, bool]] = None,
    revision: Optional[str] = None,
    **kwargs,
) -> Pipeline:
    """
    Utility method to build a pipeline for various RyzenAI tasks.

    This function creates a pipeline for a specified task, utilizing a given model or loading the default model for the task.
    The pipeline includes components such as a image processor and model.

    Args:
        task (`str`):
            The task defining which pipeline will be returned. Available tasks include:
            - "image-classification"
            - "object-detection"
        model (`Optional[Any]`, defaults to `None`):
            The model that will be used by the pipeline to make predictions. This can be a model identifier or an
            actual instance of a pretrained model. If not provided, the default model for the specified task will be loaded.
        vaip_config (`Optional[str]`, defaults to `None`):
            Runtime configuration file for inference with Ryzen IPU. A default config file can be found in the Ryzen AI VOE package,
            extracted during installation under the name `vaip_config.json`.
        model_type (`Optional[str]`, defaults to `None`):
            Model type for the model
        feature_extractor (`Union[str, "PreTrainedFeatureExtractor"]`, defaults to `None`):
            The feature extractor that will be used by the pipeline to encode data for the model. This can be a model
            identifier or an actual pretrained feature extractor.
        image_processor (`Union[str, BaseImageProcessor]`, defaults to `None`):
            The image processor that will be used by the pipeline for image-related tasks.
        use_fast (`bool`, defaults to `True`):
            Whether or not to use a Fast tokenizer if possible.
        token (`Union[str, bool`], 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`, defaults to `None`):
            The specific model version to use, specified as a branch name, tag name, or commit id.
        **kwargs:
            Additional keyword arguments passed to the underlying pipeline class.

    Returns:
        Pipeline:
            An instance of the specified pipeline for the given task and model.
    """
    model, model_id, model_type = load_model(
        model,
        task=task,
        model_type=model_type,
        vaip_config=vaip_config,
        SUPPORTED_TASKS=RYZENAI_SUPPORTED_TASKS,
        token=token,
        revision=revision,
    )

    if model.config is None:
        if model_type is None:
            raise ValueError(
                "Could not automatically find a `model_type` for the model, you must specify the `model_type` argument explicitly."
            )

        if model_type not in pipeline_map:
            raise ValueError(
                f"Model type: {model_type} is not supported by Ryzen pipelines. Please open an issue or submit a PR to add the support."
            )

        image_processor = pipeline_map[model_type]["preprocessor"]()
        ryzen_pipeline = pipeline_map[model_type]["impl"]

        model.config = PretrainedConfig.from_dict({})
    else:
        library_name = TasksManager._infer_library_from_model(model)

        if library_name != "timm" and image_processor is None:
            if model_id:
                if feature_extractor is None and image_processor is None:
                    image_processor = get_preprocessor(model_id)
            else:
                for preprocessor in model.preprocessors:
                    if isinstance(preprocessor, BaseImageProcessor):
                        image_processor = preprocessor
                        break
                if image_processor is None:
                    raise ValueError(
                        "Could not automatically find an image processor for the model, you must specifiy the argument `image_processor` explicitly."
                    )

        if task == "image-classification" and library_name == "timm":
            ryzen_pipeline = TimmImageClassificationPipeline
        else:
            ryzen_pipeline = transformers_pipeline

    return ryzen_pipeline(
        task=task,
        model=model,
        feature_extractor=feature_extractor,
        image_processor=image_processor,
        use_fast=use_fast,
        **kwargs,
    )