Initialisation du repository de Beta

This commit is contained in:
Mathis 2026-02-06 22:23:20 +01:00
commit 14985f6dbb
9469 changed files with 1903273 additions and 0 deletions

View file

@ -0,0 +1,267 @@
from typing import Literal, Optional, Union
from huggingface_hub.inference._providers.featherless_ai import (
FeatherlessConversationalTask,
FeatherlessTextGenerationTask,
)
from huggingface_hub.utils import logging
from ._common import AutoRouterConversationalTask, TaskProviderHelper, _fetch_inference_provider_mapping
from .black_forest_labs import BlackForestLabsTextToImageTask
from .cerebras import CerebrasConversationalTask
from .clarifai import ClarifaiConversationalTask
from .cohere import CohereConversationalTask
from .fal_ai import (
FalAIAutomaticSpeechRecognitionTask,
FalAIImageSegmentationTask,
FalAIImageToImageTask,
FalAIImageToVideoTask,
FalAITextToImageTask,
FalAITextToSpeechTask,
FalAITextToVideoTask,
)
from .fireworks_ai import FireworksAIConversationalTask
from .groq import GroqConversationalTask
from .hf_inference import (
HFInferenceBinaryInputTask,
HFInferenceConversational,
HFInferenceFeatureExtractionTask,
HFInferenceTask,
)
from .hyperbolic import HyperbolicTextGenerationTask, HyperbolicTextToImageTask
from .nebius import (
NebiusConversationalTask,
NebiusFeatureExtractionTask,
NebiusTextGenerationTask,
NebiusTextToImageTask,
)
from .novita import NovitaConversationalTask, NovitaTextGenerationTask, NovitaTextToVideoTask
from .nscale import NscaleConversationalTask, NscaleTextToImageTask
from .openai import OpenAIConversationalTask
from .ovhcloud import OVHcloudConversationalTask
from .publicai import PublicAIConversationalTask
from .replicate import (
ReplicateAutomaticSpeechRecognitionTask,
ReplicateImageToImageTask,
ReplicateTask,
ReplicateTextToImageTask,
ReplicateTextToSpeechTask,
)
from .sambanova import SambanovaConversationalTask, SambanovaFeatureExtractionTask
from .scaleway import ScalewayConversationalTask, ScalewayFeatureExtractionTask
from .together import TogetherConversationalTask, TogetherTextGenerationTask, TogetherTextToImageTask
from .wavespeed import (
WavespeedAIImageToImageTask,
WavespeedAIImageToVideoTask,
WavespeedAITextToImageTask,
WavespeedAITextToVideoTask,
)
from .zai_org import ZaiConversationalTask, ZaiTextToImageTask
logger = logging.get_logger(__name__)
PROVIDER_T = Literal[
"black-forest-labs",
"cerebras",
"clarifai",
"cohere",
"fal-ai",
"featherless-ai",
"fireworks-ai",
"groq",
"hf-inference",
"hyperbolic",
"nebius",
"novita",
"nscale",
"openai",
"ovhcloud",
"publicai",
"replicate",
"sambanova",
"scaleway",
"together",
"wavespeed",
"zai-org",
]
PROVIDER_OR_POLICY_T = Union[PROVIDER_T, Literal["auto"]]
CONVERSATIONAL_AUTO_ROUTER = AutoRouterConversationalTask()
PROVIDERS: dict[PROVIDER_T, dict[str, TaskProviderHelper]] = {
"black-forest-labs": {
"text-to-image": BlackForestLabsTextToImageTask(),
},
"cerebras": {
"conversational": CerebrasConversationalTask(),
},
"clarifai": {
"conversational": ClarifaiConversationalTask(),
},
"cohere": {
"conversational": CohereConversationalTask(),
},
"fal-ai": {
"automatic-speech-recognition": FalAIAutomaticSpeechRecognitionTask(),
"text-to-image": FalAITextToImageTask(),
"text-to-speech": FalAITextToSpeechTask(),
"text-to-video": FalAITextToVideoTask(),
"image-to-video": FalAIImageToVideoTask(),
"image-to-image": FalAIImageToImageTask(),
"image-segmentation": FalAIImageSegmentationTask(),
},
"featherless-ai": {
"conversational": FeatherlessConversationalTask(),
"text-generation": FeatherlessTextGenerationTask(),
},
"fireworks-ai": {
"conversational": FireworksAIConversationalTask(),
},
"groq": {
"conversational": GroqConversationalTask(),
},
"hf-inference": {
"text-to-image": HFInferenceTask("text-to-image"),
"conversational": HFInferenceConversational(),
"text-generation": HFInferenceTask("text-generation"),
"text-classification": HFInferenceTask("text-classification"),
"question-answering": HFInferenceTask("question-answering"),
"audio-classification": HFInferenceBinaryInputTask("audio-classification"),
"automatic-speech-recognition": HFInferenceBinaryInputTask("automatic-speech-recognition"),
"fill-mask": HFInferenceTask("fill-mask"),
"feature-extraction": HFInferenceFeatureExtractionTask(),
"image-classification": HFInferenceBinaryInputTask("image-classification"),
"image-segmentation": HFInferenceBinaryInputTask("image-segmentation"),
"document-question-answering": HFInferenceTask("document-question-answering"),
"image-to-text": HFInferenceBinaryInputTask("image-to-text"),
"object-detection": HFInferenceBinaryInputTask("object-detection"),
"audio-to-audio": HFInferenceBinaryInputTask("audio-to-audio"),
"zero-shot-image-classification": HFInferenceBinaryInputTask("zero-shot-image-classification"),
"zero-shot-classification": HFInferenceTask("zero-shot-classification"),
"image-to-image": HFInferenceBinaryInputTask("image-to-image"),
"sentence-similarity": HFInferenceTask("sentence-similarity"),
"table-question-answering": HFInferenceTask("table-question-answering"),
"tabular-classification": HFInferenceTask("tabular-classification"),
"text-to-speech": HFInferenceTask("text-to-speech"),
"token-classification": HFInferenceTask("token-classification"),
"translation": HFInferenceTask("translation"),
"summarization": HFInferenceTask("summarization"),
"visual-question-answering": HFInferenceBinaryInputTask("visual-question-answering"),
},
"hyperbolic": {
"text-to-image": HyperbolicTextToImageTask(),
"conversational": HyperbolicTextGenerationTask("conversational"),
"text-generation": HyperbolicTextGenerationTask("text-generation"),
},
"nebius": {
"text-to-image": NebiusTextToImageTask(),
"conversational": NebiusConversationalTask(),
"text-generation": NebiusTextGenerationTask(),
"feature-extraction": NebiusFeatureExtractionTask(),
},
"novita": {
"text-generation": NovitaTextGenerationTask(),
"conversational": NovitaConversationalTask(),
"text-to-video": NovitaTextToVideoTask(),
},
"nscale": {
"conversational": NscaleConversationalTask(),
"text-to-image": NscaleTextToImageTask(),
},
"openai": {
"conversational": OpenAIConversationalTask(),
},
"ovhcloud": {
"conversational": OVHcloudConversationalTask(),
},
"publicai": {
"conversational": PublicAIConversationalTask(),
},
"replicate": {
"automatic-speech-recognition": ReplicateAutomaticSpeechRecognitionTask(),
"image-to-image": ReplicateImageToImageTask(),
"text-to-image": ReplicateTextToImageTask(),
"text-to-speech": ReplicateTextToSpeechTask(),
"text-to-video": ReplicateTask("text-to-video"),
},
"sambanova": {
"conversational": SambanovaConversationalTask(),
"feature-extraction": SambanovaFeatureExtractionTask(),
},
"scaleway": {
"conversational": ScalewayConversationalTask(),
"feature-extraction": ScalewayFeatureExtractionTask(),
},
"together": {
"text-to-image": TogetherTextToImageTask(),
"conversational": TogetherConversationalTask(),
"text-generation": TogetherTextGenerationTask(),
},
"wavespeed": {
"text-to-image": WavespeedAITextToImageTask(),
"text-to-video": WavespeedAITextToVideoTask(),
"image-to-image": WavespeedAIImageToImageTask(),
"image-to-video": WavespeedAIImageToVideoTask(),
},
"zai-org": {
"conversational": ZaiConversationalTask(),
"text-to-image": ZaiTextToImageTask(),
},
}
def get_provider_helper(
provider: Optional[PROVIDER_OR_POLICY_T], task: str, model: Optional[str]
) -> TaskProviderHelper:
"""Get provider helper instance by name and task.
Args:
provider (`str`, *optional*): name of the provider, or "auto" to automatically select the provider for the model.
task (`str`): Name of the task
model (`str`, *optional*): Name of the model
Returns:
TaskProviderHelper: Helper instance for the specified provider and task
Raises:
ValueError: If provider or task is not supported
"""
if (model is None and provider in (None, "auto")) or (
model is not None and model.startswith(("http://", "https://"))
):
provider = "hf-inference"
if provider is None:
logger.info(
"No provider specified for task `conversational`. Defaulting to server-side auto routing."
if task == "conversational"
else "Defaulting to 'auto' which will select the first provider available for the model, sorted by the user's order in https://hf.co/settings/inference-providers."
)
provider = "auto"
if provider == "auto":
if model is None:
raise ValueError("Specifying a model is required when provider is 'auto'")
if task == "conversational":
# Special case: we have a dedicated auto-router for conversational models. No need to fetch provider mapping.
return CONVERSATIONAL_AUTO_ROUTER
provider_mapping = _fetch_inference_provider_mapping(model)
provider = next(iter(provider_mapping)).provider
provider_tasks = PROVIDERS.get(provider) # type: ignore
if provider_tasks is None:
raise ValueError(
f"Provider '{provider}' not supported. Available values: 'auto' or any provider from {list(PROVIDERS.keys())}."
"Passing 'auto' (default value) will automatically select the first provider available for the model, sorted "
"by the user's order in https://hf.co/settings/inference-providers."
)
if task not in provider_tasks:
raise ValueError(
f"Task '{task}' not supported for provider '{provider}'. Available tasks: {list(provider_tasks.keys())}"
)
return provider_tasks[task]

View file

@ -0,0 +1,363 @@
from functools import lru_cache
from typing import Any, Optional, Union, overload
from huggingface_hub import constants
from huggingface_hub.hf_api import InferenceProviderMapping
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
logger = logging.get_logger(__name__)
# Dev purposes only.
# If you want to try to run inference for a new model locally before it's registered on huggingface.co
# for a given Inference Provider, you can add it to the following dictionary.
HARDCODED_MODEL_INFERENCE_MAPPING: dict[str, dict[str, InferenceProviderMapping]] = {
# "HF model ID" => InferenceProviderMapping object initialized with "Model ID on Inference Provider's side"
#
# Example:
# "Qwen/Qwen2.5-Coder-32B-Instruct": InferenceProviderMapping(hf_model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
# provider_id="Qwen2.5-Coder-32B-Instruct",
# task="conversational",
# status="live")
"cerebras": {},
"cohere": {},
"clarifai": {},
"fal-ai": {},
"fireworks-ai": {},
"groq": {},
"hf-inference": {},
"hyperbolic": {},
"nebius": {},
"nscale": {},
"ovhcloud": {},
"replicate": {},
"sambanova": {},
"scaleway": {},
"together": {},
"wavespeed": {},
"zai-org": {},
}
@overload
def filter_none(obj: dict[str, Any]) -> dict[str, Any]: ...
@overload
def filter_none(obj: list[Any]) -> list[Any]: ...
def filter_none(obj: Union[dict[str, Any], list[Any]]) -> Union[dict[str, Any], list[Any]]:
if isinstance(obj, dict):
cleaned: dict[str, Any] = {}
for k, v in obj.items():
if v is None:
continue
if isinstance(v, (dict, list)):
v = filter_none(v)
cleaned[k] = v
return cleaned
if isinstance(obj, list):
return [filter_none(v) if isinstance(v, (dict, list)) else v for v in obj]
raise ValueError(f"Expected dict or list, got {type(obj)}")
class TaskProviderHelper:
"""Base class for task-specific provider helpers."""
def __init__(self, provider: str, base_url: str, task: str) -> None:
self.provider = provider
self.task = task
self.base_url = base_url
def prepare_request(
self,
*,
inputs: Any,
parameters: dict[str, Any],
headers: dict,
model: Optional[str],
api_key: Optional[str],
extra_payload: Optional[dict[str, Any]] = None,
) -> RequestParameters:
"""
Prepare the request to be sent to the provider.
Each step (api_key, model, headers, url, payload) can be customized in subclasses.
"""
# api_key from user, or local token, or raise error
api_key = self._prepare_api_key(api_key)
# mapped model from HF model ID
provider_mapping_info = self._prepare_mapping_info(model)
# default HF headers + user headers (to customize in subclasses)
headers = self._prepare_headers(headers, api_key)
# routed URL if HF token, or direct URL (to customize in '_prepare_route' in subclasses)
url = self._prepare_url(api_key, provider_mapping_info.provider_id)
# prepare payload (to customize in subclasses)
payload = self._prepare_payload_as_dict(inputs, parameters, provider_mapping_info=provider_mapping_info)
if payload is not None:
payload = recursive_merge(payload, filter_none(extra_payload or {}))
# body data (to customize in subclasses)
data = self._prepare_payload_as_bytes(inputs, parameters, provider_mapping_info, extra_payload)
# check if both payload and data are set and return
if payload is not None and data is not None:
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=normalized_headers,
)
def get_response(
self,
response: Union[bytes, dict],
request_params: Optional[RequestParameters] = None,
) -> Any:
"""
Return the response in the expected format.
Override this method in subclasses for customized response handling."""
return response
def _prepare_api_key(self, api_key: Optional[str]) -> str:
"""Return the API key to use for the request.
Usually not overwritten in subclasses."""
if api_key is None:
api_key = get_token()
if api_key is None:
raise ValueError(
f"You must provide an api_key to work with {self.provider} API or log in with `hf auth login`."
)
return api_key
def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
"""Return the mapped model ID to use for the request.
Usually not overwritten in subclasses."""
if model is None:
raise ValueError(f"Please provide an HF model ID supported by {self.provider}.")
# hardcoded mapping for local testing
if HARDCODED_MODEL_INFERENCE_MAPPING.get(self.provider, {}).get(model):
return HARDCODED_MODEL_INFERENCE_MAPPING[self.provider][model]
provider_mapping = None
for mapping in _fetch_inference_provider_mapping(model):
if mapping.provider == self.provider:
provider_mapping = mapping
break
if provider_mapping is None:
raise ValueError(f"Model {model} is not supported by provider {self.provider}.")
if provider_mapping.task != self.task:
raise ValueError(
f"Model {model} is not supported for task {self.task} and provider {self.provider}. "
f"Supported task: {provider_mapping.task}."
)
if provider_mapping.status == "staging":
logger.warning(
f"Model {model} is in staging mode for provider {self.provider}. Meant for test purposes only."
)
if provider_mapping.status == "error":
logger.warning(
f"Our latest automated health check on model '{model}' for provider '{self.provider}' did not complete successfully. "
"Inference call might fail."
)
return provider_mapping
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.
"""
return {**build_hf_headers(token=api_key), **headers}
def _prepare_url(self, api_key: str, mapped_model: str) -> str:
"""Return the URL to use for the request.
Usually not overwritten in subclasses."""
base_url = self._prepare_base_url(api_key)
route = self._prepare_route(mapped_model, api_key)
return f"{base_url.rstrip('/')}/{route.lstrip('/')}"
def _prepare_base_url(self, api_key: str) -> str:
"""Return the base URL to use for the request.
Usually not overwritten in subclasses."""
# Route to the proxy if the api_key is a HF TOKEN
if api_key.startswith("hf_"):
logger.info(f"Calling '{self.provider}' provider through Hugging Face router.")
return constants.INFERENCE_PROXY_TEMPLATE.format(provider=self.provider)
else:
logger.info(f"Calling '{self.provider}' provider directly.")
return self.base_url
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.
"""
return ""
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
"""Return the payload to use for the request, as a dict.
Override this method in subclasses for customized payloads.
Only one of `_prepare_payload_as_dict` and `_prepare_payload_as_bytes` should return a value.
"""
return None
def _prepare_payload_as_bytes(
self,
inputs: Any,
parameters: dict,
provider_mapping_info: InferenceProviderMapping,
extra_payload: Optional[dict],
) -> Optional[MimeBytes]:
"""Return the body to use for the request, as bytes.
Override this method in subclasses for customized body data.
Only one of `_prepare_payload_as_dict` and `_prepare_payload_as_bytes` should return a value.
"""
return None
class BaseConversationalTask(TaskProviderHelper):
"""
Base class for conversational (chat completion) tasks.
The schema follows the OpenAI API format defined here: https://platform.openai.com/docs/api-reference/chat
"""
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, api_key: str) -> str:
return "/v1/chat/completions"
def _prepare_payload_as_dict(
self,
inputs: list[Union[dict, ChatCompletionInputMessage]],
parameters: dict,
provider_mapping_info: InferenceProviderMapping,
) -> Optional[dict]:
return filter_none({"messages": inputs, **parameters, "model": provider_mapping_info.provider_id})
class AutoRouterConversationalTask(BaseConversationalTask):
"""
Auto-router for conversational tasks.
We let the Hugging Face router select the best provider for the model, based on availability and user preferences.
This is a special case since the selection is done server-side (avoid 1 API call to fetch provider mapping).
"""
def __init__(self):
super().__init__(provider="auto", base_url="https://router.huggingface.co")
def _prepare_base_url(self, api_key: str) -> str:
"""Return the base URL to use for the request.
Usually not overwritten in subclasses."""
# Route to the proxy if the api_key is a HF TOKEN
if not api_key.startswith("hf_"):
raise ValueError("Cannot select auto-router when using non-Hugging Face API key.")
else:
return self.base_url # No `/auto` suffix in the URL
def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
"""
In auto-router, we don't need to fetch provider mapping info.
We just return a dummy mapping info with provider_id set to the HF model ID.
"""
if model is None:
raise ValueError("Please provide an HF model ID.")
return InferenceProviderMapping(
provider="auto",
hf_model_id=model,
providerId=model,
status="live",
task="conversational",
)
class BaseTextGenerationTask(TaskProviderHelper):
"""
Base class for text-generation (completion) tasks.
The schema follows the OpenAI API format defined here: https://platform.openai.com/docs/api-reference/completions
"""
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, api_key: str) -> str:
return "/v1/completions"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
return filter_none({"prompt": inputs, **parameters, "model": provider_mapping_info.provider_id})
@lru_cache(maxsize=None)
def _fetch_inference_provider_mapping(model: str) -> list["InferenceProviderMapping"]:
"""
Fetch provider mappings for a model from the Hub.
"""
from huggingface_hub.hf_api import HfApi
info = HfApi().model_info(model, expand=["inferenceProviderMapping"])
provider_mapping = info.inference_provider_mapping
if provider_mapping is None:
raise ValueError(f"No provider mapping found for model {model}")
return provider_mapping
def recursive_merge(dict1: dict, dict2: dict) -> dict:
return {
**dict1,
**{
key: recursive_merge(dict1[key], value)
if (key in dict1 and isinstance(dict1[key], dict) and isinstance(value, dict))
else value
for key, value in dict2.items()
},
}

View file

@ -0,0 +1,69 @@
import time
from typing import Any, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
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
logger = logging.get_logger(__name__)
MAX_POLLING_ATTEMPTS = 6
POLLING_INTERVAL = 1.0
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[str, Any]:
headers = super()._prepare_headers(headers, api_key)
if not api_key.startswith("hf_"):
_ = headers.pop("authorization")
headers["X-Key"] = api_key
return headers
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, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
parameters = filter_none(parameters)
if "num_inference_steps" in parameters:
parameters["steps"] = parameters.pop("num_inference_steps")
if "guidance_scale" in parameters:
parameters["guidance"] = parameters.pop("guidance_scale")
return {"prompt": inputs, **parameters}
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.
"""
url = _as_dict(response).get("polling_url")
session = get_session()
for _ in range(MAX_POLLING_ATTEMPTS):
time.sleep(POLLING_INTERVAL)
response = session.get(url, headers={"Content-Type": "application/json"}) # type: ignore
response.raise_for_status() # type: ignore
response_json: dict = response.json() # type: ignore
status = response_json.get("status")
logger.info(
f"Polling generation result from {url}. Current status: {status}. "
f"Will retry after {POLLING_INTERVAL} seconds if not ready."
)
if (
status == "Ready"
and isinstance(response_json.get("result"), dict)
and (sample_url := response_json["result"].get("sample"))
):
image_resp = session.get(sample_url)
image_resp.raise_for_status()
return image_resp.content
raise TimeoutError(f"Failed to get the image URL after {MAX_POLLING_ATTEMPTS} attempts.")

View file

@ -0,0 +1,6 @@
from ._common import BaseConversationalTask
class CerebrasConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="cerebras", base_url="https://api.cerebras.ai")

View file

@ -0,0 +1,13 @@
from ._common import BaseConversationalTask
_PROVIDER = "clarifai"
_BASE_URL = "https://api.clarifai.com"
class ClarifaiConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v2/ext/openai/v1/chat/completions"

View file

@ -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

View file

@ -0,0 +1,299 @@
import base64
import time
from abc import ABC
from typing import Any, Optional, Union
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: 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, 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
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: 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)}.")
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
) -> 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"),
}
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], request_params: Optional[RequestParameters] = 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
) -> Optional[dict]:
return {"text": inputs, **filter_none(parameters)}
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(FalAIQueueTask):
def __init__(self):
super().__init__("text-to-video")
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],
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,
}
]

View file

@ -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)

View file

@ -0,0 +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")
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

View file

@ -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"

View file

@ -0,0 +1,228 @@
import json
from functools import lru_cache
from pathlib import Path
from typing import Any, Optional, Union
from urllib.parse import urlparse, urlunparse
from huggingface_hub import constants
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
class HFInferenceTask(TaskProviderHelper):
"""Base class for HF Inference API tasks."""
def __init__(self, task: str):
super().__init__(
provider="hf-inference",
base_url=constants.INFERENCE_PROXY_TEMPLATE.format(provider="hf-inference"),
task=task,
)
def _prepare_api_key(self, api_key: Optional[str]) -> str:
# special case: for HF Inference we allow not providing an API key
return api_key or get_token() # type: ignore[return-value]
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."
)
_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)
if mapped_model.startswith(("http://", "https://")):
return mapped_model
return (
# Feature-extraction and sentence-similarity are the only cases where we handle models with several tasks.
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, 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 filter_none({"inputs": inputs, "parameters": parameters})
class HFInferenceBinaryInputTask(HFInferenceTask):
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,
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
# Raise if not a binary object or a local path or a URL.
if not isinstance(inputs, (bytes, Path)) and not isinstance(inputs, str):
raise ValueError(f"Expected binary inputs or a local path or a URL. Got {inputs}")
# Send inputs as raw content when no parameters are provided
if not has_parameters:
return _open_as_mime_bytes(inputs)
# Otherwise encode as b64
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__("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 = (
mapped_model
if mapped_model.startswith(("http://", "https://"))
else f"{constants.INFERENCE_PROXY_TEMPLATE.format(provider='hf-inference')}/models/{mapped_model}"
)
return _build_chat_completion_url(base_url)
def _build_chat_completion_url(model_url: str) -> str:
parsed = urlparse(model_url)
path = parsed.path.rstrip("/")
# 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
else:
new_path = path + "/v1/chat/completions"
# 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]]:
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

View file

@ -0,0 +1,47 @@
import base64
from typing import Any, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
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, api_key: str) -> str:
return "/v1/images/generations"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
mapped_model = provider_mapping_info.provider_id
parameters = filter_none(parameters)
if "num_inference_steps" in parameters:
parameters["steps"] = parameters.pop("num_inference_steps")
if "guidance_scale" in parameters:
parameters["cfg_scale"] = parameters.pop("guidance_scale")
# For Hyperbolic, the width and height are required parameters
if "width" not in parameters:
parameters["width"] = 512
if "height" not in parameters:
parameters["height"] = 512
return {"prompt": inputs, "model_name": mapped_model, **parameters}
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"])
class HyperbolicTextGenerationTask(BaseConversationalTask):
"""
Special case for Hyperbolic, where text-generation task is handled as a conversational task.
"""
def __init__(self, task: str):
super().__init__(
provider="hyperbolic",
base_url="https://api.hyperbolic.xyz",
)
self.task = task

View file

@ -0,0 +1,83 @@
import base64
from typing import Any, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import (
BaseConversationalTask,
BaseTextGenerationTask,
TaskProviderHelper,
filter_none,
)
class NebiusTextGenerationTask(BaseTextGenerationTask):
def __init__(self):
super().__init__(provider="nebius", base_url="https://api.studio.nebius.ai")
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 NebiusConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="nebius", base_url="https://api.studio.nebius.ai")
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["guided_json"] = json_schema_details["schema"] # type: ignore [index]
return payload
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, api_key: str) -> str:
return "/v1/images/generations"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
mapped_model = provider_mapping_info.provider_id
parameters = filter_none(parameters)
if "guidance_scale" in parameters:
parameters.pop("guidance_scale")
if parameters.get("response_format") not in ("b64_json", "url"):
parameters["response_format"] = "b64_json"
return {"prompt": inputs, **parameters, "model": mapped_model}
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"])
class NebiusFeatureExtractionTask(TaskProviderHelper):
def __init__(self):
super().__init__(task="feature-extraction", provider="nebius", base_url="https://api.studio.nebius.ai")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v1/embeddings"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
return {"input": inputs, "model": provider_mapping_info.provider_id}
def get_response(self, response: Union[bytes, dict], request_params: Optional[RequestParameters] = None) -> Any:
embeddings = _as_dict(response)["data"]
return [embedding["embedding"] for embedding in embeddings]

View file

@ -0,0 +1,69 @@
from typing import Any, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import (
BaseConversationalTask,
BaseTextGenerationTask,
TaskProviderHelper,
filter_none,
)
from huggingface_hub.utils import get_session
_PROVIDER = "novita"
_BASE_URL = "https://api.novita.ai"
class NovitaTextGenerationTask(BaseTextGenerationTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
# there is no v1/ route for novita
return "/v3/openai/completions"
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 NovitaConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
# there is no v1/ route for novita
return "/v3/openai/chat/completions"
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, api_key: str) -> str:
return f"/v3/hf/{mapped_model}"
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], request_params: Optional[RequestParameters] = None) -> Any:
response_dict = _as_dict(response)
if not (
isinstance(response_dict, dict)
and "video" in response_dict
and isinstance(response_dict["video"], dict)
and "video_url" in response_dict["video"]
):
raise ValueError("Expected response format: { 'video': { 'video_url': string } }")
video_url = response_dict["video"]["video_url"]
return get_session().get(video_url).content

View file

@ -0,0 +1,44 @@
import base64
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, TaskProviderHelper, filter_none
class NscaleConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="nscale", base_url="https://inference.api.nscale.com")
class NscaleTextToImageTask(TaskProviderHelper):
def __init__(self):
super().__init__(provider="nscale", base_url="https://inference.api.nscale.com", task="text-to-image")
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, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
mapped_model = provider_mapping_info.provider_id
# Combine all parameters except inputs and parameters
parameters = filter_none(parameters)
if "width" in parameters and "height" in parameters:
parameters["size"] = f"{parameters.pop('width')}x{parameters.pop('height')}"
if "num_inference_steps" in parameters:
parameters.pop("num_inference_steps")
if "cfg_scale" in parameters:
parameters.pop("cfg_scale")
payload = {
"response_format": "b64_json",
"prompt": inputs,
"model": mapped_model,
**parameters,
}
return payload
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"])

View file

@ -0,0 +1,25 @@
from typing import Optional
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._providers._common import BaseConversationalTask
class OpenAIConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="openai", base_url="https://api.openai.com")
def _prepare_api_key(self, api_key: Optional[str]) -> str:
if api_key is None:
raise ValueError("You must provide an api_key to work with OpenAI API.")
if api_key.startswith("hf_"):
raise ValueError(
"OpenAI provider is not available through Hugging Face routing, please use your own OpenAI API key."
)
return api_key
def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
if model is None:
raise ValueError("Please provide an OpenAI model ID, e.g. `gpt-4o` or `o1`.")
return InferenceProviderMapping(
provider="openai", providerId=model, task="conversational", status="live", hf_model_id=model
)

View file

@ -0,0 +1,10 @@
from huggingface_hub.inference._providers._common import BaseConversationalTask
_PROVIDER = "ovhcloud"
_BASE_URL = "https://oai.endpoints.kepler.ai.cloud.ovh.net"
class OVHcloudConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)

View file

@ -0,0 +1,6 @@
from ._common import BaseConversationalTask
class PublicAIConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="publicai", base_url="https://api.publicai.co")

View file

@ -0,0 +1,151 @@
from typing import Any, Optional, Union
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
_PROVIDER = "replicate"
_BASE_URL = "https://api.replicate.com"
class ReplicateTask(TaskProviderHelper):
def __init__(self, task: str):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task=task)
def _prepare_headers(self, headers: dict, api_key: str) -> dict[str, Any]:
headers = super()._prepare_headers(headers, api_key)
headers["Prefer"] = "wait"
return headers
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
if ":" in mapped_model:
return "/v1/predictions"
return f"/v1/models/{mapped_model}/predictions"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
mapped_model = provider_mapping_info.provider_id
payload: dict[str, Any] = {"input": {"prompt": inputs, **filter_none(parameters)}}
if ":" in mapped_model:
version = mapped_model.split(":", 1)[1]
payload["version"] = version
return payload
def get_response(self, response: Union[bytes, dict], request_params: Optional[RequestParameters] = None) -> Any:
response_dict = _as_dict(response)
if response_dict.get("output") is None:
raise TimeoutError(
f"Inference request timed out after 60 seconds. No output generated for model {response_dict.get('model')}"
"The model might be in cold state or starting up. Please try again later."
)
output_url = (
response_dict["output"] if isinstance(response_dict["output"], str) else response_dict["output"][0]
)
return get_session().get(output_url).content
class ReplicateTextToImageTask(ReplicateTask):
def __init__(self):
super().__init__("text-to-image")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
payload: dict = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info) # type: ignore[assignment]
if provider_mapping_info.adapter_weights_path is not None:
payload["input"]["lora_weights"] = f"https://huggingface.co/{provider_mapping_info.hf_model_id}"
return payload
class ReplicateTextToSpeechTask(ReplicateTask):
def __init__(self):
super().__init__("text-to-speech")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
payload: dict = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info) # type: ignore[assignment]
payload["input"]["text"] = payload["input"].pop("prompt") # rename "prompt" to "text" for TTS
return payload
class ReplicateAutomaticSpeechRecognitionTask(ReplicateTask):
def __init__(self) -> None:
super().__init__("automatic-speech-recognition")
def _prepare_payload_as_dict(
self,
inputs: Any,
parameters: dict,
provider_mapping_info: InferenceProviderMapping,
) -> Optional[dict]:
mapped_model = provider_mapping_info.provider_id
audio_url = _as_url(inputs, default_mime_type="audio/wav")
payload: dict[str, Any] = {
"input": {
**{"audio": audio_url},
**filter_none(parameters),
}
}
if ":" in mapped_model:
payload["version"] = mapped_model.split(":", 1)[1]
return payload
def get_response(self, response: Union[bytes, dict], request_params: Optional[RequestParameters] = None) -> Any:
response_dict = _as_dict(response)
output = response_dict.get("output")
if isinstance(output, str):
return {"text": output}
if isinstance(output, list) and output:
first_item = output[0]
if isinstance(first_item, str):
return {"text": first_item}
if isinstance(first_item, dict):
output = first_item
text: Optional[str] = None
if isinstance(output, dict):
transcription = output.get("transcription")
if isinstance(transcription, str):
text = transcription
translation = output.get("translation")
if isinstance(translation, str):
text = translation
txt_file = output.get("txt_file")
if isinstance(txt_file, str):
text_response = get_session().get(txt_file)
text_response.raise_for_status()
text = text_response.text
if text is not None:
return {"text": text}
raise ValueError("Received malformed response from Replicate automatic-speech-recognition API")
class ReplicateImageToImageTask(ReplicateTask):
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")
payload: dict[str, Any] = {"input": {"input_image": image_url, **filter_none(parameters)}}
mapped_model = provider_mapping_info.provider_id
if ":" in mapped_model:
version = mapped_model.split(":", 1)[1]
payload["version"] = version
return payload

View file

@ -0,0 +1,42 @@
from typing import Any, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none
class SambanovaConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="sambanova", base_url="https://api.sambanova.ai")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
response_format_config = parameters.get("response_format")
if isinstance(response_format_config, dict):
if response_format_config.get("type") == "json_schema":
json_schema_config = response_format_config.get("json_schema", {})
strict = json_schema_config.get("strict")
if isinstance(json_schema_config, dict) and (strict is True or strict is None):
json_schema_config["strict"] = False
payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
return payload
class SambanovaFeatureExtractionTask(TaskProviderHelper):
def __init__(self):
super().__init__(provider="sambanova", base_url="https://api.sambanova.ai", task="feature-extraction")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v1/embeddings"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
parameters = filter_none(parameters)
return {"input": inputs, "model": provider_mapping_info.provider_id, **parameters}
def get_response(self, response: Union[bytes, dict], request_params: Optional[RequestParameters] = None) -> Any:
embeddings = _as_dict(response)["data"]
return [embedding["embedding"] for embedding in embeddings]

View file

@ -0,0 +1,28 @@
from typing import Any, Dict, Optional, Union
from huggingface_hub.inference._common import RequestParameters, _as_dict
from ._common import BaseConversationalTask, InferenceProviderMapping, TaskProviderHelper, filter_none
class ScalewayConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider="scaleway", base_url="https://api.scaleway.ai")
class ScalewayFeatureExtractionTask(TaskProviderHelper):
def __init__(self):
super().__init__(provider="scaleway", base_url="https://api.scaleway.ai", task="feature-extraction")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/v1/embeddings"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[Dict]:
parameters = filter_none(parameters)
return {"input": inputs, "model": provider_mapping_info.provider_id, **parameters}
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
embeddings = _as_dict(response)["data"]
return [embedding["embedding"] for embedding in embeddings]

View file

@ -0,0 +1,88 @@
import base64
from abc import ABC
from typing import Any, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import (
BaseConversationalTask,
BaseTextGenerationTask,
TaskProviderHelper,
filter_none,
)
_PROVIDER = "together"
_BASE_URL = "https://api.together.xyz"
class TogetherTask(TaskProviderHelper, ABC):
"""Base class for Together API tasks."""
def __init__(self, task: str):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task=task)
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
if self.task == "text-to-image":
return "/v1/images/generations"
elif self.task == "conversational":
return "/v1/chat/completions"
elif self.task == "text-generation":
return "/v1/completions"
raise ValueError(f"Unsupported task '{self.task}' for Together API.")
class TogetherTextGenerationTask(BaseTextGenerationTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
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 TogetherConversationalTask(BaseConversationalTask):
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]:
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
class TogetherTextToImageTask(TogetherTask):
def __init__(self):
super().__init__("text-to-image")
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
mapped_model = provider_mapping_info.provider_id
parameters = filter_none(parameters)
if "num_inference_steps" in parameters:
parameters["steps"] = parameters.pop("num_inference_steps")
if "guidance_scale" in parameters:
parameters["guidance"] = parameters.pop("guidance_scale")
return {"prompt": inputs, "response_format": "base64", **parameters, "model": mapped_model}
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"])

View file

@ -0,0 +1,138 @@
import base64
import time
from abc import ABC
from typing import Any, Optional, Union
from urllib.parse import urlparse
from huggingface_hub.hf_api import InferenceProviderMapping
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, hf_raise_for_status
from huggingface_hub.utils.logging import get_logger
logger = get_logger(__name__)
# Polling interval (in seconds)
_POLLING_INTERVAL = 0.5
class WavespeedAITask(TaskProviderHelper, ABC):
def __init__(self, task: str):
super().__init__(provider="wavespeed", base_url="https://api.wavespeed.ai", task=task)
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return f"/api/v3/{mapped_model}"
def get_response(
self,
response: Union[bytes, dict],
request_params: Optional[RequestParameters] = None,
) -> Any:
response_dict = _as_dict(response)
data = response_dict.get("data", {})
result_path = data.get("urls", {}).get("get")
if not result_path:
raise ValueError("No result URL found in the response")
if request_params is None:
raise ValueError("A `RequestParameters` object should be provided to get responses with WaveSpeed AI.")
# Parse the request URL to determine base URL
parsed_url = urlparse(request_params.url)
# Add /wavespeed to base URL if going through HF router
if parsed_url.netloc == "router.huggingface.co":
base_url = f"{parsed_url.scheme}://{parsed_url.netloc}/wavespeed"
else:
base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
# Extract path from result_path URL
if isinstance(result_path, str):
result_url_path = urlparse(result_path).path
else:
result_url_path = result_path
result_url = f"{base_url}{result_url_path}"
logger.info("Processing request, polling for results...")
# Poll until task is completed
while True:
time.sleep(_POLLING_INTERVAL)
result_response = get_session().get(result_url, headers=request_params.headers)
hf_raise_for_status(result_response)
result = result_response.json()
task_result = result.get("data", {})
status = task_result.get("status")
if status == "completed":
# Get content from the first output URL
if not task_result.get("outputs") or len(task_result["outputs"]) == 0:
raise ValueError("No output URL in completed response")
output_url = task_result["outputs"][0]
return get_session().get(output_url).content
elif status == "failed":
error_msg = task_result.get("error", "Task failed with no specific error message")
raise ValueError(f"WaveSpeed AI task failed: {error_msg}")
elif status in ["processing", "created"]:
continue
else:
raise ValueError(f"Unknown status: {status}")
class WavespeedAITextToImageTask(WavespeedAITask):
def __init__(self):
super().__init__("text-to-image")
def _prepare_payload_as_dict(
self,
inputs: Any,
parameters: dict,
provider_mapping_info: InferenceProviderMapping,
) -> Optional[dict]:
return {"prompt": inputs, **filter_none(parameters)}
class WavespeedAITextToVideoTask(WavespeedAITextToImageTask):
def __init__(self):
WavespeedAITask.__init__(self, "text-to-video")
class WavespeedAIImageToImageTask(WavespeedAITask):
def __init__(self):
super().__init__("image-to-image")
def _prepare_payload_as_dict(
self,
inputs: Any,
parameters: dict,
provider_mapping_info: InferenceProviderMapping,
) -> Optional[dict]:
# Convert inputs to image (URL or base64)
if isinstance(inputs, str) and inputs.startswith(("http://", "https://")):
image = inputs
elif isinstance(inputs, str):
# If input is a file path, read it first
with open(inputs, "rb") as f:
file_content = f.read()
image_b64 = base64.b64encode(file_content).decode("utf-8")
image = f"data:image/jpeg;base64,{image_b64}"
else:
# If input is binary data
image_b64 = base64.b64encode(inputs).decode("utf-8")
image = f"data:image/jpeg;base64,{image_b64}"
# Extract prompt from parameters if present
prompt = parameters.pop("prompt", None)
payload = {"image": image, **filter_none(parameters)}
if prompt is not None:
payload["prompt"] = prompt
return payload
class WavespeedAIImageToVideoTask(WavespeedAIImageToImageTask):
def __init__(self):
WavespeedAITask.__init__(self, "image-to-video")

View file

@ -0,0 +1,127 @@
import time
from abc import ABC
from typing import Any, Optional, Union
from huggingface_hub.hf_api import InferenceProviderMapping
from huggingface_hub.inference._common import RequestParameters, _as_dict
from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none
from huggingface_hub.utils import get_session
_PROVIDER = "zai-org"
_BASE_URL = "https://api.z.ai"
_POLLING_INTERVAL = 5 # seconds
_MAX_POLL_ATTEMPTS = 60
class ZaiTask(TaskProviderHelper, ABC):
def __init__(self, task: str):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task=task)
def _prepare_headers(self, headers: dict, api_key: str) -> dict[str, Any]:
headers = super()._prepare_headers(headers, api_key)
headers["Accept-Language"] = "en-US,en"
headers["x-source-channel"] = "hugging_face"
return headers
class ZaiConversationalTask(BaseConversationalTask):
def __init__(self):
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
def _prepare_headers(self, headers: dict, api_key: str) -> dict[str, Any]:
headers = super()._prepare_headers(headers, api_key)
headers["Accept-Language"] = "en-US,en"
headers["x-source-channel"] = "hugging_face"
return headers
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/api/paas/v4/chat/completions"
class ZaiTextToImageTask(ZaiTask):
"""Text-to-image task for ZAI provider using async API."""
def __init__(self):
super().__init__("text-to-image")
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
return "/api/paas/v4/async/images/generations"
def _prepare_payload_as_dict(
self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
) -> Optional[dict]:
width = parameters.pop("width", None)
height = parameters.pop("height", None)
size = None
if width is not None and height is not None:
size = f"{width}x{height}"
payload: dict[str, Any] = {
"model": provider_mapping_info.provider_id,
"prompt": inputs,
}
if size is not None:
payload["size"] = size
payload.update(filter_none(parameters))
return payload
def get_response(
self,
response: Union[bytes, dict],
request_params: Optional[RequestParameters] = None,
) -> Any:
"""Handle async response by polling for results."""
response_dict = _as_dict(response)
task_id = response_dict.get("id")
if task_id is None:
raise ValueError("No task_id in response from ZAI API")
task_status = response_dict.get("task_status")
if task_status == "FAIL":
raise ValueError(f"ZAI image generation failed for request {task_id}")
if task_status == "PROCESSING" and request_params is not None:
return self._poll_for_result(task_id, request_params)
return self._extract_image(response_dict)
def _poll_for_result(self, task_id: str, request_params: RequestParameters) -> bytes:
"""Poll the async-result endpoint until completion."""
session = get_session()
base_url = request_params.url.rsplit("/api/paas/v4/async/images/generations", 1)[0]
poll_url = f"{base_url}/api/paas/v4/async-result/{task_id}"
for _ in range(_MAX_POLL_ATTEMPTS):
poll_response = session.get(poll_url, headers=request_params.headers)
poll_response.raise_for_status()
result = poll_response.json()
task_status = result.get("task_status")
if task_status == "SUCCESS":
return self._extract_image(result)
elif task_status == "FAIL":
raise ValueError(f"Zai text-to-image generation failed for request {task_id}")
time.sleep(_POLLING_INTERVAL)
raise ValueError(
f"Timed out while waiting for the result from Zai API - aborting after {_MAX_POLL_ATTEMPTS} attempts"
)
def _extract_image(self, result: dict) -> bytes:
"""Extract and download the image from the result."""
image_result = result.get("image_result")
if not image_result or not isinstance(image_result, list) or len(image_result) == 0:
raise ValueError("No image_result in response from ZAI API")
image_url = image_result[0].get("url")
if not image_url:
raise ValueError("No image URL in response from ZAI API")
session = get_session()
image_response = session.get(image_url)
image_response.raise_for_status()
return image_response.content