Repository: huggingface/huggingface_hub
Pull request #1567 (2023-07-21 09:37:45)

Send Accept: image/png as header for image tasks
cc @philschmid 

Tasks that expect an image to be returned by the server (i.e. image-to-image and text-to-image) must send `Accept: image/png` as header by default. This value can be overwritten by the user if needed. This PR fixes a bug happening with Inference Endpoints which returns a base64-encoded image in case Accept header is not set. This doesn't change anything for InferenceAPI calls.

**EDIT:** also fixes a bug in AsyncInferenceClient where headers set by the user where not taken into account.

**NOTE:** failing tests are unrelated

HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/huggingface_hub/pr_1567). All of your documentation changes will be reflected on that endpoint.
StephenHodgson: I also have not been able to get this to work myself. I think some tasks needed to be updated with the changes I made last month in https://github.com/huggingface/api-inference-community/pull/287. 
Wauplin: @StephenHodgson If I understand correctly this PR and [the one](https://github.com/huggingface/api-inference-community/pull/287) you contributed last month are not exactly related. Here it is to fix InferenceClient to make it work with Inference Endpoints by force a default image mime type. But agree it could be safer to for a mime type for audio as well (though it can be done in another PR) 
StephenHodgson: @Wauplin correct me if I'm wrong, but I'm pretty sure it uses the same infra, last I checked
 That's the main reason why I ended up making the PR there because I was trying to prove it out in my Unity client.
Wauplin: @StephenHodgson To clarify things:
- `InferenceAPI` is a service to quickly test and evaluate models hosted on the Hub. It's free of charge but only meant for quick prototyping. Docs are [here](https://huggingface.co/docs/api-inference/index). Part of the source code is public under [api-inference-community](https://github.com/huggingface/api-inference-community).
- `InferenceEndpoint` is a product to deploy models for production. This is a paid service to setup by the client (1 endpoint = 1 model = 1 GPU to pay). The entire implementation is private.
- `InferenceClient` is a client in `huggingface_hub` that can run inference with both `InferenceAPI` and `InferenceEndpoint`. The goal is to be compatible with both, and also endpoints hosted on other providers (AWS Sagemaker, Azure Endpoints,...). At the moment, we've more focused our attention on the HF solutions. 

Patch:
diff --git a/src/huggingface_hub/inference/_client.py b/src/huggingface_hub/inference/_client.py
index 33c704f3..7620f85c 100644
--- a/src/huggingface_hub/inference/_client.py
+++ b/src/huggingface_hub/inference/_client.py
@@ -54,6 +54,7 @@ from requests.structures import CaseInsensitiveDict
 
 from huggingface_hub.constants import INFERENCE_ENDPOINT
 from huggingface_hub.inference._common import (
+    TASKS_EXPECTING_IMAGES,
     ContentT,
     InferenceTimeoutError,
     _b64_encode,
@@ -204,6 +205,11 @@ class InferenceClient:
         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 @@ class InferenceClient:
                         url,
                         json=json,
                         data=data_as_binary,
-                        headers=self.headers,
+                        headers=headers,
                         cookies=self.cookies,
                         timeout=self.timeout,
                         stream=stream,
diff --git a/src/huggingface_hub/inference/_common.py b/src/huggingface_hub/inference/_common.py
index 45385fa4..88b930ce 100644
--- a/src/huggingface_hub/inference/_common.py
+++ b/src/huggingface_hub/inference/_common.py
@@ -62,6 +62,9 @@ PathT = Union[str, Path]
 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__)
 
 
diff --git a/src/huggingface_hub/inference/_generated/_async_client.py b/src/huggingface_hub/inference/_generated/_async_client.py
index b17fb4c9..3d95c25d 100644
--- a/src/huggingface_hub/inference/_generated/_async_client.py
+++ b/src/huggingface_hub/inference/_generated/_async_client.py
@@ -38,6 +38,7 @@ from requests.structures import CaseInsensitiveDict
 
 from huggingface_hub.constants import INFERENCE_ENDPOINT
 from huggingface_hub.inference._common import (
+    TASKS_EXPECTING_IMAGES,
     ContentT,
     InferenceTimeoutError,
     _async_stream_text_generation_response,
@@ -190,6 +191,11 @@ class AsyncInferenceClient:
         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:
@@ -197,11 +203,11 @@ class AsyncInferenceClient:
                 # 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:
diff --git a/utils/generate_async_inference_client.py b/utils/generate_async_inference_client.py
index e2b17851..43e185c2 100644
--- a/utils/generate_async_inference_client.py
+++ b/utils/generate_async_inference_client.py
@@ -164,6 +164,11 @@ ASYNC_POST_CODE = """
         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:
@@ -171,11 +176,11 @@ ASYNC_POST_CODE = """
                 # 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:


Test patch:
diff --git a/tests/test_inference_client.py b/tests/test_inference_client.py
index f7c6117c..fd635374 100644
--- a/tests/test_inference_client.py
+++ b/tests/test_inference_client.py
@@ -330,3 +330,15 @@ class TestHeadersAndCookies(unittest.TestCase):
             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")