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

xinference/model/image/scheduler/flux.py

@@ -0,0 +1,533 @@
+# Copyright 2022-2024 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 asyncio
+import logging
+import os
+import re
+import typing
+from collections import deque
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import xoscar as xo
+
+from ..utils import handle_image_result
+
+if TYPE_CHECKING:
+    from ..stable_diffusion.core import DiffusionModel
+
+
+logger = logging.getLogger(__name__)
+DEFAULT_MAX_SEQUENCE_LENGTH = 512
+
+
+class Text2ImageRequest:
+    def __init__(
+        self,
+        unique_id,
+        future,
+        prompt: str,
+        n: int,
+        size: str,
+        response_format: str,
+        *args,
+        **kwargs,
+    ):
+        self._unique_id = unique_id
+        self.future = future
+        self._prompt = prompt
+        self._n = n
+        self._size = size
+        self._response_format = response_format
+        self._args = args
+        self._kwargs = kwargs
+        self._width = -1
+        self._height = -1
+        self._generate_kwargs: Dict[str, Any] = {}
+        self._set_width_and_height()
+        self.is_encode = True
+        self.scheduler = None
+        self.done_steps = 0
+        self.total_steps = 0
+        self.static_tensors: Dict[str, torch.Tensor] = {}
+        self.timesteps = None
+        self.dtype = None
+        self.output = None
+        self.error_msg: Optional[str] = None
+        self.aborted = False
+
+    def _set_width_and_height(self):
+        self._width, self._height = map(int, re.split(r"[^\d]+", self._size))
+
+    def set_generate_kwargs(self, generate_kwargs: Dict):
+        self._generate_kwargs = {k: v for k, v in generate_kwargs.items()}
+
+    @property
+    def prompt(self):
+        return self._prompt
+
+    @property
+    def n(self):
+        return self._n
+
+    @property
+    def size(self):
+        return self._size
+
+    @property
+    def response_format(self):
+        return self._response_format
+
+    @property
+    def kwargs(self):
+        return self._kwargs
+
+    @property
+    def width(self):
+        return self._width
+
+    @property
+    def height(self):
+        return self._height
+
+    @property
+    def generate_kwargs(self):
+        return self._generate_kwargs
+
+    @property
+    def request_id(self):
+        return self._unique_id
+
+
+class FluxBatchSchedulerActor(xo.StatelessActor):
+    @classmethod
+    def gen_uid(cls, model_uid: str):
+        return f"{model_uid}-scheduler-actor"
+
+    def __init__(self):
+        from ....device_utils import get_available_device
+
+        super().__init__()
+        self._waiting_queue: deque[Text2ImageRequest] = deque()  # type: ignore
+        self._running_queue: deque[Text2ImageRequest] = deque()  # type: ignore
+        self._model = None
+        self._available_device = get_available_device()
+        self._id_to_req: Dict[str, Text2ImageRequest] = {}
+
+    def set_model(self, model):
+        """
+        Must use `set_model`. Otherwise, the model will be copied once.
+        """
+        self._model = model
+
+    async def __post_create__(self):
+        from ....isolation import Isolation
+
+        self._isolation = Isolation(
+            asyncio.new_event_loop(), threaded=True, daemon=True
+        )
+        self._isolation.start()
+        asyncio.run_coroutine_threadsafe(self.run(), loop=self._isolation.loop)
+
+    async def __pre_destroy__(self):
+        try:
+            assert self._isolation is not None
+            self._isolation.stop()
+            del self._isolation
+        except Exception as e:
+            logger.debug(
+                f"Destroy scheduler actor failed, address: {self.address}, error: {e}"
+            )
+
+    async def add_request(self, unique_id: str, future, *args, **kwargs):
+        req = Text2ImageRequest(unique_id, future, *args, **kwargs)
+        rid = req.request_id
+        if rid is not None:
+            if rid in self._id_to_req:
+                raise KeyError(f"Request id: {rid} has already existed!")
+            self._id_to_req[rid] = req
+        self._waiting_queue.append(req)
+
+    async def abort_request(self, req_id: str) -> str:
+        """
+        Abort a request.
+        Abort a submitted request. If the request is finished or not found, this method will be a no-op.
+        """
+        from ....core.utils import AbortRequestMessage
+
+        if req_id not in self._id_to_req:
+            logger.info(f"Request id: {req_id} not found. No-op for xinference.")
+            return AbortRequestMessage.NOT_FOUND.name
+        else:
+            self._id_to_req[req_id].aborted = True
+            logger.info(f"Request id: {req_id} found to be aborted.")
+            return AbortRequestMessage.DONE.name
+
+    def _handle_request(
+        self,
+    ) -> Optional[Tuple[List[Text2ImageRequest], List[Text2ImageRequest]]]:
+        """
+        Every request may generate `n>=1` images.
+        Here we need to decide whether to wait or not based on the value of `n` of each request.
+        """
+        if self._model is None:
+            return None
+        max_num_images = self._model.get_max_num_images_for_batching()
+        cur_num_images = 0
+        abort_list: List[Text2ImageRequest] = []
+        # currently, FCFS strategy
+        running_list: List[Text2ImageRequest] = []
+        while len(self._running_queue) > 0:
+            req = self._running_queue.popleft()
+            if req.aborted:
+                abort_list.append(req)
+            else:
+                running_list.append(req)
+                cur_num_images += req.n
+
+        # Remove all the aborted requests in the waiting queue
+        waiting_tmp_list: List[Text2ImageRequest] = []
+        while len(self._waiting_queue) > 0:
+            req = self._waiting_queue.popleft()
+            if req.aborted:
+                abort_list.append(req)
+            else:
+                waiting_tmp_list.append(req)
+        self._waiting_queue.extend(waiting_tmp_list)
+
+        waiting_list: List[Text2ImageRequest] = []
+        while len(self._waiting_queue) > 0:
+            req = self._waiting_queue[0]
+            if req.n + cur_num_images <= max_num_images:
+                waiting_list.append(self._waiting_queue.popleft())
+                cur_num_images += req.n
+            else:
+                logger.warning(
+                    f"Current queue is full, with an upper limit of max_num_images: {max_num_images}. "
+                    f"Requests will continue to wait."
+                )
+                break
+
+        return waiting_list + running_list, abort_list
+
+    @staticmethod
+    def _empty_cache():
+        from ....device_utils import empty_cache
+
+        empty_cache()
+
+    async def step(self):
+        res = self._handle_request()
+        if res is None:
+            return
+        req_list, abort_list = res
+        # handle abort
+        if abort_list:
+            for r in abort_list:
+                r.future.set_exception(
+                    RuntimeError(
+                        f"Request: {r.request_id} has been cancelled by another `abort_request` request."
+                    )
+                )
+                self._id_to_req.pop(r.request_id, None)
+        if not req_list:
+            return
+        _batch_text_to_image(self._model, req_list, self._available_device)
+        # handle results
+        for r in req_list:
+            if r.error_msg is not None:
+                r.future.set_exception(ValueError(r.error_msg))
+                self._id_to_req.pop(r.request_id, None)
+                continue
+            if r.output is not None:
+                r.future.set_result(
+                    handle_image_result(r.response_format, r.output.images)
+                )
+                self._id_to_req.pop(r.request_id, None)
+            else:
+                self._running_queue.append(r)
+        self._empty_cache()
+
+    async def run(self):
+        try:
+            while True:
+                # wait 10ms
+                await asyncio.sleep(0.01)
+                await self.step()
+        except Exception as e:
+            logger.exception(
+                f"Scheduler actor uid: {self.uid}, address: {self.address} run with error: {e}"
+            )
+
+
+def _cat_tensors(infos: List[Dict]) -> Dict:
+    keys = infos[0].keys()
+    res = {}
+    for k in keys:
+        tmp = [info[k] for info in infos]
+        res[k] = torch.cat(tmp)
+    return res
+
+
+@typing.no_type_check
+@torch.inference_mode()
+def _batch_text_to_image_internal(
+    model_cls: "DiffusionModel",
+    req_list: List[Text2ImageRequest],
+    available_device: str,
+):
+    from diffusers.pipelines.flux.pipeline_flux import (
+        calculate_shift,
+        retrieve_timesteps,
+    )
+    from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+    from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
+        FlowMatchEulerDiscreteScheduler,
+    )
+
+    device = model_cls._model._execution_device
+    height, width = req_list[0].height, req_list[0].width
+    cur_batch_max_sequence_length = [
+        r.generate_kwargs.get("max_sequence_length", DEFAULT_MAX_SEQUENCE_LENGTH)
+        for r in req_list
+        if not r.is_encode
+    ]
+    for r in req_list:
+        if r.is_encode:
+            generate_kwargs = model_cls._model_spec.default_generate_config.copy()
+            generate_kwargs.update({k: v for k, v in r.kwargs.items() if v is not None})
+            model_cls._filter_kwargs(model_cls._model, generate_kwargs)
+            r.set_generate_kwargs(generate_kwargs)
+
+            # check max_sequence_length
+            max_sequence_length = r.generate_kwargs.get(
+                "max_sequence_length", DEFAULT_MAX_SEQUENCE_LENGTH
+            )
+            if (
+                cur_batch_max_sequence_length
+                and max_sequence_length != cur_batch_max_sequence_length[0]
+            ):
+                r.is_encode = False
+                r.error_msg = (
+                    f"The max_sequence_length of the current request: {max_sequence_length} is "
+                    f"different from the setting in the running batch: {cur_batch_max_sequence_length[0]}, "
+                    f"please be consistent."
+                )
+                continue
+
+            num_images_per_prompt = r.n
+            callback_on_step_end_tensor_inputs = r.generate_kwargs.get(
+                "callback_on_step_end_tensor_inputs", ["latents"]
+            )
+            num_inference_steps = r.generate_kwargs.get("num_inference_steps", 28)
+            guidance_scale = r.generate_kwargs.get("guidance_scale", 7.0)
+            generator = None
+            seed = r.generate_kwargs.get("seed", None)
+            if seed is not None:
+                generator = torch.Generator(device=available_device)  # type: ignore
+                if seed != -1:
+                    generator = generator.manual_seed(seed)
+            latents = None
+            timesteps = None
+
+            # Each request must build its own scheduler instance,
+            # otherwise the mixing of variables at `scheduler.STEP` will result in an error.
+            r.scheduler = FlowMatchEulerDiscreteScheduler(
+                model_cls._model.scheduler.config.num_train_timesteps,
+                model_cls._model.scheduler.config.shift,
+                model_cls._model.scheduler.config.use_dynamic_shifting,
+                model_cls._model.scheduler.config.base_shift,
+                model_cls._model.scheduler.config.max_shift,
+                model_cls._model.scheduler.config.base_image_seq_len,
+                model_cls._model.scheduler.config.max_image_seq_len,
+            )
+
+            # check inputs
+            model_cls._model.check_inputs(
+                r.prompt,
+                None,
+                height,
+                width,
+                prompt_embeds=None,
+                pooled_prompt_embeds=None,
+                callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+                max_sequence_length=max_sequence_length,
+            )
+
+            # handle prompt
+            (
+                prompt_embeds,
+                pooled_prompt_embeds,
+                text_ids,
+            ) = model_cls._model.encode_prompt(
+                prompt=r.prompt,
+                prompt_2=None,
+                prompt_embeds=None,
+                pooled_prompt_embeds=None,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                lora_scale=None,
+            )
+
+            # Prepare latent variables
+            num_channels_latents = model_cls._model.transformer.config.in_channels // 4
+            latents, latent_image_ids = model_cls._model.prepare_latents(
+                num_images_per_prompt,
+                num_channels_latents,
+                height,
+                width,
+                prompt_embeds.dtype,
+                device,
+                generator,
+                latents,
+            )
+
+            # Prepare timesteps
+            sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+            image_seq_len = latents.shape[1]
+
+            mu = calculate_shift(
+                image_seq_len,
+                r.scheduler.config["base_image_seq_len"],
+                r.scheduler.config["max_image_seq_len"],
+                r.scheduler.config["base_shift"],
+                r.scheduler.config["max_shift"],
+            )
+            timesteps, num_inference_steps = retrieve_timesteps(
+                r.scheduler,
+                num_inference_steps,
+                device,
+                timesteps,
+                sigmas,
+                mu=mu,
+            )
+
+            # handle guidance
+            if model_cls._model.transformer.config.guidance_embeds:
+                guidance = torch.full(
+                    [1], guidance_scale, device=device, dtype=torch.float32
+                )
+                guidance = guidance.expand(latents.shape[0])
+            else:
+                guidance = None
+
+            r.static_tensors["latents"] = latents
+            r.static_tensors["guidance"] = guidance
+            r.static_tensors["pooled_prompt_embeds"] = pooled_prompt_embeds
+            r.static_tensors["prompt_embeds"] = prompt_embeds
+            r.static_tensors["text_ids"] = text_ids
+            r.static_tensors["latent_image_ids"] = latent_image_ids
+            r.timesteps = timesteps
+            r.dtype = latents.dtype
+            r.total_steps = len(timesteps)
+            r.is_encode = False
+
+    running_req_list = [r for r in req_list if r.error_msg is None]
+    static_tensors = _cat_tensors([r.static_tensors for r in running_req_list])
+
+    # Do a step
+    timestep_tmp = []
+    for r in running_req_list:
+        timestep_tmp.append(r.timesteps[r.done_steps].expand(r.n).to(r.dtype))
+        r.done_steps += 1
+    timestep = torch.cat(timestep_tmp)
+    noise_pred = model_cls._model.transformer(
+        hidden_states=static_tensors["latents"],
+        # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transformer model (we should not keep it but I want to keep the inputs same for the model for testing)
+        timestep=timestep / 1000,
+        guidance=static_tensors["guidance"],
+        pooled_projections=static_tensors["pooled_prompt_embeds"],
+        encoder_hidden_states=static_tensors["prompt_embeds"],
+        txt_ids=static_tensors["text_ids"],
+        img_ids=static_tensors["latent_image_ids"],
+        joint_attention_kwargs=None,
+        return_dict=False,
+    )[0]
+
+    # update latents
+    start_idx = 0
+    for r in running_req_list:
+        n = r.n
+        # handle diffusion scheduler step
+        _noise_pred = noise_pred[start_idx : start_idx + n, ::]
+        _timestep = timestep[start_idx]
+        latents_out = r.scheduler.step(
+            _noise_pred, _timestep, r.static_tensors["latents"], return_dict=False
+        )[0]
+        r.static_tensors["latents"] = latents_out
+        start_idx += n
+
+        logger.info(
+            f"Request {r.request_id} has done {r.done_steps} / {r.total_steps} steps."
+        )
+
+        # process result
+        if r.done_steps == r.total_steps:
+            output_type = r.generate_kwargs.get("output_type", "pil")
+            _latents = r.static_tensors["latents"]
+            if output_type == "latent":
+                image = _latents
+            else:
+                _latents = model_cls._model._unpack_latents(
+                    _latents, height, width, model_cls._model.vae_scale_factor
+                )
+                _latents = (
+                    _latents / model_cls._model.vae.config.scaling_factor
+                ) + model_cls._model.vae.config.shift_factor
+                image = model_cls._model.vae.decode(_latents, return_dict=False)[0]
+                image = model_cls._model.image_processor.postprocess(
+                    image, output_type=output_type
+                )
+
+            is_padded = r.generate_kwargs.get("is_padded", None)
+            origin_size = r.generate_kwargs.get("origin_size", None)
+
+            if is_padded and origin_size:
+                new_images = []
+                x, y = origin_size
+                for img in image:
+                    new_images.append(img.crop((0, 0, x, y)))
+                image = new_images
+
+            r.output = FluxPipelineOutput(images=image)
+            logger.info(
+                f"Request {r.request_id} has completed total {r.total_steps} steps."
+            )
+
+
+def _batch_text_to_image(
+    model_cls: "DiffusionModel",
+    req_list: List[Text2ImageRequest],
+    available_device: str,
+):
+    from ....core.model import OutOfMemoryError
+
+    try:
+        _batch_text_to_image_internal(model_cls, req_list, available_device)
+    except OutOfMemoryError:
+        logger.exception(
+            f"Batch text_to_image out of memory. "
+            f"Xinference will restart the model: {model_cls._model_uid}. "
+            f"Please be patient for a few moments."
+        )
+        # Just kill the process and let xinference auto-recover the model
+        os._exit(1)
+    except Exception as e:
+        logger.exception(f"Internal error for batch text_to_image: {e}.")
+        # If internal error happens, just skip all the requests in this batch.
+        # If not handle here, the client will hang.
+        for r in req_list:
+            r.error_msg = str(e)

xinference/model/image/stable_diffusion/core.py

@@ -12,31 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import base64
 import contextlib
 import gc
 import inspect
 import itertools
 import logging
-import os
 import re
 import sys
-import time
-import uuid
 import warnings
-from concurrent.futures import ThreadPoolExecutor
-from functools import partial
-from io import BytesIO
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import PIL.Image
 import torch
 from PIL import ImageOps
 
-from ....constants import XINFERENCE_IMAGE_DIR
 from ....device_utils import get_available_device, move_model_to_available_device
-from ....types import Image, ImageList, LoRA
+from ....types import LoRA
 from ..sdapi import SDAPIDiffusionModelMixin
+from ..utils import handle_image_result
 
 if TYPE_CHECKING:
     from ....core.progress_tracker import Progressor
@@ -297,6 +290,9 @@ class DiffusionModel(SDAPIDiffusionModelMixin):
         if self._kwargs.get("vae_tiling", False):
             model.enable_vae_tiling()
 
+    def get_max_num_images_for_batching(self):
+        return self._kwargs.get("max_num_images", 16)
+
     @staticmethod
     def _get_scheduler(model: Any, sampler_name: str):
         if not sampler_name or sampler_name == "default":
@@ -476,28 +472,7 @@ class DiffusionModel(SDAPIDiffusionModelMixin):
         if return_images:
             return images
 
-        if response_format == "url":
-            os.makedirs(XINFERENCE_IMAGE_DIR, exist_ok=True)
-            image_list = []
-            with ThreadPoolExecutor() as executor:
-                for img in images:
-                    path = os.path.join(XINFERENCE_IMAGE_DIR, uuid.uuid4().hex + ".jpg")
-                    image_list.append(Image(url=path, b64_json=None))
-                    executor.submit(img.save, path, "jpeg")
-            return ImageList(created=int(time.time()), data=image_list)
-        elif response_format == "b64_json":
-
-            def _gen_base64_image(_img):
-                buffered = BytesIO()
-                _img.save(buffered, format="jpeg")
-                return base64.b64encode(buffered.getvalue()).decode()
-
-            with ThreadPoolExecutor() as executor:
-                results = list(map(partial(executor.submit, _gen_base64_image), images))  # type: ignore
-                image_list = [Image(url=None, b64_json=s.result()) for s in results]  # type: ignore
-            return ImageList(created=int(time.time()), data=image_list)
-        else:
-            raise ValueError(f"Unsupported response format: {response_format}")
+        return handle_image_result(response_format, images)
 
     @classmethod
     def _filter_kwargs(cls, model, kwargs: dict):

xinference/model/image/utils.py

@@ -11,16 +11,52 @@
 # 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.
-from typing import Optional
+import base64
+import os
+import time
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+from io import BytesIO
+from typing import TYPE_CHECKING, Optional
 
-from .core import ImageModelFamilyV1
+from ...constants import XINFERENCE_IMAGE_DIR
+from ...types import Image, ImageList
+
+if TYPE_CHECKING:
+    from .core import ImageModelFamilyV1
 
 
 def get_model_version(
-    image_model: ImageModelFamilyV1, controlnet: Optional[ImageModelFamilyV1]
+    image_model: "ImageModelFamilyV1", controlnet: Optional["ImageModelFamilyV1"]
 ) -> str:
     return (
         image_model.model_name
         if controlnet is None
         else f"{image_model.model_name}--{controlnet.model_name}"
     )
+
+
+def handle_image_result(response_format: str, images) -> ImageList:
+    if response_format == "url":
+        os.makedirs(XINFERENCE_IMAGE_DIR, exist_ok=True)
+        image_list = []
+        with ThreadPoolExecutor() as executor:
+            for img in images:
+                path = os.path.join(XINFERENCE_IMAGE_DIR, uuid.uuid4().hex + ".jpg")
+                image_list.append(Image(url=path, b64_json=None))
+                executor.submit(img.save, path, "jpeg")
+        return ImageList(created=int(time.time()), data=image_list)
+    elif response_format == "b64_json":
+
+        def _gen_base64_image(_img):
+            buffered = BytesIO()
+            _img.save(buffered, format="jpeg")
+            return base64.b64encode(buffered.getvalue()).decode()
+
+        with ThreadPoolExecutor() as executor:
+            results = list(map(partial(executor.submit, _gen_base64_image), images))  # type: ignore
+            image_list = [Image(url=None, b64_json=s.result()) for s in results]  # type: ignore
+        return ImageList(created=int(time.time()), data=image_list)
+    else:
+        raise ValueError(f"Unsupported response format: {response_format}")

xinference/model/llm/__init__.py

@@ -146,6 +146,7 @@ def _install():
     from .transformers.internlm2 import Internlm2PytorchChatModel
     from .transformers.minicpmv25 import MiniCPMV25Model
     from .transformers.minicpmv26 import MiniCPMV26Model
+    from .transformers.opt import OptPytorchModel
     from .transformers.qwen2_audio import Qwen2AudioChatModel
     from .transformers.qwen2_vl import Qwen2VLChatModel
     from .transformers.qwen_vl import QwenVLChatModel
@@ -190,6 +191,7 @@ def _install():
             Glm4VModel,
             DeepSeekV2PytorchModel,
             DeepSeekV2PytorchChatModel,
+            OptPytorchModel,
         ]
     )
     if OmniLMMModel:  # type: ignore