huggingface-hub 0.29.0rc2__py3-none-any.whl → 1.1.3__py3-none-any.whl

huggingface_hub/inference/_providers/cohere.py

@@ -0,0 +1,32 @@
+from typing import Any, Optional
+
+from huggingface_hub.hf_api import InferenceProviderMapping
+
+from ._common import BaseConversationalTask
+
+
+_PROVIDER = "cohere"
+_BASE_URL = "https://api.cohere.com"
+
+
+class CohereConversationalTask(BaseConversationalTask):
+    def __init__(self):
+        super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
+
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
+        return "/compatibility/v1/chat/completions"
+
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
+        payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
+        response_format = parameters.get("response_format")
+        if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
+            json_schema_details = response_format.get("json_schema")
+            if isinstance(json_schema_details, dict) and "schema" in json_schema_details:
+                payload["response_format"] = {  # type: ignore [index]
+                    "type": "json_object",
+                    "schema": json_schema_details["schema"],
+                }
+
+        return payload

huggingface_hub/inference/_providers/fal_ai.py

@@ -1,31 +1,98 @@
 import base64
+import time
 from abc import ABC
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
+from urllib.parse import urlparse
 
-from huggingface_hub.inference._common import _as_dict
+from huggingface_hub import constants
+from huggingface_hub.hf_api import InferenceProviderMapping
+from huggingface_hub.inference._common import RequestParameters, _as_dict, _as_url
 from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
-from huggingface_hub.utils import get_session
+from huggingface_hub.utils import get_session, hf_raise_for_status
+from huggingface_hub.utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+# Arbitrary polling interval
+_POLLING_INTERVAL = 0.5
 
 
 class FalAITask(TaskProviderHelper, ABC):
     def __init__(self, task: str):
         super().__init__(provider="fal-ai", base_url="https://fal.run", task=task)
 
-    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["authorization"] = f"Key {api_key}"
         return headers
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         return f"/{mapped_model}"
 
 
+class FalAIQueueTask(TaskProviderHelper, ABC):
+    def __init__(self, task: str):
+        super().__init__(provider="fal-ai", base_url="https://queue.fal.run", task=task)
+
+    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["authorization"] = f"Key {api_key}"
+        return headers
+
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
+        if api_key.startswith("hf_"):
+            # Use the queue subdomain for HF routing
+            return f"/{mapped_model}?_subdomain=queue"
+        return f"/{mapped_model}"
+
+    def get_response(
+        self,
+        response: Union[bytes, dict],
+        request_params: Optional[RequestParameters] = None,
+    ) -> Any:
+        response_dict = _as_dict(response)
+
+        request_id = response_dict.get("request_id")
+        if not request_id:
+            raise ValueError("No request ID found in the response")
+        if request_params is None:
+            raise ValueError(
+                f"A `RequestParameters` object should be provided to get {self.task} responses with Fal AI."
+            )
+
+        # extract the base url and query params
+        parsed_url = urlparse(request_params.url)
+        # a bit hacky way to concatenate the provider name without parsing `parsed_url.path`
+        base_url = f"{parsed_url.scheme}://{parsed_url.netloc}{'/fal-ai' if parsed_url.netloc == 'router.huggingface.co' else ''}"
+        query_param = f"?{parsed_url.query}" if parsed_url.query else ""
+
+        # extracting the provider model id for status and result urls
+        # from the response as it might be different from the mapped model in `request_params.url`
+        model_id = urlparse(response_dict.get("response_url")).path
+        status_url = f"{base_url}{str(model_id)}/status{query_param}"
+        result_url = f"{base_url}{str(model_id)}{query_param}"
+
+        status = response_dict.get("status")
+        logger.info("Generating the output.. this can take several minutes.")
+        while status != "COMPLETED":
+            time.sleep(_POLLING_INTERVAL)
+            status_response = get_session().get(status_url, headers=request_params.headers)
+            hf_raise_for_status(status_response)
+            status = status_response.json().get("status")
+
+        return get_session().get(result_url, headers=request_params.headers).json()
+
+
 class FalAIAutomaticSpeechRecognitionTask(FalAITask):
     def __init__(self):
         super().__init__("automatic-speech-recognition")
 
-    def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
         if isinstance(inputs, str) and inputs.startswith(("http://", "https://")):
             # If input is a URL, pass it directly
             audio_url = inputs
@@ -41,7 +108,7 @@ class FalAIAutomaticSpeechRecognitionTask(FalAITask):
 
         return {"audio_url": audio_url, **filter_none(parameters)}
 
-    def get_response(self, response: Union[bytes, Dict]) -> Any:
+    def get_response(self, response: Union[bytes, dict], request_params: Optional[RequestParameters] = None) -> Any:
         text = _as_dict(response)["text"]
         if not isinstance(text, str):
             raise ValueError(f"Unexpected output format from FalAI API. Expected string, got {type(text)}.")
@@ -52,16 +119,33 @@ class FalAITextToImageTask(FalAITask):
     def __init__(self):
         super().__init__("text-to-image")
 
-    def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
-        parameters = filter_none(parameters)
-        if "width" in parameters and "height" in parameters:
-            parameters["image_size"] = {
-                "width": parameters.pop("width"),
-                "height": parameters.pop("height"),
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
+        payload: dict[str, Any] = {
+            "prompt": inputs,
+            **filter_none(parameters),
+        }
+        if "width" in payload and "height" in payload:
+            payload["image_size"] = {
+                "width": payload.pop("width"),
+                "height": payload.pop("height"),
             }
-        return {"prompt": inputs, **parameters}
+        if provider_mapping_info.adapter_weights_path is not None:
+            lora_path = constants.HUGGINGFACE_CO_URL_TEMPLATE.format(
+                repo_id=provider_mapping_info.hf_model_id,
+                revision="main",
+                filename=provider_mapping_info.adapter_weights_path,
+            )
+            payload["loras"] = [{"path": lora_path, "scale": 1}]
+            if provider_mapping_info.provider_id == "fal-ai/lora":
+                # little hack: fal requires the base model for stable-diffusion-based loras but not for flux-based
+                # See payloads in https://fal.ai/models/fal-ai/lora/api vs https://fal.ai/models/fal-ai/flux-lora/api
+                payload["model_name"] = "stabilityai/stable-diffusion-xl-base-1.0"
+
+        return payload
 
-    def get_response(self, response: Union[bytes, Dict]) -> Any:
+    def get_response(self, response: Union[bytes, dict], request_params: Optional[RequestParameters] = None) -> Any:
         url = _as_dict(response)["images"][0]["url"]
         return get_session().get(url).content
 
@@ -70,21 +154,146 @@ class FalAITextToSpeechTask(FalAITask):
     def __init__(self):
         super().__init__("text-to-speech")
 
-    def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
-        return {"lyrics": inputs, **filter_none(parameters)}
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
+        return {"text": inputs, **filter_none(parameters)}
 
-    def get_response(self, response: Union[bytes, Dict]) -> Any:
+    def get_response(self, response: Union[bytes, dict], request_params: Optional[RequestParameters] = None) -> Any:
         url = _as_dict(response)["audio"]["url"]
         return get_session().get(url).content
 
 
-class FalAITextToVideoTask(FalAITask):
+class FalAITextToVideoTask(FalAIQueueTask):
     def __init__(self):
         super().__init__("text-to-video")
 
-    def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
         return {"prompt": inputs, **filter_none(parameters)}
 
-    def get_response(self, response: Union[bytes, Dict]) -> Any:
-        url = _as_dict(response)["video"]["url"]
+    def get_response(
+        self,
+        response: Union[bytes, dict],
+        request_params: Optional[RequestParameters] = None,
+    ) -> Any:
+        output = super().get_response(response, request_params)
+        url = _as_dict(output)["video"]["url"]
         return get_session().get(url).content
+
+
+class FalAIImageToImageTask(FalAIQueueTask):
+    def __init__(self):
+        super().__init__("image-to-image")
+
+    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/jpeg")
+        if "target_size" in parameters:
+            parameters["image_size"] = parameters.pop("target_size")
+        payload: dict[str, Any] = {
+            "image_url": image_url,
+            **filter_none(parameters),
+        }
+        if provider_mapping_info.adapter_weights_path is not None:
+            lora_path = constants.HUGGINGFACE_CO_URL_TEMPLATE.format(
+                repo_id=provider_mapping_info.hf_model_id,
+                revision="main",
+                filename=provider_mapping_info.adapter_weights_path,
+            )
+            payload["loras"] = [{"path": lora_path, "scale": 1}]
+
+        return payload
+
+    def get_response(
+        self,
+        response: Union[bytes, dict],
+        request_params: Optional[RequestParameters] = None,
+    ) -> Any:
+        output = super().get_response(response, request_params)
+        url = _as_dict(output)["images"][0]["url"]
+        return get_session().get(url).content
+
+
+class FalAIImageToVideoTask(FalAIQueueTask):
+    def __init__(self):
+        super().__init__("image-to-video")
+
+    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/jpeg")
+        payload: dict[str, Any] = {
+            "image_url": image_url,
+            **filter_none(parameters),
+        }
+        if provider_mapping_info.adapter_weights_path is not None:
+            lora_path = constants.HUGGINGFACE_CO_URL_TEMPLATE.format(
+                repo_id=provider_mapping_info.hf_model_id,
+                revision="main",
+                filename=provider_mapping_info.adapter_weights_path,
+            )
+            payload["loras"] = [{"path": lora_path, "scale": 1}]
+        return payload
+
+    def get_response(
+        self,
+        response: Union[bytes, dict],
+        request_params: Optional[RequestParameters] = None,
+    ) -> Any:
+        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,
+            }
+        ]

huggingface_hub/inference/_providers/featherless_ai.py

@@ -0,0 +1,38 @@
+from typing import Any, Optional, Union
+
+from huggingface_hub.hf_api import InferenceProviderMapping
+from huggingface_hub.inference._common import RequestParameters, _as_dict
+
+from ._common import BaseConversationalTask, BaseTextGenerationTask, filter_none
+
+
+_PROVIDER = "featherless-ai"
+_BASE_URL = "https://api.featherless.ai"
+
+
+class FeatherlessTextGenerationTask(BaseTextGenerationTask):
+    def __init__(self):
+        super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
+
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
+        params = filter_none(parameters.copy())
+        params["max_tokens"] = params.pop("max_new_tokens", None)
+
+        return {"prompt": inputs, **params, "model": provider_mapping_info.provider_id}
+
+    def get_response(self, response: Union[bytes, dict], request_params: Optional[RequestParameters] = None) -> Any:
+        output = _as_dict(response)["choices"][0]
+        return {
+            "generated_text": output["text"],
+            "details": {
+                "finish_reason": output.get("finish_reason"),
+                "seed": output.get("seed"),
+            },
+        }
+
+
+class FeatherlessConversationalTask(BaseConversationalTask):
+    def __init__(self):
+        super().__init__(provider=_PROVIDER, base_url=_BASE_URL)

huggingface_hub/inference/_providers/fireworks_ai.py

@@ -1,6 +1,27 @@
+from typing import Any, Optional
+
+from huggingface_hub.hf_api import InferenceProviderMapping
+
 from ._common import BaseConversationalTask
 
 
 class FireworksAIConversationalTask(BaseConversationalTask):
     def __init__(self):
-        super().__init__(provider="fireworks-ai", base_url="https://api.fireworks.ai/inference")
+        super().__init__(provider="fireworks-ai", base_url="https://api.fireworks.ai")
+
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
+        return "/inference/v1/chat/completions"
+
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
+        payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
+        response_format = parameters.get("response_format")
+        if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
+            json_schema_details = response_format.get("json_schema")
+            if isinstance(json_schema_details, dict) and "schema" in json_schema_details:
+                payload["response_format"] = {  # type: ignore [index]
+                    "type": "json_object",
+                    "schema": json_schema_details["schema"],
+                }
+        return payload

huggingface_hub/inference/_providers/groq.py

@@ -0,0 +1,9 @@
+from ._common import BaseConversationalTask
+
+
+class GroqConversationalTask(BaseConversationalTask):
+    def __init__(self):
+        super().__init__(provider="groq", base_url="https://api.groq.com")
+
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
+        return "/openai/v1/chat/completions"

huggingface_hub/inference/_providers/hf_inference.py

@@ -1,10 +1,18 @@
 import json
 from functools import lru_cache
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Optional, Union
+from urllib.parse import urlparse, urlunparse
 
 from huggingface_hub import constants
-from huggingface_hub.inference._common import _b64_encode, _open_as_binary
+from huggingface_hub.hf_api import InferenceProviderMapping
+from huggingface_hub.inference._common import (
+    MimeBytes,
+    RequestParameters,
+    _b64_encode,
+    _bytes_to_dict,
+    _open_as_mime_bytes,
+)
 from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
 from huggingface_hub.utils import build_hf_headers, get_session, get_token, hf_raise_for_status
 
@@ -23,16 +31,21 @@ class HFInferenceTask(TaskProviderHelper):
         # special case: for HF Inference we allow not providing an API key
         return api_key or get_token()  # type: ignore[return-value]
 
-    def _prepare_mapped_model(self, model: Optional[str]) -> str:
-        if model is not None:
-            return model
-        model = _fetch_recommended_models().get(self.task)
-        if model is None:
+    def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
+        if model is not None and model.startswith(("http://", "https://")):
+            return InferenceProviderMapping(
+                provider="hf-inference", providerId=model, hf_model_id=model, task=self.task, status="live"
+            )
+        model_id = model if model is not None else _fetch_recommended_models().get(self.task)
+        if model_id is None:
             raise ValueError(
                 f"Task {self.task} has no recommended model for HF Inference. Please specify a model"
                 " explicitly. Visit https://huggingface.co/tasks for more info."
             )
-        return model
+        _check_supported_task(model_id, self.task)
+        return InferenceProviderMapping(
+            provider="hf-inference", providerId=model_id, hf_model_id=model_id, task=self.task, status="live"
+        )
 
     def _prepare_url(self, api_key: str, mapped_model: str) -> str:
         # hf-inference provider can handle URLs (e.g. Inference Endpoints or TGI deployment)
@@ -40,28 +53,36 @@ class HFInferenceTask(TaskProviderHelper):
             return mapped_model
         return (
             # Feature-extraction and sentence-similarity are the only cases where we handle models with several tasks.
-            f"{self.base_url}/pipeline/{self.task}/{mapped_model}"
+            f"{self.base_url}/models/{mapped_model}/pipeline/{self.task}"
             if self.task in ("feature-extraction", "sentence-similarity")
             # Otherwise, we use the default endpoint
             else f"{self.base_url}/models/{mapped_model}"
         )
 
-    def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
         if isinstance(inputs, bytes):
             raise ValueError(f"Unexpected binary input for task {self.task}.")
         if isinstance(inputs, Path):
             raise ValueError(f"Unexpected path input for task {self.task} (got {inputs})")
-        return {"inputs": inputs, "parameters": filter_none(parameters)}
+        return filter_none({"inputs": inputs, "parameters": parameters})
 
 
 class HFInferenceBinaryInputTask(HFInferenceTask):
-    def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
         return None
 
     def _prepare_payload_as_bytes(
-        self, inputs: Any, parameters: Dict, mapped_model: str, extra_payload: Optional[Dict]
-    ) -> Optional[bytes]:
-        parameters = filter_none({k: v for k, v in parameters.items() if v is not None})
+        self,
+        inputs: Any,
+        parameters: dict,
+        provider_mapping_info: InferenceProviderMapping,
+        extra_payload: Optional[dict],
+    ) -> Optional[MimeBytes]:
+        parameters = filter_none(parameters)
         extra_payload = extra_payload or {}
         has_parameters = len(parameters) > 0 or len(extra_payload) > 0
 
@@ -71,21 +92,36 @@ class HFInferenceBinaryInputTask(HFInferenceTask):
 
         # Send inputs as raw content when no parameters are provided
         if not has_parameters:
-            with _open_as_binary(inputs) as data:
-                data_as_bytes = data if isinstance(data, bytes) else data.read()
-                return data_as_bytes
+            return _open_as_mime_bytes(inputs)
 
         # Otherwise encode as b64
-        return json.dumps({"inputs": _b64_encode(inputs), "parameters": parameters, **extra_payload}).encode("utf-8")
+        return MimeBytes(
+            json.dumps({"inputs": _b64_encode(inputs), "parameters": parameters, **extra_payload}).encode("utf-8"),
+            mime_type="application/json",
+        )
 
 
 class HFInferenceConversational(HFInferenceTask):
     def __init__(self):
-        super().__init__("text-generation")
-
-    def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
-        payload_model = "tgi" if mapped_model.startswith(("http://", "https://")) else mapped_model
-        return {**filter_none(parameters), "model": payload_model, "messages": inputs}
+        super().__init__("conversational")
+
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
+        payload = filter_none(parameters)
+        mapped_model = provider_mapping_info.provider_id
+        payload_model = parameters.get("model") or mapped_model
+
+        if payload_model is None or payload_model.startswith(("http://", "https://")):
+            payload_model = "dummy"
+
+        response_format = parameters.get("response_format")
+        if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
+            payload["response_format"] = {
+                "type": "json_object",
+                "value": response_format["json_schema"]["schema"],
+            }
+        return {**payload, "model": payload_model, "messages": inputs}
 
     def _prepare_url(self, api_key: str, mapped_model: str) -> str:
         base_url = (
@@ -97,22 +133,96 @@ class HFInferenceConversational(HFInferenceTask):
 
 
 def _build_chat_completion_url(model_url: str) -> str:
-    # Strip trailing /
-    model_url = model_url.rstrip("/")
+    parsed = urlparse(model_url)
+    path = parsed.path.rstrip("/")
 
-    # Append /chat/completions if not already present
-    if model_url.endswith("/v1"):
-        model_url += "/chat/completions"
+    # If the path already ends with /chat/completions, we're done!
+    if path.endswith("/chat/completions"):
+        return model_url
 
+    # Append /chat/completions if not already present
+    if path.endswith("/v1"):
+        new_path = path + "/chat/completions"
+    # If path was empty or just "/", set the full path
+    elif not path:
+        new_path = "/v1/chat/completions"
     # Append /v1/chat/completions if not already present
-    if not model_url.endswith("/chat/completions"):
-        model_url += "/v1/chat/completions"
+    else:
+        new_path = path + "/v1/chat/completions"
 
-    return model_url
+    # Reconstruct the URL with the new path and original query parameters.
+    new_parsed = parsed._replace(path=new_path)
+    return str(urlunparse(new_parsed))
 
 
 @lru_cache(maxsize=1)
-def _fetch_recommended_models() -> Dict[str, Optional[str]]:
+def _fetch_recommended_models() -> dict[str, Optional[str]]:
     response = get_session().get(f"{constants.ENDPOINT}/api/tasks", headers=build_hf_headers())
     hf_raise_for_status(response)
     return {task: next(iter(details["widgetModels"]), None) for task, details in response.json().items()}
+
+
+@lru_cache(maxsize=None)
+def _check_supported_task(model: str, task: str) -> None:
+    from huggingface_hub.hf_api import HfApi
+
+    model_info = HfApi().model_info(model)
+    pipeline_tag = model_info.pipeline_tag
+    tags = model_info.tags or []
+    is_conversational = "conversational" in tags
+    if task in ("text-generation", "conversational"):
+        if pipeline_tag == "text-generation":
+            # text-generation + conversational tag -> both tasks allowed
+            if is_conversational:
+                return
+            # text-generation without conversational tag -> only text-generation allowed
+            if task == "text-generation":
+                return
+            raise ValueError(f"Model '{model}' doesn't support task '{task}'.")
+
+    if pipeline_tag == "text2text-generation":
+        if task == "text-generation":
+            return
+        raise ValueError(f"Model '{model}' doesn't support task '{task}'.")
+
+    if pipeline_tag == "image-text-to-text":
+        if is_conversational and task == "conversational":
+            return  # Only conversational allowed if tagged as conversational
+        raise ValueError("Non-conversational image-text-to-text task is not supported.")
+
+    if (
+        task in ("feature-extraction", "sentence-similarity")
+        and pipeline_tag in ("feature-extraction", "sentence-similarity")
+        and task in tags
+    ):
+        # feature-extraction and sentence-similarity are interchangeable for HF Inference
+        return
+
+    # For all other tasks, just check pipeline tag
+    if pipeline_tag != task:
+        raise ValueError(
+            f"Model '{model}' doesn't support task '{task}'. Supported tasks: '{pipeline_tag}', got: '{task}'"
+        )
+    return
+
+
+class HFInferenceFeatureExtractionTask(HFInferenceTask):
+    def __init__(self):
+        super().__init__("feature-extraction")
+
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[dict]:
+        if isinstance(inputs, bytes):
+            raise ValueError(f"Unexpected binary input for task {self.task}.")
+        if isinstance(inputs, Path):
+            raise ValueError(f"Unexpected path input for task {self.task} (got {inputs})")
+
+        # Parameters are sent at root-level for feature-extraction task
+        # See specs: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/feature-extraction/spec/input.json
+        return {"inputs": inputs, **filter_none(parameters)}
+
+    def get_response(self, response: Union[bytes, dict], request_params: Optional[RequestParameters] = None) -> Any:
+        if isinstance(response, bytes):
+            return _bytes_to_dict(response)
+        return response