o i ] @sddlZddlmZddlmZddlZddlmZddlmm Z ddl m Z ddl mZddlmZddlmZddlmZd d lmZmZGd d d ejZGd ddejZGdddejZGdddejZGdddejZdejdede eje eefffddZ!dejdede eefde eefdejf ddZ"ded ed!ejdejfd"d#Z#d$ejd%ejd&ejde eefd e eefdejf d'd(Z$Gd)d*d*ejZ%d+d,Z&d-d.Z'Gd/d0d0eZ(Gd1d2d2ejZ)dS)3N)Iterable)partial)CLIPVisionConfig)MMEncoderAttention) Conv2dLayer)QuantizationConfig)default_weight_loader) CLIPEncoderCLIPVisionEmbeddingsc sPeZdZejfdededeejddffdd Zde j de j fd d Z Z S) MLPBlock embedding_dimmlp_dimactreturnNcs2tt|||_t|||_||_dSN)super__init__nnLinearlin1lin2r)selfr rr __class__\/home/ubuntu/vllm_env/lib/python3.10/site-packages/vllm/model_executor/models/deepencoder.pyrs  zMLPBlock.__init__xcCs||||Sr)rrrrrrrrforward)szMLPBlock.forward) __name__ __module__ __qualname__rGELUinttypeModulertorchTensorr __classcell__rrrrr s r csBeZdZd dededdffdd Zdejdejfd d ZZ S) LayerNorm2dư> num_channelsepsrNcs8ttt||_tt||_||_ dSr) rrr Parameterr'onesweightzerosbiasr-)rr,r-rrrr0s  zLayerNorm2d.__init__rcCsn|jddd}||djddd}||t||j}|jddddf||jddddf}|S)Nr T)keepdim)meanpowr'sqrtr-r0r2)rrusrrrr6s ,zLayerNorm2d.forward)r+) r r!r"r$floatrr'r(rr)rrrrr*/sr*c%seZdZdddddddddejejdd dd d dfd ed edededededededede ej de ej dedededede edfdeddf$fd d! Z d"e jd#efd$d%Zd&e jde jfd'd(ZZS))ImageEncoderViT @TFrrimg_size patch_sizein_chans embed_dimdepth num_heads mlp_ratio out_chansqkv_bias norm_layer act_layer use_abs_pos use_rel_posrel_pos_zero_init window_sizeglobal_attn_indexes.last_conv_outputrNcst||_t||f||f||d|_d|_| r*tt d||||||_t |_ t |D]"}t |||| | | | |||vrD|nd||||fd }|j |q3tt||dddt|t||dddd t||_td d dd ddd |_td |dd ddd |_dS)a Args: img_size (int): Input image size. patch_size (int): Patch size. in_chans (int): Number of input image channels. embed_dim (int): Patch embedding dimension. depth (int): Depth of ViT. num_heads (int): Number of attention heads in each ViT block. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. qkv_bias (bool): If True, add a learnable bias to query, key, value. norm_layer (nn.Module): Normalization layer. act_layer (nn.Module): Activation layer. use_abs_pos (bool): If True, use absolute positional embeddings. use_rel_pos (bool): If True, add relative positional embeddings to the attention map. rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. window_size (int): Window size for window attention blocks. global_attn_indexes (list): Indexes for blocks using global attention. ) kernel_sizestriderErFNr r) dimrHrIrKrLrMrOrPrQ input_sizeF)rTr2r>)rTpaddingr2rBir4)rTrUrXr2)rrrC PatchEmbed patch_embed pos_embedrr.r'r1 ModuleListblocksrangeBlockappend Sequentialrr*necknet_2net_3)rrCrDrErFrGrHrIrJrKrLrMrNrOrPrQrRrSiblockrrrr@sj &      zImageEncoderViT.__init__abs_postgt_sizecCsj|j}|d}||kr3|dddd}|tj}tj|||fdddd|}|dddd}|S|S) Nr rr>r4bicubicTFsizemode antialias align_corners)dtyperkpermutetor'float32F interpolate)rrgrhrosrc_size old_pos_embed new_pos_embedrrr get_abs_poss"  zImageEncoderViT.get_abs_posrcCsn||}|jdur|||j|d}|jD]}||}q||dddd}||}||}|S)Nr rr>r4) rZr[rxrkr]rbrprcrd)rrblk neck_output conv2_output conv3_outputrrrrs      zImageEncoderViT.forward)r r!r"r LayerNormr#r$r:boolr%r&tuplerr'r(rxrr)rrrrr;?sr      br;cseZdZdZddejejddddfdeded ed e d e ej d e ej d e de dede eefdBddffdd Z dejdejfddZZS)r_zWTransformer blocks with support of window attention and residual propagation blocksrATFrNrVrHrIrKrLrMrOrPrQrWrc sft|||_t|||||| dkr| n| | fd|_|||_t|t|||d|_| |_ dS)ai Args: dim (int): Number of input channels. num_heads (int): Number of attention heads in each ViT block. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. qkv_bias (bool): If True, add a learnable bias to query, key, value. norm_layer (nn.Module): Normalization layer. act_layer (nn.Module): Activation layer. use_rel_pos (bool): If True, add relative positional embeddings to the attention map. rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. window_size (int): Window size for window attention blocks. If it equals 0, then use global attention. input_size (tuple(int, int) or None): Input resolution for calculating the relative positional parameter size. r)rHrKrOrPrW)r rrN) rrnorm1RelPosAttentionattnnorm2r r$mlprQ) rrVrHrIrKrLrMrOrPrQrWrrrrs    zBlock.__init__rcCs|}||}|jdkr|jd|jd}}t||j\}}||}|jdkr3t||j|||f}||}||||}|S)Nrr r4)rrQshapewindow_partitionrwindow_unpartitionrr)rrshortcutHWpad_hwrrrrs    z Block.forward)r r!r"__doc__rr}r#r$r:r~r%r&rrr'r(rr)rrrrr_sD    .r_csleZdZdZ     ddededed ed ed eeefdBd dffd d Zdej d ej fddZ Z S)rz=Multi-head Attention block with relative position embeddings.TFNrVrHrKrOrPrWrcst||_||}|d|_tj||d|d|_t|||_||_|jrS|dus1Jdt t d|dd||_ t t d|dd||_ dSdS) a Args: dim (int): Number of input channels. num_heads (int): Number of attention heads. qkv_bias (bool): If True, add a learnable bias to query, key, value. rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. input_size (tuple(int, int) or None): Input resolution for calculating the relative positional parameter size. gr>)r2NzBInput size must be provided if using relative positional encoding.r4rr )rrrHscalerrqkvprojrOr.r'r1 rel_pos_h rel_pos_w)rrVrHrKrOrPrWhead_dimrrrr s    $zRelPosAttention.__init__rc Cs|j\}}}}|||||d|jdddddd}|d||j||dd\}}} d\} } |jrGt||j|j ||f||f\} } | ||j||d}| ||j||d}| ||j||d} |jr| ||j| d| d| d} | ||j| d| d| d} | |  ||j| d| d| d} t j jj||| | d}n t j j||| }| ||j||dddddd|||d}||}|S) Nr>r4rr )NN) attn_mask)rrreshaperHrpunbindrOadd_decomposed_rel_posrrviewrkr'r functionalscaled_dot_product_attentionr) rrBrr_rqkvrel_hrel_w attn_biasrrrr0s@*&   zRelPosAttention.forward)rTFTN) r r!r"rr$r~rrr'r(rr)rrrrr s.#rrrQrc Cs|j\}}}}||||}||||}|dks|dkr+t|ddd|d|f}||||}} |||||| |||}|ddddddd|||} | || ffS)aU Partition into non-overlapping windows with padding if needed. Args: x (tensor): input tokens with [B, H, W, C]. window_size (int): window size. Returns: windows: windows after partition with [B * num_windows, window_size, window_size, C]. (Hp, Wp): padded height and width before partition rr r>r4rr)rrspadrrp contiguous) rrQrrrCpad_hpad_wHpWpwindowsrrrr^s " rrrhwc Cs|\}}|\}}|jd||||}||||||||d} | dddddd|||d} ||ks=||krO| ddd|d|ddf} | S) a Window unpartition into original sequences and removing padding. Args: windows (tensor): input tokens with [B * num_windows, window_size, window_size, C]. window_size (int): window size. pad_hw (Tuple): padded height and width (Hp, Wp). hw (Tuple): original height and width (H, W) before padding. Returns: x: unpartitioned sequences with [B, H, W, C]. rrr r>r4rrN)rrrpr) rrQrrrrrrrrrrrrzs$$rq_sizek_sizerel_posc Cstdt||d}|jd|kr>|j}|tj}tj| d|jdd ddd|dd|}| d| dd}n|}tj ||j ddddft||d }tj ||j ddddft||d }|||dt||d }|| S) a\ Get relative positional embeddings according to the relative positions of query and key sizes. Args: q_size (int): size of query q. k_size (int): size of key k. rel_pos (Tensor): relative position embeddings (L, C). Returns: Extracted positional embeddings according to relative positions. r4r rrlinear)rkrl)deviceNg?)r$maxrrorqr'rrrsrtrrparangerlong) rrr max_rel_distrorel_pos_resizedq_coordsk_coordsrelative_coordsrrr get_rel_poss*   rrrrcCs|\}}|\}}t|||} t|||} |j\} } } || ||| }td|| }td|| }|d}|d}|| |||d}|| ||d|}||fS)a Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py Args: q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C). rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis. rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis. q_size (Tuple): spatial sequence size of query q with (q_h, q_w). k_size (Tuple): spatial sequence size of key k with (k_h, k_w). Returns: attn (Tensor): attention map with added relative positional embeddings. zbhwc,hkc->bhwkzbhwc,wkc->bhwkrr )rrrr'einsum unsqueeze)rrrrrq_hq_wk_hk_wRhRwrrrVr_qrrrrrrs     rc steZdZdZ     ddeeefdeeefdeeefd ed ed d f fd d Zdejd ejfddZ Z S)rYz# Image to Patch Embedding. r=r=rrr>r?rTrUrXrErFrNcs"tt|||||d|_dS)aP Args: kernel_size (Tuple): kernel size of the projection layer. stride (Tuple): stride of the projection layer. padding (Tuple): padding size of the projection layer. in_chans (int): Number of input image channels. embed_dim (int): Patch embedding dimension. )rTrUrXN)rrrr)rrTrUrXrErFrrrrs   zPatchEmbed.__init__rcCs||}|dddd}|S)Nrr4r>r )rrprrrrrs zPatchEmbed.forward)rrrr>r?) r r!r"rrr$rr'r(rr)rrrrrYs*   rYcCstdddgddS)Nr?r@)r4rr )encoder_embed_dim encoder_depthencoder_num_headsencoder_global_attn_indexes) _build_samrrrrbuild_sam_vit_bs rcCs<d}d}d}t|||dttjjdd||dd|d|d }|S) NrBr<r=rr+r-T) rGrFrCrIrLrHrDrKrOrRrQrJ)r;rr'rr})rrrrprompt_embed_dim image_sizevit_patch_size image_encoderrrrrs$rc@sDeZdZdejdefddZ d dejdejdBdejfd d ZdS) DeepCLIPVisionEmbeddingsrgrhc Cs|d}|d}|dd|dd}}tt|jdd}tt|}|j}||kr||d|||dddd }| t j }t j|||fdddd  |} | dddd} | |||} t j|| gdd } | d||d|} | S|S) Nrrr r>r4riTFrjrV)rksqueezer$mathr7rrorrprrqr'rrrsrtcat) rrgrhrV abs_pos_new cls_tokenrvrurorwvision_pos_embedrrrrx+s6    z$DeepCLIPVisionEmbeddings.get_abs_posN pixel_values patch_embedsrcCsx|jd}|dur |}n||}|ddd}|j|dd}tj||gdd}||| |j | d}|S)Nrr4r rr) rpatch_embeddingflatten transposeclass_embeddingexpandr'rrxposition_embedding position_idsrk)rrr batch_size class_embeds embeddingsrrrrMs  z DeepCLIPVisionEmbeddings.forwardr)r r!r"r'r(r$rxrrrrrr*s#rc seZdZ dddddededBdedBdeddf fd d Zed d Z ed dZ dddde j de j dBde edBde j fddZdeeee j fdeefddZZS)DeepCLIPVisionTransformerN)num_hidden_layers_overrideprefixconfig quant_configrrrcst||_|j}t||_tj||jd|_ t ||||dt d|_ |j }t|j j|j krAtd|dt|j jddS)Nrz.encoder)rrrrattn_clszThe original encoder only has z layers, but you requested z layers.)rrr hidden_sizerrrr}layer_norm_eps pre_layrnormr r transformernum_hidden_layerslenlayers ValueError)rrrrrrFrrrrr`s(   z"DeepCLIPVisionTransformer.__init__cCt|jSr)next parametersrorrrrrozDeepCLIPVisionTransformer.dtypecCrr)rrrrrrrrrz DeepCLIPVisionTransformer.device) select_layersrrrcCs,|||}||}|j||dud}|S)N) inputs_embedsreturn_all_hidden_states)rrr)rrrr hidden_statesencoder_outputsrrrrs  z!DeepCLIPVisionTransformer.forwardweightscCsLt|}t}|D]\}}||}t|dt}|||||q |S)N weight_loader)dictnamed_parameterssetgetattrradd)rr params_dict loaded_paramsname loaded_weightparamrrrr load_weightss     z&DeepCLIPVisionTransformer.load_weightsr)r r!r"rrr$strrpropertyrorr'r(listrrrrrr)rrrrr_sB"    ,r)*rcollections.abcr functoolsrr'torch.nnrtorch.nn.functionalrrs transformersr$vllm.model_executor.layers.attentionrvllm.model_executor.layers.convr'vllm.model_executor.layers.quantizationr-vllm.model_executor.model_loader.weight_utilsrclipr r r&r r*r;r_rr(r$rrrrrrYrrrrrrrrsn        ET    '   %# 5