o
    vZh                     @   st   d dl Z d dlmZmZmZmZ d dlmZ d dlm	Z	m
Z
 d dlmZmZmZ G dd deZG dd	 d	eZdS )
    N)AnyDictOptionalUnion)InferenceProviderMapping)RequestParameters_as_dict)BaseConversationalTaskTaskProviderHelperfilter_nonec                	       st   e Zd Z fddZdededefddZded	ed
ede	e fddZ
ddeeef de	e defddZ  ZS )HyperbolicTextToImageTaskc                    s   t  jdddd d S )N
hyperbolichttps://api.hyperbolic.xyzztext-to-image)providerbase_urltask)super__init__)self	__class__ ^/var/www/auris/lib/python3.10/site-packages/huggingface_hub/inference/_providers/hyperbolic.pyr   
   s   z"HyperbolicTextToImageTask.__init__mapped_modelapi_keyreturnc                 C   s   dS )Nz/v1/images/generationsr   )r   r   r   r   r   r   _prepare_route   s   z(HyperbolicTextToImageTask._prepare_routeinputs
parametersprovider_mapping_infoc                 C   sh   |j }t|}d|v r|d|d< d|v r|d|d< d|vr%d|d< d|vr-d|d< ||d|S )	NZnum_inference_stepsZstepsZguidance_scaleZ	cfg_scalewidthi   height)promptZ
model_name)Zprovider_idr   pop)r   r   r   r   r   r   r   r   _prepare_payload_as_dict   s   z2HyperbolicTextToImageTask._prepare_payload_as_dictNresponserequest_paramsc                 C   s   t |}t|d d d S )NZimagesr   image)r   base64	b64decode)r   r%   r&   Zresponse_dictr   r   r   get_response    s   z&HyperbolicTextToImageTask.get_response)N)__name__
__module____qualname__r   strr   r   r   r   r   r$   r   bytesr   r*   __classcell__r   r   r   r   r   	   s    
,r   c                       s&   e Zd ZdZdef fddZ  ZS )HyperbolicTextGenerationTaskzf
    Special case for Hyperbolic, where text-generation task is handled as a conversational task.
    r   c                    s   t  jddd || _d S )Nr   r   )r   r   )r   r   r   )r   r   r   r   r   r   *   s
   
z%HyperbolicTextGenerationTask.__init__)r+   r,   r-   __doc__r.   r   r0   r   r   r   r   r1   %   s    r1   )r(   typingr   r   r   r   Zhuggingface_hub.hf_apir   Z!huggingface_hub.inference._commonr   r   Z,huggingface_hub.inference._providers._commonr	   r
   r   r   r1   r   r   r   r   <module>   s    