evalscope 0.10.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

evalscope/models/text2image_model.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import importlib
+import time
+import torch
+from logging import getLogger
+from typing import Any, Dict, List, Literal, Optional, Protocol, Tuple, Union, cast
+
+from evalscope.api.messages import (
+    ChatMessage,
+    ChatMessageAssistant,
+    ContentAudio,
+    ContentImage,
+    ContentText,
+    ContentVideo,
+)
+from evalscope.api.model import (
+    ChatCompletionChoice,
+    GenerateConfig,
+    Logprob,
+    Logprobs,
+    ModelAPI,
+    ModelOutput,
+    ModelUsage,
+    TopLogprob,
+)
+from evalscope.api.tool import ToolChoice, ToolInfo
+from evalscope.utils.io_utils import PIL_to_base64
+from evalscope.utils.model_utils import get_device
+
+logger = getLogger()
+
+
+class Text2ImageAPI(ModelAPI):
+
+    def __init__(
+        self,
+        model_name: str,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        config: GenerateConfig = GenerateConfig(),
+        **model_args: Any,
+    ):
+        super().__init__(
+            model_name=model_name,
+            base_url=base_url,
+            api_key=api_key,
+            config=config,
+        )
+
+        # collect known model_args (then delete them so we can pass the rest on)
+        def collect_model_arg(name: str) -> Optional[Any]:
+            nonlocal model_args
+            value = model_args.get(name, None)
+            if value is not None:
+                model_args.pop(name)
+            return value
+
+        model_path = collect_model_arg('model_path')
+        torch_dtype = collect_model_arg('precision') or collect_model_arg('torch_dtype')
+        device_map = collect_model_arg('device_map')
+        # torch dtype
+        DTYPE_MAP = {'float16': torch.float16, 'float32': torch.float32, 'bfloat16': torch.bfloat16, 'auto': 'auto'}
+
+        if isinstance(torch_dtype, str) and torch_dtype != 'auto':
+            torch_dtype = DTYPE_MAP.get(torch_dtype, torch.float32)
+        self.torch_dtype = torch_dtype
+        self.device = device_map or get_device()
+
+        self.pipeline_cls = collect_model_arg('pipeline_cls')
+        # default to DiffusionPipeline if not specified
+        if self.pipeline_cls is None:
+            if 'flux' in model_name.lower():
+                self.pipeline_cls = 'FluxPipeline'
+            else:
+                self.pipeline_cls = 'DiffusionPipeline'
+
+        model_name_or_path = model_path or model_name
+
+        # from modelscope import pipeline_cls
+        module = getattr(importlib.import_module('modelscope'), self.pipeline_cls)
+        logger.info(f'Loading model {model_name_or_path} with {self.pipeline_cls} ...')
+
+        self.model = module.from_pretrained(
+            model_name_or_path,
+            torch_dtype=self.torch_dtype,
+            **model_args,
+        )
+
+        self.model.to(self.device)
+
+    def generate(
+        self,
+        input: List[ChatMessage],
+        tools: List[ToolInfo],
+        tool_choice: ToolChoice,
+        config: GenerateConfig,
+    ) -> ModelOutput:
+
+        # prepare generator
+        kwargs: Dict[str, Any] = {}
+        if config.height is not None:
+            kwargs['height'] = config.height
+        if config.width is not None:
+            kwargs['width'] = config.width
+        if config.num_inference_steps is not None:
+            kwargs['num_inference_steps'] = config.num_inference_steps
+        if config.guidance_scale is not None:
+            kwargs['guidance_scale'] = config.guidance_scale
+        # update with extra model parameters
+        kwargs.update(config.model_extra)
+
+        # assume the first text as prompt
+        prompt = input[0].text
+        # get the first image as output
+        image = self.model(prompt=prompt, **kwargs).images[0]
+
+        image_base64 = PIL_to_base64(image)
+
+        return ModelOutput(
+            model=self.model_name,
+            choices=[ChatCompletionChoice.from_content(content=[ContentImage(image=image_base64)])],
+            time=time.time(),
+        )