import base64
import time
from abc import ABC
from typing import Any
from urllib.parse import urlparse
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, 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[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:
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: bytes | dict,
request_params: RequestParameters | None = 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, provider_mapping_info: InferenceProviderMapping
) -> dict | None:
if isinstance(inputs, str) and inputs.startswith(("http://", "https://")):
# If input is a URL, pass it directly
audio_url = inputs
else:
# If input is a file path, read it first
if isinstance(inputs, str):
with open(inputs, "rb") as f:
inputs = f.read()
audio_b64 = base64.b64encode(inputs).decode()
content_type = "audio/mpeg"
audio_url = f"data:{content_type};base64,{audio_b64}"
return {"audio_url": audio_url, **filter_none(parameters)}
def get_response(self, response: bytes | dict, request_params: RequestParameters | None = 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)}.")
return {"text": text}
class FalAITextToImageTask(FalAITask):
def __init__(self):
super().__init__("text-to-image")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> dict | None:
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"),
}
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: bytes | dict, request_params: RequestParameters | None = None) -> Any:
url = _as_dict(response)["images"][0]["url"]
return get_session().get(url).content
class FalAITextToSpeechTask(FalAITask):
def __init__(self):
super().__init__("text-to-speech")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> dict | None:
return {"text": inputs, **filter_none(parameters)}
def get_response(self, response: bytes | dict, request_params: RequestParameters | None = None) -> Any:
url = _as_dict(response)["audio"]["url"]
return get_session().get(url).content
class FalAITextToVideoTask(FalAIQueueTask):
def __init__(self):
super().__init__("text-to-video")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> dict | None:
return {"prompt": inputs, **filter_none(parameters)}
def get_response(
self,
response: bytes | dict,
request_params: RequestParameters | None = 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
) -> dict | None:
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,
"image_urls": [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: bytes | dict,
request_params: RequestParameters | None = 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
) -> dict | None:
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: bytes | dict,
request_params: RequestParameters | None = 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
) -> dict | None:
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: bytes | dict,
request_params: RequestParameters | None = 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,
}
]