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

xinference/__init__.py

@@ -26,13 +26,9 @@ except:
 def _install():
     from xoscar.backends.router import Router
 
-    from .model import _install as install_model
-
     default_router = Router.get_instance_or_empty()
     Router.set_instance(default_router)
 
-    install_model()
-
 
 _install()
 del _install

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-10-12T18:28:41+0800",
+ "date": "2024-10-25T12:51:06+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "c0be11504c70f6c392cbdb67c86cf12153353f70",
- "version": "0.15.4"
+ "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/constants.py

@@ -27,8 +27,8 @@ XINFERENCE_ENV_HEALTH_CHECK_INTERVAL = "XINFERENCE_HEALTH_CHECK_INTERVAL"
 XINFERENCE_ENV_HEALTH_CHECK_TIMEOUT = "XINFERENCE_HEALTH_CHECK_TIMEOUT"
 XINFERENCE_ENV_DISABLE_HEALTH_CHECK = "XINFERENCE_DISABLE_HEALTH_CHECK"
 XINFERENCE_ENV_DISABLE_METRICS = "XINFERENCE_DISABLE_METRICS"
-XINFERENCE_ENV_TRANSFORMERS_ENABLE_BATCHING = "XINFERENCE_TRANSFORMERS_ENABLE_BATCHING"
 XINFERENCE_ENV_DOWNLOAD_MAX_ATTEMPTS = "XINFERENCE_DOWNLOAD_MAX_ATTEMPTS"
+XINFERENCE_ENV_TEXT_TO_IMAGE_BATCHING_SIZE = "XINFERENCE_TEXT_TO_IMAGE_BATCHING_SIZE"
 
 
 def get_xinference_home() -> str:
@@ -80,9 +80,9 @@ XINFERENCE_DISABLE_HEALTH_CHECK = bool(
 XINFERENCE_DISABLE_METRICS = bool(
     int(os.environ.get(XINFERENCE_ENV_DISABLE_METRICS, 0))
 )
-XINFERENCE_TRANSFORMERS_ENABLE_BATCHING = bool(
-    int(os.environ.get(XINFERENCE_ENV_TRANSFORMERS_ENABLE_BATCHING, 0))
-)
 XINFERENCE_DOWNLOAD_MAX_ATTEMPTS = int(
     os.environ.get(XINFERENCE_ENV_DOWNLOAD_MAX_ATTEMPTS, 3)
 )
+XINFERENCE_TEXT_TO_IMAGE_BATCHING_SIZE = os.environ.get(
+    XINFERENCE_ENV_TEXT_TO_IMAGE_BATCHING_SIZE, None
+)

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,
@@ -41,7 +40,7 @@ from typing import (
 import sse_starlette.sse
 import xoscar as xo
 
-from ..constants import XINFERENCE_TRANSFORMERS_ENABLE_BATCHING
+from ..constants import XINFERENCE_TEXT_TO_IMAGE_BATCHING_SIZE
 
 if TYPE_CHECKING:
     from .progress_tracker import ProgressTrackerActor
@@ -74,6 +73,8 @@ XINFERENCE_BATCHING_ALLOWED_VISION_MODELS = [
     "MiniCPM-V-2.6",
 ]
 
+XINFERENCE_TEXT_TO_IMAGE_BATCHING_ALLOWED_MODELS = ["FLUX.1-dev", "FLUX.1-schnell"]
+
 
 def request_limit(fn):
     """
@@ -153,6 +154,16 @@ class ModelActor(xo.StatelessActor):
                     f"Destroy scheduler actor failed, address: {self.address}, error: {e}"
                 )
 
+        if self.allow_batching_for_text_to_image():
+            try:
+                assert self._text_to_image_scheduler_ref is not None
+                await xo.destroy_actor(self._text_to_image_scheduler_ref)
+                del self._text_to_image_scheduler_ref
+            except Exception as e:
+                logger.debug(
+                    f"Destroy text_to_image scheduler actor failed, address: {self.address}, error: {e}"
+                )
+
         if hasattr(self._model, "stop") and callable(self._model.stop):
             self._model.stop()
 
@@ -197,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(
@@ -220,10 +230,15 @@ class ModelActor(xo.StatelessActor):
         self._loop: Optional[asyncio.AbstractEventLoop] = None
 
         self._scheduler_ref = None
+        self._text_to_image_scheduler_ref = None
 
     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
 
@@ -233,6 +248,15 @@ class ModelActor(xo.StatelessActor):
                 uid=SchedulerActor.gen_uid(self.model_uid(), self._model.rep_id),
             )
 
+        if self.allow_batching_for_text_to_image():
+            from ..model.image.scheduler.flux import FluxBatchSchedulerActor
+
+            self._text_to_image_scheduler_ref = await xo.create_actor(
+                FluxBatchSchedulerActor,
+                address=self.address,
+                uid=FluxBatchSchedulerActor.gen_uid(self.model_uid()),
+            )
+
     async def _record_completion_metrics(
         self, duration, completion_tokens, prompt_tokens
     ):
@@ -311,10 +335,8 @@ class ModelActor(xo.StatelessActor):
 
         model_ability = self._model_description.get("model_ability", [])
 
-        condition = XINFERENCE_TRANSFORMERS_ENABLE_BATCHING and isinstance(
-            self._model, PytorchModel
-        )
-        if condition and "vision" in model_ability:
+        condition = isinstance(self._model, PytorchModel)
+        if condition and ("vision" in model_ability or "audio" in model_ability):
             if (
                 self._model.model_family.model_name
                 in XINFERENCE_BATCHING_ALLOWED_VISION_MODELS
@@ -331,6 +353,26 @@ class ModelActor(xo.StatelessActor):
                 return False
         return condition
 
+    def allow_batching_for_text_to_image(self) -> bool:
+        from ..model.image.stable_diffusion.core import DiffusionModel
+
+        condition = XINFERENCE_TEXT_TO_IMAGE_BATCHING_SIZE is not None and isinstance(
+            self._model, DiffusionModel
+        )
+
+        if condition:
+            model_name = self._model._model_spec.model_name  # type: ignore
+            if model_name in XINFERENCE_TEXT_TO_IMAGE_BATCHING_ALLOWED_MODELS:
+                return True
+            else:
+                logger.warning(
+                    f"Currently for image models with text_to_image ability, "
+                    f"xinference only supports {', '.join(XINFERENCE_TEXT_TO_IMAGE_BATCHING_ALLOWED_MODELS)} for batching. "
+                    f"Your model {model_name} is disqualified."
+                )
+                return False
+        return condition
+
     async def load(self):
         self._model.load()
         if self.allow_batching():
@@ -338,6 +380,11 @@ class ModelActor(xo.StatelessActor):
             logger.debug(
                 f"Batching enabled for model: {self.model_uid()}, max_num_seqs: {self._model.get_max_num_seqs()}"
             )
+        if self.allow_batching_for_text_to_image():
+            await self._text_to_image_scheduler_ref.set_model(self._model)
+            logger.debug(
+                f"Batching enabled for model: {self.model_uid()}, max_num_images: {self._model.get_max_num_images_for_batching()}"
+            )
 
     def model_uid(self):
         return (
@@ -429,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)
 
@@ -442,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):
@@ -449,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:
@@ -547,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
@@ -617,12 +730,16 @@ class ModelActor(xo.StatelessActor):
                 )
 
     async def abort_request(self, request_id: str) -> str:
-        from .scheduler import AbortRequestMessage
+        from .utils import AbortRequestMessage
 
         if self.allow_batching():
             if self._scheduler_ref is None:
                 return AbortRequestMessage.NOT_FOUND.name
             return await self._scheduler_ref.abort_request(request_id)
+        elif self.allow_batching_for_text_to_image():
+            if self._text_to_image_scheduler_ref is None:
+                return AbortRequestMessage.NOT_FOUND.name
+            return await self._text_to_image_scheduler_ref.abort_request(request_id)
         return AbortRequestMessage.NO_OP.name
 
     @request_limit
@@ -747,6 +864,22 @@ class ModelActor(xo.StatelessActor):
             f"Model {self._model.model_spec} is not for creating speech."
         )
 
+    async def handle_image_batching_request(self, unique_id, *args, **kwargs):
+        size = args[2]
+        if XINFERENCE_TEXT_TO_IMAGE_BATCHING_SIZE != size:
+            raise RuntimeError(
+                f"The image size: {size} of text_to_image for batching "
+                f"must be the same as the environment variable: {XINFERENCE_TEXT_TO_IMAGE_BATCHING_SIZE} you set."
+            )
+        assert self._loop is not None
+        future = ConcurrentFuture()
+        await self._text_to_image_scheduler_ref.add_request(
+            unique_id, future, *args, **kwargs
+        )
+        fut = asyncio.wrap_future(future, loop=self._loop)
+        result = await fut
+        return await asyncio.to_thread(json_dumps, result)
+
     @request_limit
     @log_async(logger=logger)
     async def text_to_image(
@@ -759,19 +892,25 @@ class ModelActor(xo.StatelessActor):
         **kwargs,
     ):
         if hasattr(self._model, "text_to_image"):
-            progressor = kwargs["progressor"] = await self._get_progressor(
-                kwargs.pop("request_id", None)
-            )
-            with progressor:
-                return await self._call_wrapper_json(
-                    self._model.text_to_image,
-                    prompt,
-                    n,
-                    size,
-                    response_format,
-                    *args,
-                    **kwargs,
+            if self.allow_batching_for_text_to_image():
+                unique_id = kwargs.pop("request_id", None)
+                return await self.handle_image_batching_request(
+                    unique_id, prompt, n, size, response_format, *args, **kwargs
+                )
+            else:
+                progressor = kwargs["progressor"] = await self._get_progressor(
+                    kwargs.pop("request_id", None)
                 )
+                with progressor:
+                    return await self._call_wrapper_json(
+                        self._model.text_to_image,
+                        prompt,
+                        n,
+                        size,
+                        response_format,
+                        *args,
+                        **kwargs,
+                    )
         raise AttributeError(
             f"Model {self._model.model_spec} is not for creating image."
         )
@@ -882,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(
@@ -923,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

@@ -17,11 +17,12 @@ import functools
 import logging
 import uuid
 from collections import deque
-from enum import Enum
 from typing import Dict, List, Optional, Set, Tuple, Union
 
 import xoscar as xo
 
+from .utils import AbortRequestMessage
+
 logger = logging.getLogger(__name__)
 
 XINFERENCE_STREAMING_DONE_FLAG = "<XINFERENCE_STREAMING_DONE>"
@@ -30,12 +31,6 @@ XINFERENCE_STREAMING_ABORT_FLAG = "<XINFERENCE_STREAMING_ABORT>"
 XINFERENCE_NON_STREAMING_ABORT_FLAG = "<XINFERENCE_NON_STREAMING_ABORT>"
 
 
-class AbortRequestMessage(Enum):
-    NOT_FOUND = 1
-    DONE = 2
-    NO_OP = 3
-
-
 class InferenceRequest:
     def __init__(
         self,
@@ -81,6 +76,10 @@ class InferenceRequest:
         self.padding_len = 0
         # Use in stream mode
         self.last_output_length = 0
+        # For tool call
+        self.tools = None
+        # Currently, for storing tool call streaming results.
+        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.
@@ -112,6 +111,10 @@ class InferenceRequest:
         """
         return self._prompt
 
+    @prompt.setter
+    def prompt(self, value: str):
+        self._prompt = value
+
     @property
     def call_ability(self):
         return self._call_ability

xinference/core/utils.py

@@ -16,6 +16,7 @@ import os
 import random
 import string
 import uuid
+from enum import Enum
 from typing import Dict, Generator, List, Optional, Tuple, Union
 
 import orjson
@@ -27,6 +28,12 @@ from ..constants import XINFERENCE_LOG_ARG_MAX_LENGTH
 logger = logging.getLogger(__name__)
 
 
+class AbortRequestMessage(Enum):
+    NOT_FOUND = 1
+    DONE = 2
+    NO_OP = 3
+
+
 def truncate_log_arg(arg) -> str:
     s = str(arg)
     if len(s) > XINFERENCE_LOG_ARG_MAX_LENGTH:
@@ -51,6 +58,8 @@ def log_async(
             request_id_str = kwargs.get("request_id", "")
             if not request_id_str:
                 request_id_str = uuid.uuid1()
+                if func_name == "text_to_image":
+                    kwargs["request_id"] = request_id_str
             request_id_str = f"[request {request_id_str}]"
             formatted_args = ",".join(map(truncate_log_arg, args))
             formatted_kwargs = ",".join(

xinference/model/__init__.py

@@ -29,3 +29,7 @@ def _install():
     image_install()
     rerank_install()
     video_install()
+
+
+_install()
+del _install