FEAT: support getting progress for image model (#2395)

Author: qinxuye

xinference/api/restful_api.py

@@ -524,6 +524,16 @@ async def internal_exception_handler(request: Request, exc: Exception):
                 else None
             ),
         )
+        self._router.add_api_route(
+            "/v1/requests/{request_id}/progress",
+            self.get_progress,
+            methods=["get"],
+            dependencies=(
+                [Security(self._auth_service, scopes=["models:read"])]
+                if self.is_authenticated()
+                else None
+            ),
+        )
         self._router.add_api_route(
             "/v1/images/generations",
             self.create_images,
@@ -1486,6 +1496,17 @@ async def create_speech(
             await self._report_error_event(model_uid, str(e))
             raise HTTPException(status_code=500, detail=str(e))
 
+    async def get_progress(self, request_id: str) -> JSONResponse:
+        try:
+            supervisor_ref = await self._get_supervisor_ref()
+            result = {"progress": await supervisor_ref.get_progress(request_id)}
+            return JSONResponse(content=result)
+        except KeyError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+        except Exception as e:
+            logger.error(e, exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
     async def create_images(self, request: Request) -> Response:
         body = TextToImageRequest.parse_obj(await request.json())
         model_uid = body.model

xinference/client/restful/restful_client.py

@@ -1385,6 +1385,16 @@ def get_supervisor_info(self):
         response_json = response.json()
         return response_json
 
+    def get_progress(self, request_id: str):
+        url = f"{self.base_url}/v1/requests/{request_id}/progress"
+        response = requests.get(url, headers=self._headers)
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"Failed to get progress, detail: {_get_error_string(response)}"
+            )
+        response_json = response.json()
+        return response_json
+
     def abort_cluster(self):
         url = f"{self.base_url}/v1/clusters"
         response = requests.delete(url, headers=self._headers)

xinference/core/image_interface.py

@@ -16,6 +16,9 @@
 import io
 import logging
 import os
+import threading
+import time
+import uuid
 from typing import Dict, List, Optional, Union
 
 import gradio as gr
@@ -84,6 +87,7 @@ def text_generate_image(
             num_inference_steps: int,
             negative_prompt: Optional[str] = None,
             sampler_name: Optional[str] = None,
+            progress=gr.Progress(),
         ) -> PIL.Image.Image:
             from ..client import RESTfulClient
 
@@ -99,19 +103,43 @@ def text_generate_image(
             )
             sampler_name = None if sampler_name == "default" else sampler_name
 
-            response = model.text_to_image(
-                prompt=prompt,
-                n=n,
-                size=size,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=guidance_scale,
-                negative_prompt=negative_prompt,
-                sampler_name=sampler_name,
-                response_format="b64_json",
-            )
+            response = None
+            exc = None
+            request_id = str(uuid.uuid4())
+
+            def run_in_thread():
+                nonlocal exc, response
+                try:
+                    response = model.text_to_image(
+                        request_id=request_id,
+                        prompt=prompt,
+                        n=n,
+                        size=size,
+                        num_inference_steps=num_inference_steps,
+                        guidance_scale=guidance_scale,
+                        negative_prompt=negative_prompt,
+                        sampler_name=sampler_name,
+                        response_format="b64_json",
+                    )
+                except Exception as e:
+                    exc = e
+
+            t = threading.Thread(target=run_in_thread)
+            t.start()
+            while t.is_alive():
+                try:
+                    cur_progress = client.get_progress(request_id)["progress"]
+                except (KeyError, RuntimeError):
+                    cur_progress = 0.0
+
+                progress(cur_progress, desc="Generating images")
+                time.sleep(1)
+
+            if exc:
+                raise exc
 
             images = []
-            for image_dict in response["data"]:
+            for image_dict in response["data"]:  # type: ignore
                 assert image_dict["b64_json"] is not None
                 image_data = base64.b64decode(image_dict["b64_json"])
                 image = PIL.Image.open(io.BytesIO(image_data))
@@ -184,6 +212,7 @@ def image_generate_image(
             num_inference_steps: int,
             padding_image_to_multiple: int,
             sampler_name: Optional[str] = None,
+            progress=gr.Progress(),
         ) -> PIL.Image.Image:
             from ..client import RESTfulClient
 
@@ -205,20 +234,44 @@ def image_generate_image(
             bio = io.BytesIO()
             image.save(bio, format="png")
 
-            response = model.image_to_image(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                n=n,
-                image=bio.getvalue(),
-                size=size,
-                response_format="b64_json",
-                num_inference_steps=num_inference_steps,
-                padding_image_to_multiple=padding_image_to_multiple,
-                sampler_name=sampler_name,
-            )
+            response = None
+            exc = None
+            request_id = str(uuid.uuid4())
+
+            def run_in_thread():
+                nonlocal exc, response
+                try:
+                    response = model.image_to_image(
+                        request_id=request_id,
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        n=n,
+                        image=bio.getvalue(),
+                        size=size,
+                        response_format="b64_json",
+                        num_inference_steps=num_inference_steps,
+                        padding_image_to_multiple=padding_image_to_multiple,
+                        sampler_name=sampler_name,
+                    )
+                except Exception as e:
+                    exc = e
+
+            t = threading.Thread(target=run_in_thread)
+            t.start()
+            while t.is_alive():
+                try:
+                    cur_progress = client.get_progress(request_id)["progress"]
+                except (KeyError, RuntimeError):
+                    cur_progress = 0.0
+
+                progress(cur_progress, desc="Generating images")
+                time.sleep(1)
+
+            if exc:
+                raise exc
 
             images = []
-            for image_dict in response["data"]:
+            for image_dict in response["data"]:  # type: ignore
                 assert image_dict["b64_json"] is not None
                 image_data = base64.b64decode(image_dict["b64_json"])
                 image = PIL.Image.open(io.BytesIO(image_data))

xinference/core/model.py

@@ -44,6 +44,7 @@
 from ..constants import XINFERENCE_TRANSFORMERS_ENABLE_BATCHING
 
 if TYPE_CHECKING:
+    from .progress_tracker import ProgressTrackerActor
     from .worker import WorkerActor
     from ..model.llm.core import LLM
     from ..model.core import ModelDescription
@@ -177,6 +178,7 @@ async def __pre_destroy__(self):
 
     def __init__(
         self,
+        supervisor_address: str,
         worker_address: str,
         model: "LLM",
         model_description: Optional["ModelDescription"] = None,
@@ -188,6 +190,7 @@ def __init__(
         from ..model.llm.transformers.core import PytorchModel
         from ..model.llm.vllm.core import VLLMModel
 
+        self._supervisor_address = supervisor_address
         self._worker_address = worker_address
         self._model = model
         self._model_description = (
@@ -205,6 +208,7 @@ def __init__(
             else asyncio.locks.Lock()
         )
         self._worker_ref = None
+        self._progress_tracker_ref = None
         self._serve_count = 0
         self._metrics_labels = {
             "type": self._model_description.get("model_type", "unknown"),
@@ -275,6 +279,28 @@ async def _get_worker_ref(self) -> xo.ActorRefType["WorkerActor"]:
             )
         return self._worker_ref
 
+    async def _get_progress_tracker_ref(
+        self,
+    ) -> xo.ActorRefType["ProgressTrackerActor"]:
+        from .progress_tracker import ProgressTrackerActor
+
+        if self._progress_tracker_ref is None:
+            self._progress_tracker_ref = await xo.actor_ref(
+                address=self._supervisor_address, uid=ProgressTrackerActor.default_uid()
+            )
+        return self._progress_tracker_ref
+
+    async def _get_progressor(self, request_id: str):
+        from .progress_tracker import Progressor
+
+        progressor = Progressor(
+            request_id,
+            await self._get_progress_tracker_ref(),
+            asyncio.get_running_loop(),
+        )
+        await progressor.start()
+        return progressor
+
     def is_vllm_backend(self) -> bool:
         from ..model.llm.vllm.core import VLLMModel
 
@@ -732,17 +758,20 @@ async def text_to_image(
         *args,
         **kwargs,
     ):
-        kwargs.pop("request_id", None)
         if hasattr(self._model, "text_to_image"):
-            return await self._call_wrapper_json(
-                self._model.text_to_image,
-                prompt,
-                n,
-                size,
-                response_format,
-                *args,
-                **kwargs,
+            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."
         )
@@ -753,12 +782,15 @@ async def txt2img(
         self,
         **kwargs,
     ):
-        kwargs.pop("request_id", None)
         if hasattr(self._model, "txt2img"):
-            return await self._call_wrapper_json(
-                self._model.txt2img,
-                **kwargs,
+            progressor = kwargs["progressor"] = await self._get_progressor(
+                kwargs.pop("request_id", None)
             )
+            with progressor:
+                return await self._call_wrapper_json(
+                    self._model.txt2img,
+                    **kwargs,
+                )
         raise AttributeError(f"Model {self._model.model_spec} is not for txt2img.")
 
     @log_async(
@@ -776,19 +808,22 @@ async def image_to_image(
         *args,
         **kwargs,
     ):
-        kwargs.pop("request_id", None)
         kwargs["negative_prompt"] = negative_prompt
         if hasattr(self._model, "image_to_image"):
-            return await self._call_wrapper_json(
-                self._model.image_to_image,
-                image,
-                prompt,
-                n,
-                size,
-                response_format,
-                *args,
-                **kwargs,
+            progressor = kwargs["progressor"] = await self._get_progressor(
+                kwargs.pop("request_id", None)
             )
+            with progressor:
+                return await self._call_wrapper_json(
+                    self._model.image_to_image,
+                    image,
+                    prompt,
+                    n,
+                    size,
+                    response_format,
+                    *args,
+                    **kwargs,
+                )
         raise AttributeError(
             f"Model {self._model.model_spec} is not for creating image."
         )
@@ -799,12 +834,15 @@ async def img2img(
         self,
         **kwargs,
     ):
-        kwargs.pop("request_id", None)
         if hasattr(self._model, "img2img"):
-            return await self._call_wrapper_json(
-                self._model.img2img,
-                **kwargs,
+            progressor = kwargs["progressor"] = await self._get_progressor(
+                kwargs.pop("request_id", None)
             )
+            with progressor:
+                return await self._call_wrapper_json(
+                    self._model.img2img,
+                    **kwargs,
+                )
         raise AttributeError(f"Model {self._model.model_spec} is not for img2img.")
 
     @log_async(
@@ -823,20 +861,23 @@ async def inpainting(
         *args,
         **kwargs,
     ):
-        kwargs.pop("request_id", None)
+        kwargs["negative_prompt"] = negative_prompt
         if hasattr(self._model, "inpainting"):
-            return await self._call_wrapper_json(
-                self._model.inpainting,
-                image,
-                mask_image,
-                prompt,
-                negative_prompt,
-                n,
-                size,
-                response_format,
-                *args,
-                **kwargs,
+            progressor = kwargs["progressor"] = await self._get_progressor(
+                kwargs.pop("request_id", None)
             )
+            with progressor:
+                return await self._call_wrapper_json(
+                    self._model.inpainting,
+                    image,
+                    mask_image,
+                    prompt,
+                    n,
+                    size,
+                    response_format,
+                    *args,
+                    **kwargs,
+                )
         raise AttributeError(
             f"Model {self._model.model_spec} is not for creating image."
         )