[Inference Providers] implement image-segmentation for fal (#3521)

* implement image-segmentation for fal.ai

* Update src/huggingface_hub/inference/_providers/fal_ai.py

Co-authored-by: Lucain <lucain@huggingface.co>

* update tests

---------

Co-authored-by: Lucain <lucain@huggingface.co>

Author: hanouticelina

docs/source/en/guides/inference.md

@@ -202,7 +202,7 @@ For more details, refer to the [Inference Providers pricing documentation](https
 | [`~InferenceClient.feature_extraction`]             | ❌                 | ❌        | ❌        | ❌      | ❌      | ❌              | ❌            | ❌    | ✅            | ❌          | ✅                | ❌         | ❌      | ❌          | ❌         | ✅         | ✅         | ❌        | ❌         | ❌   |
 | [`~InferenceClient.fill_mask`]                      | ❌                 | ❌        | ❌        | ❌      | ❌      | ❌              | ❌            | ❌    | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        | ❌      | ❌          | ❌         | ❌         | ❌   |
 | [`~InferenceClient.image_classification`]           | ❌                 | ❌        | ❌        | ❌      | ❌      | ❌              | ❌            | ❌    | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        | ❌      | ❌          | ❌         | ❌         | ❌   |
-| [`~InferenceClient.image_segmentation`]             | ❌                 | ❌        | ❌        | ❌      | ❌      | ❌              | ❌            | ❌    | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        | ❌      | ❌          | ❌         | ❌         | ❌   |
+| [`~InferenceClient.image_segmentation`]             | ❌                 | ❌        | ❌        | ❌      | ✅      | ❌              | ❌            | ❌    | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        | ❌      | ❌          | ❌         | ❌         | ❌   |
 | [`~InferenceClient.image_to_image`]                 | ❌                 | ❌        | ❌        | ❌      | ✅      | ❌              | ❌            | ❌    | ✅            | ❌          | ❌                | ❌         | ❌      | ❌          | ✅         | ❌         | ❌         | ❌        | ✅         | ❌   |
 | [`~InferenceClient.image_to_video`]                 | ❌                 | ❌        | ❌        | ❌      | ✅      | ❌              | ❌            | ❌    | ❌            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        | ❌      | ❌          | ❌         | ✅         | ❌   |
 | [`~InferenceClient.image_to_text`]                  | ❌                 | ❌        | ❌        | ❌      | ❌      | ❌              | ❌            | ❌    | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        | ❌      | ❌          | ❌         | ❌         | ❌   |

src/huggingface_hub/inference/_client.py

@@ -1262,6 +1262,7 @@ def image_segmentation(
             api_key=self.token,
         )
         response = self._inner_post(request_parameters)
+        response = provider_helper.get_response(response, request_parameters)
         output = ImageSegmentationOutputElement.parse_obj_as_list(response)
         for item in output:
             item.mask = _b64_to_image(item.mask)  # type: ignore [assignment]

src/huggingface_hub/inference/_generated/_async_client.py

@@ -1293,6 +1293,7 @@ async def image_segmentation(
             api_key=self.token,
         )
         response = await self._inner_post(request_parameters)
+        response = provider_helper.get_response(response, request_parameters)
         output = ImageSegmentationOutputElement.parse_obj_as_list(response)
         for item in output:
             item.mask = _b64_to_image(item.mask)  # type: ignore [assignment]

src/huggingface_hub/inference/_providers/__init__.py

@@ -13,6 +13,7 @@
 from .cohere import CohereConversationalTask
 from .fal_ai import (
     FalAIAutomaticSpeechRecognitionTask,
+    FalAIImageSegmentationTask,
     FalAIImageToImageTask,
     FalAIImageToVideoTask,
     FalAITextToImageTask,
@@ -102,6 +103,7 @@
         "text-to-video": FalAITextToVideoTask(),
         "image-to-video": FalAIImageToVideoTask(),
         "image-to-image": FalAIImageToImageTask(),
+        "image-segmentation": FalAIImageSegmentationTask(),
     },
     "featherless-ai": {
         "conversational": FeatherlessConversationalTask(),

src/huggingface_hub/inference/_providers/fal_ai.py

@@ -246,3 +246,54 @@ def get_response(
         output = super().get_response(response, request_params)
         url = _as_dict(output)["video"]["url"]
         return get_session().get(url).content
+
+
+class FalAIImageSegmentationTask(FalAIQueueTask):
+    def __init__(self):
+        super().__init__("image-segmentation")
+
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
+        image_url = _as_url(inputs, default_mime_type="image/png")
+        payload: dict[str, Any] = {
+            "image_url": image_url,
+            **filter_none(parameters),
+            "sync_mode": True,
+        }
+        return payload
+
+    def get_response(
+        self,
+        response: Union[bytes, dict],
+        request_params: Optional[RequestParameters] = None,
+    ) -> Any:
+        result = super().get_response(response, request_params)
+        result_dict = _as_dict(result)
+
+        if "image" not in result_dict:
+            raise ValueError(f"Response from fal ai image-segmentation API does not contain an image: {result_dict}")
+
+        image_data = result_dict["image"]
+        if "url" not in image_data:
+            raise ValueError(f"Image data from fal ai image-segmentation API does not contain a URL: {image_data}")
+
+        image_url = image_data["url"]
+
+        if isinstance(image_url, str) and image_url.startswith("data:"):
+            if "," in image_url:
+                mask_base64 = image_url.split(",", 1)[1]
+            else:
+                raise ValueError(f"Invalid data URL format: {image_url}")
+        else:
+            # or it's a regular URL, fetch it
+            mask_response = get_session().get(image_url)
+            hf_raise_for_status(mask_response)
+            mask_base64 = base64.b64encode(mask_response.content).decode()
+
+        return [
+            {
+                "label": "mask",
+                "mask": mask_base64,
+            }
+        ]

tests/test_inference_providers.py

@@ -22,6 +22,7 @@
 from huggingface_hub.inference._providers.fal_ai import (
     _POLLING_INTERVAL,
     FalAIAutomaticSpeechRecognitionTask,
+    FalAIImageSegmentationTask,
     FalAIImageToImageTask,
     FalAIImageToVideoTask,
     FalAITextToImageTask,
@@ -629,6 +630,124 @@ def test_image_to_video_response(self, mocker):
         mock_sleep.assert_called_once_with(_POLLING_INTERVAL)
         assert response == b"video_content"
 
+    def test_image_segmentation_payload(self):
+        helper = FalAIImageSegmentationTask()
+        mapping_info = InferenceProviderMapping(
+            provider="fal-ai",
+            hf_model_id="briaai/RMBG-2.0",
+            providerId="fal-ai/rmbg-2.0",
+            task="image-segmentation",
+            status="live",
+        )
+        payload = helper._prepare_payload_as_dict("https://example.com/image.png", {"threshold": 0.5}, mapping_info)
+        assert payload == {"image_url": "https://example.com/image.png", "threshold": 0.5, "sync_mode": True}
+
+        payload = helper._prepare_payload_as_dict(b"dummy_image_data", {"mask_threshold": 0.8}, mapping_info)
+        assert payload == {
+            "image_url": f"data:image/png;base64,{base64.b64encode(b'dummy_image_data').decode()}",
+            "mask_threshold": 0.8,
+            "sync_mode": True,
+        }
+
+    def test_image_segmentation_response_with_data_url(self, mocker):
+        """Test image segmentation response when image URL is a data URL."""
+        helper = FalAIImageSegmentationTask()
+        mock_session = mocker.patch("huggingface_hub.inference._providers.fal_ai.get_session")
+        mock_sleep = mocker.patch("huggingface_hub.inference._providers.fal_ai.time.sleep")
+        dummy_mask_base64 = base64.b64encode(b"mask_content").decode()
+        data_url = f"data:image/png;base64,{dummy_mask_base64}"
+        mock_session.return_value.get.side_effect = [
+            # First call: status
+            mocker.Mock(json=lambda: {"status": "COMPLETED"}, headers={"Content-Type": "application/json"}),
+            # Second call: get result
+            mocker.Mock(json=lambda: {"image": {"url": data_url}}, headers={"Content-Type": "application/json"}),
+        ]
+        api_key = helper._prepare_api_key("hf_token")
+        headers = helper._prepare_headers({}, api_key)
+        url = helper._prepare_url(api_key, "username/repo_name")
+
+        request_params = RequestParameters(
+            url=url,
+            headers=headers,
+            task="image-segmentation",
+            model="username/repo_name",
+            data=None,
+            json=None,
+        )
+        response = helper.get_response(
+            b'{"request_id": "test_request_id", "status": "PROCESSING", "response_url": "https://queue.fal.run/username_provider/repo_name_provider/requests/test_request_id", "status_url": "https://queue.fal.run/username_provider/repo_name_provider/requests/test_request_id/status"}',
+            request_params,
+        )
+
+        # Verify the correct URLs were called (only status and result, no fetch needed for data URL)
+        assert mock_session.return_value.get.call_count == 2
+        mock_session.return_value.get.assert_has_calls(
+            [
+                mocker.call(
+                    "https://router.huggingface.co/fal-ai/username_provider/repo_name_provider/requests/test_request_id/status?_subdomain=queue",
+                    headers=request_params.headers,
+                ),
+                mocker.call(
+                    "https://router.huggingface.co/fal-ai/username_provider/repo_name_provider/requests/test_request_id?_subdomain=queue",
+                    headers=request_params.headers,
+                ),
+            ]
+        )
+        mock_sleep.assert_called_once_with(_POLLING_INTERVAL)
+        assert response == [{"label": "mask", "mask": dummy_mask_base64}]
+
+    def test_image_segmentation_response_with_regular_url(self, mocker):
+        """Test image segmentation response when image URL is a regular HTTP URL."""
+        helper = FalAIImageSegmentationTask()
+        mock_session = mocker.patch("huggingface_hub.inference._providers.fal_ai.get_session")
+        mock_sleep = mocker.patch("huggingface_hub.inference._providers.fal_ai.time.sleep")
+        dummy_mask_base64 = base64.b64encode(b"mask_content").decode()
+        mock_session.return_value.get.side_effect = [
+            # First call: status
+            mocker.Mock(json=lambda: {"status": "COMPLETED"}, headers={"Content-Type": "application/json"}),
+            # Second call: get result
+            mocker.Mock(
+                json=lambda: {"image": {"url": "https://example.com/mask.png"}},
+                headers={"Content-Type": "application/json"},
+            ),
+            # Third call: get mask content
+            mocker.Mock(content=b"mask_content", raise_for_status=lambda: None),
+        ]
+        api_key = helper._prepare_api_key("hf_token")
+        headers = helper._prepare_headers({}, api_key)
+        url = helper._prepare_url(api_key, "username/repo_name")
+
+        request_params = RequestParameters(
+            url=url,
+            headers=headers,
+            task="image-segmentation",
+            model="username/repo_name",
+            data=None,
+            json=None,
+        )
+        response = helper.get_response(
+            b'{"request_id": "test_request_id", "status": "PROCESSING", "response_url": "https://queue.fal.run/username_provider/repo_name_provider/requests/test_request_id", "status_url": "https://queue.fal.run/username_provider/repo_name_provider/requests/test_request_id/status"}',
+            request_params,
+        )
+
+        # Verify the correct URLs were called (status, result, and mask fetch)
+        assert mock_session.return_value.get.call_count == 3
+        mock_session.return_value.get.assert_has_calls(
+            [
+                mocker.call(
+                    "https://router.huggingface.co/fal-ai/username_provider/repo_name_provider/requests/test_request_id/status?_subdomain=queue",
+                    headers=request_params.headers,
+                ),
+                mocker.call(
+                    "https://router.huggingface.co/fal-ai/username_provider/repo_name_provider/requests/test_request_id?_subdomain=queue",
+                    headers=request_params.headers,
+                ),
+                mocker.call("https://example.com/mask.png"),
+            ]
+        )
+        mock_sleep.assert_called_once_with(_POLLING_INTERVAL)
+        assert response == [{"label": "mask", "mask": dummy_mask_base64}]
+
 
 class TestFeatherlessAIProvider:
     def test_prepare_route_chat_completionurl(self):