xinference 0.13.0__py3-none-any.whl → 0.13.2__py3-none-any.whl

xinference/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-07-05T18:19:09+0800",
+ "date": "2024-07-19T19:15:54+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "007408c55272bc343821dd152df780de5dc9c037",
- "version": "0.13.0"
+ "full-revisionid": "880929cbbc73e5206ca069591b03d9d16dd858bf",
+ "version": "0.13.2"
 }
 '''  # END VERSION_JSON
 

xinference/api/restful_api.py

@@ -129,10 +129,12 @@ class SpeechRequest(BaseModel):
     voice: Optional[str]
     response_format: Optional[str] = "mp3"
     speed: Optional[float] = 1.0
+    stream: Optional[bool] = False
 
 
 class RegisterModelRequest(BaseModel):
     model: str
+    worker_ip: Optional[str]
     persist: bool
 
 
@@ -490,6 +492,17 @@ class RESTfulAPI:
                 else None
             ),
         )
+        self._router.add_api_route(
+            "/v1/images/inpainting",
+            self.create_inpainting,
+            methods=["POST"],
+            response_model=ImageList,
+            dependencies=(
+                [Security(self._auth_service, scopes=["models:read"])]
+                if self.is_authenticated()
+                else None
+            ),
+        )
         self._router.add_api_route(
             "/v1/chat/completions",
             self.create_chat_completion,
@@ -501,6 +514,16 @@ class RESTfulAPI:
                 else None
             ),
         )
+        self._router.add_api_route(
+            "/v1/flexible/infers",
+            self.create_flexible_infer,
+            methods=["POST"],
+            dependencies=(
+                [Security(self._auth_service, scopes=["models:read"])]
+                if self.is_authenticated()
+                else None
+            ),
+        )
 
         # for custom models
         self._router.add_api_route(
@@ -772,6 +795,7 @@ class RESTfulAPI:
         peft_model_config = payload.get("peft_model_config", None)
         worker_ip = payload.get("worker_ip", None)
         gpu_idx = payload.get("gpu_idx", None)
+        download_hub = payload.get("download_hub", None)
 
         exclude_keys = {
             "model_uid",
@@ -787,6 +811,7 @@ class RESTfulAPI:
             "peft_model_config",
             "worker_ip",
             "gpu_idx",
+            "download_hub",
         }
 
         kwargs = {
@@ -834,9 +859,9 @@ class RESTfulAPI:
                 peft_model_config=peft_model_config,
                 worker_ip=worker_ip,
                 gpu_idx=gpu_idx,
+                download_hub=download_hub,
                 **kwargs,
             )
-
         except ValueError as ve:
             logger.error(str(ve), exc_info=True)
             raise HTTPException(status_code=400, detail=str(ve))
@@ -1304,8 +1329,14 @@ class RESTfulAPI:
                 voice=body.voice,
                 response_format=body.response_format,
                 speed=body.speed,
+                stream=body.stream,
             )
-            return Response(media_type="application/octet-stream", content=out)
+            if body.stream:
+                return EventSourceResponse(
+                    media_type="application/octet-stream", content=out
+                )
+            else:
+                return Response(media_type="application/octet-stream", content=out)
         except RuntimeError as re:
             logger.error(re, exc_info=True)
             await self._report_error_event(model_uid, str(re))
@@ -1397,6 +1428,94 @@ class RESTfulAPI:
             await self._report_error_event(model_uid, str(e))
             raise HTTPException(status_code=500, detail=str(e))
 
+    async def create_inpainting(
+        self,
+        model: str = Form(...),
+        image: UploadFile = File(media_type="application/octet-stream"),
+        mask_image: UploadFile = File(media_type="application/octet-stream"),
+        prompt: Optional[Union[str, List[str]]] = Form(None),
+        negative_prompt: Optional[Union[str, List[str]]] = Form(None),
+        n: Optional[int] = Form(1),
+        response_format: Optional[str] = Form("url"),
+        size: Optional[str] = Form(None),
+        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)
+            mask_im = Image.open(mask_image.file)
+            if not size:
+                w, h = im.size
+                size = f"{w}*{h}"
+            image_list = await model_ref.inpainting(
+                image=im,
+                mask_image=mask_im,
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                n=n,
+                size=size,
+                response_format=response_format,
+                **parsed_kwargs,
+            )
+            return Response(content=image_list, media_type="application/json")
+        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()
+
+        model_uid = payload.get("model")
+
+        exclude = {
+            "model",
+        }
+        kwargs = {key: value for key, value in payload.items() if key not in exclude}
+
+        try:
+            model = 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:
+            result = await model.infer(**kwargs)
+            return Response(result, media_type="application/json")
+        except RuntimeError as re:
+            logger.error(re, exc_info=True)
+            await self._report_error_event(model_uid, str(re))
+            self.handle_request_limit_error(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_chat_completion(self, request: Request) -> Response:
         raw_body = await request.json()
         body = CreateChatCompletion.parse_obj(raw_body)
@@ -1593,11 +1712,12 @@ class RESTfulAPI:
     async def register_model(self, model_type: str, request: Request) -> JSONResponse:
         body = RegisterModelRequest.parse_obj(await request.json())
         model = body.model
+        worker_ip = body.worker_ip
         persist = body.persist
 
         try:
             await (await self._get_supervisor_ref()).register_model(
-                model_type, model, persist
+                model_type, model, persist, worker_ip
             )
         except ValueError as re:
             logger.error(re, exc_info=True)

xinference/client/restful/restful_client.py

@@ -294,6 +294,81 @@ class RESTfulImageModelHandle(RESTfulModelHandle):
         response_data = response.json()
         return response_data
 
+    def inpainting(
+        self,
+        image: Union[str, bytes],
+        mask_image: Union[str, bytes],
+        prompt: str,
+        negative_prompt: Optional[str] = None,
+        n: int = 1,
+        size: Optional[str] = None,
+        response_format: str = "url",
+        **kwargs,
+    ) -> "ImageList":
+        """
+        Inpaint an image by the input text.
+
+        Parameters
+        ----------
+        image: `Union[str, bytes]`
+            an image batch to be inpainted (which parts of the image to
+            be masked out with `mask_image` and repainted according to `prompt`). For both numpy array and pytorch
+            tensor, the expected value range is between `[0, 1]` If it's a tensor or a list or tensors, the
+            expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a list of arrays, the
+            expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image latents as `image`, but
+            if passing latents directly it is not encoded again.
+        mask_image: `Union[str, bytes]`
+            representing an image batch to mask `image`. White pixels in the mask
+            are repainted while black pixels are preserved. If `mask_image` is a PIL image, it is converted to a
+            single channel (luminance) before use. If it's a numpy array or pytorch tensor, it should contain one
+            color channel (L) instead of 3, so the expected shape for pytorch tensor would be `(B, 1, H, W)`, `(B,
+            H, W)`, `(1, H, W)`, `(H, W)`. And for numpy array would be for `(B, H, W, 1)`, `(B, H, W)`, `(H, W,
+            1)`, or `(H, W)`.
+        prompt: `str` or `List[str]`
+            The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+        negative_prompt (`str` or `List[str]`, *optional*):
+            The prompt or prompts not to guide the image generation. If not defined, one has to pass
+            `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+            less than `1`).
+        n: `int`, defaults to 1
+            The number of images to generate per prompt. Must be between 1 and 10.
+        size: `str`, defaults to None
+            The width*height in pixels of the generated image.
+        response_format: `str`, defaults to `url`
+            The format in which the generated images are returned. Must be one of url or b64_json.
+        Returns
+        -------
+        ImageList
+            A list of image objects.
+            :param prompt:
+            :param image:
+        """
+        url = f"{self._base_url}/v1/images/inpainting"
+        params = {
+            "model": self._model_uid,
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "n": n,
+            "size": size,
+            "response_format": response_format,
+            "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")))
+        files.append(
+            ("mask_image", ("mask_image", mask_image, "application/octet-stream"))
+        )
+        response = requests.post(url, files=files, headers=self.auth_headers)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to inpaint the images, detail: {_get_error_string(response)}"
+            )
+
+        response_data = response.json()
+        return response_data
+
 
 class RESTfulGenerateModelHandle(RESTfulModelHandle):
     def generate(
@@ -692,6 +767,7 @@ class RESTfulAudioModelHandle(RESTfulModelHandle):
         voice: str = "",
         response_format: str = "mp3",
         speed: float = 1.0,
+        stream: bool = False,
     ):
         """
         Generates audio from the input text.
@@ -707,6 +783,8 @@ class RESTfulAudioModelHandle(RESTfulModelHandle):
             The format to audio in.
         speed: str
             The speed of the generated audio.
+        stream: bool
+            Use stream or not.
 
         Returns
         -------
@@ -720,6 +798,7 @@ class RESTfulAudioModelHandle(RESTfulModelHandle):
             "voice": voice,
             "response_format": response_format,
             "speed": speed,
+            "stream": stream,
         }
         response = requests.post(url, json=params, headers=self.auth_headers)
         if response.status_code != 200:
@@ -727,6 +806,44 @@ class RESTfulAudioModelHandle(RESTfulModelHandle):
                 f"Failed to speech the text, detail: {_get_error_string(response)}"
             )
 
+        if stream:
+            return response.iter_content(chunk_size=1024)
+
+        return response.content
+
+
+class RESTfulFlexibleModelHandle(RESTfulModelHandle):
+    def infer(
+        self,
+        **kwargs,
+    ):
+        """
+        Call flexible model.
+
+        Parameters
+        ----------
+
+        kwargs: dict
+            The inference arguments.
+
+
+        Returns
+        -------
+        bytes
+            The inference result.
+        """
+        url = f"{self._base_url}/v1/flexible/infers"
+        params = {
+            "model": self._model_uid,
+        }
+        params.update(kwargs)
+
+        response = requests.post(url, json=params, headers=self.auth_headers)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to predict, detail: {_get_error_string(response)}"
+            )
+
         return response.content
 
 
@@ -1009,6 +1126,10 @@ class Client:
             return RESTfulAudioModelHandle(
                 model_uid, self.base_url, auth_headers=self._headers
             )
+        elif desc["model_type"] == "flexible":
+            return RESTfulFlexibleModelHandle(
+                model_uid, self.base_url, auth_headers=self._headers
+            )
         else:
             raise ValueError(f"Unknown model type:{desc['model_type']}")
 
@@ -1062,7 +1183,13 @@ class Client:
             )
         return response.json()
 
-    def register_model(self, model_type: str, model: str, persist: bool):
+    def register_model(
+        self,
+        model_type: str,
+        model: str,
+        persist: bool,
+        worker_ip: Optional[str] = None,
+    ):
         """
         Register a custom model.
 
@@ -1072,6 +1199,8 @@ class Client:
             The type of model.
         model: str
             The model definition. (refer to: https://inference.readthedocs.io/en/latest/models/custom.html)
+        worker_ip: Optional[str]
+            The IP address of the worker on which the model is running.
         persist: bool
 
 
@@ -1081,7 +1210,7 @@ class Client:
             Report failure to register the custom model. Provide details of failure through error message.
         """
         url = f"{self.base_url}/v1/model_registrations/{model_type}"
-        request_body = {"model": model, "persist": persist}
+        request_body = {"model": model, "worker_ip": worker_ip, "persist": persist}
         response = requests.post(url, json=request_body, headers=self._headers)
         if response.status_code != 200:
             raise RuntimeError(

xinference/core/model.py

@@ -310,7 +310,7 @@ class ModelActor(xo.StatelessActor):
             )
         )
 
-    def _to_json_generator(self, gen: types.GeneratorType):
+    def _to_generator(self, output_type: str, gen: types.GeneratorType):
         start_time = time.time()
         time_to_first_token = None
         final_usage = None
@@ -318,8 +318,13 @@ class ModelActor(xo.StatelessActor):
             for v in gen:
                 if time_to_first_token is None:
                     time_to_first_token = (time.time() - start_time) * 1000
-                final_usage = v.get("usage", None)
-                v = dict(data=json.dumps(v, ensure_ascii=False))
+                if output_type == "json":
+                    final_usage = v.get("usage", None)
+                    v = dict(data=json.dumps(v, ensure_ascii=False))
+                else:
+                    assert (
+                        output_type == "binary"
+                    ), f"Unknown output type '{output_type}'"
                 yield sse_starlette.sse.ensure_bytes(v, None)
         except OutOfMemoryError:
             logger.exception(
@@ -342,7 +347,7 @@ class ModelActor(xo.StatelessActor):
                 )
                 asyncio.run_coroutine_threadsafe(coro, loop=self._loop)
 
-    async def _to_json_async_gen(self, gen: types.AsyncGeneratorType):
+    async def _to_async_gen(self, output_type: str, gen: types.AsyncGeneratorType):
         start_time = time.time()
         time_to_first_token = None
         final_usage = None
@@ -351,8 +356,13 @@ class ModelActor(xo.StatelessActor):
                 if time_to_first_token is None:
                     time_to_first_token = (time.time() - start_time) * 1000
                 final_usage = v.get("usage", None)
-                v = await asyncio.to_thread(json.dumps, v)
-                v = dict(data=v)  # noqa: F821
+                if output_type == "json":
+                    v = await asyncio.to_thread(json.dumps, v, ensure_ascii=False)
+                    v = dict(data=v)  # noqa: F821
+                else:
+                    assert (
+                        output_type == "binary"
+                    ), f"Unknown output type '{output_type}'"
                 yield await asyncio.to_thread(sse_starlette.sse.ensure_bytes, v, None)
         except OutOfMemoryError:
             logger.exception(
@@ -379,8 +389,14 @@ class ModelActor(xo.StatelessActor):
                 )
             await asyncio.gather(*coros)
 
+    async def _call_wrapper_json(self, fn: Callable, *args, **kwargs):
+        return await self._call_wrapper("json", fn, *args, **kwargs)
+
+    async def _call_wrapper_binary(self, fn: Callable, *args, **kwargs):
+        return await self._call_wrapper("binary", fn, *args, **kwargs)
+
     @oom_check
-    async def _call_wrapper(self, fn: Callable, *args, **kwargs):
+    async def _call_wrapper(self, output_type: str, fn: Callable, *args, **kwargs):
         if self._lock is None:
             if inspect.iscoroutinefunction(fn):
                 ret = await fn(*args, **kwargs)
@@ -397,16 +413,18 @@ class ModelActor(xo.StatelessActor):
             raise Exception("Parallel generation is not supported by ggml.")
 
         if inspect.isgenerator(ret):
-            gen = self._to_json_generator(ret)
+            gen = self._to_generator(output_type, ret)
             self._current_generator = weakref.ref(gen)
             return gen
         if inspect.isasyncgen(ret):
-            gen = self._to_json_async_gen(ret)
+            gen = self._to_async_gen(output_type, ret)
             self._current_generator = weakref.ref(gen)
             return gen
-        if isinstance(ret, bytes):
+        if output_type == "json":
+            return await asyncio.to_thread(json_dumps, ret)
+        else:
+            assert output_type == "binary", f"Unknown output type '{output_type}'"
             return ret
-        return await asyncio.to_thread(json_dumps, ret)
 
     @log_async(logger=logger)
     @request_limit
@@ -419,11 +437,11 @@ class ModelActor(xo.StatelessActor):
         else:
             kwargs.pop("raw_params", None)
             if hasattr(self._model, "generate"):
-                return await self._call_wrapper(
+                return await self._call_wrapper_json(
                     self._model.generate, prompt, *args, **kwargs
                 )
             if hasattr(self._model, "async_generate"):
-                return await self._call_wrapper(
+                return await self._call_wrapper_json(
                     self._model.async_generate, prompt, *args, **kwargs
                 )
             raise AttributeError(f"Model {self._model.model_spec} is not for generate.")
@@ -471,7 +489,7 @@ class ModelActor(xo.StatelessActor):
             queue: Queue[Any] = Queue()
             ret = self._queue_consumer(queue)
             await self._scheduler_ref.add_request(prompt, queue, *args, **kwargs)
-            gen = self._to_json_async_gen(ret)
+            gen = self._to_async_gen("json", ret)
             self._current_generator = weakref.ref(gen)
             return gen
         else:
@@ -502,12 +520,12 @@ class ModelActor(xo.StatelessActor):
             else:
                 kwargs.pop("raw_params", None)
                 if hasattr(self._model, "chat"):
-                    response = await self._call_wrapper(
+                    response = await self._call_wrapper_json(
                         self._model.chat, prompt, *args, **kwargs
                     )
                     return response
                 if hasattr(self._model, "async_chat"):
-                    response = await self._call_wrapper(
+                    response = await self._call_wrapper_json(
                         self._model.async_chat, prompt, *args, **kwargs
                     )
                     return response
@@ -543,7 +561,7 @@ class ModelActor(xo.StatelessActor):
     @request_limit
     async def create_embedding(self, input: Union[str, List[str]], *args, **kwargs):
         if hasattr(self._model, "create_embedding"):
-            return await self._call_wrapper(
+            return await self._call_wrapper_json(
                 self._model.create_embedding, input, *args, **kwargs
             )
 
@@ -565,7 +583,7 @@ class ModelActor(xo.StatelessActor):
         **kwargs,
     ):
         if hasattr(self._model, "rerank"):
-            return await self._call_wrapper(
+            return await self._call_wrapper_json(
                 self._model.rerank,
                 documents,
                 query,
@@ -590,7 +608,7 @@ class ModelActor(xo.StatelessActor):
         timestamp_granularities: Optional[List[str]] = None,
     ):
         if hasattr(self._model, "transcriptions"):
-            return await self._call_wrapper(
+            return await self._call_wrapper_json(
                 self._model.transcriptions,
                 audio,
                 language,
@@ -615,7 +633,7 @@ class ModelActor(xo.StatelessActor):
         timestamp_granularities: Optional[List[str]] = None,
     ):
         if hasattr(self._model, "translations"):
-            return await self._call_wrapper(
+            return await self._call_wrapper_json(
                 self._model.translations,
                 audio,
                 language,
@@ -630,16 +648,23 @@ class ModelActor(xo.StatelessActor):
 
     @log_async(logger=logger)
     @request_limit
+    @xo.generator
     async def speech(
-        self, input: str, voice: str, response_format: str = "mp3", speed: float = 1.0
+        self,
+        input: str,
+        voice: str,
+        response_format: str = "mp3",
+        speed: float = 1.0,
+        stream: bool = False,
     ):
         if hasattr(self._model, "speech"):
-            return await self._call_wrapper(
+            return await self._call_wrapper_binary(
                 self._model.speech,
                 input,
                 voice,
                 response_format,
                 speed,
+                stream,
             )
         raise AttributeError(
             f"Model {self._model.model_spec} is not for creating speech."
@@ -657,7 +682,7 @@ class ModelActor(xo.StatelessActor):
         **kwargs,
     ):
         if hasattr(self._model, "text_to_image"):
-            return await self._call_wrapper(
+            return await self._call_wrapper_json(
                 self._model.text_to_image,
                 prompt,
                 n,
@@ -682,7 +707,7 @@ class ModelActor(xo.StatelessActor):
         **kwargs,
     ):
         if hasattr(self._model, "image_to_image"):
-            return await self._call_wrapper(
+            return await self._call_wrapper_json(
                 self._model.image_to_image,
                 image,
                 prompt,
@@ -697,6 +722,50 @@ class ModelActor(xo.StatelessActor):
             f"Model {self._model.model_spec} is not for creating image."
         )
 
+    async def inpainting(
+        self,
+        image: "PIL.Image",
+        mask_image: "PIL.Image",
+        prompt: str,
+        negative_prompt: str,
+        n: int = 1,
+        size: str = "1024*1024",
+        response_format: str = "url",
+        *args,
+        **kwargs,
+    ):
+        if hasattr(self._model, "inpainting"):
+            return await self._call_wrapper(
+                self._model.inpainting,
+                image,
+                mask_image,
+                prompt,
+                negative_prompt,
+                n,
+                size,
+                response_format,
+                *args,
+                **kwargs,
+            )
+        raise AttributeError(
+            f"Model {self._model.model_spec} is not for creating image."
+        )
+
+    @log_async(logger=logger)
+    @request_limit
+    async def infer(
+        self,
+        **kwargs,
+    ):
+        if hasattr(self._model, "infer"):
+            return await self._call_wrapper(
+                self._model.infer,
+                **kwargs,
+            )
+        raise AttributeError(
+            f"Model {self._model.model_spec} is not for flexible infer."
+        )
+
     async def record_metrics(self, name, op, kwargs):
         worker_ref = await self._get_worker_ref()
         await worker_ref.record_metrics(name, op, kwargs)