o
    Û·ir  ã                   @   s¸   d dl Z ddlmZmZmZ ddlmZmZ eeefZeefZ	dZ
dZdZeh e
£e£e£ƒZe jje jje jje jje jje jje jje jjfZde jjd	ed
e jjdB fdd„ZdS )é    Né   )ÚAttentionModuleMixinÚFeedForwardÚLuminaFeedForward)Ú	AttentionÚMochiAttention)ÚblocksÚtransformer_blocksÚsingle_transformer_blocksÚlayersÚvisual_transformer_blocks)Útemporal_transformer_blocks)r   r	   r   ÚmoduleÚfqnÚreturnc                 C   s&   |   ¡ D ]\}}||kr|  S qd S )N)Únamed_modules)r   r   Úsubmodule_nameÚ	submodule© r   úM/home/ubuntu/vllm_env/lib/python3.10/site-packages/diffusers/hooks/_common.pyÚ_get_submodule_from_fqn9   s
   ÿr   )ÚtorchÚmodels.attentionr   r   r   Úmodels.attention_processorr   r   Ú_ATTENTION_CLASSESÚ_FEEDFORWARD_CLASSESÚ&_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERSÚ'_TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERSÚ$_CROSS_TRANSFORMER_BLOCK_IDENTIFIERSÚtupleÚ"_ALL_TRANSFORMER_BLOCK_IDENTIFIERSÚnnÚConv1dÚConv2dÚConv3dÚConvTranspose1dÚConvTranspose2dÚConvTranspose3dÚLinearÚ	EmbeddingÚ_GO_LC_SUPPORTED_PYTORCH_LAYERSÚModuleÚstrr   r   r   r   r   Ú<module>   s6   
ÿþýÿ
ø&