o
    h
                     @   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                       s   e Zd Z fddZ  ZS )NscaleConversationalTaskc                    s   t  jddd d S )Nnscale https://inference.api.nscale.com)providerbase_urlsuper__init__self	__class__ a/var/www/html/ai/venv/lib/python3.10/site-packages/huggingface_hub/inference/_providers/nscale.pyr      s   z!NscaleConversationalTask.__init__)__name__
__module____qualname__r   __classcell__r   r   r   r   r   
   s    r   c                	       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 )NscaleTextToImageTaskc                    s   t  jdddd d S )Nr   r   ztext-to-image)r   r   taskr   r   r   r   r   r      s   zNscaleTextToImageTask.__init__mapped_modelapi_keyreturnc                 C   s   dS )Nz/v1/images/generationsr   )r   r!   r"   r   r   r   _prepare_route   s   z$NscaleTextToImageTask._prepare_routeinputs
parametersprovider_mapping_infoc                 C   st   |j }t|}d|v rd|v r|d d|d |d< d|v r'|d d|v r0|d d||d|}|S )	Nwidthheightxsizenum_inference_steps	cfg_scaleb64_json)response_formatpromptmodel)provider_idr   pop)r   r%   r&   r'   r!   payloadr   r   r   _prepare_payload_as_dict   s   

z.NscaleTextToImageTask._prepare_payload_as_dictNresponserequest_paramsc                 C   s   t |}t|d d d S )Ndatar   r.   )r   base64	b64decode)r   r6   r7   response_dictr   r   r   get_response*   s   z"NscaleTextToImageTask.get_response)N)r   r   r   r   strr$   r   r   r   r   r5   r   bytesr   r<   r   r   r   r   r   r      s    
,r   )r9   typingr   r   r   r   huggingface_hub.hf_apir   !huggingface_hub.inference._commonr   r   _commonr
   r   r   r   r   r   r   r   r   <module>   s    