FEAT: OpenAI image edit api support (#4110)

Author: OliverBryant

.github/workflows/python.yaml

@@ -45,7 +45,7 @@ jobs:
           sortPaths: "xinference"
           configuration: "--check-only --diff --sp setup.cfg"
       - name: mypy
-        run: pip install 'mypy<1.16.0' && mypy --install-types --non-interactive xinference
+        run: pip install 'mypy==1.18.1' && mypy --install-types --non-interactive xinference
       - name: codespell
         run: pip install codespell && codespell --ignore-words-list thirdparty xinference
       - name: Set up Node.js

xinference/api/restful_api.py

@@ -74,6 +74,7 @@ def __init__(self, model_uid: str, base_url: str, auth_headers: Dict):
         self._model_uid = model_uid
         self._base_url = base_url
         self.auth_headers = auth_headers
+        self.timeout = aiohttp.ClientTimeout(total=1800)
         self.session = aiohttp.ClientSession(
             connector=aiohttp.TCPConnector(force_close=True)
         )
@@ -356,7 +357,7 @@ async def image_to_image(
         else:
             # Single image
             files.append(("image", ("image", image, "application/octet-stream")))
-        response = await self.session.post(url, files=files, headers=self.auth_headers)
+        response = await self.session.post(url, data=files, headers=self.auth_headers)
         if response.status != 200:
             raise RuntimeError(
                 f"Failed to variants the images, detail: {await _get_error_string(response)}"
@@ -366,6 +367,157 @@ async def image_to_image(
         await _release_response(response)
         return response_data
 
+    async def image_edit(
+        self,
+        image: Union[Union[str, bytes], List[Union[str, bytes]]],
+        prompt: str,
+        mask: Optional[Union[str, bytes]] = None,
+        n: int = 1,
+        size: Optional[str] = None,
+        response_format: str = "url",
+        **kwargs,
+    ) -> "ImageList":
+        """
+        Edit image(s) by the input text and optional mask.
+
+        Parameters
+        ----------
+        image: `Union[Union[str, bytes], List[Union[str, bytes]]]`
+            The input image(s) to edit. Can be:
+            - Single image: file path, URL, or binary image data
+            - Multiple images: list of file paths, URLs, or binary image data
+            When multiple images are provided, the first image is used as the primary image
+            and subsequent images are used as reference images for better editing results.
+        prompt: `str`
+            The prompt or prompts to guide image editing. If not defined, you need to pass `prompt_embeds`.
+        mask: `Optional[Union[str, bytes]]`, optional
+            An optional mask image. White pixels in the mask are repainted while black pixels are preserved.
+            If provided, this will trigger inpainting mode. If not provided, this will trigger image-to-image mode.
+        n: `int`, defaults to 1
+            The number of images to generate per prompt. Must be between 1 and 10.
+        size: `Optional[str]`, optional
+            The width*height in pixels of the generated image. If not specified, uses the original image size.
+        response_format: `str`, defaults to `url`
+            The format in which the generated images are returned. Must be one of url or b64_json.
+        **kwargs
+            Additional parameters to pass to the model.
+
+        Returns
+        -------
+        ImageList
+            A list of edited image objects.
+
+        Raises
+        ------
+        RuntimeError
+            If the image editing request fails.
+
+        Examples
+        --------
+        # Single image editing
+        result = await model.image_edit(
+            image="path/to/image.png",
+            prompt="make this image look like a painting"
+        )
+
+        # Multiple image editing with reference images
+        result = await model.image_edit(
+            image=["primary_image.png", "reference1.jpg", "reference2.png"],
+            prompt="edit the main image using the style from reference images"
+        )
+        """
+        url = f"{self._base_url}/v1/images/edits"
+        params = {
+            "model": self._model_uid,
+            "prompt": prompt,
+            "n": n,
+            "size": size,
+            "response_format": response_format,
+            "kwargs": json.dumps(kwargs),
+        }
+        params = _filter_params(params)
+        files: List[Any] = []
+        for key, value in params.items():
+            files.append((key, (None, value)))
+
+        # Handle single image or multiple images
+        import aiohttp
+
+        data = aiohttp.FormData()
+
+        # Add all parameters as form fields
+        for key, value in params.items():
+            if value is not None:
+                data.add_field(key, str(value))
+
+        # Handle single image or multiple images
+        if isinstance(image, list):
+            # Validate image list is not empty
+            if len(image) == 0:
+                raise ValueError("Image list cannot be empty")
+            # Multiple images - send as image[] array
+            for i, img in enumerate(image):
+                if isinstance(img, str):
+                    # File path - read file content
+                    with open(img, "rb") as f:
+                        content = f.read()
+                    data.add_field(
+                        f"image[]",
+                        content,
+                        filename=f"image_{i}.png",
+                        content_type="image/png",
+                    )
+                else:
+                    # Binary data
+                    data.add_field(
+                        f"image[]",
+                        img,
+                        filename=f"image_{i}.png",
+                        content_type="image/png",
+                    )
+        else:
+            # Single image
+            if isinstance(image, str):
+                # File path - read file content
+                with open(image, "rb") as f:
+                    content = f.read()
+                data.add_field(
+                    "image", content, filename="image.png", content_type="image/png"
+                )
+            else:
+                # Binary data
+                data.add_field(
+                    "image", image, filename="image.png", content_type="image/png"
+                )
+
+        if mask is not None:
+            if isinstance(mask, str):
+                # File path - read file content
+                with open(mask, "rb") as f:
+                    content = f.read()
+                data.add_field(
+                    "mask", content, filename="mask.png", content_type="image/png"
+                )
+            else:
+                # Binary data
+                data.add_field(
+                    "mask", mask, filename="mask.png", content_type="image/png"
+                )
+
+        try:
+            response = await self.session.post(
+                url, data=data, headers=self.auth_headers
+            )
+            if response.status != 200:
+                raise RuntimeError(
+                    f"Failed to edit the images, detail: {await _get_error_string(response)}"
+                )
+
+            response_data = await response.json()
+            return response_data
+        finally:
+            await _release_response(response) if "response" in locals() else None
+
     async def inpainting(
         self,
         image: Union[str, bytes],
@@ -436,7 +588,7 @@ async def inpainting(
                 ("mask_image", mask_image, "application/octet-stream"),
             )
         )
-        response = await self.session.post(url, files=files, headers=self.auth_headers)
+        response = await self.session.post(url, data=files, headers=self.auth_headers)
         if response.status != 200:
             raise RuntimeError(
                 f"Failed to inpaint the images, detail: {await _get_error_string(response)}"
@@ -457,7 +609,7 @@ async def ocr(self, image: Union[str, bytes], **kwargs):
         for key, value in params.items():
             files.append((key, (None, value)))
         files.append(("image", ("image", image, "application/octet-stream")))
-        response = await self.session.post(url, files=files, headers=self.auth_headers)
+        response = await self.session.post(url, data=files, headers=self.auth_headers)
         if response.status != 200:
             raise RuntimeError(
                 f"Failed to ocr the images, detail: {await _get_error_string(response)}"
@@ -547,7 +699,7 @@ async def image_to_video(
         for key, value in params.items():
             files.append((key, (None, value)))
         files.append(("image", ("image", image, "application/octet-stream")))
-        response = await self.session.post(url, files=files, headers=self.auth_headers)
+        response = await self.session.post(url, data=files, headers=self.auth_headers)
         if response.status != 200:
             raise RuntimeError(
                 f"Failed to create the video from image, detail: {await _get_error_string(response)}"
@@ -987,8 +1139,9 @@ def __init__(self, base_url, api_key: Optional[str] = None):
         self.base_url = base_url
         self._headers: Dict[str, str] = {}
         self._cluster_authed = False
+        self.timeout = aiohttp.ClientTimeout(total=1800)
         self.session = aiohttp.ClientSession(
-            connector=aiohttp.TCPConnector(force_close=True)
+            connector=aiohttp.TCPConnector(force_close=True), timeout=self.timeout
         )
         self._check_cluster_authenticated()
         if api_key is not None and self._cluster_authed:

xinference/client/restful/async_restful_client.py

@@ -329,6 +329,137 @@ def image_to_image(
         response_data = response.json()
         return response_data
 
+    def image_edit(
+        self,
+        image: Union[Union[str, bytes], List[Union[str, bytes]]],
+        prompt: str,
+        mask: Optional[Union[str, bytes]] = None,
+        n: int = 1,
+        size: Optional[str] = None,
+        response_format: str = "url",
+        **kwargs,
+    ) -> "ImageList":
+        """
+        Edit image(s) by the input text and optional mask.
+
+        Parameters
+        ----------
+        image: `Union[Union[str, bytes], List[Union[str, bytes]]]`
+            The input image(s) to edit. Can be:
+            - Single image: file path, URL, or binary image data
+            - Multiple images: list of file paths, URLs, or binary image data
+            When multiple images are provided, the first image is used as the primary image
+            and subsequent images are used as reference images for better editing results.
+        prompt: `str`
+            The prompt or prompts to guide image editing. If not defined, you need to pass `prompt_embeds`.
+        mask: `Optional[Union[str, bytes]]`, optional
+            An optional mask image. White pixels in the mask are repainted while black pixels are preserved.
+            If provided, this will trigger inpainting mode. If not provided, this will trigger image-to-image mode.
+        n: `int`, defaults to 1
+            The number of images to generate per prompt. Must be between 1 and 10.
+        size: `Optional[str]`, optional
+            The width*height in pixels of the generated image. If not specified, uses the original image size.
+        response_format: `str`, defaults to `url`
+            The format in which the generated images are returned. Must be one of url or b64_json.
+        **kwargs
+            Additional parameters to pass to the model.
+
+        Returns
+        -------
+        ImageList
+            A list of edited image objects.
+
+        Raises
+        ------
+        RuntimeError
+            If the image editing request fails.
+
+        Examples
+        --------
+        # Single image editing
+        result = model.image_edit(
+            image="path/to/image.png",
+            prompt="make this image look like a painting"
+        )
+
+        # Multiple image editing with reference images
+        result = model.image_edit(
+            image=["primary_image.png", "reference1.jpg", "reference2.png"],
+            prompt="edit the main image using the style from reference images"
+        )
+        """
+        url = f"{self._base_url}/v1/images/edits"
+        params = {
+            "model": self._model_uid,
+            "prompt": prompt,
+            "n": n,
+            "size": size,
+            "response_format": response_format,
+            "kwargs": json.dumps(kwargs),
+        }
+        files: List[Any] = []
+        for key, value in params.items():
+            if value is not None:
+                files.append((key, (None, value)))
+
+        # Handle single image or multiple images using requests format
+        if isinstance(image, list):
+            # Validate image list is not empty
+            if len(image) == 0:
+                raise ValueError("Image list cannot be empty")
+            # Multiple images - send as image[] array
+            for i, img in enumerate(image):
+                if isinstance(img, str):
+                    # File path - open file
+                    f = open(img, "rb")
+                    files.append(
+                        (f"image[]", (f"image_{i}", f, "application/octet-stream"))
+                    )
+                else:
+                    # Binary data
+                    files.append(
+                        (f"image[]", (f"image_{i}", img, "application/octet-stream"))
+                    )
+        else:
+            # Single image
+            if isinstance(image, str):
+                # File path - open file
+                f = open(image, "rb")
+                files.append(("image", ("image", f, "application/octet-stream")))
+            else:
+                # Binary data
+                files.append(("image", ("image", image, "application/octet-stream")))
+
+        if mask is not None:
+            if isinstance(mask, str):
+                # File path - open file
+                f = open(mask, "rb")
+                files.append(("mask", ("mask", f, "application/octet-stream")))
+            else:
+                # Binary data
+                files.append(("mask", ("mask", mask, "application/octet-stream")))
+
+        try:
+            response = self.session.post(url, files=files, headers=self.auth_headers)
+            if response.status_code != 200:
+                raise RuntimeError(
+                    f"Failed to edit the images, detail: {_get_error_string(response)}"
+                )
+
+            response_data = response.json()
+            return response_data
+        finally:
+            # Close all opened files
+            for file_item in files:
+                if (
+                    len(file_item) >= 2
+                    and hasattr(file_item[1], "__len__")
+                    and len(file_item[1]) >= 2
+                ):
+                    file_obj = file_item[1][1]
+                    if hasattr(file_obj, "close"):
+                        file_obj.close()
+
     def inpainting(
         self,
         image: Union[str, bytes],

xinference/client/restful/restful_client.py

@@ -1661,6 +1661,9 @@ async def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]:
         if isinstance(worker_ref, list):
             # get first worker to fetch information if model across workers
             worker_ref = worker_ref[0]
+        assert not isinstance(
+            worker_ref, (list, tuple)
+        ), "worker_ref must be a single worker"
         return await worker_ref.get_model(model_uid=replica_model_uid)
 
     @log_async(logger=logger)
@@ -1673,6 +1676,9 @@ async def get_model_status(self, replica_model_uid: str):
         if isinstance(worker_ref, list):
             # get status from first shard if model has multiple shards across workers
             worker_ref = worker_ref[0]
+        assert not isinstance(
+            worker_ref, (list, tuple)
+        ), "worker_ref must be a single worker"
         return await worker_ref.get_model_status(replica_model_uid)
 
     @log_async(logger=logger)
@@ -1691,6 +1697,9 @@ async def describe_model(self, model_uid: str) -> Dict[str, Any]:
         if isinstance(worker_ref, list):
             # get status from first shard if model has multiple shards across workers
             worker_ref = worker_ref[0]
+        assert not isinstance(
+            worker_ref, (list, tuple)
+        ), "worker_ref must be a single worker"
         info = await worker_ref.describe_model(model_uid=replica_model_uid)
         info["replica"] = replica_info.replica
         return info
@@ -1766,6 +1775,9 @@ async def abort_request(
             if isinstance(worker_ref, list):
                 # get status from first shard if model has multiple shards across workers
                 worker_ref = worker_ref[0]
+            assert not isinstance(
+                worker_ref, (list, tuple)
+            ), "worker_ref must be a single worker"
             model_ref = await worker_ref.get_model(model_uid=rep_mid)
             result_info = await model_ref.abort_request(request_id, block_duration)
             res["msg"] = result_info