-
Notifications
You must be signed in to change notification settings - Fork 839
[inference provider] Add wavespeed.ai as an inference provider #3474
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[inference provider] Add wavespeed.ai as an inference provider #3474
Conversation
|
Related to huggingface/huggingface.js#1424 |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
@bot /style |
|
Style bot fixed some files and pushed the changes. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
thanks @arabot777 for the PR! 🤗
could you please update this line
| response = provider_helper.get_response(response) |
request_parameters into get_response():
response = provider_helper.get_response(response, request_parameters)also, could you add Wavespeed to the Inference providers table in the documentation here?
I left a few comments to avoid duplicating the payload dictionary construction and use built-in types!
| def _prepare_payload_as_dict( | ||
| self, | ||
| inputs: Any, | ||
| parameters: Dict, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| parameters: Dict, | |
| parameters: dict, |
|
|
||
| def get_response( | ||
| self, | ||
| response: Union[bytes, Dict], |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| response: Union[bytes, Dict], | |
| response: Union[bytes, dict], |
| def __init__(self, task: str): | ||
| super().__init__(provider="wavespeed", base_url="https://api.wavespeed.ai", task=task) | ||
|
|
||
| def _prepare_headers(self, headers: Dict, api_key: str) -> Dict: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
let's use built-in types everywhere
| def _prepare_headers(self, headers: Dict, api_key: str) -> Dict: | |
| def _prepare_headers(self, headers: dict, api_key: str) -> dict: |
| inputs: Any, | ||
| parameters: Dict, | ||
| provider_mapping_info: InferenceProviderMapping, | ||
| ) -> Optional[Dict]: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| ) -> Optional[Dict]: | |
| ) -> Optional[dict]: |
| class WavespeedAITextToVideoTask(WavespeedAITask): | ||
| 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)} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
WavespeedAITextToVideoTask can inherit from WavespeedAITextToImageTask since it's the same payload preparation when the input modality is text
| class WavespeedAITextToVideoTask(WavespeedAITask): | |
| 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)} | |
| class WavespeedAITextToVideoTask(WavespeedAITextToImageTask): | |
| def __init__(self): | |
| WavespeedAITask.__init__(self, "text-to-video") |
| class WavespeedAIImageToVideoTask(WavespeedAITask): | ||
| def __init__(self): | ||
| super().__init__("image-to-video") | ||
|
|
||
| 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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
since the payload is built the same for both image-to-image and image-to-video
| class WavespeedAIImageToVideoTask(WavespeedAITask): | |
| def __init__(self): | |
| super().__init__("image-to-video") | |
| 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") |
|
@hanouticelina thanks for the review! I've addressed all your feedback: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
almost ready to be merged! 🤗 thanks for the iteration @arabot777
| def _prepare_headers(self, headers: dict, api_key: str) -> dict: | ||
| headers = super()._prepare_headers(headers, api_key) | ||
| if not api_key.startswith("hf_"): | ||
| headers["Authorization"] = f"Bearer {api_key}" | ||
| return headers |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
no need to redefine _prepare_headers, TaskProviderHelper. _prepare_headers already sets the bearer prefix in the authorization header
| def _prepare_headers(self, headers: dict, api_key: str) -> dict: | |
| headers = super()._prepare_headers(headers, api_key) | |
| if not api_key.startswith("hf_"): | |
| headers["Authorization"] = f"Bearer {api_key}" | |
| return headers |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@hanouticelina Done! Thanks for the thorough review!
|
@bot /style |
|
Style bot fixed some files and pushed the changes. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
looks good, thank you @arabot777 for the contribution!
What’s in this PR
This PR adds wavespeed AI as a new inference provider for conversational tasks.