
    wi'                        d dl Z d dlZd dlmZ d dlmZmZmZmZ d dl	m
Z
 d dlmZ d dlmZ d dlmZmZmZ d dlmZmZ d d	lmZmZ d d
lmZ  ee          ZdZ G d dee          Z G d dee          Z G d de          Z  G d de          Z! G d de          Z" G d de          Z# G d de          Z$ G d de          Z%dS )    N)ABC)AnyDictOptionalUnion)urlparse)	constants)InferenceProviderMapping)RequestParameters_as_dict_as_url)TaskProviderHelperfilter_none)get_sessionhf_raise_for_status)
get_loggerg      ?c                   `     e Zd Zdef fdZdededeeef         f fdZdededefdZ xZ	S )		FalAITasktaskc                 P    t                                          dd|           d S )Nfal-aizhttps://fal.runproviderbase_urlr   super__init__selfr   	__class__s     /var/www/development/aibuddy-work/election-extract/venv/lib/python3.11/site-packages/huggingface_hub/inference/_providers/fal_ai.pyr   zFalAITask.__init__   s)    (5FTRRRRR    headersapi_keyreturnc                     t                                          ||          }|                    d          sd| |d<   |S Nhf_zKey authorizationr   _prepare_headers
startswithr   r#   r$   r    s      r!   r+   zFalAITask._prepare_headers   I    ''**7G<<!!%(( 	8'7g'7'7GO$r"   mapped_modelc                     d| S )N/ r   r/   r$   s      r!   _prepare_routezFalAITask._prepare_route   s    !<!!!r"   )
__name__
__module____qualname__strr   r   r   r+   r4   __classcell__r    s   @r!   r   r      s        SS S S S S S S s tCH~      "3 " " " " " " " " " "r"   r   c                        e Zd Zdef fdZdededeeef         f fdZdededefdZ	 dd
e	e
ef         dee         defdZ xZS )FalAIQueueTaskr   c                 P    t                                          dd|           d S )Nr   zhttps://queue.fal.runr   r   r   s     r!   r   zFalAIQueueTask.__init__$   s*    (5LSWXXXXXr"   r#   r$   r%   c                     t                                          ||          }|                    d          sd| |d<   |S r'   r*   r-   s      r!   r+   zFalAIQueueTask._prepare_headers'   r.   r"   r/   c                 B    |                     d          rd| dS d| S )Nr(   r1   z?_subdomain=queue)r,   r3   s      r!   r4   zFalAIQueueTask._prepare_route-   s7    e$$ 	76|6666!<!!!r"   Nresponserequest_paramsc                    t          |          }|                    d          }|st          d          |t          d| j         d          t	          |j                  }|j         d|j         |j        dk    rdnd }|j        r
d	|j         nd}t	          |                    d
                    j	        }| t          |           d| }	| t          |           | }
|                    d          }t                              d           |dk    r}t          j        t                     t!                                          |	|j                  }t%          |           |                                                    d          }|dk    }t!                                          |
|j                                                  S )N
request_idz#No request ID found in the responsez7A `RequestParameters` object should be provided to get z responses with Fal AI.z://zrouter.huggingface.coz/fal-ai ?response_urlz/statusstatusz6Generating the output.. this can take several minutes.	COMPLETED)r#   )r   get
ValueErrorr   r   urlschemenetlocquerypathr8   loggerinfotimesleep_POLLING_INTERVALr   r#   r   json)r   r@   rA   response_dictrC   
parsed_urlr   query_parammodel_id
status_url
result_urlrG   status_responses                r!   get_responsezFalAIQueueTask.get_response3   s   
 !**"&&|44
 	DBCCC!l$)lll  
 n011
 '  C  CJ,=  CJL]axLxLxyy  A  C  C0:0@H,**,,,b M--n==>>C E#h--EEEE
 >#h-->>>
""8,,LMMM##J())))mm//
NDZ/[[O000$))++//99F	 ## }}  ^5K LLQQSSSr"   N)r5   r6   r7   r8   r   r   r   r+   r4   r   bytesr   r   r]   r9   r:   s   @r!   r<   r<   #   s       YS Y Y Y Y Y Y s tCH~      "3 " " " " " " 7;#T #Tt$#T !!23#T 
	#T #T #T #T #T #T #T #Tr"   r<   c            	       t     e Zd Z fdZdedededee         fdZdde	e
ef         d	ee         defd
Z xZS )#FalAIAutomaticSpeechRecognitionTaskc                 J    t                                          d           d S )Nzautomatic-speech-recognitionr   r   r    s    r!   r   z,FalAIAutomaticSpeechRecognitionTask.__init__Z   s"    788888r"   inputs
parametersprovider_mapping_infor%   c                    t          |t                    r|                    d          r|}nt          |t                    r<t          |d          5 }|                                }d d d            n# 1 swxY w Y   t          j        |                                          }d}d| d| }d|it          |          S )N)zhttp://zhttps://rbz
audio/mpegzdata:z;base64,	audio_url)	
isinstancer8   r,   openreadbase64	b64encodedecoder   )r   rd   re   rf   ri   f	audio_b64content_types           r!   _prepare_payload_as_dictz<FalAIAutomaticSpeechRecognitionTask._prepare_payload_as_dict]   s    fc"" 	Bv'8'89P'Q'Q 	BII &#&& &&$'' &1VVXXF& & & & & & & & & & & & & & & (007799I'LAAAiAAIYB+j*A*ABBs   A44A8;A8Nr@   rA   c                     t          |          d         }t          |t                    s t          dt	          |           d          |S )Ntextz>Unexpected output format from FalAI API. Expected string, got .)r   rj   r8   rJ   type)r   r@   rA   ru   s       r!   r]   z0FalAIAutomaticSpeechRecognitionTask.get_responseo   sP    !!&)$$$ 	mk^bcg^h^hkkklllr"   r^   r5   r6   r7   r   r   r   r
   r   rs   r   r_   r   r]   r9   r:   s   @r!   ra   ra   Y   s        9 9 9 9 9CC'+CD\C	$C C C C$ U5$;%7 RcId ps        r"   ra   c            	       t     e Zd Z fdZdedededee         fdZdde	e
ef         d	ee         defd
Z xZS )FalAITextToImageTaskc                 J    t                                          d           d S )Nztext-to-imager   rc   s    r!   r   zFalAITextToImageTask.__init__w   !    )))))r"   rd   re   rf   r%   c                 0   d|it          |          }d|v r2d|v r.|                    d          |                    d          d|d<   |j        Et          j                            |j        d|j                  }|dd	g|d
<   |j        dk    rd|d<   |S )Npromptwidthheight)r   r   
image_sizemainrepo_idrevisionfilename   rO   scaleloraszfal-ai/loraz(stabilityai/stable-diffusion-xl-base-1.0
model_name)r   popadapter_weights_pathr	   HUGGINGFACE_CO_URL_TEMPLATEformathf_model_idprovider_id)r   rd   re   rf   payload	lora_paths         r!   rs   z-FalAITextToImageTask._prepare_payload_as_dictz   s     f#
*%%#
 g(g"5"5 W--!++h//% %GL! !5A!=DD-9.C E  I
 *3Q ? ?@GG$0MAA )S%r"   Nr@   rA   c                     t          |          d         d         d         }t                                          |          j        S Nimagesr   rK   r   r   rI   contentr   r@   rA   rK   s       r!   r]   z!FalAITextToImageTask.get_response   s9    x  *1-e4}}  %%--r"   r^   rx   r:   s   @r!   rz   rz   v   s        * * * * *'+D\	$   4. .U5$;%7 .RcId .ps . . . . . . . .r"   rz   c            	       t     e Zd Z fdZdedededee         fdZdde	e
ef         d	ee         defd
Z xZS )FalAITextToSpeechTaskc                 J    t                                          d           d S )Nztext-to-speechr   rc   s    r!   r   zFalAITextToSpeechTask.__init__   "    )*****r"   rd   re   rf   r%   c                 (    d|it          |          S )Nru   r   r   rd   re   rf   s       r!   rs   z.FalAITextToSpeechTask._prepare_payload_as_dict   s     :+j"9"9::r"   Nr@   rA   c                     t          |          d         d         }t                                          |          j        S )NaudiorK   r   r   s       r!   r]   z"FalAITextToSpeechTask.get_response   s4    x  )%0}}  %%--r"   r^   rx   r:   s   @r!   r   r      s        + + + + +;;'+;D\;	$; ; ; ;
. .U5$;%7 .RcId .ps . . . . . . . .r"   r   c            	       z     e Zd Z fdZdedededee         fdZ	 dde	e
ef         d	ee         def fd
Z xZS )FalAITextToVideoTaskc                 J    t                                          d           d S )Nztext-to-videor   rc   s    r!   r   zFalAITextToVideoTask.__init__   r|   r"   rd   re   rf   r%   c                 (    d|it          |          S )Nr~   r   r   s       r!   rs   z-FalAITextToVideoTask._prepare_payload_as_dict   s     &<K
$;$;<<r"   Nr@   rA   c                     t                                          ||          }t          |          d         d         }t                                          |          j        S NvideorK   r   r]   r   r   rI   r   r   r@   rA   outputrK   r    s        r!   r]   z!FalAITextToVideoTask.get_response   P    
 %%h??vw'.}}  %%--r"   r^   rx   r:   s   @r!   r   r      s        * * * * *=='+=D\=	$= = = = 7;. .t$. !!23. 
	. . . . . . . . . .r"   r   c            	       z     e Zd Z fdZdedededee         fdZ	 dde	e
ef         d	ee         def fd
Z xZS )FalAIImageToImageTaskc                 J    t                                          d           d S )Nzimage-to-imager   rc   s    r!   r   zFalAIImageToImageTask.__init__   r   r"   rd   re   rf   r%   c                     t          |d          }d|v r|                    d          |d<   d|it          |          }|j        5t          j                            |j        d|j                  }|dd	g|d
<   |S )N
image/jpegdefault_mime_typetarget_sizer   	image_urlr   r   r   r   r   )r   r   r   r   r	   r   r   r   r   rd   re   rf   r   r   r   s          r!   rs   z.FalAIImageToImageTask._prepare_payload_as_dict   s     FlCCC	J&&'1~~m'D'DJ|$#
*%%#
 !5A!=DD-9.C E  I
 *3Q ? ?@GGr"   Nr@   rA   c                     t                                          ||          }t          |          d         d         d         }t                                          |          j        S r   r   r   s        r!   r]   z"FalAIImageToImageTask.get_response   sU    
 %%h??vx(+E2}}  %%--r"   r^   rx   r:   s   @r!   r   r      s        + + + + +'+D\	$   . 7;. .t$. !!23. 
	. . . . . . . . . .r"   r   c            	       z     e Zd Z fdZdedededee         fdZ	 dde	e
ef         d	ee         def fd
Z xZS )FalAIImageToVideoTaskc                 J    t                                          d           d S )Nzimage-to-videor   rc   s    r!   r   zFalAIImageToVideoTask.__init__   r   r"   rd   re   rf   r%   c                     t          |d          }d|it          |          }|j        5t          j                            |j        d|j                  }|ddg|d<   |S )	Nr   r   r   r   r   r   r   r   )r   r   r   r	   r   r   r   r   s          r!   rs   z.FalAIImageToVideoTask._prepare_payload_as_dict   s     FlCCC	#
*%%#
 !5A!=DD-9.C E  I
 *3Q ? ?@GGr"   Nr@   rA   c                     t                                          ||          }t          |          d         d         }t                                          |          j        S r   r   r   s        r!   r]   z"FalAIImageToVideoTask.get_response   r   r"   r^   rx   r:   s   @r!   r   r      s        + + + + +'+D\	$   ( 7;. .t$. !!23. 
	. . . . . . . . . .r"   r   )&rm   rR   abcr   typingr   r   r   r   urllib.parser   huggingface_hubr	   huggingface_hub.hf_apir
   !huggingface_hub.inference._commonr   r   r   ,huggingface_hub.inference._providers._commonr   r   huggingface_hub.utilsr   r   huggingface_hub.utils.loggingr   r5   rP   rT   r   r<   ra   rz   r   r   r   r   r2   r"   r!   <module>r      sg           - - - - - - - - - - - - ! ! ! ! ! ! % % % % % % ; ; ; ; ; ; R R R R R R R R R R X X X X X X X X B B B B B B B B 4 4 4 4 4 4 
H		  " " " " ""C " " "3T 3T 3T 3T 3T' 3T 3T 3Tl    )   : .  .  .  .  .9  .  .  .F. . . . .I . . .. . . . .> . . .&. . . . .N . . .D. . . . .N . . . . .r"   