ENH: support JoyAI image editing models

doc/source/models/builtin/image/index.rst

@@ -48,6 +48,10 @@ The following is a list of built-in image models in Xinference:
    hunyuanocr
 
    ideogram4
+
+   joyai-image-edit
+
+   joyai-image-edit-plus
 
    kolors
 

doc/source/models/builtin/image/joyai-image-edit-plus.rst

@@ -0,0 +1,19 @@
+.. _models_builtin_joyai-image-edit-plus:
+
+=====================
+joyai-image-edit-plus
+=====================
+
+- **Model Name:** joyai-image-edit-plus
+- **Model Family:** stable_diffusion
+- **Abilities:** image2image
+- **Available ControlNet:** None
+
+Specifications
+^^^^^^^^^^^^^^
+
+- **Model ID:** jdopensource/JoyAI-Image-Edit-Plus-Diffusers
+
+Execute the following command to launch the model::
+
+   xinference launch --model-name joyai-image-edit-plus --model-type image

doc/source/models/builtin/image/joyai-image-edit.rst

@@ -0,0 +1,19 @@
+.. _models_builtin_joyai-image-edit:
+
+================
+joyai-image-edit
+================
+
+- **Model Name:** joyai-image-edit
+- **Model Family:** stable_diffusion
+- **Abilities:** image2image
+- **Available ControlNet:** None
+
+Specifications
+^^^^^^^^^^^^^^
+
+- **Model ID:** jdopensource/JoyAI-Image-Edit-Diffusers
+
+Execute the following command to launch the model::
+
+   xinference launch --model-name joyai-image-edit --model-type image

xinference/model/image/model_spec.json

@@ -2321,5 +2321,89 @@
     },
     "featured": false,
     "updated_at": 1786731484
+  },
+  {
+    "version": 2,
+    "model_name": "joyai-image-edit",
+    "model_family": "stable_diffusion",
+    "model_description": "JoyAI-Image-Edit is a single-image instruction-guided editing model with spatial understanding for precise object, viewpoint, and scene edits.",
+    "model_ability": [
+      "image2image"
+    ],
+    "model_src": {
+      "huggingface": {
+        "model_id": "jdopensource/JoyAI-Image-Edit-Diffusers",
+        "model_revision": "main"
+      },
+      "modelscope": {
+        "model_id": "jd-opensource/JoyAI-Image-Edit-Diffusers",
+        "model_revision": "master"
+      }
+    },
+    "default_model_config": {
+      "torch_dtype": "bfloat16"
+    },
+    "default_generate_config": {
+      "num_inference_steps": 40,
+      "guidance_scale": 4.0
+    },
+    "virtualenv": {
+      "packages": [
+        "diffusers>=0.40.0 ; #engine# == \"diffusers\"",
+        "transformers>=5.2.0,<6 ; #engine# == \"diffusers\"",
+        "accelerate>=1.0.0 ; #engine# == \"diffusers\"",
+        "sentencepiece ; #engine# == \"diffusers\"",
+        "huggingface-hub>=1.23.0,<2.0 ; #engine# == \"diffusers\"",
+        "#system_torch# ; #engine# == \"diffusers\"",
+        "#system_torchvision# ; #engine# == \"diffusers\"",
+        "#system_numpy# ; #engine# == \"diffusers\""
+      ],
+      "no_build_isolation": false,
+      "index_strategy": "unsafe-best-match"
+    },
+    "featured": false,
+    "updated_at": 1788101544
+  },
+  {
+    "version": 2,
+    "model_name": "joyai-image-edit-plus",
+    "model_family": "stable_diffusion",
+    "model_description": "JoyAI-Image-Edit-Plus is a multi-image instruction-guided editing model that combines content from one to six reference images.",
+    "model_ability": [
+      "image2image"
+    ],
+    "model_src": {
+      "huggingface": {
+        "model_id": "jdopensource/JoyAI-Image-Edit-Plus-Diffusers",
+        "model_revision": "main"
+      },
+      "modelscope": {
+        "model_id": "jd-opensource/JoyAI-Image-Edit-Plus-Diffusers",
+        "model_revision": "master"
+      }
+    },
+    "default_model_config": {
+      "torch_dtype": "bfloat16"
+    },
+    "default_generate_config": {
+      "num_inference_steps": 30,
+      "guidance_scale": 4.0
+    },
+    "virtualenv": {
+      "packages": [
+        "diffusers>=0.40.0 ; #engine# == \"diffusers\"",
+        "transformers>=5.2.0,<6 ; #engine# == \"diffusers\"",
+        "accelerate>=1.0.0 ; #engine# == \"diffusers\"",
+        "sentencepiece ; #engine# == \"diffusers\"",
+        "huggingface-hub>=1.23.0,<2.0 ; #engine# == \"diffusers\"",
+        "#system_torch# ; #engine# == \"diffusers\"",
+        "#system_torchvision# ; #engine# == \"diffusers\"",
+        "#system_numpy# ; #engine# == \"diffusers\""
+      ],
+      "no_build_isolation": false,
+      "index_strategy": "unsafe-best-match"
+    },
+    "featured": false,
+    "updated_at": 1788101545
   }
 ]

xinference/model/image/stable_diffusion/core.py

@@ -148,6 +148,18 @@ def _is_glm_image_model(self) -> bool:
         model_name = self._model_spec.model_name.lower().replace("_", "-")
         return model_name.startswith("glm-image")
 
+    def _is_joyai_image_model(self) -> bool:
+        if self._model_spec is None:
+            return False
+        model_name = self._model_spec.model_name.lower().replace("_", "-")
+        return model_name.startswith("joyai-image-edit")
+
+    def _is_joyai_image_edit_plus_model(self) -> bool:
+        if self._model_spec is None:
+            return False
+        model_name = self._model_spec.model_name.lower().replace("_", "-")
+        return model_name == "joyai-image-edit-plus"
+
     @staticmethod
     def _get_pipeline_type(ability: str) -> type:
         if ability == "text2image":
@@ -264,7 +276,11 @@ def _get_layer_cls(self, layer: str):
             return getattr(module, class_name)
 
     def load(self):
-        if self._is_glm_image_model():
+        if self._is_joyai_image_model():
+            # JoyAI pipelines are image-edit-only and are not registered in
+            # Diffusers' AutoPipelineForText2Image mapping.
+            from diffusers import DiffusionPipeline as AutoPipelineModel
+        elif self._is_glm_image_model():
             from diffusers import GlmImagePipeline as AutoPipelineModel
         elif "text2image" in self._abilities or "image2image" in self._abilities:
             from diffusers import AutoPipelineForText2Image as AutoPipelineModel
@@ -753,10 +769,17 @@ def _wrap_deepcache(self, model: Any):
                 self._deepcache_helper.pipe = None
 
     @staticmethod
-    def _process_progressor(kwargs: dict):
+    def _process_progressor(
+        kwargs: dict,
+        *,
+        progressor: Optional["Progressor"] = None,
+        pipeline_call_index: int = 0,
+        pipeline_call_count: int = 1,
+    ):
         import diffusers
 
-        progressor: Progressor = kwargs.pop("progressor", None)
+        if progressor is None:
+            progressor = kwargs.pop("progressor", None)
 
         def report_status_callback(
             pipe: diffusers.DiffusionPipeline,
@@ -765,7 +788,10 @@ def report_status_callback(
             callback_kwargs: dict,
         ):
             num_steps = pipe.num_timesteps
-            progressor.set_progress((step + 1) / num_steps)
+            local_progress = (step + 1) / num_steps
+            progressor.set_progress(
+                (pipeline_call_index + local_progress) / pipeline_call_count
+            )
 
             return callback_kwargs
 
@@ -776,16 +802,20 @@ def _call_model(
         self,
         response_format: str,
         model=None,
+        _num_pipeline_calls: int = 1,
         **kwargs,
     ):
         model = model if model is not None else self._model
         is_padded = kwargs.pop("is_padded", None)
         origin_size = kwargs.pop("origin_size", None)
         seed = kwargs.pop("seed", None)
         return_images = kwargs.pop("_return_images", None)
-        seeds = resolve_image_seed_list(
-            seed, int(kwargs.get("num_images_per_prompt", 1))
+        seed_count = (
+            _num_pipeline_calls
+            if _num_pipeline_calls > 1
+            else int(kwargs.get("num_images_per_prompt", 1))
         )
+        seeds = resolve_image_seed_list(seed, seed_count)
         if seeds is not None:
             kwargs["generator"] = [
                 torch.Generator(  # type: ignore
@@ -797,7 +827,7 @@ def _call_model(
             kwargs["generator"] = generator = torch.Generator(device=get_available_device())  # type: ignore
             kwargs["generator"] = generator.manual_seed(seed)
         sampler_name = kwargs.pop("sampler_name", None)
-        self._process_progressor(kwargs)
+        progressor = kwargs.pop("progressor", None)
         if self._is_ideogram4_model() and kwargs.get("guidance_scale") is not None:
             # Ideogram4 defaults to a per-step guidance schedule. Its pipeline
             # rejects passing that default together with a constant scale.
@@ -812,8 +842,26 @@ def _call_model(
             # Some pipelines (e.g., Z-Image img2img) can't handle guidance_scale=None.
             if kwargs.get("guidance_scale", "unset") is None:
                 kwargs.pop("guidance_scale", None)
-            self._filter_kwargs(model, kwargs)
-            images = model(**kwargs).images
+            if _num_pipeline_calls == 1:
+                self._process_progressor(kwargs, progressor=progressor)
+                self._filter_kwargs(model, kwargs)
+                images = model(**kwargs).images
+            else:
+                self._filter_kwargs(model, kwargs)
+                images = []
+                generators = kwargs.get("generator")
+                for call_index in range(_num_pipeline_calls):
+                    per_call_kwargs = kwargs.copy()
+                    if isinstance(generators, list):
+                        per_call_kwargs["generator"] = generators[call_index]
+                    self._process_progressor(
+                        per_call_kwargs,
+                        progressor=progressor,
+                        pipeline_call_index=call_index,
+                        pipeline_call_count=_num_pipeline_calls,
+                    )
+                    self._filter_kwargs(model, per_call_kwargs)
+                    images.extend(model(**per_call_kwargs).images)
 
         if images and isinstance(images[0], (list, tuple)):
             images = list(itertools.chain.from_iterable(images))
@@ -995,12 +1043,14 @@ def image_to_image(
                 raise RuntimeError(f"{self._model_uid} does not support image2image")
             model = self._get_model(ability)
 
-        # These pipelines consume all reference images through their ``image``
-        # argument. The OpenAI-compatible endpoint exposes the first upload as
-        # ``image`` and the remaining uploads as ``reference_images``.
+        # Multi-reference pipelines consume all uploaded images through one
+        # pipeline argument. The OpenAI-compatible endpoint exposes the first
+        # upload as ``image`` and the rest as ``reference_images``.
+        is_joyai_image_edit_plus = self._is_joyai_image_edit_plus_model()
         if kwargs.get("reference_images") and (
             type(model).__name__ == "QwenImageEditPlusPipeline"
             or self._is_glm_image_model()
+            or is_joyai_image_edit_plus
         ):
             reference_images = kwargs.pop("reference_images")
             primary_images = image if isinstance(image, list) else [image]
@@ -1011,6 +1061,10 @@ def image_to_image(
             )
             image = primary_images + reference_images
 
+        # JoyImageEditPlusPipeline always expects a list, including one image.
+        if is_joyai_image_edit_plus and not isinstance(image, list):
+            image = [image]
+
         # GlmImagePipeline expects a list even for one conditioning image.
         if self._is_glm_image_model() and not isinstance(image, list):
             image = [image]
@@ -1038,7 +1092,7 @@ def image_to_image(
         else:
             # SD3 image2image cannot accept width and height
             allow_width_height = model_accept_param(["width", "height"], model)
-            if allow_width_height:
+            if allow_width_height and not is_joyai_image_edit_plus:
                 if isinstance(image, list):
                     kwargs["width"], kwargs["height"] = image[0].size
                 else:
@@ -1052,12 +1106,18 @@ def image_to_image(
         # generate config for lightning
         self._gen_config_for_lightning(kwargs)
 
+        if is_joyai_image_edit_plus:
+            kwargs["images"] = image
+            image = None
+
+        # JoyAI Image Edit Plus produces one image per pipeline invocation.
         return self._call_model(
             image=image,
             prompt=prompt,
-            num_images_per_prompt=n,
+            num_images_per_prompt=1 if is_joyai_image_edit_plus else n,
             response_format=response_format,
             model=model,
+            _num_pipeline_calls=n if is_joyai_image_edit_plus else 1,
             **kwargs,
         )