xinference 0.16.0__py3-none-any.whl → 0.16.2__py3-none-any.whl

xinference/model/image/model_spec.json

@@ -8,7 +8,11 @@
       "text2image",
       "image2image",
       "inpainting"
-    ]
+    ],
+    "default_model_config": {
+      "quantize": true,
+      "quantize_text_encoder": "text_encoder_2"
+    }
   },
   {
     "model_name": "FLUX.1-dev",
@@ -19,7 +23,11 @@
       "text2image",
       "image2image",
       "inpainting"
-    ]
+    ],
+    "default_model_config": {
+      "quantize": true,
+      "quantize_text_encoder": "text_encoder_2"
+    }
   },
   {
     "model_name": "sd3-medium",
@@ -30,7 +38,11 @@
       "text2image",
       "image2image",
       "inpainting"
-    ]
+    ],
+    "default_model_config": {
+      "quantize": true,
+      "quantize_text_encoder": "text_encoder_3"
+    }
   },
   {
     "model_name": "sd-turbo",
@@ -178,5 +190,14 @@
     "model_ability": [
       "inpainting"
     ]
+  },
+  {
+    "model_name": "GOT-OCR2_0",
+    "model_family": "ocr",
+    "model_id": "stepfun-ai/GOT-OCR2_0",
+    "model_revision": "cf6b7386bc89a54f09785612ba74cb12de6fa17c",
+    "model_ability": [
+      "ocr"
+    ]
   }
 ]

xinference/model/image/model_spec_modelscope.json

@@ -9,7 +9,11 @@
       "text2image",
       "image2image",
       "inpainting"
-    ]
+    ],
+    "default_model_config": {
+      "quantize": true,
+      "quantize_text_encoder": "text_encoder_2"
+    }
   },
   {
     "model_name": "FLUX.1-dev",
@@ -21,7 +25,11 @@
       "text2image",
       "image2image",
       "inpainting"
-    ]
+    ],
+    "default_model_config": {
+      "quantize": true,
+      "quantize_text_encoder": "text_encoder_2"
+    }
   },
   {
     "model_name": "sd3-medium",
@@ -33,7 +41,11 @@
       "text2image",
       "image2image",
       "inpainting"
-    ]
+    ],
+    "default_model_config": {
+      "quantize": true,
+      "quantize_text_encoder": "text_encoder_3"
+    }
   },
   {
     "model_name": "sd-turbo",
@@ -148,5 +160,15 @@
         "model_revision": "62134b9d8e703b5d6f74f1534457287a8bba77ef"
       }
     ]
+  },
+  {
+    "model_name": "GOT-OCR2_0",
+    "model_family": "ocr",
+    "model_id": "stepfun-ai/GOT-OCR2_0",
+    "model_revision": "master",
+    "model_hub": "modelscope",
+    "model_ability": [
+      "ocr"
+    ]
   }
 ]

xinference/model/image/ocr/__init__.py

@@ -0,0 +1,13 @@
+# 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.

xinference/model/image/ocr/got_ocr2.py

@@ -0,0 +1,79 @@
+# 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 logging
+from typing import TYPE_CHECKING, Optional
+
+import PIL.Image
+
+if TYPE_CHECKING:
+    from ..core import ImageModelFamilyV1
+
+logger = logging.getLogger(__name__)
+
+
+class GotOCR2Model:
+    def __init__(
+        self,
+        model_uid: str,
+        model_path: Optional[str] = None,
+        device: Optional[str] = None,
+        model_spec: Optional["ImageModelFamilyV1"] = None,
+        **kwargs,
+    ):
+        self._model_uid = model_uid
+        self._model_path = model_path
+        self._device = device
+        # model info when loading
+        self._model = None
+        self._tokenizer = None
+        # info
+        self._model_spec = model_spec
+        self._abilities = model_spec.model_ability or []  # type: ignore
+        self._kwargs = kwargs
+
+    @property
+    def model_ability(self):
+        return self._abilities
+
+    def load(self):
+        from transformers import AutoModel, AutoTokenizer
+
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            self._model_path, trust_remote_code=True
+        )
+        model = AutoModel.from_pretrained(
+            self._model_path,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+            device_map="cuda",
+            use_safetensors=True,
+            pad_token_id=self._tokenizer.eos_token_id,
+        )
+        self._model = model.eval().cuda()
+
+    def ocr(
+        self,
+        image: PIL.Image,
+        **kwargs,
+    ):
+        logger.info("Got OCR 2.0 kwargs: %s", kwargs)
+        if "ocr_type" not in kwargs:
+            kwargs["ocr_type"] = "ocr"
+        if image.mode == "RGBA" or image.mode == "CMYK":
+            # convert to RGB
+            image = image.convert("RGB")
+        assert self._model is not None
+        # This chat API limits the max new tokens inside.
+        return self._model.chat(self._tokenizer, image, gradio_input=True, **kwargs)

xinference/model/image/scheduler/flux.py

@@ -124,7 +124,7 @@ class FluxBatchSchedulerActor(xo.StatelessActor):
         self._running_queue: deque[Text2ImageRequest] = deque()  # type: ignore
         self._model = None
         self._available_device = get_available_device()
-        self._id_to_req: Dict[str, Text2ImageRequest] = {}
+        self._id_to_req: Dict[str, Text2ImageRequest] = {}  # type: ignore
 
     def set_model(self, model):
         """

xinference/model/image/stable_diffusion/core.py

@@ -283,9 +283,8 @@ class DiffusionModel(SDAPIDiffusionModelMixin):
             model.enable_sequential_cpu_offload()
         elif not self._kwargs.get("device_map"):
             logger.debug("Loading model to available device")
-            model = move_model_to_available_device(self._model)
-        # Recommended if your computer has < 64 GB of RAM
-        if self._kwargs.get("attention_slicing", True):
+            model = move_model_to_available_device(model)
+        if self._kwargs.get("attention_slicing", False):
             model.enable_attention_slicing()
         if self._kwargs.get("vae_tiling", False):
             model.enable_vae_tiling()

xinference/model/image/stable_diffusion/mlx.py

@@ -0,0 +1,221 @@
+# 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 contextlib
+import gc
+import logging
+import re
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+import numpy as np
+from PIL import Image
+from xoscar.utils import classproperty
+
+from ....types import LoRA
+from ..sdapi import SDAPIDiffusionModelMixin
+from ..utils import handle_image_result
+
+if TYPE_CHECKING:
+    from ....core.progress_tracker import Progressor
+    from ..core import ImageModelFamilyV1
+
+
+logger = logging.getLogger(__name__)
+
+
+def quantization_predicate(name: str, m) -> bool:
+    return hasattr(m, "to_quantized") and m.weight.shape[1] % 512 == 0
+
+
+def to_latent_size(image_size: Tuple[int, int]):
+    h, w = image_size
+    h = ((h + 15) // 16) * 16
+    w = ((w + 15) // 16) * 16
+
+    if (h, w) != image_size:
+        print(
+            "Warning: The image dimensions need to be divisible by 16px. "
+            f"Changing size to {h}x{w}."
+        )
+
+    return (h // 8, w // 8)
+
+
+class MLXDiffusionModel(SDAPIDiffusionModelMixin):
+    def __init__(
+        self,
+        model_uid: str,
+        model_path: Optional[str] = None,
+        device: Optional[str] = None,
+        lora_model: Optional[List[LoRA]] = None,
+        lora_load_kwargs: Optional[Dict] = None,
+        lora_fuse_kwargs: Optional[Dict] = None,
+        model_spec: Optional["ImageModelFamilyV1"] = None,
+        **kwargs,
+    ):
+        self._model_uid = model_uid
+        self._model_path = model_path
+        self._device = device
+        # model info when loading
+        self._model = None
+        self._lora_model = lora_model
+        self._lora_load_kwargs = lora_load_kwargs or {}
+        self._lora_fuse_kwargs = lora_fuse_kwargs or {}
+        # info
+        self._model_spec = model_spec
+        self._abilities = model_spec.model_ability or []  # type: ignore
+        self._kwargs = kwargs
+
+    @property
+    def model_ability(self):
+        return self._abilities
+
+    @classproperty
+    def supported_models(self):
+        return ["FLUX.1-schnell", "FLUX.1-dev"]
+
+    def load(self):
+        try:
+            import mlx.nn as nn
+        except ImportError:
+            error_message = "Failed to import module 'mlx'"
+            installation_guide = [
+                "Please make sure 'mlx' is installed. ",
+                "You can install it by `pip install mlx`\n",
+            ]
+
+            raise ImportError(f"{error_message}\n\n{''.join(installation_guide)}")
+
+        from ....thirdparty.mlx.flux import FluxPipeline
+
+        logger.debug(
+            "Loading model from %s, kwargs: %s", self._model_path, self._kwargs
+        )
+        flux = self._model = FluxPipeline(
+            "flux-" + self._model_spec.model_name.split("-")[1],
+            model_path=self._model_path,
+            t5_padding=self._kwargs.get("t5_padding", True),
+        )
+        self._apply_lora()
+
+        quantize = self._kwargs.get("quantize", True)
+        if quantize:
+            nn.quantize(flux.flow, class_predicate=quantization_predicate)
+            nn.quantize(flux.t5, class_predicate=quantization_predicate)
+            nn.quantize(flux.clip, class_predicate=quantization_predicate)
+
+    def _apply_lora(self):
+        if self._lora_model is not None:
+            import mlx.core as mx
+
+            for lora_model in self._lora_model:
+                weights, lora_config = mx.load(
+                    lora_model.local_path, return_metadata=True
+                )
+                rank = int(lora_config.get("lora_rank", 8))
+                num_blocks = int(lora_config.get("lora_blocks", -1))
+                flux = self._model
+                flux.linear_to_lora_layers(rank, num_blocks)
+                flux.flow.load_weights(list(weights.items()), strict=False)
+                flux.fuse_lora_layers()
+            logger.info(f"Successfully loaded the LoRA for model {self._model_uid}.")
+
+    @staticmethod
+    @contextlib.contextmanager
+    def _release_after():
+        import mlx.core as mx
+
+        try:
+            yield
+        finally:
+            gc.collect()
+            mx.metal.clear_cache()
+
+    def text_to_image(
+        self,
+        prompt: str,
+        n: int = 1,
+        size: str = "1024*1024",
+        response_format: str = "url",
+        **kwargs,
+    ):
+        import mlx.core as mx
+
+        flux = self._model
+        width, height = map(int, re.split(r"[^\d]+", size))
+
+        # Make the generator
+        latent_size = to_latent_size((height, width))
+        gen_latent_kwargs = {}
+        if (num_steps := kwargs.get("num_inference_steps")) is None:
+            num_steps = 50 if "dev" in self._model_spec.model_name else 2  # type: ignore
+        gen_latent_kwargs["num_steps"] = num_steps
+        if guidance := kwargs.get("guidance_scale"):
+            gen_latent_kwargs["guidance"] = guidance
+        if seed := kwargs.get("seed"):
+            gen_latent_kwargs["seed"] = seed
+
+        with self._release_after():
+            latents = flux.generate_latents(  # type: ignore
+                prompt, n_images=n, latent_size=latent_size, **gen_latent_kwargs
+            )
+
+            # First we get and eval the conditioning
+            conditioning = next(latents)
+            mx.eval(conditioning)
+            peak_mem_conditioning = mx.metal.get_peak_memory() / 1024**3
+            mx.metal.reset_peak_memory()
+
+            progressor: Progressor = kwargs.pop("progressor", None)
+            # Actual denoising loop
+            for i, x_t in enumerate(latents):
+                mx.eval(x_t)
+                progressor.set_progress((i + 1) / num_steps)
+
+            peak_mem_generation = mx.metal.get_peak_memory() / 1024**3
+            mx.metal.reset_peak_memory()
+
+            # Decode them into images
+            decoded = []
+            for i in range(n):
+                decoded.append(flux.decode(x_t[i : i + 1], latent_size))  # type: ignore
+                mx.eval(decoded[-1])
+            peak_mem_decoding = mx.metal.get_peak_memory() / 1024**3
+            peak_mem_overall = max(
+                peak_mem_conditioning, peak_mem_generation, peak_mem_decoding
+            )
+
+            images = []
+            x = mx.concatenate(decoded, axis=0)
+            x = (x * 255).astype(mx.uint8)
+            for i in range(len(x)):
+                im = Image.fromarray(np.array(x[i]))
+                images.append(im)
+
+        logger.debug(
+            f"Peak memory used for the text:       {peak_mem_conditioning:.3f}GB"
+        )
+        logger.debug(
+            f"Peak memory used for the generation: {peak_mem_generation:.3f}GB"
+        )
+        logger.debug(f"Peak memory used for the decoding:   {peak_mem_decoding:.3f}GB")
+        logger.debug(f"Peak memory used overall:            {peak_mem_overall:.3f}GB")
+
+        return handle_image_result(response_format, images)
+
+    def image_to_image(self, **kwargs):
+        raise NotImplementedError
+
+    def inpainting(self, **kwargs):
+        raise NotImplementedError

xinference/model/llm/__init__.py

@@ -32,6 +32,7 @@ from .llm_family import (
     BUILTIN_LLM_MODEL_TOOL_CALL_FAMILIES,
     BUILTIN_LLM_PROMPT_STYLE,
     BUILTIN_MODELSCOPE_LLM_FAMILIES,
+    BUILTIN_OPENMIND_HUB_LLM_FAMILIES,
     LLAMA_CLASSES,
     LLM_ENGINES,
     LMDEPLOY_CLASSES,
@@ -258,6 +259,36 @@ def _install():
         if "tools" in model_spec.model_ability:
             BUILTIN_LLM_MODEL_TOOL_CALL_FAMILIES.add(model_spec.model_name)
 
+    openmind_hub_json_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "llm_family_openmind_hub.json"
+    )
+    for json_obj in json.load(
+        codecs.open(openmind_hub_json_path, "r", encoding="utf-8")
+    ):
+        model_spec = LLMFamilyV1.parse_obj(json_obj)
+        BUILTIN_OPENMIND_HUB_LLM_FAMILIES.append(model_spec)
+
+        # register prompt style, in case that we have something missed
+        # if duplicated with huggingface json, keep it as the huggingface style
+
+        if (
+            "chat" in model_spec.model_ability
+            and isinstance(model_spec.chat_template, str)
+            and model_spec.model_name not in BUILTIN_LLM_PROMPT_STYLE
+        ):
+            BUILTIN_LLM_PROMPT_STYLE[model_spec.model_name] = {
+                "chat_template": model_spec.chat_template,
+                "stop_token_ids": model_spec.stop_token_ids,
+                "stop": model_spec.stop,
+            }
+        # register model family
+        if "chat" in model_spec.model_ability:
+            BUILTIN_LLM_MODEL_CHAT_FAMILIES.add(model_spec.model_name)
+        else:
+            BUILTIN_LLM_MODEL_GENERATE_FAMILIES.add(model_spec.model_name)
+        if "tools" in model_spec.model_ability:
+            BUILTIN_LLM_MODEL_TOOL_CALL_FAMILIES.add(model_spec.model_name)
+
     csghub_json_path = os.path.join(
         os.path.dirname(os.path.abspath(__file__)), "llm_family_csghub.json"
     )
@@ -288,6 +319,7 @@ def _install():
     for llm_specs in [
         BUILTIN_LLM_FAMILIES,
         BUILTIN_MODELSCOPE_LLM_FAMILIES,
+        BUILTIN_OPENMIND_HUB_LLM_FAMILIES,
         BUILTIN_CSGHUB_LLM_FAMILIES,
     ]:
         for llm_spec in llm_specs:
@@ -298,6 +330,7 @@ def _install():
     for families in [
         BUILTIN_LLM_FAMILIES,
         BUILTIN_MODELSCOPE_LLM_FAMILIES,
+        BUILTIN_OPENMIND_HUB_LLM_FAMILIES,
         BUILTIN_CSGHUB_LLM_FAMILIES,
     ]:
         for family in families:

xinference/model/llm/core.py

@@ -193,7 +193,9 @@ def create_llm_model_instance(
     model_size_in_billions: Optional[Union[int, str]] = None,
     quantization: Optional[str] = None,
     peft_model_config: Optional[PeftModelConfig] = None,
-    download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
+    download_hub: Optional[
+        Literal["huggingface", "modelscope", "openmind_hub", "csghub"]
+    ] = None,
     model_path: Optional[str] = None,
     **kwargs,
 ) -> Tuple[LLM, LLMDescription]:

xinference/model/llm/llm_family.json

@@ -8176,6 +8176,15 @@
         ],
         "model_id": "Qwen/Qwen2.5-Coder-7B-Instruct"
       },
+      {
+        "model_format": "gptq",
+        "model_size_in_billions": "7",
+        "quantizations": [
+            "Int4",
+            "Int8"
+        ],
+        "model_id": "Qwen/Qwen2.5-Coder-7B-Instruct-GPTQ-{quantization}"
+      },
       {
         "model_format": "ggufv2",
         "model_size_in_billions": "1_5",

xinference/model/llm/llm_family.py

@@ -41,6 +41,7 @@ from ..utils import (
     create_symlink,
     download_from_csghub,
     download_from_modelscope,
+    download_from_openmind_hub,
     is_valid_model_uri,
     parse_uri,
     retry_download,
@@ -239,6 +240,7 @@ LLAMA_CLASSES: List[Type[LLM]] = []
 
 BUILTIN_LLM_FAMILIES: List["LLMFamilyV1"] = []
 BUILTIN_MODELSCOPE_LLM_FAMILIES: List["LLMFamilyV1"] = []
+BUILTIN_OPENMIND_HUB_LLM_FAMILIES: List["LLMFamilyV1"] = []
 BUILTIN_CSGHUB_LLM_FAMILIES: List["LLMFamilyV1"] = []
 
 SGLANG_CLASSES: List[Type[LLM]] = []
@@ -301,6 +303,9 @@ def cache(
             elif llm_spec.model_hub == "modelscope":
                 logger.info(f"Caching from Modelscope: {llm_spec.model_id}")
                 return cache_from_modelscope(llm_family, llm_spec, quantization)
+            elif llm_spec.model_hub == "openmind_hub":
+                logger.info(f"Caching from openmind_hub: {llm_spec.model_id}")
+                return cache_from_openmind_hub(llm_family, llm_spec, quantization)
             elif llm_spec.model_hub == "csghub":
                 logger.info(f"Caching from CSGHub: {llm_spec.model_id}")
                 return cache_from_csghub(llm_family, llm_spec, quantization)
@@ -474,7 +479,7 @@ def _skip_download(
     model_revision: Optional[str],
     quantization: Optional[str] = None,
 ) -> bool:
-    if model_format == "pytorch":
+    if model_format in ["pytorch", "mindspore"]:
         model_hub_to_meta_path = {
             "huggingface": _get_meta_path(
                 cache_dir, model_format, "huggingface", quantization
@@ -482,6 +487,9 @@ def _skip_download(
             "modelscope": _get_meta_path(
                 cache_dir, model_format, "modelscope", quantization
             ),
+            "openmind_hub": _get_meta_path(
+                cache_dir, model_format, "openmind_hub", quantization
+            ),
             "csghub": _get_meta_path(cache_dir, model_format, "csghub", quantization),
         }
         if valid_model_revision(model_hub_to_meta_path[model_hub], model_revision):
@@ -702,6 +710,50 @@ def cache_from_modelscope(
     return cache_dir
 
 
+def cache_from_openmind_hub(
+    llm_family: LLMFamilyV1,
+    llm_spec: "LLMSpecV1",
+    quantization: Optional[str] = None,
+) -> str:
+    """
+    Cache model from openmind_hub. Return the cache directory.
+    """
+    from openmind_hub import snapshot_download
+
+    cache_dir = _get_cache_dir(llm_family, llm_spec)
+    if _skip_download(
+        cache_dir,
+        llm_spec.model_format,
+        llm_spec.model_hub,
+        llm_spec.model_revision,
+        quantization,
+    ):
+        return cache_dir
+
+    if llm_spec.model_format in ["pytorch", "mindspore"]:
+        download_dir = retry_download(
+            snapshot_download,
+            llm_family.model_name,
+            {
+                "model_size": llm_spec.model_size_in_billions,
+                "model_format": llm_spec.model_format,
+            },
+            llm_spec.model_id,
+            revision=llm_spec.model_revision,
+        )
+        create_symlink(download_dir, cache_dir)
+
+    else:
+        raise ValueError(f"Unsupported format: {llm_spec.model_format}")
+
+    meta_path = _get_meta_path(
+        cache_dir, llm_spec.model_format, llm_spec.model_hub, quantization
+    )
+    _generate_meta_file(meta_path, llm_family, llm_spec, quantization)
+
+    return cache_dir
+
+
 def cache_from_huggingface(
     llm_family: LLMFamilyV1,
     llm_spec: "LLMSpecV1",
@@ -893,7 +945,9 @@ def match_llm(
     model_format: Optional[str] = None,
     model_size_in_billions: Optional[Union[int, str]] = None,
     quantization: Optional[str] = None,
-    download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
+    download_hub: Optional[
+        Literal["huggingface", "modelscope", "openmind_hub", "csghub"]
+    ] = None,
 ) -> Optional[Tuple[LLMFamilyV1, LLMSpecV1, str]]:
     """
     Find an LLM family, spec, and quantization that satisfy given criteria.
@@ -924,6 +978,12 @@ def match_llm(
             + BUILTIN_LLM_FAMILIES
             + user_defined_llm_families
         )
+    elif download_hub == "openmind_hub":
+        all_families = (
+            BUILTIN_OPENMIND_HUB_LLM_FAMILIES
+            + BUILTIN_LLM_FAMILIES
+            + user_defined_llm_families
+        )
     elif download_hub == "csghub":
         all_families = (
             BUILTIN_CSGHUB_LLM_FAMILIES
@@ -938,6 +998,12 @@ def match_llm(
             + BUILTIN_LLM_FAMILIES
             + user_defined_llm_families
         )
+    elif download_from_openmind_hub():
+        all_families = (
+            BUILTIN_OPENMIND_HUB_LLM_FAMILIES
+            + BUILTIN_LLM_FAMILIES
+            + user_defined_llm_families
+        )
     elif download_from_csghub():
         all_families = (
             BUILTIN_CSGHUB_LLM_FAMILIES

xinference/model/llm/llm_family_modelscope.json

@@ -5880,6 +5880,17 @@
         "model_revision": "master",
         "model_hub": "modelscope"
       },
+      {
+        "model_format": "gptq",
+        "model_size_in_billions": 7,
+        "quantizations": [
+          "Int4",
+          "Int8"
+        ],
+        "model_id": "qwen/Qwen2.5-Coder-7B-Instruct-GPTQ-{quantization}",
+        "model_revision": "master",
+        "model_hub": "modelscope"
+      },
       {
         "model_format": "ggufv2",
         "model_size_in_billions": "1_5",