FEAT: Support sd-turbo

.github/workflows/python.yaml

@@ -126,6 +126,7 @@ jobs:
           if [ "$MODULE" == "gpu" ]; then
             ${{ env.SELF_HOST_PYTHON }} -m pip install "openai>1"
             ${{ env.SELF_HOST_PYTHON }} -m pip install -U modelscope
+            ${{ env.SELF_HOST_PYTHON }} -m pip install -U sse_starlette
             ${{ env.SELF_HOST_PYTHON }} -m pytest --timeout=1500 \
               -W ignore::PendingDeprecationWarning \
               --cov-config=setup.cfg --cov-report=xml --cov=xinference xinference/model/image/tests/test_stable_diffusion.py

xinference/api/restful_api.py

@@ -104,6 +104,7 @@ class TextToImageRequest(BaseModel):
     n: Optional[int] = 1
     response_format: Optional[str] = "url"
     size: Optional[str] = "1024*1024"
+    kwargs: Optional[str] = None
     user: Optional[str] = None
 
 
@@ -645,8 +646,14 @@ async def create_images(self, request: TextToImageRequest) -> JSONResponse:
             raise HTTPException(status_code=500, detail=str(e))
 
         try:
+            if request.kwargs:
+                kwargs = json.loads(request.kwargs)
             image_list = await model.text_to_image(
-                request.prompt, request.n, request.size, request.response_format
+                prompt=request.prompt,
+                n=request.n,
+                size=request.size,
+                response_format=request.response_format,
+                **kwargs,
             )
             return JSONResponse(content=image_list)
         except RuntimeError as re:

xinference/client/restful/restful_client.py

@@ -154,6 +154,7 @@ def text_to_image(
         n: int = 1,
         size: str = "1024*1024",
         response_format: str = "url",
+        **kwargs,
     ) -> "ImageList":
         """
         Creates an image by the input text.
@@ -180,6 +181,7 @@ def text_to_image(
             "n": n,
             "size": size,
             "response_format": response_format,
+            "kwargs": json.dumps(kwargs),
         }
         response = requests.post(url, json=request_body)
         if response.status_code != 200:

xinference/model/image/model_spec.json

@@ -1,4 +1,10 @@
 [
+  {
+    "model_name": "sd-turbo",
+    "model_family": "stable_diffusion",
+    "model_id": "stabilityai/sd-turbo",
+    "model_revision": "1681ed09e0cff58eeb41e878a49893228b78b94c"
+  },
   {
     "model_name": "stable-diffusion-v1.5",
     "model_family": "stable_diffusion",

xinference/model/image/stable_diffusion/core.py

@@ -71,7 +71,15 @@ def _call_model(
         response_format: str,
         **kwargs,
     ):
-        logger.debug("stable diffusion kwargs: %s", kwargs)
+        logger.debug(
+            "stable diffusion args: %s",
+            dict(
+                kwargs,
+                height=height,
+                width=width,
+                num_images_per_prompt=num_images_per_prompt,
+            ),
+        )
         assert callable(self._model)
         images = self._model(
             height=height,

xinference/model/image/tests/test_stable_diffusion.py

@@ -144,6 +144,32 @@ def test_restful_api_for_image_with_mlsd_controlnet(setup):
     logger.info("test result %s", r)
 
 
+def test_restful_api_for_sd_turbo(setup):
+    endpoint, _ = setup
+    from ....client import Client
+
+    client = Client(endpoint)
+
+    model_uid = client.launch_model(
+        model_uid="my_controlnet",
+        model_name="sd-turbo",
+        model_type="image",
+    )
+    model = client.get_model(model_uid)
+
+    r = model.text_to_image(
+        prompt="A cinematic shot of a baby raccoon wearing an intricate italian priest robe.",
+        size="512*512",
+        num_inference_steps=10,
+    )
+    logger.info("test result %s", r)
+    from PIL import Image
+
+    with open(r["data"][0]["url"], "rb") as f:
+        img = Image.open(f)
+        assert img.size == (512, 512)
+
+
 def test_get_cache_status():
     from ..core import get_cache_status