From 54c418e745992a8623596577dc466d2b12f7bc15 Mon Sep 17 00:00:00 2001
From: "Xuye (Chris) Qin" <qinxuye@gmail.com>
Date: Wed, 7 Aug 2024 00:15:38 +0800
Subject: [PATCH] FEAT: support kolors image model (#2028)

---
 doc/source/models/builtin/image/kolors.rst    | 19 +++++++++++++++++++
 xinference/core/image_interface.py            |  9 +++++++++
 xinference/model/image/model_spec.json        | 10 ++++++++++
 .../model/image/stable_diffusion/core.py      |  1 +
 4 files changed, 39 insertions(+)
 create mode 100644 doc/source/models/builtin/image/kolors.rst

diff --git a/doc/source/models/builtin/image/kolors.rst b/doc/source/models/builtin/image/kolors.rst
new file mode 100644
index 0000000000..19d11c4201
--- /dev/null
+++ b/doc/source/models/builtin/image/kolors.rst
@@ -0,0 +1,19 @@
+.. _models_builtin_kolors:
+
+======
+kolors
+======
+
+- **Model Name:** kolors
+- **Model Family:** stable_diffusion
+- **Abilities:** text2image, image2image
+- **Available ControlNet:** None
+
+Specifications
+^^^^^^^^^^^^^^
+
+- **Model ID:** Kwai-Kolors/Kolors-diffusers
+
+Execute the following command to launch the model::
+
+   xinference launch --model-name kolors --model-type image
\ No newline at end of file
diff --git a/xinference/core/image_interface.py b/xinference/core/image_interface.py
index 791a0c5515..3d80b72df4 100644
--- a/xinference/core/image_interface.py
+++ b/xinference/core/image_interface.py
@@ -162,6 +162,7 @@ def image_generate_image(
             n: int,
             size_width: int,
             size_height: int,
+            num_inference_steps: int,
         ) -> PIL.Image.Image:
             from ..client import RESTfulClient
 
@@ -174,6 +175,9 @@ def image_generate_image(
                 size = f"{int(size_width)}*{int(size_height)}"
             else:
                 size = None
+            num_inference_steps = (
+                None if num_inference_steps == -1 else num_inference_steps  # type: ignore
+            )
 
             bio = io.BytesIO()
             image.save(bio, format="png")
@@ -185,6 +189,7 @@ def image_generate_image(
                 image=bio.getvalue(),
                 size=size,
                 response_format="b64_json",
+                num_inference_steps=num_inference_steps,
             )
 
             images = []
@@ -217,6 +222,9 @@ def image_generate_image(
                     n = gr.Number(label="Number of image", value=1)
                     size_width = gr.Number(label="Width", value=-1)
                     size_height = gr.Number(label="Height", value=-1)
+                    num_inference_steps = gr.Number(
+                        label="Inference Step Number", value=-1
+                    )
 
                 with gr.Row():
                     with gr.Column(scale=1):
@@ -233,6 +241,7 @@ def image_generate_image(
                     n,
                     size_width,
                     size_height,
+                    num_inference_steps,
                 ],
                 outputs=output_gallery,
             )
diff --git a/xinference/model/image/model_spec.json b/xinference/model/image/model_spec.json
index 847bffbaaf..db062dace7 100644
--- a/xinference/model/image/model_spec.json
+++ b/xinference/model/image/model_spec.json
@@ -129,6 +129,16 @@
       }
     ]
   },
+  {
+    "model_name": "kolors",
+    "model_family": "stable_diffusion",
+    "model_id": "Kwai-Kolors/Kolors-diffusers",
+    "model_revision": "7e091c75199e910a26cd1b51ed52c28de5db3711",
+    "model_ability": [
+      "text2image",
+      "image2image"
+    ]
+  },
   {
     "model_name": "stable-diffusion-inpainting",
     "model_family": "stable_diffusion",
diff --git a/xinference/model/image/stable_diffusion/core.py b/xinference/model/image/stable_diffusion/core.py
index 14cb5dbcd3..17c83e9f18 100644
--- a/xinference/model/image/stable_diffusion/core.py
+++ b/xinference/model/image/stable_diffusion/core.py
@@ -88,6 +88,7 @@ def load(self):
         if sys.platform != "darwin" and torch_dtype is None:
             # The following params crashes on Mac M2
             self._kwargs["torch_dtype"] = torch.float16
+            self._kwargs["variant"] = "fp16"
             self._kwargs["use_safetensors"] = True
         if isinstance(torch_dtype, str):
             self._kwargs["torch_dtype"] = getattr(torch, torch_dtype)