Send Accept: image/png as header for image tasks

src/huggingface_hub/inference/_client.py

@@ -54,6 +54,7 @@
 
 from huggingface_hub.constants import INFERENCE_ENDPOINT
 from huggingface_hub.inference._common import (
+    TASKS_EXPECTING_IMAGES,
     ContentT,
     InferenceTimeoutError,
     _b64_encode,
@@ -204,6 +205,11 @@ def post(
         if data is not None and json is not None:
             warnings.warn("Ignoring `json` as `data` is passed as binary.")
 
+        # Set Accept header if relevant
+        headers = self.headers.copy()
+        if task in TASKS_EXPECTING_IMAGES and "Accept" not in headers:
+            headers["Accept"] = "image/png"
+
         t0 = time.time()
         timeout = self.timeout
         while True:
@@ -213,7 +219,7 @@ def post(
                         url,
                         json=json,
                         data=data_as_binary,
-                        headers=self.headers,
+                        headers=headers,
                         cookies=self.cookies,
                         timeout=self.timeout,
                         stream=stream,

src/huggingface_hub/inference/_common.py

@@ -62,6 +62,9 @@
 BinaryT = Union[bytes, BinaryIO]
 ContentT = Union[BinaryT, PathT, UrlT]
 
+# Use to set a Accept: image/png header
+TASKS_EXPECTING_IMAGES = {"text-to-image", "image-to-image"}
+
 logger = logging.getLogger(__name__)
 
 

src/huggingface_hub/inference/_generated/_async_client.py

@@ -38,6 +38,7 @@
 
 from huggingface_hub.constants import INFERENCE_ENDPOINT
 from huggingface_hub.inference._common import (
+    TASKS_EXPECTING_IMAGES,
     ContentT,
     InferenceTimeoutError,
     _async_stream_text_generation_response,
@@ -190,18 +191,23 @@ async def post(
         if data is not None and json is not None:
             warnings.warn("Ignoring `json` as `data` is passed as binary.")
 
+        # Set Accept header if relevant
+        headers = self.headers.copy()
+        if task in TASKS_EXPECTING_IMAGES and "Accept" not in headers:
+            headers["Accept"] = "image/png"
+
         t0 = time.time()
         timeout = self.timeout
         while True:
             with _open_as_binary(data) as data_as_binary:
                 # Do not use context manager as we don't want to close the connection immediately when returning
                 # a stream
                 client = aiohttp.ClientSession(
-                    headers=self.headers, cookies=self.cookies, timeout=aiohttp.ClientTimeout(self.timeout)
+                    headers=headers, cookies=self.cookies, timeout=aiohttp.ClientTimeout(self.timeout)
                 )
 
                 try:
-                    response = await client.post(url, headers=build_hf_headers(), json=json, data=data_as_binary)
+                    response = await client.post(url, json=json, data=data_as_binary)
                     response_error_payload = None
                     if response.status != 200:
                         try:

tests/test_inference_client.py

@@ -330,3 +330,15 @@ def test_mocked_post(self, get_session_mock: MagicMock) -> None:
             timeout=None,
             stream=False,
         )
+
+    @patch("huggingface_hub.inference._client._bytes_to_image")
+    @patch("huggingface_hub.inference._client.get_session")
+    def test_accept_header_image(self, get_session_mock: MagicMock, bytes_to_image_mock: MagicMock) -> None:
+        """Test that Accept: image/png header is set for image tasks."""
+        client = InferenceClient()
+
+        response = client.text_to_image("An astronaut riding a horse")
+        self.assertEqual(response, bytes_to_image_mock.return_value)
+
+        headers = get_session_mock().post.call_args_list[0].kwargs["headers"]
+        self.assertEqual(headers["Accept"], "image/png")

utils/generate_async_inference_client.py

@@ -164,18 +164,23 @@ def _rename_to_AsyncInferenceClient(code: str) -> str:
         if data is not None and json is not None:
             warnings.warn("Ignoring `json` as `data` is passed as binary.")
 
+        # Set Accept header if relevant
+        headers = self.headers.copy()
+        if task in TASKS_EXPECTING_IMAGES and "Accept" not in headers:
+            headers["Accept"] = "image/png"
+
         t0 = time.time()
         timeout = self.timeout
         while True:
             with _open_as_binary(data) as data_as_binary:
                 # Do not use context manager as we don't want to close the connection immediately when returning
                 # a stream
                 client = aiohttp.ClientSession(
-                    headers=self.headers, cookies=self.cookies, timeout=aiohttp.ClientTimeout(self.timeout)
+                    headers=headers, cookies=self.cookies, timeout=aiohttp.ClientTimeout(self.timeout)
                 )
 
                 try:
-                    response = await client.post(url, headers=build_hf_headers(), json=json, data=data_as_binary)
+                    response = await client.post(url, json=json, data=data_as_binary)
                     response_error_payload = None
                     if response.status != 200:
                         try: