[Inference Providers] Async calls for fal.ai (#2927)

* add async calls for fal-ai

* fix

* fix test

* nit

* pass request params to get response

* fixes

* nit

* fix quality

* nit

* fixes post-review

* fix

Co-authored-by: Lucain <[email protected]>

* remove unnecessary type: ignore

Co-authored-by: Lucain <[email protected]>

* fix

Co-authored-by: Lucain <[email protected]>

---------

Co-authored-by: Lucain <[email protected]>

Author: hanouticelina

src/huggingface_hub/inference/_client.py

@@ -2617,7 +2617,7 @@ def text_to_video(
             api_key=self.token,
         )
         response = self._inner_post(request_parameters)
-        response = provider_helper.get_response(response)
+        response = provider_helper.get_response(response, request_parameters)
         return response
 
     def text_to_speech(

src/huggingface_hub/inference/_generated/_async_client.py

@@ -2674,7 +2674,7 @@ async def text_to_video(
             api_key=self.token,
         )
         response = await self._inner_post(request_parameters)
-        response = provider_helper.get_response(response)
+        response = provider_helper.get_response(response, request_parameters)
         return response
 
     async def text_to_speech(

src/huggingface_hub/inference/_providers/_common.py

@@ -84,7 +84,11 @@ def prepare_request(
             raise ValueError("Either payload or data must be set in the request.")
         return RequestParameters(url=url, task=self.task, model=mapped_model, json=payload, data=data, headers=headers)
 
-    def get_response(self, response: Union[bytes, Dict]) -> Any:
+    def get_response(
+        self,
+        response: Union[bytes, Dict],
+        request_params: Optional[RequestParameters] = None,
+    ) -> Any:
         """
         Return the response in the expected format.
 
@@ -142,7 +146,7 @@ def _prepare_url(self, api_key: str, mapped_model: str) -> str:
 
         Usually not overwritten in subclasses."""
         base_url = self._prepare_base_url(api_key)
-        route = self._prepare_route(mapped_model)
+        route = self._prepare_route(mapped_model, api_key)
         return f"{base_url.rstrip('/')}/{route.lstrip('/')}"
 
     def _prepare_base_url(self, api_key: str) -> str:
@@ -157,7 +161,7 @@ def _prepare_base_url(self, api_key: str) -> str:
             logger.info(f"Calling '{self.provider}' provider directly.")
             return self.base_url
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         """Return the route to use for the request.
 
         Override this method in subclasses for customized routes.
@@ -192,7 +196,7 @@ class BaseConversationalTask(TaskProviderHelper):
     def __init__(self, provider: str, base_url: str):
         super().__init__(provider=provider, base_url=base_url, task="conversational")
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         return "/v1/chat/completions"
 
     def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
@@ -208,7 +212,7 @@ class BaseTextGenerationTask(TaskProviderHelper):
     def __init__(self, provider: str, base_url: str):
         super().__init__(provider=provider, base_url=base_url, task="text-generation")
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         return "/v1/completions"
 
     def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:

src/huggingface_hub/inference/_providers/black_forest_labs.py

@@ -1,7 +1,7 @@
 import time
 from typing import Any, Dict, Optional, Union
 
-from huggingface_hub.inference._common import _as_dict
+from huggingface_hub.inference._common import RequestParameters, _as_dict
 from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
 from huggingface_hub.utils import logging
 from huggingface_hub.utils._http import get_session
@@ -24,7 +24,7 @@ def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
             headers["X-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"/v1/{mapped_model}"
 
     def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
@@ -36,7 +36,7 @@ def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model:
 
         return {"prompt": inputs, **parameters}
 
-    def get_response(self, response: Union[bytes, Dict]) -> Any:
+    def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
         """
         Polling mechanism for Black Forest Labs since the API is asynchronous.
         """

src/huggingface_hub/inference/_providers/cohere.py

@@ -11,5 +11,5 @@ class CohereConversationalTask(BaseConversationalTask):
     def __init__(self):
         super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         return "/compatibility/v1/chat/completions"

src/huggingface_hub/inference/_providers/fal_ai.py

@@ -1,10 +1,18 @@
 import base64
+import time
 from abc import ABC
 from typing import Any, Dict, Optional, Union
 
-from huggingface_hub.inference._common import _as_dict
+from huggingface_hub.inference._common import RequestParameters, _as_dict
 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 = 2.0
 
 
 class FalAITask(TaskProviderHelper, ABC):
@@ -17,7 +25,7 @@ def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
             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}"
 
 
@@ -41,7 +49,7 @@ def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model:
 
         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)}.")
@@ -61,7 +69,7 @@ def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model:
             }
         return {"prompt": inputs, **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)["images"][0]["url"]
         return get_session().get(url).content
 
@@ -73,7 +81,7 @@ def __init__(self):
     def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
         return {"lyrics": 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
 
@@ -82,9 +90,52 @@ class FalAITextToVideoTask(FalAITask):
     def __init__(self):
         super().__init__("text-to-video")
 
+    def _prepare_base_url(self, api_key: str) -> str:
+        if api_key.startswith("hf_"):
+            return super()._prepare_base_url(api_key)
+        else:
+            logger.info(f"Calling '{self.provider}' provider directly.")
+            return "https://queue.fal.run"
+
+    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 _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
         return {"prompt": 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:
+        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(
+                "A `RequestParameters` object should be provided to get text-to-video responses with Fal AI."
+            )
+
+        # extract the base url and query params
+        base_url = request_params.url.split("?")[0]  # or parsed.scheme + "://" + parsed.netloc + parsed.path ?
+        query = "?_subdomain=queue" if request_params.url.endswith("_subdomain=queue") else ""
+
+        status_url = f"{base_url}/requests/{request_id}/status{query}"
+        result_url = f"{base_url}/requests/{request_id}{query}"
+
+        status = response_dict.get("status")
+        logger.info("Generating the video.. 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")
+
+        response = get_session().get(result_url, headers=request_params.headers).json()
         url = _as_dict(response)["video"]["url"]
         return get_session().get(url).content

src/huggingface_hub/inference/_providers/fireworks_ai.py

@@ -5,5 +5,5 @@ class FireworksAIConversationalTask(BaseConversationalTask):
     def __init__(self):
         super().__init__(provider="fireworks-ai", base_url="https://api.fireworks.ai")
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         return "/inference/v1/chat/completions"

src/huggingface_hub/inference/_providers/hyperbolic.py

@@ -1,15 +1,15 @@
 import base64
 from typing import Any, Dict, Optional, Union
 
-from huggingface_hub.inference._common import _as_dict
+from huggingface_hub.inference._common import RequestParameters, _as_dict
 from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none
 
 
 class HyperbolicTextToImageTask(TaskProviderHelper):
     def __init__(self):
         super().__init__(provider="hyperbolic", base_url="https://api.hyperbolic.xyz", task="text-to-image")
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         return "/v1/images/generations"
 
     def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
@@ -25,7 +25,7 @@ def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model:
             parameters["height"] = 512
         return {"prompt": inputs, "model_name": mapped_model, **parameters}
 
-    def get_response(self, response: Union[bytes, Dict]) -> Any:
+    def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
         response_dict = _as_dict(response)
         return base64.b64decode(response_dict["images"][0]["image"])
 

src/huggingface_hub/inference/_providers/nebius.py

@@ -1,7 +1,7 @@
 import base64
 from typing import Any, Dict, Optional, Union
 
-from huggingface_hub.inference._common import _as_dict
+from huggingface_hub.inference._common import RequestParameters, _as_dict
 from huggingface_hub.inference._providers._common import (
     BaseConversationalTask,
     BaseTextGenerationTask,
@@ -24,7 +24,7 @@ class NebiusTextToImageTask(TaskProviderHelper):
     def __init__(self):
         super().__init__(task="text-to-image", provider="nebius", base_url="https://api.studio.nebius.ai")
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         return "/v1/images/generations"
 
     def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
@@ -36,6 +36,6 @@ def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model:
 
         return {"prompt": inputs, **parameters, "model": mapped_model}
 
-    def get_response(self, response: Union[bytes, Dict]) -> Any:
+    def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
         response_dict = _as_dict(response)
         return base64.b64decode(response_dict["data"][0]["b64_json"])

src/huggingface_hub/inference/_providers/novita.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, Optional, Union
 
-from huggingface_hub.inference._common import _as_dict
+from huggingface_hub.inference._common import RequestParameters, _as_dict
 from huggingface_hub.inference._providers._common import (
     BaseConversationalTask,
     BaseTextGenerationTask,
@@ -18,7 +18,7 @@ class NovitaTextGenerationTask(BaseTextGenerationTask):
     def __init__(self):
         super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         # there is no v1/ route for novita
         return "/v3/openai/completions"
 
@@ -27,7 +27,7 @@ class NovitaConversationalTask(BaseConversationalTask):
     def __init__(self):
         super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         # there is no v1/ route for novita
         return "/v3/openai/chat/completions"
 
@@ -36,13 +36,13 @@ class NovitaTextToVideoTask(TaskProviderHelper):
     def __init__(self):
         super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task="text-to-video")
 
-    def _prepare_route(self, mapped_model: str) -> str:
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
         return f"/v3/hf/{mapped_model}"
 
     def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
         return {"prompt": 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:
         response_dict = _as_dict(response)
         if not (
             isinstance(response_dict, dict)