[InferenceClient] Add content-type header whenever possible + refacto (#3321)

* [InferenceClient] Add content-type header whenever possible + refacto

* address comments

* Update src/huggingface_hub/inference/_common.py

Co-authored-by: célina <hanouticelina@gmail.com>

---------

Co-authored-by: célina <hanouticelina@gmail.com>

Author: Wauplin

src/huggingface_hub/inference/_client.py

@@ -53,7 +53,6 @@
     _bytes_to_list,
     _get_unsupported_text_generation_kwargs,
     _import_numpy,
-    _open_as_binary,
     _set_unsupported_text_generation_kwargs,
     _stream_chat_completion_response,
     _stream_text_generation_response,
@@ -257,21 +256,20 @@ def _inner_post(
         if request_parameters.task in TASKS_EXPECTING_IMAGES and "Accept" not in request_parameters.headers:
             request_parameters.headers["Accept"] = "image/png"
 
-        with _open_as_binary(request_parameters.data) as data_as_binary:
-            try:
-                response = get_session().post(
-                    request_parameters.url,
-                    json=request_parameters.json,
-                    data=data_as_binary,
-                    headers=request_parameters.headers,
-                    cookies=self.cookies,
-                    timeout=self.timeout,
-                    stream=stream,
-                    proxies=self.proxies,
-                )
-            except TimeoutError as error:
-                # Convert any `TimeoutError` to a `InferenceTimeoutError`
-                raise InferenceTimeoutError(f"Inference call timed out: {request_parameters.url}") from error  # type: ignore
+        try:
+            response = get_session().post(
+                request_parameters.url,
+                json=request_parameters.json,
+                data=request_parameters.data,
+                headers=request_parameters.headers,
+                cookies=self.cookies,
+                timeout=self.timeout,
+                stream=stream,
+                proxies=self.proxies,
+            )
+        except TimeoutError as error:
+            # Convert any `TimeoutError` to a `InferenceTimeoutError`
+            raise InferenceTimeoutError(f"Inference call timed out: {request_parameters.url}") from error  # type: ignore
 
         try:
             hf_raise_for_status(response)

src/huggingface_hub/inference/_common.py

@@ -19,17 +19,14 @@
 import json
 import logging
 import mimetypes
-from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
     AsyncIterable,
     BinaryIO,
-    ContextManager,
     Dict,
-    Generator,
     Iterable,
     List,
     Literal,
@@ -61,8 +58,7 @@
 # TYPES
 UrlT = str
 PathT = Union[str, Path]
-BinaryT = Union[bytes, BinaryIO]
-ContentT = Union[BinaryT, PathT, UrlT, "Image"]
+ContentT = Union[bytes, BinaryIO, PathT, UrlT, "Image", bytearray, memoryview]
 
 # Use to set a Accept: image/png header
 TASKS_EXPECTING_IMAGES = {"text-to-image", "image-to-image"}
@@ -76,10 +72,35 @@ class RequestParameters:
     task: str
     model: Optional[str]
     json: Optional[Union[str, Dict, List]]
-    data: Optional[ContentT]
+    data: Optional[bytes]
     headers: Dict[str, Any]
 
 
+class MimeBytes(bytes):
+    """
+    A bytes object with a mime type.
+    To be returned by `_prepare_payload_open_as_mime_bytes` in subclasses.
+
+    Example:
+    ```python
+        >>> b = MimeBytes(b"hello", "text/plain")
+        >>> isinstance(b, bytes)
+        True
+        >>> b.mime_type
+        'text/plain'
+    ```
+    """
+
+    mime_type: Optional[str]
+
+    def __new__(cls, data: bytes, mime_type: Optional[str] = None):
+        obj = super().__new__(cls, data)
+        obj.mime_type = mime_type
+        if isinstance(data, MimeBytes) and mime_type is None:
+            obj.mime_type = data.mime_type
+        return obj
+
+
 ## IMPORT UTILS
 
 
@@ -117,31 +138,49 @@ def _import_pil_image():
 
 
 @overload
-def _open_as_binary(
-    content: ContentT,
-) -> ContextManager[BinaryT]: ...  # means "if input is not None, output is not None"
+def _open_as_mime_bytes(content: ContentT) -> MimeBytes: ...  # means "if input is not None, output is not None"
 
 
 @overload
-def _open_as_binary(
-    content: Literal[None],
-) -> ContextManager[Literal[None]]: ...  # means "if input is None, output is None"
+def _open_as_mime_bytes(content: Literal[None]) -> Literal[None]: ...  # means "if input is None, output is None"
 
 
-@contextmanager  # type: ignore
-def _open_as_binary(content: Optional[ContentT]) -> Generator[Optional[BinaryT], None, None]:
+def _open_as_mime_bytes(content: Optional[ContentT]) -> Optional[MimeBytes]:
     """Open `content` as a binary file, either from a URL, a local path, raw bytes, or a PIL Image.
 
     Do nothing if `content` is None.
-
-    TODO: handle base64 as input
     """
+    # If content is None, yield None
+    if content is None:
+        return None
+
+    # If content is bytes, return it
+    if isinstance(content, bytes):
+        return MimeBytes(content)
+
+    # If content is raw binary data (bytearray, memoryview)
+    if isinstance(content, (bytearray, memoryview)):
+        return MimeBytes(bytes(content))
+
+    # If content is a binary file-like object
+    if hasattr(content, "read"):  # duck-typing instead of isinstance(content, BinaryIO)
+        logger.debug("Reading content from BinaryIO")
+        data = content.read()
+        mime_type = mimetypes.guess_type(content.name)[0] if hasattr(content, "name") else None
+        if isinstance(data, str):
+            raise TypeError("Expected binary stream (bytes), but got text stream")
+        return MimeBytes(data, mime_type=mime_type)
+
     # If content is a string => must be either a URL or a path
     if isinstance(content, str):
         if content.startswith("https://") or content.startswith("http://"):
             logger.debug(f"Downloading content from {content}")
-            yield get_session().get(content).content  # TODO: retrieve as stream and pipe to post request ?
-            return
+            response = get_session().get(content)
+            mime_type = response.headers.get("Content-Type")
+            if mime_type is None:
+                mime_type = mimetypes.guess_type(content)[0]
+            return MimeBytes(response.content, mime_type=mime_type)
+
         content = Path(content)
         if not content.exists():
             raise FileNotFoundError(
@@ -152,9 +191,7 @@ def _open_as_binary(content: Optional[ContentT]) -> Generator[Optional[BinaryT],
     # If content is a Path => open it
     if isinstance(content, Path):
         logger.debug(f"Opening content from {content}")
-        with content.open("rb") as f:
-            yield f
-        return
+        return MimeBytes(content.read_bytes(), mime_type=mimetypes.guess_type(content)[0])
 
     # If content is a PIL Image => convert to bytes
     if is_pillow_available():
@@ -163,38 +200,37 @@ def _open_as_binary(content: Optional[ContentT]) -> Generator[Optional[BinaryT],
         if isinstance(content, Image.Image):
             logger.debug("Converting PIL Image to bytes")
             buffer = io.BytesIO()
-            content.save(buffer, format=content.format or "PNG")
-            yield buffer.getvalue()
-            return
+            format = content.format or "PNG"
+            content.save(buffer, format=format)
+            return MimeBytes(buffer.getvalue(), mime_type=f"image/{format.lower()}")
 
-    # Otherwise: already a file-like object or None
-    yield content  # type: ignore
+    # If nothing matched, raise error
+    raise TypeError(
+        f"Unsupported content type: {type(content)}. "
+        "Expected one of: bytes, bytearray, BinaryIO, memoryview, Path, str (URL or file path), or PIL.Image.Image."
+    )
 
 
 def _b64_encode(content: ContentT) -> str:
     """Encode a raw file (image, audio) into base64. Can be bytes, an opened file, a path or a URL."""
-    with _open_as_binary(content) as data:
-        data_as_bytes = data if isinstance(data, bytes) else data.read()
-        return base64.b64encode(data_as_bytes).decode()
+    raw_bytes = _open_as_mime_bytes(content)
+    return base64.b64encode(raw_bytes).decode()
 
 
 def _as_url(content: ContentT, default_mime_type: str) -> str:
-    if isinstance(content, str) and (content.startswith("https://") or content.startswith("http://")):
+    if isinstance(content, str) and content.startswith(("http://", "https://", "data:")):
         return content
 
-    # Handle MIME type detection for different content types
-    mime_type = None
-    if isinstance(content, (str, Path)):
-        mime_type = mimetypes.guess_type(content, strict=False)[0]
-    elif is_pillow_available():
-        from PIL import Image
+    # Convert content to bytes
+    raw_bytes = _open_as_mime_bytes(content)
 
-        if isinstance(content, Image.Image):
-            # Determine MIME type from PIL Image format, in sync with `_open_as_binary`
-            mime_type = f"image/{(content.format or 'PNG').lower()}"
+    # Get MIME type
+    mime_type = raw_bytes.mime_type or default_mime_type
 
-    mime_type = mime_type or default_mime_type
-    encoded_data = _b64_encode(content)
+    # Encode content to base64
+    encoded_data = base64.b64encode(raw_bytes).decode()
+
+    # Build data URL
     return f"data:{mime_type};base64,{encoded_data}"
 
 
@@ -239,9 +275,6 @@ def _as_dict(response: Union[bytes, Dict]) -> Dict:
     return json.loads(response) if isinstance(response, bytes) else response
 
 
-## PAYLOAD UTILS
-
-
 ## STREAMING UTILS
 
 

src/huggingface_hub/inference/_generated/_async_client.py

@@ -40,7 +40,6 @@
     _bytes_to_list,
     _get_unsupported_text_generation_kwargs,
     _import_numpy,
-    _open_as_binary,
     _set_unsupported_text_generation_kwargs,
     raise_text_generation_error,
 )
@@ -255,39 +254,38 @@ async def _inner_post(
         if request_parameters.task in TASKS_EXPECTING_IMAGES and "Accept" not in request_parameters.headers:
             request_parameters.headers["Accept"] = "image/png"
 
-        with _open_as_binary(request_parameters.data) as data_as_binary:
-            # Do not use context manager as we don't want to close the connection immediately when returning
-            # a stream
-            session = self._get_client_session(headers=request_parameters.headers)
+        # Do not use context manager as we don't want to close the connection immediately when returning
+        # a stream
+        session = self._get_client_session(headers=request_parameters.headers)
 
-            try:
-                response = await session.post(
-                    request_parameters.url, json=request_parameters.json, data=data_as_binary, proxy=self.proxies
-                )
-                response_error_payload = None
-                if response.status != 200:
-                    try:
-                        response_error_payload = await response.json()  # get payload before connection closed
-                    except Exception:
-                        pass
-                response.raise_for_status()
-                if stream:
-                    return _async_yield_from(session, response)
-                else:
-                    content = await response.read()
-                    await session.close()
-                    return content
-            except asyncio.TimeoutError as error:
-                await session.close()
-                # Convert any `TimeoutError` to a `InferenceTimeoutError`
-                raise InferenceTimeoutError(f"Inference call timed out: {request_parameters.url}") from error  # type: ignore
-            except aiohttp.ClientResponseError as error:
-                error.response_error_payload = response_error_payload
-                await session.close()
-                raise error
-            except Exception:
+        try:
+            response = await session.post(
+                request_parameters.url, json=request_parameters.json, data=request_parameters.data, proxy=self.proxies
+            )
+            response_error_payload = None
+            if response.status != 200:
+                try:
+                    response_error_payload = await response.json()  # get payload before connection closed
+                except Exception:
+                    pass
+            response.raise_for_status()
+            if stream:
+                return _async_yield_from(session, response)
+            else:
+                content = await response.read()
                 await session.close()
-                raise
+                return content
+        except asyncio.TimeoutError as error:
+            await session.close()
+            # Convert any `TimeoutError` to a `InferenceTimeoutError`
+            raise InferenceTimeoutError(f"Inference call timed out: {request_parameters.url}") from error  # type: ignore
+        except aiohttp.ClientResponseError as error:
+            error.response_error_payload = response_error_payload
+            await session.close()
+            raise error
+        except Exception:
+            await session.close()
+            raise
 
     async def __aenter__(self):
         return self

src/huggingface_hub/inference/_providers/_common.py

@@ -3,7 +3,7 @@
 
 from huggingface_hub import constants
 from huggingface_hub.hf_api import InferenceProviderMapping
-from huggingface_hub.inference._common import RequestParameters
+from huggingface_hub.inference._common import MimeBytes, RequestParameters
 from huggingface_hub.inference._generated.types.chat_completion import ChatCompletionInputMessage
 from huggingface_hub.utils import build_hf_headers, get_token, logging
 
@@ -108,8 +108,17 @@ def prepare_request(
             raise ValueError("Both payload and data cannot be set in the same request.")
         if payload is None and data is None:
             raise ValueError("Either payload or data must be set in the request.")
+
+        # normalize headers to lowercase and add content-type if not present
+        normalized_headers = self._normalize_headers(headers, payload, data)
+
         return RequestParameters(
-            url=url, task=self.task, model=provider_mapping_info.provider_id, json=payload, data=data, headers=headers
+            url=url,
+            task=self.task,
+            model=provider_mapping_info.provider_id,
+            json=payload,
+            data=data,
+            headers=normalized_headers,
         )
 
     def get_response(
@@ -172,7 +181,22 @@ def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMappin
             )
         return provider_mapping
 
-    def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
+    def _normalize_headers(
+        self, headers: Dict[str, Any], payload: Optional[Dict[str, Any]], data: Optional[MimeBytes]
+    ) -> Dict[str, Any]:
+        """Normalize the headers to use for the request.
+
+        Override this method in subclasses for customized headers.
+        """
+        normalized_headers = {key.lower(): value for key, value in headers.items() if value is not None}
+        if normalized_headers.get("content-type") is None:
+            if data is not None and data.mime_type is not None:
+                normalized_headers["content-type"] = data.mime_type
+            elif payload is not None:
+                normalized_headers["content-type"] = "application/json"
+        return normalized_headers
+
+    def _prepare_headers(self, headers: Dict, api_key: str) -> Dict[str, Any]:
         """Return the headers to use for the request.
 
         Override this method in subclasses for customized headers.
@@ -222,7 +246,7 @@ def _prepare_payload_as_bytes(
         parameters: Dict,
         provider_mapping_info: InferenceProviderMapping,
         extra_payload: Optional[Dict],
-    ) -> Optional[bytes]:
+    ) -> Optional[MimeBytes]:
         """Return the body to use for the request, as bytes.
 
         Override this method in subclasses for customized body data.

src/huggingface_hub/inference/_providers/black_forest_labs.py

@@ -18,7 +18,7 @@ class BlackForestLabsTextToImageTask(TaskProviderHelper):
     def __init__(self):
         super().__init__(provider="black-forest-labs", base_url="https://api.us1.bfl.ai", task="text-to-image")
 
-    def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
+    def _prepare_headers(self, headers: Dict, api_key: str) -> Dict[str, Any]:
         headers = super()._prepare_headers(headers, api_key)
         if not api_key.startswith("hf_"):
             _ = headers.pop("authorization")