xinference 0.14.4.post1__py3-none-any.whl → 0.15.1__py3-none-any.whl

xinference/model/embedding/core.py

@@ -19,6 +19,7 @@ from collections import defaultdict
 from typing import Dict, List, Literal, Optional, Tuple, Union, no_type_check
 
 import numpy as np
+import torch
 
 from ...device_utils import empty_cache
 from ...types import Embedding, EmbeddingData, EmbeddingUsage
@@ -34,7 +35,11 @@ EMBEDDING_MODEL_DESCRIPTIONS: Dict[str, List[Dict]] = defaultdict(list)
 EMBEDDING_EMPTY_CACHE_COUNT = int(
     os.getenv("XINFERENCE_EMBEDDING_EMPTY_CACHE_COUNT", "10")
 )
+EMBEDDING_EMPTY_CACHE_TOKENS = int(
+    os.getenv("XINFERENCE_EMBEDDING_EMPTY_CACHE_TOKENS", "8192")
+)
 assert EMBEDDING_EMPTY_CACHE_COUNT > 0
+assert EMBEDDING_EMPTY_CACHE_TOKENS > 0
 
 
 def get_embedding_model_descriptions():
@@ -149,6 +154,25 @@ class EmbeddingModel:
             def to(self, *args, **kwargs):
                 pass
 
+        torch_dtype = None
+        if torch_dtype_str := self._kwargs.get("torch_dtype"):
+            try:
+                torch_dtype = getattr(torch, torch_dtype_str)
+                if torch_dtype not in [
+                    torch.float16,
+                    torch.float32,
+                    torch.bfloat16,
+                ]:
+                    logger.warning(
+                        f"Load embedding model with unsupported torch dtype :  {torch_dtype_str}. Using default torch dtype: fp32."
+                    )
+                    torch_dtype = torch.float32
+            except AttributeError:
+                logger.warning(
+                    f"Load embedding model with  unknown torch dtype '{torch_dtype_str}'. Using default torch dtype: fp32."
+                )
+                torch_dtype = torch.float32
+
         from ..utils import patch_trust_remote_code
 
         patch_trust_remote_code()
@@ -156,42 +180,21 @@ class EmbeddingModel:
             "gte" in self._model_spec.model_name.lower()
             and "qwen2" in self._model_spec.model_name.lower()
         ):
-            import torch
-
-            torch_dtype_str = self._kwargs.get("torch_dtype")
-            if torch_dtype_str is not None:
-                try:
-                    torch_dtype = getattr(torch, torch_dtype_str)
-                    if torch_dtype not in [
-                        torch.float16,
-                        torch.float32,
-                        torch.bfloat16,
-                    ]:
-                        logger.warning(
-                            f"Load embedding model with unsupported torch dtype :  {torch_dtype_str}. Using default torch dtype: fp32."
-                        )
-                        torch_dtype = torch.float32
-                except AttributeError:
-                    logger.warning(
-                        f"Load embedding model with  unknown torch dtype '{torch_dtype_str}'. Using default torch dtype: fp32."
-                    )
-                    torch_dtype = torch.float32
-            else:
-                torch_dtype = "auto"
+            model_kwargs = {"device_map": "auto"}
+            if torch_dtype:
+                model_kwargs["torch_dtype"] = torch_dtype
             self._model = XSentenceTransformer(
                 self._model_path,
                 device=self._device,
-                model_kwargs={"device_map": "auto", "torch_dtype": torch_dtype},
+                model_kwargs=model_kwargs,
             )
         else:
-            self._model = SentenceTransformer(self._model_path, device=self._device)
+            model_kwargs = {"torch_dtype": torch_dtype} if torch_dtype else None
+            self._model = SentenceTransformer(
+                self._model_path, device=self._device, model_kwargs=model_kwargs
+            )
 
     def create_embedding(self, sentences: Union[str, List[str]], **kwargs):
-        self._counter += 1
-        if self._counter % EMBEDDING_EMPTY_CACHE_COUNT == 0:
-            logger.debug("Empty embedding cache.")
-            gc.collect()
-            empty_cache()
         from sentence_transformers import SentenceTransformer
 
         kwargs.setdefault("normalize_embeddings", True)
@@ -309,7 +312,9 @@ class EmbeddingModel:
                 features = model.tokenize(sentences_batch)
                 features = batch_to_device(features, device)
                 features.update(extra_features)
-                all_token_nums += sum([len(f) for f in features])
+                # when batching, the attention mask 1 means there is a token
+                # thus we just sum up it to get the total number of tokens
+                all_token_nums += features["attention_mask"].sum().item()
 
                 with torch.no_grad():
                     out_features = model.forward(features)
@@ -393,13 +398,29 @@ class EmbeddingModel:
         usage = EmbeddingUsage(
             prompt_tokens=all_token_nums, total_tokens=all_token_nums
         )
-        return Embedding(
+        result = Embedding(
             object="list",
             model=self._model_uid,
             data=embedding_list,
             usage=usage,
         )
 
+        # clean cache if possible
+        self._counter += 1
+        if (
+            self._counter % EMBEDDING_EMPTY_CACHE_COUNT == 0
+            or all_token_nums >= EMBEDDING_EMPTY_CACHE_TOKENS
+        ):
+            logger.debug(
+                "Empty embedding cache, calling count %s, all_token_nums %s",
+                self._counter,
+                all_token_nums,
+            )
+            gc.collect()
+            empty_cache()
+
+        return result
+
 
 def match_embedding(
     model_name: str,

xinference/model/image/core.py

@@ -47,6 +47,7 @@ class ImageModelFamilyV1(CacheableModelSpec):
     model_hub: str = "huggingface"
     model_ability: Optional[List[str]]
     controlnet: Optional[List["ImageModelFamilyV1"]]
+    default_generate_config: Optional[dict] = {}
 
 
 class ImageModelDescription(ModelDescription):
@@ -238,7 +239,7 @@ def create_image_model_instance(
         lora_model_paths=lora_model,
         lora_load_kwargs=lora_load_kwargs,
         lora_fuse_kwargs=lora_fuse_kwargs,
-        abilities=model_spec.model_ability,
+        model_spec=model_spec,
         **kwargs,
     )
     model_description = ImageModelDescription(

xinference/model/image/model_spec.json

@@ -5,7 +5,9 @@
     "model_id": "black-forest-labs/FLUX.1-schnell",
     "model_revision": "768d12a373ed5cc9ef9a9dea7504dc09fcc14842",
     "model_ability": [
-      "text2image"
+      "text2image",
+      "image2image",
+      "inpainting"
     ]
   },
   {
@@ -14,7 +16,9 @@
     "model_id": "black-forest-labs/FLUX.1-dev",
     "model_revision": "01aa605f2c300568dd6515476f04565a954fcb59",
     "model_ability": [
-      "text2image"
+      "text2image",
+      "image2image",
+      "inpainting"
     ]
   },
   {
@@ -35,7 +39,11 @@
     "model_revision": "1681ed09e0cff58eeb41e878a49893228b78b94c",
     "model_ability": [
       "text2image"
-    ]
+    ],
+    "default_generate_config": {
+      "guidance_scale": 0.0,
+      "num_inference_steps": 1
+    }
   },
   {
     "model_name": "sdxl-turbo",
@@ -44,7 +52,11 @@
     "model_revision": "f4b0486b498f84668e828044de1d0c8ba486e05b",
     "model_ability": [
       "text2image"
-    ]
+    ],
+    "default_generate_config": {
+      "guidance_scale": 0.0,
+      "num_inference_steps": 1
+    }
   },
   {
     "model_name": "stable-diffusion-v1.5",

xinference/model/image/model_spec_modelscope.json

@@ -6,7 +6,9 @@
     "model_id": "AI-ModelScope/FLUX.1-schnell",
     "model_revision": "master",
     "model_ability": [
-      "text2image"
+      "text2image",
+      "image2image",
+      "inpainting"
     ]
   },
   {
@@ -16,7 +18,9 @@
     "model_id": "AI-ModelScope/FLUX.1-dev",
     "model_revision": "master",
     "model_ability": [
-      "text2image"
+      "text2image",
+      "image2image",
+      "inpainting"
     ]
   },
   {
@@ -39,7 +43,11 @@
     "model_revision": "master",
     "model_ability": [
       "text2image"
-    ]
+    ],
+    "default_generate_config": {
+      "guidance_scale": 0.0,
+      "num_inference_steps": 1
+    }
   },
   {
     "model_name": "sdxl-turbo",
@@ -49,7 +57,11 @@
     "model_revision": "master",
     "model_ability": [
       "text2image"
-    ]
+    ],
+    "default_generate_config": {
+      "guidance_scale": 0.0,
+      "num_inference_steps": 1
+    }
   },
   {
     "model_name": "stable-diffusion-v1.5",

xinference/model/image/sdapi.py

@@ -0,0 +1,136 @@
+# Copyright 2022-2023 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import base64
+import io
+import warnings
+
+from PIL import Image
+
+
+class SDAPIToDiffusersConverter:
+    txt2img_identical_args = {
+        "prompt",
+        "negative_prompt",
+        "seed",
+        "width",
+        "height",
+        "sampler_name",
+    }
+    txt2img_arg_mapping = {
+        "steps": "num_inference_steps",
+        "cfg_scale": "guidance_scale",
+        # "denoising_strength": "strength",
+    }
+    img2img_identical_args = {
+        "prompt",
+        "negative_prompt",
+        "seed",
+        "width",
+        "height",
+        "sampler_name",
+    }
+    img2img_arg_mapping = {
+        "init_images": "image",
+        "steps": "num_inference_steps",
+        "cfg_scale": "guidance_scale",
+        "denoising_strength": "strength",
+    }
+
+    @staticmethod
+    def convert_to_diffusers(sd_type: str, params: dict) -> dict:
+        diffusers_params = {}
+
+        identical_args = getattr(SDAPIToDiffusersConverter, f"{sd_type}_identical_args")
+        mapping_args = getattr(SDAPIToDiffusersConverter, f"{sd_type}_arg_mapping")
+        for param, value in params.items():
+            if param in identical_args:
+                diffusers_params[param] = value
+            elif param in mapping_args:
+                diffusers_params[mapping_args[param]] = value
+            else:
+                raise ValueError(f"Unknown arg: {param}")
+
+        return diffusers_params
+
+    @staticmethod
+    def get_available_args(sd_type: str) -> set:
+        identical_args = getattr(SDAPIToDiffusersConverter, f"{sd_type}_identical_args")
+        mapping_args = getattr(SDAPIToDiffusersConverter, f"{sd_type}_arg_mapping")
+        return identical_args.union(mapping_args)
+
+
+class SDAPIDiffusionModelMixin:
+    @staticmethod
+    def _check_kwargs(sd_type: str, kwargs: dict):
+        available_args = SDAPIToDiffusersConverter.get_available_args(sd_type)
+        unknown_args = []
+        available_kwargs = {}
+        for arg, value in kwargs.items():
+            if arg in available_args:
+                available_kwargs[arg] = value
+            else:
+                unknown_args.append(arg)
+        if unknown_args:
+            warnings.warn(
+                f"Some args are not supported for now and will be ignored: {unknown_args}"
+            )
+
+        converted_kwargs = SDAPIToDiffusersConverter.convert_to_diffusers(
+            sd_type, available_kwargs
+        )
+
+        width, height = converted_kwargs.pop("width", None), converted_kwargs.pop(
+            "height", None
+        )
+        if width and height:
+            converted_kwargs["size"] = f"{width}*{height}"
+
+        return converted_kwargs
+
+    def txt2img(self, **kwargs):
+        converted_kwargs = self._check_kwargs("txt2img", kwargs)
+        result = self.text_to_image(response_format="b64_json", **converted_kwargs)  # type: ignore
+
+        # convert to SD API result
+        return {
+            "images": [r["b64_json"] for r in result["data"]],
+            "info": {"created": result["created"]},
+            "parameters": {},
+        }
+
+    @staticmethod
+    def _decode_b64_img(img_str: str) -> Image:
+        # img_str in a format: "data:image/png;base64," + raw_b64_img(image)
+        f, data = img_str.split(",", 1)
+        f, encode_type = f.split(";", 1)
+        assert encode_type == "base64"
+        f = f.split("/", 1)[1]
+        b = base64.b64decode(data)
+        return Image.open(io.BytesIO(b), formats=[f])
+
+    def img2img(self, **kwargs):
+        init_images = kwargs.pop("init_images", [])
+        kwargs["init_images"] = [self._decode_b64_img(i) for i in init_images]
+        clip_skip = kwargs.get("override_settings", {}).get("clip_skip")
+        converted_kwargs = self._check_kwargs("img2img", kwargs)
+        if clip_skip:
+            converted_kwargs["clip_skip"] = clip_skip
+        result = self.image_to_image(response_format="b64_json", **converted_kwargs)  # type: ignore
+
+        # convert to SD API result
+        return {
+            "images": [r["b64_json"] for r in result["data"]],
+            "info": {"created": result["created"]},
+            "parameters": {},
+        }

xinference/model/image/stable_diffusion/core.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import base64
+import contextlib
+import inspect
 import logging
 import os
 import re
@@ -22,19 +24,43 @@ import uuid
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from io import BytesIO
-from typing import Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import PIL.Image
+import torch
 from PIL import ImageOps
 
 from ....constants import XINFERENCE_IMAGE_DIR
 from ....device_utils import move_model_to_available_device
 from ....types import Image, ImageList, LoRA
+from ..sdapi import SDAPIDiffusionModelMixin
 
-logger = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from ..core import ImageModelFamilyV1
 
+logger = logging.getLogger(__name__)
 
-class DiffusionModel:
+SAMPLING_METHODS = [
+    "default",
+    "DPM++ 2M",
+    "DPM++ 2M Karras",
+    "DPM++ 2M SDE",
+    "DPM++ 2M SDE Karras",
+    "DPM++ SDE",
+    "DPM++ SDE Karras",
+    "DPM2",
+    "DPM2 Karras",
+    "DPM2 a",
+    "DPM2 a Karras",
+    "Euler",
+    "Euler a",
+    "Heun",
+    "LMS",
+    "LMS Karras",
+]
+
+
+class DiffusionModel(SDAPIDiffusionModelMixin):
     def __init__(
         self,
         model_uid: str,
@@ -43,7 +69,7 @@ class DiffusionModel:
         lora_model: Optional[List[LoRA]] = None,
         lora_load_kwargs: Optional[Dict] = None,
         lora_fuse_kwargs: Optional[Dict] = None,
-        abilities: Optional[List[str]] = None,
+        model_spec: Optional["ImageModelFamilyV1"] = None,
         **kwargs,
     ):
         self._model_uid = model_uid
@@ -59,7 +85,8 @@ class DiffusionModel:
         self._lora_model = lora_model
         self._lora_load_kwargs = lora_load_kwargs or {}
         self._lora_fuse_kwargs = lora_fuse_kwargs or {}
-        self._abilities = abilities or []
+        self._model_spec = model_spec
+        self._abilities = model_spec.model_ability or []  # type: ignore
         self._kwargs = kwargs
 
     @property
@@ -80,8 +107,6 @@ class DiffusionModel:
             logger.info(f"Successfully loaded the LoRA for model {self._model_uid}.")
 
     def load(self):
-        import torch
-
         if "text2image" in self._abilities or "image2image" in self._abilities:
             from diffusers import AutoPipelineForText2Image as AutoPipelineModel
         elif "inpainting" in self._abilities:
@@ -143,7 +168,9 @@ class DiffusionModel:
                 self._kwargs[text_encoder_name] = text_encoder
                 self._kwargs["device_map"] = "balanced"
 
-        logger.debug("Loading model %s", AutoPipelineModel)
+        logger.debug(
+            "Loading model from %s, kwargs: %s", self._model_path, self._kwargs
+        )
         self._model = AutoPipelineModel.from_pretrained(
             self._model_path,
             **self._kwargs,
@@ -158,6 +185,89 @@ class DiffusionModel:
         self._model.enable_attention_slicing()
         self._apply_lora()
 
+    @staticmethod
+    def _get_scheduler(model: Any, sampler_name: str):
+        if not sampler_name:
+            return
+
+        assert model is not None
+
+        import diffusers
+
+        # see https://github.com/huggingface/diffusers/issues/4167
+        # to get A1111 <> Diffusers Scheduler mapping
+        if sampler_name == "DPM++ 2M":
+            return diffusers.DPMSolverMultistepScheduler.from_config(
+                model.scheduler.config
+            )
+        elif sampler_name == "DPM++ 2M Karras":
+            return diffusers.DPMSolverMultistepScheduler.from_config(
+                model.scheduler.config, use_karras_sigmas=True
+            )
+        elif sampler_name == "DPM++ 2M SDE":
+            return diffusers.DPMSolverMultistepScheduler.from_config(
+                model.scheduler.config, algorithm_type="sde-dpmsolver++"
+            )
+        elif sampler_name == "DPM++ 2M SDE Karras":
+            return diffusers.DPMSolverMultistepScheduler.from_config(
+                model.scheduler.config,
+                algorithm_type="sde-dpmsolver++",
+                use_karras_sigmas=True,
+            )
+        elif sampler_name == "DPM++ SDE":
+            return diffusers.DPMSolverSinglestepScheduler.from_config(
+                model.scheduler.config
+            )
+        elif sampler_name == "DPM++ SDE Karras":
+            return diffusers.DPMSolverSinglestepScheduler.from_config(
+                model.scheduler.config, use_karras_sigmas=True
+            )
+        elif sampler_name == "DPM2":
+            return diffusers.KDPM2DiscreteScheduler.from_config(model.scheduler.config)
+        elif sampler_name == "DPM2 Karras":
+            return diffusers.KDPM2DiscreteScheduler.from_config(
+                model.scheduler.config, use_karras_sigmas=True
+            )
+        elif sampler_name == "DPM2 a":
+            return diffusers.KDPM2AncestralDiscreteScheduler.from_config(
+                model.scheduler.config
+            )
+        elif sampler_name == "DPM2 a Karras":
+            return diffusers.KDPM2AncestralDiscreteScheduler.from_config(
+                model.scheduler.config, use_karras_sigmas=True
+            )
+        elif sampler_name == "Euler":
+            return diffusers.EulerDiscreteScheduler.from_config(model.scheduler.config)
+        elif sampler_name == "Euler a":
+            return diffusers.EulerAncestralDiscreteScheduler.from_config(
+                model.scheduler.config
+            )
+        elif sampler_name == "Heun":
+            return diffusers.HeunDiscreteScheduler.from_config(model.scheduler.config)
+        elif sampler_name == "LMS":
+            return diffusers.LMSDiscreteScheduler.from_config(model.scheduler.config)
+        elif sampler_name == "LMS Karras":
+            return diffusers.LMSDiscreteScheduler.from_config(
+                model.scheduler.config, use_karras_sigmas=True
+            )
+        else:
+            raise ValueError(f"Unknown sampler: {sampler_name}")
+
+    @staticmethod
+    @contextlib.contextmanager
+    def _reset_when_done(model: Any, sampler_name: str):
+        assert model is not None
+        scheduler = DiffusionModel._get_scheduler(model, sampler_name)
+        if scheduler:
+            default_scheduler = model.scheduler
+            model.scheduler = scheduler
+            try:
+                yield
+            finally:
+                model.scheduler = default_scheduler
+        else:
+            yield
+
     def _call_model(
         self,
         response_format: str,
@@ -168,13 +278,27 @@ class DiffusionModel:
 
         from ....device_utils import empty_cache
 
-        logger.debug(
-            "stable diffusion args: %s",
-            kwargs,
-        )
         model = model if model is not None else self._model
+        is_padded = kwargs.pop("is_padded", None)
+        origin_size = kwargs.pop("origin_size", None)
+        seed = kwargs.pop("seed", None)
+        if seed is not None:
+            kwargs["generator"] = generator = torch.Generator(device=self._model.device)  # type: ignore
+            if seed != -1:
+                kwargs["generator"] = generator.manual_seed(seed)
+        sampler_name = kwargs.pop("sampler_name", None)
         assert callable(model)
-        images = model(**kwargs).images
+        with self._reset_when_done(model, sampler_name):
+            logger.debug("stable diffusion args: %s, model: %s", kwargs, model)
+            images = model(**kwargs).images
+
+        # revert padding if padded
+        if is_padded and origin_size:
+            new_images = []
+            x, y = origin_size
+            for img in images:
+                new_images.append(img.crop((0, 0, x, y)))
+            images = new_images
 
         # clean cache
         gc.collect()
@@ -198,7 +322,7 @@ class DiffusionModel:
 
             with ThreadPoolExecutor() as executor:
                 results = list(map(partial(executor.submit, _gen_base64_image), images))  # type: ignore
-                image_list = [Image(url=None, b64_json=s.result()) for s in results]
+                image_list = [Image(url=None, b64_json=s.result()) for s in results]  # type: ignore
             return ImageList(created=int(time.time()), data=image_list)
         else:
             raise ValueError(f"Unsupported response format: {response_format}")
@@ -220,14 +344,16 @@ class DiffusionModel:
         # References:
         # https://huggingface.co/docs/diffusers/main/en/api/pipelines/controlnet_sdxl
         width, height = map(int, re.split(r"[^\d]+", size))
-        self._filter_kwargs(kwargs)
+        generate_kwargs = self._model_spec.default_generate_config.copy()  # type: ignore
+        generate_kwargs.update({k: v for k, v in kwargs.items() if v is not None})
+        self._filter_kwargs(generate_kwargs)
         return self._call_model(
             prompt=prompt,
             height=height,
             width=width,
             num_images_per_prompt=n,
             response_format=response_format,
-            **kwargs,
+            **generate_kwargs,
         )
 
     @staticmethod
@@ -265,6 +391,9 @@ class DiffusionModel:
         if padding_image_to_multiple := kwargs.pop("padding_image_to_multiple", None):
             # Model like SD3 image to image requires image's height and width is times of 16
             # padding the image if specified
+            origin_x, origin_y = image.size
+            kwargs["origin_size"] = (origin_x, origin_y)
+            kwargs["is_padded"] = True
             image = self.pad_to_multiple(image, multiple=int(padding_image_to_multiple))
 
         if size:
@@ -273,12 +402,24 @@ class DiffusionModel:
                 width, height = image.size
             kwargs["width"] = width
             kwargs["height"] = height
-
+        else:
+            # SD3 image2image cannot accept width and height
+            parameters = inspect.signature(model.__call__).parameters  # type: ignore
+            allow_width_height = False
+            for param in parameters.values():
+                if param.kind == inspect.Parameter.VAR_KEYWORD:
+                    allow_width_height = True
+                    break
+            if "width" in parameters or "height" in parameters:
+                allow_width_height = True
+            if allow_width_height:
+                kwargs["width"], kwargs["height"] = image.size
+
+        kwargs["negative_prompt"] = negative_prompt
         self._filter_kwargs(kwargs)
         return self._call_model(
             image=image,
             prompt=prompt,
-            negative_prompt=negative_prompt,
             num_images_per_prompt=n,
             response_format=response_format,
             model=model,
@@ -318,6 +459,9 @@ class DiffusionModel:
         if padding_image_to_multiple := kwargs.pop("padding_image_to_multiple", None):
             # Model like SD3 inpainting requires image's height and width is times of 16
             # padding the image if specified
+            origin_x, origin_y = image.size
+            kwargs["origin_size"] = (origin_x, origin_y)
+            kwargs["is_padded"] = True
             image = self.pad_to_multiple(image, multiple=int(padding_image_to_multiple))
             mask_image = self.pad_to_multiple(
                 mask_image, multiple=int(padding_image_to_multiple)
@@ -325,11 +469,12 @@ class DiffusionModel:
             # calculate actual image size after padding
             width, height = image.size
 
+        kwargs["negative_prompt"] = negative_prompt
+        self._filter_kwargs(kwargs)
         return self._call_model(
             image=image,
             mask_image=mask_image,
             prompt=prompt,
-            negative_prompt=negative_prompt,
             height=height,
             width=width,
             num_images_per_prompt=n,