diff --git a/doc/source/models/builtin/image/index.rst b/doc/source/models/builtin/image/index.rst index e0c43fd708..4656c35407 100644 --- a/doc/source/models/builtin/image/index.rst +++ b/doc/source/models/builtin/image/index.rst @@ -26,4 +26,6 @@ The following is a list of built-in image models in Xinference: stable-diffusion-xl-base-1.0 stable-diffusion-xl-inpainting + + stable-diffusion-xl-inpainting \ No newline at end of file diff --git a/xinference/core/scheduler.py b/xinference/core/scheduler.py index da9d4bbe28..c8a0747e56 100644 --- a/xinference/core/scheduler.py +++ b/xinference/core/scheduler.py @@ -81,7 +81,7 @@ def __init__(self, prompt, future_or_queue, is_prefill, *args, **kwargs): self.future_or_queue = future_or_queue # Record error message when this request has error. # Must set stopped=True when this field is set. - self.error_msg: Optional[str] = None + self.error_msg: Optional[str] = None # type: ignore # For compatibility. Record some extra parameters for some special cases. self.extra_kwargs = {} @@ -295,11 +295,11 @@ def gen_uid(cls, model_uid: str, replica_id: str): def __init__(self): super().__init__() - self._waiting_queue: deque[InferenceRequest] = deque() - self._running_queue: deque[InferenceRequest] = deque() + self._waiting_queue: deque[InferenceRequest] = deque() # type: ignore + self._running_queue: deque[InferenceRequest] = deque() # type: ignore self._model = None self._id_to_req = {} - self._abort_req_ids: Set[str] = set() + self._abort_req_ids: Set[str] = set() # type: ignore self._isolation = None async def __post_create__(self): diff --git a/xinference/model/flexible/launchers/__init__.py b/xinference/model/flexible/launchers/__init__.py index 29559082d2..1dc7941c39 100644 --- a/xinference/model/flexible/launchers/__init__.py +++ b/xinference/model/flexible/launchers/__init__.py @@ -12,4 +12,5 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .image_process_launcher import launcher as image_process from .transformers_launcher import launcher as transformers diff --git a/xinference/model/flexible/launchers/image_process_launcher.py b/xinference/model/flexible/launchers/image_process_launcher.py new file mode 100644 index 0000000000..e1e3bece90 --- /dev/null +++ b/xinference/model/flexible/launchers/image_process_launcher.py @@ -0,0 +1,70 @@ +# Copyright 2022-2023 XProbe Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import base64 +from io import BytesIO + +import PIL.Image +import PIL.ImageOps + +from ....types import Image +from ..core import FlexibleModel, FlexibleModelSpec + + +class ImageRemoveBackgroundModel(FlexibleModel): + def infer(self, **kwargs): + invert = kwargs.get("invert", False) + b64_image: str = kwargs.get("image") # type: ignore + only_mask = kwargs.pop("only_mask", True) + image_format = kwargs.pop("image_format", "PNG") + if not b64_image: + raise ValueError("No image found to remove background") + image = base64.b64decode(b64_image) + + try: + from rembg import remove + except ImportError: + error_message = "Failed to import module 'rembg'" + installation_guide = [ + "Please make sure 'rembg' is installed. ", + "You can install it by visiting the installation section of the git repo:\n", + "https://github.com/danielgatis/rembg?tab=readme-ov-file#installation", + ] + + raise ImportError(f"{error_message}\n\n{''.join(installation_guide)}") + + im = PIL.Image.open(BytesIO(image)) + om = remove(im, only_mask=only_mask, **kwargs) + if invert: + om = PIL.ImageOps.invert(om) + + buffered = BytesIO() + om.save(buffered, format=image_format) + img_str = base64.b64encode(buffered.getvalue()).decode() + return Image(url=None, b64_json=img_str) + + +def launcher(model_uid: str, model_spec: FlexibleModelSpec, **kwargs) -> FlexibleModel: + task = kwargs.get("task") + device = kwargs.get("device") + + if task == "remove_background": + return ImageRemoveBackgroundModel( + model_uid=model_uid, + model_path=model_spec.model_uri, # type: ignore + device=device, + config=kwargs, + ) + else: + raise ValueError(f"Unknown Task for image processing: {task}") diff --git a/xinference/model/llm/pytorch/cogvlm2.py b/xinference/model/llm/pytorch/cogvlm2.py index 07f44ad47f..791ecedc98 100644 --- a/xinference/model/llm/pytorch/cogvlm2.py +++ b/xinference/model/llm/pytorch/cogvlm2.py @@ -387,7 +387,7 @@ def _get_full_prompt(self, prompt, system_prompt, chat_history, tools): prompt, system_prompt=system_prompt, chat_history=chat_history ) - input_by_model: dict = self._model.build_conversation_input_ids( + input_by_model: dict = self._model.build_conversation_input_ids( # type: ignore self._tokenizer, query=query, history=history,