o Gi@sddlZddlZddlmZddlmZmZddlmZddlm Z ddl m Z ddl m Z mZmZdd lmZdd lmZdd lmZeeZd d Zd ddZGdddejZGdddejZGdddejZGdddejZGdddeeZ GdddeeZ!dS)!N)nn) ConfigMixinregister_to_config) ModelMixin) FeedForward) Attention)TimestepEmbedding Timestepsget_2d_sincos_pos_embed)Transformer2DModelOutput) AdaLayerNorm)loggingcCsdd}||d|ks||d|krtdtB||||}||||}|d|dd|d|||td| ||j ||d|WdS1sbwYdS)NcSsdt|tddS)N?@)matherfsqrt)xra/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/pipelines/unidiffuser/modeling_uvit.pynorm_cdfsz(_no_grad_trunc_normal_..norm_cdfzjmean is more than 2 std from [a, b] in nn.init.trunc_normal_. The distribution of values may be incorrect.r)minmax) loggerwarningtorchno_graduniform_erfinv_mul_rradd_clamp_)tensormeanstdabrlurrr_no_grad_trunc_normal_s   $r,rrcCst|||||S)aFills the input Tensor with values drawn from a truncated normal distribution. The values are effectively drawn from the normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)` with values outside :math:`[a, b]` redrawn until they are within the bounds. The method used for generating the random values works best when :math:`a \leq \text{mean} \leq b`. Args: tensor: an n-dimensional `torch.Tensor` mean: the mean of the normal distribution std: the standard deviation of the normal distribution a: the minimum cutoff value b: the maximum cutoff value Examples: >>> w = torch.empty(3, 5) >>> nn.init.trunc_normal_(w) )r,)r%r&r'r(r)rrr trunc_normal_8sr/cs<eZdZdZ         d fdd Zd d ZZS) PatchEmbedz2D Image to Patch EmbeddingrFTc st||||} ||_||_tj||||f||d|_|r,tj|ddd|_nd|_| |_ |j rOt |t | ddd} |j d| d dd dSdS) N) kernel_sizestridebiasFgư>)elementwise_affineeps?pt) output_type pos_embedr) persistent)super__init__flatten layer_normrConv2dproj LayerNormnorm use_pos_embedr intregister_bufferfloat unsqueeze) selfheightwidth patch_size in_channels embed_dimrAr@r6rF num_patchesr< __class__rrr?Ns zPatchEmbed.__init__cCsF||}|jr|ddd}|jr||}|jr!||jS|S)Nrr)rCr@ transposerArErFr<)rKlatentrrrforwardms   zPatchEmbed.forward) r1r1r2rr3FTTT)__name__ __module__ __qualname____doc__r?rV __classcell__rrrRrr0Ksr0cs*eZdZdeffdd ZddZZS) SkipBlockdimcs,ttd|||_t||_dS)Nr)r>r?rLinear skip_linearrDrE)rKr]rRrrr?zs zSkipBlock.__init__cCs&|tj||gdd}||}|S)Nr])r_rcatrE)rKrskiprrrrVs zSkipBlock.forward)rWrXrYrGr?rVr[rrrRrr\ysr\cseZdZdZ            dded ed ed edBd ed edBdedededededededeffdd Z      dddZZ S)UTransformerBlockaS A modification of BasicTransformerBlock which supports pre-LayerNorm and post-LayerNorm configurations. Parameters: dim (`int`): The number of channels in the input and output. num_attention_heads (`int`): The number of heads to use for multi-head attention. attention_head_dim (`int`): The number of channels in each head. dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. num_embeds_ada_norm (:obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. attention_bias (:obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. only_cross_attention (`bool`, *optional*): Whether to use only cross-attention layers. In this case two cross attention layers are used. double_self_attention (`bool`, *optional*): Whether to use two self-attention layers. In this case no cross attention layers are used. upcast_attention (`bool`, *optional*): Whether to upcast the query and key to float32 when performing the attention calculation. norm_elementwise_affine (`bool`, *optional*): Whether to use learnable per-element affine parameters during layer normalization. norm_type (`str`, defaults to `"layer_norm"`): The layer norm implementation to use. pre_layer_norm (`bool`, *optional*): Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"), as opposed to after ("post-LayerNorm"). Note that `BasicTransformerBlock` uses pre-LayerNorm, e.g. `pre_layer_norm = True`. final_dropout (`bool`, *optional*): Whether to use a final Dropout layer after the feedforward network. r-NgegluFTrAr]num_attention_headsattention_head_dimcross_attention_dim activation_fnnum_embeds_ada_normattention_biasonly_cross_attentiondouble_self_attentionupcast_attentionnorm_elementwise_affine norm_typepre_layer_norm final_dropoutc  t| |_|duo| dk|_||_| dvr'|dur'td| d| dt|||||| r1|nd| d|_|dus=| rNt|| sC|nd||||| d|_nd|_|jr[t |||_ nt j || d|_ |dusi| rz|jrqt ||nt j || d|_ nd|_ t j || d|_t||||d |_dS Nada_norm)ru ada_norm_zeroz`norm_type` is set to zw, but `num_embeds_ada_norm` is not defined. Please make sure to define `num_embeds_ada_norm` if setting `norm_type` to .) query_dimheadsdim_headdropoutr6rhrn)rxrhryrzr{r6rn)r7)r{rirrr>r?rluse_ada_layer_normrq ValueErrorrattn1attn2r norm1rrDnorm2norm3rffrKr]rfrgr{rhrirjrkrlrmrnrorprqrrrRrrr?R        zUTransformerBlock.__init__c Cs8|jr|jr |||}n||}n|}|dur|ni}|j|f|jr&|nd|d|} |js@|jr;|| |} n|| } | |}|jdur~|jr[|jrU|||n||}n|}|j|f||d|} |jsz|jru|| |n|| } | |}|jr||}n|}||} |js|| } | |}|SN)encoder_hidden_statesattention_mask rqr}rrrlrrrr) rK hidden_statesrrencoder_attention_masktimestepcross_attention_kwargs class_labelsnorm_hidden_states attn_output ff_outputrrrrVsV        zUTransformerBlock.forward) r-NreNFFFFTrATFNNNNNN rWrXrYrZrGstrboolr?rVr[rrrRrrdsd&     PrdcseZdZdZ            dded ed ed edBd ed edBdedededededededeffdd Z      dddZZ S)UniDiffuserBlocka@ A modification of BasicTransformerBlock which supports pre-LayerNorm and post-LayerNorm configurations and puts the LayerNorms on the residual backbone of the block. This matches the transformer block in the [original UniDiffuser implementation](https://github.com/thu-ml/unidiffuser/blob/main/libs/uvit_multi_post_ln_v1.py#L104). Parameters: dim (`int`): The number of channels in the input and output. num_attention_heads (`int`): The number of heads to use for multi-head attention. attention_head_dim (`int`): The number of channels in each head. dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. num_embeds_ada_norm (:obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. attention_bias (:obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. only_cross_attention (`bool`, *optional*): Whether to use only cross-attention layers. In this case two cross attention layers are used. double_self_attention (`bool`, *optional*): Whether to use two self-attention layers. In this case no cross attention layers are used. upcast_attention (`bool`, *optional*): Whether to upcast the query and key to float() when performing the attention calculation. norm_elementwise_affine (`bool`, *optional*): Whether to use learnable per-element affine parameters during layer normalization. norm_type (`str`, defaults to `"layer_norm"`): The layer norm implementation to use. pre_layer_norm (`bool`, *optional*): Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"), as opposed to after ("post-LayerNorm"). The original UniDiffuser implementation is post-LayerNorm (`pre_layer_norm = False`). final_dropout (`bool`, *optional*): Whether to use a final Dropout layer after the feedforward network. r-NreFTrAr]rfrgrhrirjrkrlrmrnrorprqrrc rsrtr|rrRrrr?qrzUniDiffuserBlock.__init__c Cs&|jr|jr |||}n||}|dur|ni}|j|f|jr#|nd|d|}||}|jsA|jr<|||}n||}|jdurx|jrW|jrR|||n||}|j|f||d|}||}|jsx|jrs|||n||}|jr||}||} | |}|js||}|Srr) rKrrrrrrrrrrrrrVsR        zUniDiffuserBlock.forward) r-NreNFFFFTrAFTrrrrrRrrMsd(     Prc.seZdZdZe                  d+dedededBdedBdededededBdededBdedBdedBdededBdedededed ed!ed"ed#ef,fd$d% Z      d,d&ed'ed(efd)d*Z Z S)-UTransformer2DModelay Transformer model based on the [U-ViT](https://github.com/baofff/U-ViT) architecture for image-like data. Compared to [`Transformer2DModel`], this model has skip connections between transformer blocks in a "U"-shaped fashion, similar to a U-Net. Supports only continuous (actual embeddings) inputs, which are embedded via a [`PatchEmbed`] layer and then reshaped to (b, t, d). Parameters: num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. in_channels (`int`, *optional*): Pass if the input is continuous. The number of channels in the input. out_channels (`int`, *optional*): The number of output channels; if `None`, defaults to `in_channels`. num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. norm_num_groups (`int`, *optional*, defaults to `32`): The number of groups to use when performing Group Normalization. cross_attention_dim (`int`, *optional*): The number of encoder_hidden_states dimensions to use. attention_bias (`bool`, *optional*): Configure if the TransformerBlocks' attention should contain a bias parameter. sample_size (`int`, *optional*): Pass if the input is discrete. The width of the latent images. Note that this is fixed at training time as it is used for learning a number of position embeddings. See `ImagePositionalEmbeddings`. num_vector_embeds (`int`, *optional*): Pass if the input is discrete. The number of classes of the vector embeddings of the latent pixels. Includes the class for the masked latent pixel. patch_size (`int`, *optional*, defaults to 2): The patch size to use in the patch embedding. activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. num_embeds_ada_norm ( `int`, *optional*): Pass if at least one of the norm_layers is `AdaLayerNorm`. The number of diffusion steps used during training. Note that this is fixed at training time as it is used to learn a number of embeddings that are added to the hidden states. During inference, you can denoise for up to but not more than steps than `num_embeds_ada_norm`. use_linear_projection (int, *optional*): TODO: Not used only_cross_attention (`bool`, *optional*): Whether to use only cross-attention layers. In this case two cross attention layers are used in each transformer block. upcast_attention (`bool`, *optional*): Whether to upcast the query and key to float() when performing the attention calculation. norm_type (`str`, *optional*, defaults to `"layer_norm"`): The Layer Normalization implementation to use. Defaults to `torch.nn.LayerNorm`. block_type (`str`, *optional*, defaults to `"unidiffuser"`): The transformer block implementation to use. If `"unidiffuser"`, has the LayerNorms on the residual backbone of each transformer block; otherwise has them in the attention/feedforward branches (the standard behavior in `diffusers`.) pre_layer_norm (`bool`, *optional*): Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"), as opposed to after ("post-LayerNorm"). The original UniDiffuser implementation is post-LayerNorm (`pre_layer_norm = False`). norm_elementwise_affine (`bool`, *optional*): Whether to use learnable per-element affine parameters during layer normalization. use_patch_pos_embed (`bool`, *optional*): Whether to use position embeddings inside the patch embedding layer (`PatchEmbed`). final_dropout (`bool`, *optional*): Whether to use a final Dropout layer after the feedforward network. r2XNrr- FrrerA unidiffuserTrfrgrO out_channels num_layersr{norm_num_groupsrhrk sample_sizenum_vector_embedsrNrirjuse_linear_projectionrlrnrp block_typerqroff_final_dropoutcsJt||_ |_|_ |dur| dusJd| dus&Jd| |_| |_| |_t| | | ||d|_ |dkrAt nt t  fddt|dD|_    d|_t  fd dt|dD|_|dur|n||_t |_dS) Nz0Patch input requires in_channels and patch_size.z?UTransformer2DModel over patched input must provide sample_sizerLrMrNrOrPrFrcs0g|]}    dqS) r{rhrirjrkrlrnrprqrorrr.0drirkrg block_clsrhr{r inner_dimrorprfrjrlrqrnrr s&z0UTransformer2DModel.__init__..rrcs@g|]}tt    ddqS)r)rcblock)r ModuleDictr\rrrrrs2)r>r?rrfrgrLrMrNr0r<rrdr ModuleListrangetransformer_in_blockstransformer_mid_blocktransformer_out_blocksrrDnorm_out)rKrfrgrOrrr{rrhrkrrrNrirjrrlrnrprrqrouse_patch_pos_embedrrRrrr?Lsb & & zUTransformer2DModel.__init__ return_dicthidden_states_is_embedding unpatchifyc Cs&|s|rtd|d|d|d|s||}g} |jD]} | |||||d}| |q||}|jD]} | d|| }| d|||||d}q7||}|rt|j dd } } |j d | | |j |j |j fd }t d |}|j d |j | |j | |j fd }n|}|s|fSt|d S)a Args: hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`. When continuous, `torch.Tensor` of shape `(batch size, channel, height, width)`): Input hidden_states encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*): Conditional embeddings for cross attention layer. If not given, cross-attention defaults to self-attention. timestep ( `torch.long`, *optional*): Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step. class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): Optional class labels to be applied as an embedding in AdaLayerZeroNorm. Used to indicate class labels conditioning. cross_attention_kwargs (*optional*): Keyword arguments to supply to the cross attention layers, if used. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple. hidden_states_is_embedding (`bool`, *optional*, defaults to `False`): Whether or not hidden_states is an embedding directly usable by the transformer. In this case we will ignore input handling (e.g. continuous, vectorized, etc.) and directly feed hidden_states into the transformer blocks. unpatchify (`bool`, *optional*, defaults to `True`): Whether to unpatchify the transformer output. Returns: [`~models.transformer_2d.Transformer2DModelOutput`] or `tuple`: [`~models.transformer_2d.Transformer2DModelOutput`] if `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is the sample tensor. z!Cannot both define `unpatchify`: z and `return_dict`: z since when `unpatchify` is zy the returned output is of shape (batch_size, seq_len, hidden_dim) rather than (batch_size, num_channels, height, width).)rrrrrcrrr9r`shapenhwpqc->nchpwq)sample)r~r<rappendrrpoprrGrreshaperNrreinsumr )rKrrrrrrrrskipsin_block out_blockrLrMoutputrrrrVsT*        zUTransformer2DModel.forward)r2rNNrr-rNFNNrreNFFFrArFTFF)NNNNTFT) rWrXrYrZrrGrIrrr?rVr[rrrRrrs9      rc6sFeZdZdZe                 d6dededededededBdedBdededededBdededBdedBdedBded edBd!ed"ed#ed$ed%ed&ed'ed(ed)ef4fd*d+ Z e j j d,d-Z   d7d.e jd/e jd0e jd1e jeBeBd2e jeBeBd3e jeBeBdBf d4d5ZZS)8UniDiffuserModela Transformer model for a image-text [UniDiffuser](https://huggingface.co/papers/2303.06555) model. This is a modification of [`UTransformer2DModel`] with input and output heads for the VAE-embedded latent image, the CLIP-embedded image, and the CLIP-embedded prompt (see paper for more details). Parameters: text_dim (`int`): The hidden dimension of the CLIP text model used to embed images. clip_img_dim (`int`): The hidden dimension of the CLIP vision model used to embed prompts. num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. in_channels (`int`, *optional*): Pass if the input is continuous. The number of channels in the input. out_channels (`int`, *optional*): The number of output channels; if `None`, defaults to `in_channels`. num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. norm_num_groups (`int`, *optional*, defaults to `32`): The number of groups to use when performing Group Normalization. cross_attention_dim (`int`, *optional*): The number of encoder_hidden_states dimensions to use. attention_bias (`bool`, *optional*): Configure if the TransformerBlocks' attention should contain a bias parameter. sample_size (`int`, *optional*): Pass if the input is discrete. The width of the latent images. Note that this is fixed at training time as it is used for learning a number of position embeddings. See `ImagePositionalEmbeddings`. num_vector_embeds (`int`, *optional*): Pass if the input is discrete. The number of classes of the vector embeddings of the latent pixels. Includes the class for the masked latent pixel. patch_size (`int`, *optional*, defaults to 2): The patch size to use in the patch embedding. activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. num_embeds_ada_norm ( `int`, *optional*): Pass if at least one of the norm_layers is `AdaLayerNorm`. The number of diffusion steps used during training. Note that this is fixed at training time as it is used to learn a number of embeddings that are added to the hidden states. During inference, you can denoise for up to but not more than steps than `num_embeds_ada_norm`. use_linear_projection (int, *optional*): TODO: Not used only_cross_attention (`bool`, *optional*): Whether to use only cross-attention layers. In this case two cross attention layers are used in each transformer block. upcast_attention (`bool`, *optional*): Whether to upcast the query and key to float32 when performing the attention calculation. norm_type (`str`, *optional*, defaults to `"layer_norm"`): The Layer Normalization implementation to use. Defaults to `torch.nn.LayerNorm`. block_type (`str`, *optional*, defaults to `"unidiffuser"`): The transformer block implementation to use. If `"unidiffuser"`, has the LayerNorms on the residual backbone of each transformer block; otherwise has them in the attention/feedforward branches (the standard behavior in `diffusers`.) pre_layer_norm (`bool`, *optional*): Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"), as opposed to after ("post-LayerNorm"). The original UniDiffuser implementation is post-LayerNorm (`pre_layer_norm = False`). norm_elementwise_affine (`bool`, *optional*): Whether to use learnable per-element affine parameters during layer normalization. use_patch_pos_embed (`bool`, *optional*): Whether to use position embeddings inside the patch embedding layer (`PatchEmbed`). ff_final_dropout (`bool`, *optional*): Whether to use a final Dropout layer after the feedforward network. use_data_type_embedding (`bool`, *optional*): Whether to use a data type embedding. This is only relevant for UniDiffuser-v1 style models; UniDiffuser-v1 is continue-trained from UniDiffuser-v0 on non-publically-available data and accepts a `data_type` argument, which can either be `1` to use the weights trained on non-publically-available data or `0` otherwise. This argument is subsequently embedded by the data type embedding, if used. r3Mr2rNrr-rFrerArTtext_dim clip_img_dimnum_text_tokensrfrgrOrrr{rrhrkrrrNrirjrrlrnrprrqroruse_data_type_embeddingcs`t|||_| dusJd| |_||_|dur|n||_||_|j||j||_t| | |||j|d|_ t ||j|_ t ||j|_ t|jddd|_|rbt|jd|j|jdnt |_t|jddd|_|r~t|jd|j|jdnt |_||_d|d |j|_t td |j|j|_t j| d |_t|jd d ||_|jrt d|j|_ t td d |j|_!t"d$id |d|d|d|d|d| d| d| d| d| d|d|d|d|d|d|d|d|d|d |d!|d"|d#||_#|d|}t |j||_$t |j||_%t |j||_&dS)%Nzr?rrrOrrNrQr0 vae_img_inrr^ clip_img_intext_inr timestep_img_projr Identitytimestep_img_embedtimestep_text_projtimestep_text_embedr num_tokens Parameterrzerosr<Dropoutpos_embed_dropr/r Embeddingdata_type_token_embeddingdata_type_pos_embed_tokenr transformer vae_img_out clip_img_outtext_out)rKrrrrfrgrOrrr{rrhrkrrrNrirjrrlrnrprrquse_timestep_embeddingrorrr patch_dimrRrrr?s           zUniDiffuserModel.__init__cCsdhS)Nr<r)rKrrrno_weight_decay sz UniDiffuserModel.no_weight_decaylatent_image_embeds image_embeds prompt_embeds timestep_img timestep_text data_typec  Cs|jd} ||} ||} ||} | d| d} }t|s/tj|gtj| j d}|tj | |j |j d}| |}|j |j d}||}|jdd}t|sbtj|gtj| j d}|tj | |j |j d}||}|j |j d}||}|jdd}|jr|dusJdt|stj|gtj| j d}|tj | |j |j d}||jdd}tj|||| | | gdd}n tj||| | | gdd}|jrtj|jddddddf|j|jddddddfgdd}n|j}||}||}|j||dd|d d d d d}|jr)|jddd| d|fdd\}}}}}}n|jdd| d|fdd\}}}}}||}t|jdd }}|jd |||j|j|jfd}td|}|jd |j||j||jfd}| |}|!|}|||fS)am Args: latent_image_embeds (`torch.Tensor` of shape `(batch size, latent channels, height, width)`): Latent image representation from the VAE encoder. image_embeds (`torch.Tensor` of shape `(batch size, 1, clip_img_dim)`): CLIP-embedded image representation (unsqueezed in the first dimension). prompt_embeds (`torch.Tensor` of shape `(batch size, seq_len, text_dim)`): CLIP-embedded text representation. timestep_img (`torch.long` or `float` or `int`): Current denoising step for the image. timestep_text (`torch.long` or `float` or `int`): Current denoising step for the text. data_type: (`torch.int` or `float` or `int`, *optional*, defaults to `1`): Only used in UniDiffuser-v1-style models. Can be either `1`, to use weights trained on nonpublic data, or `0` otherwise. encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*): Conditional embeddings for cross attention layer. If not given, cross-attention defaults to self-attention. cross_attention_kwargs (*optional*): Keyword arguments to supply to the cross attention layers, if used. Returns: `tuple`: Returns relevant parts of the model's noise prediction: the first element of the tuple is tbe VAE image embedding, the second element is the CLIP image embedding, and the third element is the CLIP text embedding. rr)dtypedevice)rraNzBdata_type must be supplied if the model uses a data type embeddingrFT)rrrrrrrr9r`rr)"rrrrsizer is_tensorr%longronesrrtorrJrrrrGrrbr<rrrsplitrrrNrrrr)rKrrrrrrrr batch_sizevae_hidden_statesclip_hidden_statestext_hidden_statesrnum_img_tokenstimestep_img_tokentimestep_text_tokendata_type_tokenrr<t_img_token_outt_text_token_outdata_type_token_outr img_clip_out img_vae_outrLrMrrrrV s &             <        zUniDiffuserModel.forward)r3rrr2rNNrr-rNFNNNreNFFFrArFFTFTF)rNN)rWrXrYrZrrGrIrrr?rjitignorerTensorrVr[rrrRrr@s?         r)r-rr.r)"rrrconfiguration_utilsrrmodelsrmodels.attentionrmodels.attention_processorrmodels.embeddingsr r r models.modeling_outputsr models.normalizationr utilsr get_loggerrWrr,r/Moduler0r\rdrrrrrrrs0         %.BF0