FEAT: support kolors image model (#2028)

xorbitsai · Aug 6, 2024 · 54c418e · 54c418e
1 parent 957c613
commit 54c418e
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 0 deletions.
diff --git a/doc/source/models/builtin/image/kolors.rst 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
diff --git 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
@@ -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
@@ -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)