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

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-10-18T12:49:02+0800",
+ "date": "2024-10-25T12:51:06+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "5f7dea44832a1c41f887b9a01377191894550057",
- "version": "0.16.0"
+ "full-revisionid": "d4cd7b15104c16838e3c562cf2d33337e3d38897",
+ "version": "0.16.1"
 }
 '''  # END VERSION_JSON
 

xinference/api/restful_api.py

@@ -567,6 +567,16 @@ class RESTfulAPI:
                 else None
             ),
         )
+        self._router.add_api_route(
+            "/v1/images/ocr",
+            self.create_ocr,
+            methods=["POST"],
+            dependencies=(
+                [Security(self._auth_service, scopes=["models:read"])]
+                if self.is_authenticated()
+                else None
+            ),
+        )
         # SD WebUI API
         self._router.add_api_route(
             "/sdapi/v1/options",
@@ -1754,6 +1764,44 @@ class RESTfulAPI:
             await self._report_error_event(model_uid, str(e))
             raise HTTPException(status_code=500, detail=str(e))
 
+    async def create_ocr(
+        self,
+        model: str = Form(...),
+        image: UploadFile = File(media_type="application/octet-stream"),
+        kwargs: Optional[str] = Form(None),
+    ) -> Response:
+        model_uid = model
+        try:
+            model_ref = await (await self._get_supervisor_ref()).get_model(model_uid)
+        except ValueError as ve:
+            logger.error(str(ve), exc_info=True)
+            await self._report_error_event(model_uid, str(ve))
+            raise HTTPException(status_code=400, detail=str(ve))
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            await self._report_error_event(model_uid, str(e))
+            raise HTTPException(status_code=500, detail=str(e))
+
+        try:
+            if kwargs is not None:
+                parsed_kwargs = json.loads(kwargs)
+            else:
+                parsed_kwargs = {}
+            im = Image.open(image.file)
+            text = await model_ref.ocr(
+                image=im,
+                **parsed_kwargs,
+            )
+            return Response(content=text, media_type="text/plain")
+        except RuntimeError as re:
+            logger.error(re, exc_info=True)
+            await self._report_error_event(model_uid, str(re))
+            raise HTTPException(status_code=400, detail=str(re))
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            await self._report_error_event(model_uid, str(e))
+            raise HTTPException(status_code=500, detail=str(e))
+
     async def create_flexible_infer(self, request: Request) -> Response:
         payload = await request.json()
 

xinference/client/restful/restful_client.py

@@ -369,6 +369,25 @@ class RESTfulImageModelHandle(RESTfulModelHandle):
         response_data = response.json()
         return response_data
 
+    def ocr(self, image: Union[str, bytes], **kwargs):
+        url = f"{self._base_url}/v1/images/ocr"
+        params = {
+            "model": self._model_uid,
+            "kwargs": json.dumps(kwargs),
+        }
+        files: List[Any] = []
+        for key, value in params.items():
+            files.append((key, (None, value)))
+        files.append(("image", ("image", image, "application/octet-stream")))
+        response = requests.post(url, files=files, headers=self.auth_headers)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to ocr the images, detail: {_get_error_string(response)}"
+            )
+
+        response_data = response.json()
+        return response_data
+
 
 class RESTfulVideoModelHandle(RESTfulModelHandle):
     def text_to_video(

xinference/core/chat_interface.py

@@ -74,7 +74,11 @@ class GradioInterface:
         # Gradio initiates the queue during a startup event, but since the app has already been
         # started, that event will not run, so manually invoke the startup events.
         # See: https://github.com/gradio-app/gradio/issues/5228
-        interface.startup_events()
+        try:
+            interface.run_startup_events()
+        except AttributeError:
+            # compatibility
+            interface.startup_events()
         favicon_path = os.path.join(
             os.path.dirname(os.path.abspath(__file__)),
             os.path.pardir,

xinference/core/image_interface.py

@@ -63,7 +63,11 @@ class ImageInterface:
         # Gradio initiates the queue during a startup event, but since the app has already been
         # started, that event will not run, so manually invoke the startup events.
         # See: https://github.com/gradio-app/gradio/issues/5228
-        interface.startup_events()
+        try:
+            interface.run_startup_events()
+        except AttributeError:
+            # compatibility
+            interface.startup_events()
         favicon_path = os.path.join(
             os.path.dirname(os.path.abspath(__file__)),
             os.path.pardir,

xinference/core/model.py

@@ -17,10 +17,10 @@ import functools
 import inspect
 import json
 import os
+import queue
 import time
 import types
 import uuid
-import weakref
 from asyncio.queues import Queue
 from asyncio.tasks import wait_for
 from concurrent.futures import Future as ConcurrentFuture
@@ -32,7 +32,6 @@ from typing import (
     Callable,
     Dict,
     Generator,
-    Iterator,
     List,
     Optional,
     Union,
@@ -209,9 +208,8 @@ class ModelActor(xo.StatelessActor):
             model_description.to_dict() if model_description else {}
         )
         self._request_limits = request_limits
-
-        self._generators: Dict[str, Union[Iterator, AsyncGenerator]] = {}
-        self._current_generator = lambda: None
+        self._pending_requests: asyncio.Queue = asyncio.Queue()
+        self._handle_pending_requests_task = None
         self._lock = (
             None
             if isinstance(
@@ -237,6 +235,10 @@ class ModelActor(xo.StatelessActor):
     async def __post_create__(self):
         self._loop = asyncio.get_running_loop()
 
+        self._handle_pending_requests_task = asyncio.create_task(
+            self._handle_pending_requests()
+        )
+
         if self.allow_batching():
             from .scheduler import SchedulerActor
 
@@ -474,6 +476,43 @@ class ModelActor(xo.StatelessActor):
                 )
             await asyncio.gather(*coros)
 
+    async def _handle_pending_requests(self):
+        logger.info("Start requests handler.")
+        while True:
+            gen, stream_out, stop = await self._pending_requests.get()
+
+            async def _async_wrapper(_gen):
+                try:
+                    # anext is only available for Python >= 3.10
+                    return await _gen.__anext__()  # noqa: F821
+                except StopAsyncIteration:
+                    return stop
+
+            def _wrapper(_gen):
+                # Avoid issue: https://github.com/python/cpython/issues/112182
+                try:
+                    return next(_gen)
+                except StopIteration:
+                    return stop
+
+            while True:
+                try:
+                    if inspect.isgenerator(gen):
+                        r = await asyncio.to_thread(_wrapper, gen)
+                    elif inspect.isasyncgen(gen):
+                        r = await _async_wrapper(gen)
+                    else:
+                        raise Exception(
+                            f"The generator {gen} should be a generator or an async generator, "
+                            f"but a {type(gen)} is got."
+                        )
+                    stream_out.put_nowait(r)
+                    if r is not stop:
+                        continue
+                except Exception:
+                    logger.exception("stream encountered an error.")
+                break
+
     async def _call_wrapper_json(self, fn: Callable, *args, **kwargs):
         return await self._call_wrapper("json", fn, *args, **kwargs)
 
@@ -487,6 +526,13 @@ class ModelActor(xo.StatelessActor):
                 ret = await fn(*args, **kwargs)
             else:
                 ret = await asyncio.to_thread(fn, *args, **kwargs)
+
+            if inspect.isgenerator(ret):
+                gen = self._to_generator(output_type, ret)
+                return gen
+            if inspect.isasyncgen(ret):
+                gen = self._to_async_gen(output_type, ret)
+                return gen
         else:
             async with self._lock:
                 if inspect.iscoroutinefunction(fn):
@@ -494,17 +540,40 @@ class ModelActor(xo.StatelessActor):
                 else:
                     ret = await asyncio.to_thread(fn, *args, **kwargs)
 
-        if self._lock is not None and self._current_generator():
-            raise Exception("Parallel generation is not supported by llama-cpp-python.")
+                stream_out: Union[queue.Queue, asyncio.Queue]
+
+                if inspect.isgenerator(ret):
+                    gen = self._to_generator(output_type, ret)
+                    stream_out = queue.Queue()
+                    stop = object()
+                    self._pending_requests.put_nowait((gen, stream_out, stop))
+
+                    def _stream_out_generator():
+                        while True:
+                            o = stream_out.get()
+                            if o is stop:
+                                break
+                            else:
+                                yield o
+
+                    return _stream_out_generator()
+
+                if inspect.isasyncgen(ret):
+                    gen = self._to_async_gen(output_type, ret)
+                    stream_out = asyncio.Queue()
+                    stop = object()
+                    self._pending_requests.put_nowait((gen, stream_out, stop))
+
+                    async def _stream_out_async_gen():
+                        while True:
+                            o = await stream_out.get()
+                            if o is stop:
+                                break
+                            else:
+                                yield o
+
+                    return _stream_out_async_gen()
 
-        if inspect.isgenerator(ret):
-            gen = self._to_generator(output_type, ret)
-            self._current_generator = weakref.ref(gen)
-            return gen
-        if inspect.isasyncgen(ret):
-            gen = self._to_async_gen(output_type, ret)
-            self._current_generator = weakref.ref(gen)
-            return gen
         if output_type == "json":
             return await asyncio.to_thread(json_dumps, ret)
         else:
@@ -592,7 +661,6 @@ class ModelActor(xo.StatelessActor):
                 prompt_or_messages, queue, call_ability, *args, **kwargs
             )
             gen = self._to_async_gen("json", ret)
-            self._current_generator = weakref.ref(gen)
             return gen
         else:
             from .scheduler import XINFERENCE_NON_STREAMING_ABORT_FLAG
@@ -953,6 +1021,25 @@ class ModelActor(xo.StatelessActor):
             f"Model {self._model.model_spec} is not for creating image."
         )
 
+    @log_async(
+        logger=logger,
+        ignore_kwargs=["image"],
+    )
+    async def ocr(
+        self,
+        image: "PIL.Image",
+        *args,
+        **kwargs,
+    ):
+        if hasattr(self._model, "ocr"):
+            return await self._call_wrapper_json(
+                self._model.ocr,
+                image,
+                *args,
+                **kwargs,
+            )
+        raise AttributeError(f"Model {self._model.model_spec} is not for ocr.")
+
     @request_limit
     @log_async(logger=logger, ignore_kwargs=["image"])
     async def infer(
@@ -994,3 +1081,6 @@ class ModelActor(xo.StatelessActor):
     async def record_metrics(self, name, op, kwargs):
         worker_ref = await self._get_worker_ref()
         await worker_ref.record_metrics(name, op, kwargs)
+
+    async def get_pending_requests_count(self):
+        return self._pending_requests.qsize()

xinference/core/scheduler.py

@@ -79,7 +79,7 @@ class InferenceRequest:
         # For tool call
         self.tools = None
         # Currently, for storing tool call streaming results.
-        self.outputs: List[str] = []
+        self.outputs: List[str] = []  # type: ignore
         # inference results,
         # it is a list type because when stream=True,
         # self.completion contains all the results in a decode round.

xinference/deploy/supervisor.py

@@ -31,10 +31,6 @@ from .utils import health_check
 
 logger = logging.getLogger(__name__)
 
-from ..model import _install as install_model
-
-install_model()
-
 
 async def _start_supervisor(address: str, logging_conf: Optional[Dict] = None):
     logging.config.dictConfig(logging_conf)  # type: ignore

xinference/model/audio/chattts.py

@@ -54,7 +54,11 @@ class ChatTTSModel:
         torch.set_float32_matmul_precision("high")
         self._model = ChatTTS.Chat()
         logger.info("Load ChatTTS model with kwargs: %s", self._kwargs)
-        self._model.load(source="custom", custom_path=self._model_path, **self._kwargs)
+        ok = self._model.load(
+            source="custom", custom_path=self._model_path, **self._kwargs
+        )
+        if not ok:
+            raise Exception(f"The ChatTTS model is not correct: {self._model_path}")
 
     def speech(
         self,
@@ -114,16 +118,15 @@ class ChatTTSModel:
                     last_pos = 0
                     with writer.open():
                         for it in iter:
-                            for itt in it:
-                                for chunk in itt:
-                                    chunk = np.array([chunk]).transpose()
-                                    writer.write_audio_chunk(i, torch.from_numpy(chunk))
-                                    new_last_pos = out.tell()
-                                    if new_last_pos != last_pos:
-                                        out.seek(last_pos)
-                                        encoded_bytes = out.read()
-                                        yield encoded_bytes
-                                        last_pos = new_last_pos
+                            for chunk in it:
+                                chunk = np.array([chunk]).transpose()
+                                writer.write_audio_chunk(i, torch.from_numpy(chunk))
+                                new_last_pos = out.tell()
+                                if new_last_pos != last_pos:
+                                    out.seek(last_pos)
+                                    encoded_bytes = out.read()
+                                    yield encoded_bytes
+                                    last_pos = new_last_pos
 
             return _generator()
         else:
@@ -131,7 +134,15 @@ class ChatTTSModel:
 
             # Save the generated audio
             with BytesIO() as out:
-                torchaudio.save(
-                    out, torch.from_numpy(wavs[0]), 24000, format=response_format
-                )
+                try:
+                    torchaudio.save(
+                        out,
+                        torch.from_numpy(wavs[0]).unsqueeze(0),
+                        24000,
+                        format=response_format,
+                    )
+                except:
+                    torchaudio.save(
+                        out, torch.from_numpy(wavs[0]), 24000, format=response_format
+                    )
                 return out.getvalue()

xinference/model/audio/model_spec.json

@@ -127,7 +127,7 @@
     "model_name": "ChatTTS",
     "model_family": "ChatTTS",
     "model_id": "2Noise/ChatTTS",
-    "model_revision": "ce5913842aebd78e4a01a02d47244b8d62ac4ee3",
+    "model_revision": "3b34118f6d25850440b8901cef3e71c6ef8619c8",
     "model_ability": "text-to-audio",
     "multilingual": true
   },

xinference/model/audio/model_spec_modelscope.json

@@ -42,7 +42,7 @@
     "model_name": "ChatTTS",
     "model_family": "ChatTTS",
     "model_hub": "modelscope",
-    "model_id": "pzc163/chatTTS",
+    "model_id": "AI-ModelScope/ChatTTS",
     "model_revision": "master",
     "model_ability": "text-to-audio",
     "multilingual": true

xinference/model/embedding/model_spec.json

@@ -233,7 +233,7 @@
   },
   {
     "model_name": "gte-Qwen2",
-    "dimensions": 3584,
+    "dimensions": 4096,
     "max_tokens": 32000,
     "language": ["zh", "en"],
     "model_id": "Alibaba-NLP/gte-Qwen2-7B-instruct",

xinference/model/image/core.py

@@ -11,17 +11,21 @@
 # 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 collections.abc
 import logging
 import os
+import platform
 from collections import defaultdict
-from typing import Dict, List, Literal, Optional, Tuple
+from typing import Dict, List, Literal, Optional, Tuple, Union
 
 from ...constants import XINFERENCE_CACHE_DIR
 from ...types import PeftModelConfig
 from ..core import CacheableModelSpec, ModelDescription
 from ..utils import valid_model_revision
+from .ocr.got_ocr2 import GotOCR2Model
 from .stable_diffusion.core import DiffusionModel
+from .stable_diffusion.mlx import MLXDiffusionModel
 
 logger = logging.getLogger(__name__)
 
@@ -45,6 +49,7 @@ class ImageModelFamilyV1(CacheableModelSpec):
     model_hub: str = "huggingface"
     model_ability: Optional[List[str]]
     controlnet: Optional[List["ImageModelFamilyV1"]]
+    default_model_config: Optional[dict] = {}
     default_generate_config: Optional[dict] = {}
 
 
@@ -180,6 +185,28 @@ def get_cache_status(
         return valid_model_revision(meta_path, model_spec.model_revision)
 
 
+def create_ocr_model_instance(
+    subpool_addr: str,
+    devices: List[str],
+    model_uid: str,
+    model_spec: ImageModelFamilyV1,
+    model_path: Optional[str] = None,
+    **kwargs,
+) -> Tuple[GotOCR2Model, ImageModelDescription]:
+    if not model_path:
+        model_path = cache(model_spec)
+    model = GotOCR2Model(
+        model_uid,
+        model_path,
+        model_spec=model_spec,
+        **kwargs,
+    )
+    model_description = ImageModelDescription(
+        subpool_addr, devices, model_spec, model_path=model_path
+    )
+    return model, model_description
+
+
 def create_image_model_instance(
     subpool_addr: str,
     devices: List[str],
@@ -189,8 +216,26 @@ def create_image_model_instance(
     download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
     model_path: Optional[str] = None,
     **kwargs,
-) -> Tuple[DiffusionModel, ImageModelDescription]:
+) -> Tuple[
+    Union[DiffusionModel, MLXDiffusionModel, GotOCR2Model], ImageModelDescription
+]:
     model_spec = match_diffusion(model_name, download_hub)
+    if model_spec.model_ability and "ocr" in model_spec.model_ability:
+        return create_ocr_model_instance(
+            subpool_addr=subpool_addr,
+            devices=devices,
+            model_uid=model_uid,
+            model_name=model_name,
+            model_spec=model_spec,
+            model_path=model_path,
+            **kwargs,
+        )
+
+    # use default model config
+    model_default_config = (model_spec.default_model_config or {}).copy()
+    model_default_config.update(kwargs)
+    kwargs = model_default_config
+
     controlnet = kwargs.get("controlnet")
     # Handle controlnet
     if controlnet is not None:
@@ -232,10 +277,20 @@ def create_image_model_instance(
         lora_load_kwargs = None
         lora_fuse_kwargs = None
 
-    model = DiffusionModel(
+    if (
+        platform.system() == "Darwin"
+        and "arm" in platform.machine().lower()
+        and model_name in MLXDiffusionModel.supported_models
+    ):
+        # Mac with M series silicon chips
+        model_cls = MLXDiffusionModel
+    else:
+        model_cls = DiffusionModel  # type: ignore
+
+    model = model_cls(
         model_uid,
         model_path,
-        lora_model_paths=lora_model,
+        lora_model=lora_model,
         lora_load_kwargs=lora_load_kwargs,
         lora_fuse_kwargs=lora_fuse_kwargs,
         model_spec=model_spec,

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.