hyperbolic.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647
  1. import base64
  2. from typing import Any
  3. from huggingface_hub.hf_api import InferenceProviderMapping
  4. from huggingface_hub.inference._common import RequestParameters, _as_dict
  5. from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none
  6. class HyperbolicTextToImageTask(TaskProviderHelper):
  7. def __init__(self):
  8. super().__init__(provider="hyperbolic", base_url="https://api.hyperbolic.xyz", task="text-to-image")
  9. def _prepare_route(self, mapped_model: str, api_key: str) -> str:
  10. return "/v1/images/generations"
  11. def _prepare_payload_as_dict(
  12. self, inputs: Any, parameters: dict, provider_mapping_info: InferenceProviderMapping
  13. ) -> dict | None:
  14. mapped_model = provider_mapping_info.provider_id
  15. parameters = filter_none(parameters)
  16. if "num_inference_steps" in parameters:
  17. parameters["steps"] = parameters.pop("num_inference_steps")
  18. if "guidance_scale" in parameters:
  19. parameters["cfg_scale"] = parameters.pop("guidance_scale")
  20. # For Hyperbolic, the width and height are required parameters
  21. if "width" not in parameters:
  22. parameters["width"] = 512
  23. if "height" not in parameters:
  24. parameters["height"] = 512
  25. return {"prompt": inputs, "model_name": mapped_model, **parameters}
  26. def get_response(self, response: bytes | dict, request_params: RequestParameters | None = None) -> Any:
  27. response_dict = _as_dict(response)
  28. return base64.b64decode(response_dict["images"][0]["image"])
  29. class HyperbolicTextGenerationTask(BaseConversationalTask):
  30. """
  31. Special case for Hyperbolic, where text-generation task is handled as a conversational task.
  32. """
  33. def __init__(self, task: str):
  34. super().__init__(
  35. provider="hyperbolic",
  36. base_url="https://api.hyperbolic.xyz",
  37. )
  38. self.task = task