xinference 0.15.3__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/sdapi.py

@@ -11,11 +11,12 @@
 # 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
 import io
 import warnings
 
-from PIL import Image
+from PIL import Image, ImageOps
 
 
 class SDAPIToDiffusersConverter:
@@ -26,11 +27,12 @@ class SDAPIToDiffusersConverter:
         "width",
         "height",
         "sampler_name",
+        "progressor",
     }
     txt2img_arg_mapping = {
         "steps": "num_inference_steps",
         "cfg_scale": "guidance_scale",
-        # "denoising_strength": "strength",
+        "denoising_strength": "strength",
     }
     img2img_identical_args = {
         "prompt",
@@ -39,12 +41,15 @@ class SDAPIToDiffusersConverter:
         "width",
         "height",
         "sampler_name",
+        "progressor",
     }
     img2img_arg_mapping = {
         "init_images": "image",
+        "mask": "mask_image",
         "steps": "num_inference_steps",
         "cfg_scale": "guidance_scale",
         "denoising_strength": "strength",
+        "inpaint_full_res_padding": "padding_mask_crop",
     }
 
     @staticmethod
@@ -121,12 +126,38 @@ class SDAPIDiffusionModelMixin:
 
     def img2img(self, **kwargs):
         init_images = kwargs.pop("init_images", [])
-        kwargs["init_images"] = [self._decode_b64_img(i) for i in init_images]
+        kwargs["init_images"] = init_images = [
+            self._decode_b64_img(i) for i in init_images
+        ]
+        if len(init_images) == 1:
+            kwargs["init_images"] = init_images[0]
+        mask_image = kwargs.pop("mask", None)
+        if mask_image:
+            if kwargs.pop("inpainting_mask_invert"):
+                mask_image = ImageOps.invert(mask_image)
+
+            kwargs["mask"] = self._decode_b64_img(mask_image)
+
+            # process inpaint_full_res and inpaint_full_res_padding
+            if kwargs.pop("inpaint_full_res", None):
+                kwargs["inpaint_full_res_padding"] = kwargs.pop(
+                    "inpaint_full_res_padding", 0
+                )
+            else:
+                # inpaint_full_res_padding is turned `into padding_mask_crop`
+                # in diffusers, if padding_mask_crop is passed, it will do inpaint_full_res
+                # so if not inpaint_full_rs, we need to pop this option
+                kwargs.pop("inpaint_full_res_padding", None)
+
         clip_skip = kwargs.get("override_settings", {}).get("clip_skip")
         converted_kwargs = self._check_kwargs("img2img", kwargs)
         if clip_skip:
             converted_kwargs["clip_skip"] = clip_skip
-        result = self.image_to_image(response_format="b64_json", **converted_kwargs)  # type: ignore
+
+        if not converted_kwargs.get("mask_image"):
+            result = self.image_to_image(response_format="b64_json", **converted_kwargs)  # type: ignore
+        else:
+            result = self.inpainting(response_format="b64_json", **converted_kwargs)  # type: ignore
 
         # convert to SD API result
         return {