Initialisation du repository de Beta
This commit is contained in:
commit
14985f6dbb
9469 changed files with 1903273 additions and 0 deletions
|
|
@ -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]
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -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()
|
||||
},
|
||||
}
|
||||
|
|
@ -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.")
|
||||
|
|
@ -0,0 +1,6 @@
|
|||
from ._common import BaseConversationalTask
|
||||
|
||||
|
||||
class CerebrasConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider="cerebras", base_url="https://api.cerebras.ai")
|
||||
|
|
@ -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"
|
||||
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
}
|
||||
]
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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]
|
||||
|
|
@ -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
|
||||
|
|
@ -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"])
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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)
|
||||
|
|
@ -0,0 +1,6 @@
|
|||
from ._common import BaseConversationalTask
|
||||
|
||||
|
||||
class PublicAIConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider="publicai", base_url="https://api.publicai.co")
|
||||
|
|
@ -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
|
||||
|
|
@ -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]
|
||||
|
|
@ -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]
|
||||
|
|
@ -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"])
|
||||
|
|
@ -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")
|
||||
|
|
@ -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
|
||||
Loading…
Add table
Add a link
Reference in a new issue