
    Yh                     F   d dl 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Zd dlZd dlmZ ddlmZ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 d
dlmZmZ g dZe G d d                      Z dee!         de!fdZ"dej#        de!de!de$ej#        e!f         fdZ%dej#        de!de!de!dej#        f
dZ&ej'        (                    d           ej'        (                    d            G d dej)                  Z*dej#        de!dej#        fd Z+d!ej#        d"ej#        d#e$e!e!e!f         d$e$e!e!e!f         d%ej#        d&ej#        d'ej#        dej#        fd(Z,dej#        d)ej#        d*e-fd+Z.ej'        (                    d,           ej'        (                    d-            G d. d/ej)                  Z/ G d0 d1ej)                  Z0 G d2 d3ej)                  Z1 G d4 d5ej)                  Z2d6e3e          d7e4d8e
e         d9e-d:ede2fd;Z5 G d< d=e          Z6 G d> d?e          Z7 e             ed@e6j8        fA          ddBdCd8e
e6         d9e-d:ede2fdD                        Z9 e             ed@e7j8        fA          ddBdCd8e
e7         d9e-d:ede2fdE                        Z:dS )F    N)Sequence)	dataclass)partial)AnyCallableOptional   )MLPStochasticDepth)VideoClassification)_log_api_usage_once   )register_modelWeightsWeightsEnum)_KINETICS400_CATEGORIES)_ovewrite_named_paramhandle_legacy_interface)MViTMViT_V1_B_Weights	mvit_v1_bMViT_V2_S_Weights	mvit_v2_sc                       e Zd ZU eed<   eed<   eed<   ee         ed<   ee         ed<   ee         ed<   ee         ed<   dS )	MSBlockConfig	num_headsinput_channelsoutput_channelskernel_q	kernel_kvstride_q	stride_kvN)__name__
__module____qualname__int__annotations__list     o/var/www/tools.fuzzalab.pt/emblema-extractor/venv/lib/python3.11/site-packages/torchvision/models/video/mvit.pyr   r      ss         NNN3iCy3iCyr*   r   sreturnc                     d}| D ]}||z  }|S N   r)   )r,   productvs      r+   _prodr3   '   s&    G  1Nr*   x
target_dim
expand_dimc                     |                                  }||dz
  k    r|                     |          } n||k    rt          d| j                   | |fS )Nr0   zUnsupported input dimension )dim	unsqueeze
ValueErrorshaper4   r5   r6   
tensor_dims       r+   
_unsqueezer>   .   s_    JZ!^##KK
##	z	!	!AAABBBj=r*   r=   c                 B    ||dz
  k    r|                      |          } | S r/   )squeezer<   s       r+   _squeezerA   7   s'    Z!^##IIj!!Hr*   r>   rA   c                        e Zd Z	 	 ddej        deej                 deej                 deddf
 fdZd	ej	        d
e
eeef         de
ej	        e
eeef         f         fdZ xZS )PoolNFpoolnorm
activationnorm_before_poolr-   c                     t                                                       || _        g }||                    |           ||                    |           |rt	          j        | nd | _        || _        d S )N)super__init__rD   appendnn
Sequentialnorm_actrG   )selfrD   rE   rF   rG   layers	__class__s         r+   rJ   zPool.__init__B   s}     		MM$!MM*%%%28Bv..d 0r*   r4   thwc                    t          |dd          \  }}t          j        |dd          \  }}|                    dd          }|j        d d         \  }}}|                    ||z  |f|z                                             }| j        r| j        |                     |          }| 	                    |          }|j        dd          \  }}	}
|                    |||d                              dd          }t          j
        ||fd          }| j        s| j        |                     |          }t          |dd|          }|||	|
ffS )	N   r0   r0   r   )indicesr8   r	   r8   )r>   torchtensor_split	transposer;   reshape
contiguousrG   rN   rD   catrA   )rO   r4   rR   r=   class_tokenBNCTHWs              r+   forwardzPool.forwardS   sW   "1a++: +AtCCCQKK1'"1"+1aIIq1uaj3&''2244   	!T]%>a  A IIaLL'!""+1aIIaAr"",,Q22I{A&A...$ 	!)Ba  AQ1j))1a)|r*   )NF)r#   r$   r%   rL   Moduler   boolrJ   rY   Tensortupler&   rf   __classcell__rQ   s   @r+   rC   rC   A   s        
 +/!&1 1i1 ry!1 RY'	1
 1 
1 1 1 1 1 1" E#sC-,@ U5<Y^_bdgil_lYmKmEn        r*   rC   	embeddingdc                    | j         d         |k    r| S t          j                            |                     dd                              d          |d                              d                              dd          S )Nr   r0   linear)sizemode)r;   rL   
functionalinterpolatepermuter9   r@   )rm   rn   s     r+   _interpolaterv   m   s|    qQ 	!!a##--a00 	" 	
 	

 
	Ar*   attnqq_thwk_thw	rel_pos_h	rel_pos_w	rel_pos_tc                    |\  }}}	|\  }
}}t          dt          ||          z  dz
            }t          dt          |	|          z  dz
            }t          dt          ||
          z  dz
            }t          ||z  d          }t          ||z  d          }t          j        |          d d d f         |z  t          j        |          d d d f         d|z
  z   |z  z
  }t          ||	z  d          }t          |	|z  d          }t          j        |	          d d d f         |z  t          j        |          d d d f         d|z
  z   |z  z
  }t          |
|z  d          }t          ||
z  d          }t          j        |          d d d f         |z  t          j        |
          d d d f         d|
z
  z   |z  z
  }t	          ||          }t	          ||          }t	          ||          }||                                         }||                                         }||                                         }|j        \  }}}}|d d d d dd f                             |||||	|          } t          j        d| |          }!t          j        d| |          }"| 	                    dddddd	                              |||z  |z  |	z  |          } t          j
        | |                    dd                                        dd          }#|#                    ||||	||
          	                    dddddd	          }#|!d d d d d d d d d d d d d d f         |"d d d d d d d d d d d d d d f         z   |#d d d d d d d d d d d d d d f         z                       ||||z  |	z  |
|z  |z            }$| d d d d dd dd fxx         |$z  cc<   | S )
Nr   r0         ?zbythwc,hkc->bythwkzbythwc,wkc->bythwkr   r	   rT      )r&   maxrY   arangerv   longr;   r\   einsumru   matmulr[   view)%rw   rx   ry   rz   r{   r|   r}   q_tq_hq_wk_tk_hk_wdhdwdt	q_h_ratio	k_h_ratiodist_h	q_w_ratio	k_w_ratiodist_w	q_t_ratio	k_t_ratiodist_tRhRwRtr`   n_head_r8   r_qrel_h_qrel_w_qrel_q_trel_poss%                                        r+   _add_rel_posr   |   sJ    MCcMCc	QS#"	#	#B	QS#"	#	#B	QS#"	#	#B C#Is##IC#Is##I\#qqq$w')3u|C7H7Hqqq7QUX[^U^7_cl6llFC#Is##IC#Is##I\#qqq$w')3u|C7H7Hqqq7QUX[^U^7_cl6llFC#Is##IC#Is##I\#qqq$w')3u|C7H7Hqqq7QUX[^U^7_cl6llF Y++IY++IY++I	6;;==	!B	6;;==	!B	6;;==	!BAvq#
AAAqqq!""H+

ac3
<
<Cl/b99Gl/b99G
++aAq!Q
'
'
/
/QZ#5E5KS
Q
QCl3Q 2 233==aCCGll1fc3S99AA!Q1aQRSSG 	111aaaAAAtQQQ,-
!!!QQQ111aaatQQQ.
/	0
!!!QQQ111aaaD$.
/	0 gasS#)c/::	  	AAAqrr122'!Kr*   shortcutresidual_with_cls_embedc           	          |r|                      |           n.| d d d d dd d d fxx         |d d d d dd d d f         z  cc<   | S r/   )add_)r4   r   r   s      r+   _add_shortcutr      so     0	x	!!!QQQAAA+(111aaaQQQ;//Hr*   r   r   c                       e Zd Zdej        fdee         dedededee         dee         dee         d	ee         d
edededede	dej
        f         ddf fdZdej        deeeef         deej        eeeef         f         fdZ xZS )MultiscaleAttention        
input_size	embed_dim
output_dimr   r   r    r!   r"   residual_poolr   rel_pos_embeddropout
norm_layer.r-   Nc                    t                                                       || _        || _        || _        ||z  | _        dt          j        | j                  z  | _        |	| _	        |
| _
        t          j        |d|z            | _        t          j        ||          g}|dk    r)|                    t          j        |d                     t          j        | | _        d | _        t'          |          dk    st'          |          dk    rWd |D             }t)          t          j        | j        | j        |||| j        d	           || j                            | _        d | _        d | _        t'          |          dk    st'          |          dk    rd
 |D             }t)          t          j        | j        | j        |||| j        d	           || j                            | _        t)          t          j        | j        | j        |||| j        d	           || j                            | _        d | _        d | _        d | _        |rt7          |dd                    }t9          |          dk    r||d         z  n|}t9          |          dk    r||d         z  n|}dt7          ||          z  dz
  }d|d         z  dz
  }t          j        t=          j        || j                            | _        t          j        t=          j        || j                            | _        t          j        t=          j        || j                            | _        t          j         !                    | j        d           t          j         !                    | j        d           t          j         !                    | j        d           d S d S )Nr   r	   r   Tinplacer0   c                 2    g | ]}t          |d z            S r   r&   ).0rx   s     r+   
<listcomp>z0MultiscaleAttention.__init__.<locals>.<listcomp>   s"    777Q!V777r*   F)stridepaddinggroupsbiasc                 2    g | ]}t          |d z            S r   r   )r   kvs     r+   r   z0MultiscaleAttention.__init__.<locals>.<listcomp>   s"    ;;;2#bAg,,;;;r*   r   r   {Gz?std)"rI   rJ   r   r   r   head_dimmathsqrtscalerr   r   rL   LinearqkvrK   DropoutrM   projectpool_qr3   rC   Conv3dpool_kpool_vr{   r|   r}   r   len	ParameterrY   zerosinittrunc_normal_)rO   r   r   r   r   r   r    r!   r"   r   r   r   r   r   rP   	padding_q
padding_kvrq   q_sizekv_sizespatial_dimtemporal_dimrQ   s                         r+   rJ   zMultiscaleAttention.__init__   s     	"$""i/DIdm444*'>$9YJ77#%9Z#D#D"ES==MM"*Wd;;;<<<}f-+/??Q%//A"5"577h777I	MM#%=   
4=)) DK ,0+/a5#3#3a#7#7;;;;;J	MM$&=   
4=)) DK 	MM$&=   
4=)) DK 261515 	<z!""~&&D,/MMA,=,=TXa[((4F.1)nnq.@.@dil**dGc&'222Q6Kz!},q0L\%+k4=*Q*QRRDN\%+k4=*Q*QRRDN\%+lDM*R*RSSDNG!!$.d!;;;G!!$.d!;;;G!!$.d!;;;;;	< 	<r*   r4   rR   c           	         |j         \  }}}|                     |                              ||d| j        | j                                      dd                              d          \  }}}| j        |                     ||          \  }}	n|}	| j        |                     ||          d         }| j	        | 	                    ||          \  }}t          j        | j        |z  |                    dd                    }
| j        2| j        +| j        $t!          |
|||	| j        | j        | j                  }
|
                    d          }
t          j        |
|          }| j        rt'          ||| j                   |                    dd                              |d| j                  }|                     |          }||fS )Nr	   r0   r   rX   r   rW   )r;   r   r\   r   r   r[   unbindr   r   r   rY   r   r   r{   r|   r}   r   softmaxr   r   r   r   r   )rO   r4   rR   r`   ra   rb   rx   kr2   rz   rw   s              r+   rf   zMultiscaleAttention.forward!  s   '1a((1++%%aAt~t}MMWWXY[\]]ddijdkk1a;"{{1c**HAuuE;"As##A&A;"[[C((FAs|DK!OQ[[A->->??>%$.*DIc D |||##Lq!! 	>!Q <===KK1%%aT_==LLOO#vr*   )r#   r$   r%   rL   	LayerNormr(   r&   rh   floatr   rg   rJ   rY   ri   rj   rf   rk   rl   s   @r+   r   r      s]        /1|Z< Z<IZ< Z< 	Z<
 Z< s)Z< 9Z< s)Z< 9Z< Z< "&Z< Z< Z< S")^,Z< 
Z< Z< Z< Z< Z< Z<x   E#sC-,@  U5<Y^_bdgil_lYmKmEn                r*   r   c                        e Zd Zddej        fdee         dededededede	d	e	d
e
dej        f         ddf fdZdej        deeeef         deej        eeeef         f         fdZ xZS )MultiscaleBlockr   r   cnfr   r   r   proj_after_attnr   stochastic_depth_probr   .r-   Nc
                 Z   t                                                       || _        d | _        t	          |j                  dk    rLd |j        D             }
d |
D             }t          t          j        |
|j        |          d           | _        |r|j	        n|j
        } |	|j
                  | _         |	|          | _        t          | j        t          j                  | _        t!          ||j
        ||j        |j        |j        |j        |j        |||||	          | _        t-          |d|z  |j	        gt          j        |d           | _        t3          |d          | _        d | _        |j
        |j	        k    r&t          j        |j
        |j	                  | _        d S d S )	Nr0   c                 (    g | ]}|d k    r|d z   n|S rU   r)   )r   r,   s     r+   r   z,MultiscaleBlock.__init__.<locals>.<listcomp>V  s(    GGGQAEE1q55qGGGr*   c                 2    g | ]}t          |d z            S r   r   )r   r   s     r+   r   z,MultiscaleBlock.__init__.<locals>.<listcomp>W  s"    ===ACQKK===r*   )r   r   )	r   r    r!   r"   r   r   r   r   r   rT   )activation_layerr   r   row)rI   rJ   r   	pool_skipr3   r!   rC   rL   	MaxPool3dr   r   norm1norm2
isinstanceBatchNorm1dneeds_transposalr   r   r   r    r"   rw   r
   GELUmlpr   stochastic_depthr   r   )rO   r   r   r   r   r   r   r   r   r   kernel_skippadding_skipattn_dimrQ   s                r+   rJ   zMultiscaleBlock.__init__E  s    	..2""GG#,GGGK=====L![|TTTVZ DN +:Q3&&s?QZ 233
Z))
 *4:r~ F F'M\m\m''$;!
 
 
	 \3./W
 
 
 !00Eu M M,0!4449S%79LMMDLLL 54r*   r4   rR   c                     | j         r=|                     |                    dd                                        dd          n|                     |          }|                     ||          \  }}| j        | j        s|n|                     |          }| j        |n|                     ||          d         }||                     |          z   }| j         r=|                     |                    dd                                        dd          n|                     |          }| j        | j        r|n|                     |          }||                     | 	                    |                    z   |fS )Nr0   r   r   )
r   r   r[   rw   r   r   r   r   r   r   )	rO   r4   rR   x_norm1x_attnthw_newx_skipx_norm2x_projs	            r+   rf   zMultiscaleBlock.forward  s`   CGCXk$**Q[[A..//99!Q???^b^h^hij^k^k))GS11%T-A%AAt||T[G\G\n,$..C2H2H2KT**6222CGCXk$**Q[[A..//99!Q???^b^h^hij^k^kl*d.B*U\H]H]--dhhw.?.?@@@'IIr*   )r#   r$   r%   rL   r   r(   r&   r   rh   r   r   rg   rJ   rY   ri   rj   rf   rk   rl   s   @r+   r   r   D  s/        '*/1|8N 8NI8N 8N 	8N
 "&8N 8N 8N 8N  %8N S")^,8N 
8N 8N 8N 8N 8N 8Nt
J 
JE#sC-,@ 
JU5<Y^_bdgil_lYmKmEn 
J 
J 
J 
J 
J 
J 
J 
Jr*   r   c            
       h     e Zd Zdedeeef         dededdf
 fdZdej        dej        fd	Z	 xZ
S )
PositionalEncoding
embed_sizespatial_sizetemporal_sizer   r-   Nc                 .   t                                                       || _        || _        t	          j        t          j        |                    | _        d | _	        d | _
        d | _        |st	          j        t          j        | j        d         | j        d         z  |                    | _	        t	          j        t          j        | j        |                    | _
        t	          j        t          j        |                    | _        d S d S )Nr   r0   )rI   rJ   r  r  rL   r   rY   r   r_   spatial_postemporal_pos	class_pos)rO   r  r  r  r   rQ   s        r+   rJ   zPositionalEncoding.__init__  s    (*<J(?(?@@374815 	C!|EK8I!8LtO`abOc8ceo,p,pqqD "U[9KZ-X-X Y YD\%+j*A*ABBDNNN	C 	Cr*   r4   c                    | j                             |                    d          d                              d          }t	          j        ||fd          }| j        | j        | j        | j        j	        \  }}t	          j
        | j        |d          }|                    | j                            d                              | j        dd                              d|                     t	          j        | j                            d          |fd                              d          }|                    |           |S )Nr   rW   r0   rX   )r_   expandrq   r9   rY   r^   r  r  r  r;   repeat_interleaver   r  r\   )rO   r4   r_   hw_sizer  pos_embeddings         r+   rf   zPositionalEncoding.forward  s,   &--affQii<<FFqIII{A&A...'D,=,IdnNh"&"2"8GZ!3D4EwTUVVVMt/99!<<CCDDVXZ\^__gghjlvwwxxx!It~'?'?'B'BM&RXYZZZddefggMFF=!!!r*   )r#   r$   r%   r&   rj   rh   rJ   rY   ri   rf   rk   rl   s   @r+   r   r     s        C3 CeCHo CVY Cjn Csw C C C C C C %,        r*   r   c            $       :    e Zd Z	 	 	 	 	 	 	 	 	 ddeeef         d	ed
ee         dedededededededede	e
dej        f                  de	e
dej        f                  deeeef         deeeef         deeeef         ddf" fdZdej        dej        fdZ xZS )r         ?r     Nr	      r  r   rT   rT   r0   r	   r	   r  r  block_settingr   r   r   r   r   attention_dropoutr   num_classesblock.r   patch_embed_kernelpatch_embed_stridepatch_embed_paddingr-   c                    t                                                       t          |            t          |          }|dk    rt	          d          |t
          }|t          t          j        d          }t          j	        d|d         j
        |||          | _        d t          |f|z   | j        j                  D             }t          |d         j
        |d	         |d
         f|d         |          | _        t          j                    | _        t%          |          D ]s\  }}|
|z  |dz
  z  }| j                             ||||||||	||	  	                   t          |j                  dk    rd t          ||j                  D             }t ||d         j                  | _        t          j        t          j        |d          t          j        |d         j        |                    | _        |                                 D ]J}t9          |t          j                  rmt          j                            |j        d           t9          |t          j                  r,|j         %t          j        !                    |j         d           t9          |t          j                  rY|j        %t          j        !                    |j        d           |j         %t          j        !                    |j         d           t9          |t                    r8|"                                D ]#}t          j                            |d           $LdS )a  
        MViT main class.

        Args:
            spatial_size (tuple of ints): The spacial size of the input as ``(H, W)``.
            temporal_size (int): The temporal size ``T`` of the input.
            block_setting (sequence of MSBlockConfig): The Network structure.
            residual_pool (bool): If True, use MViTv2 pooling residual connection.
            residual_with_cls_embed (bool): If True, the addition on the residual connection will include
                the class embedding.
            rel_pos_embed (bool): If True, use MViTv2's relative positional embeddings.
            proj_after_attn (bool): If True, apply the projection after the attention.
            dropout (float): Dropout rate. Default: 0.0.
            attention_dropout (float): Attention dropout rate. Default: 0.0.
            stochastic_depth_prob: (float): Stochastic depth rate. Default: 0.0.
            num_classes (int): The number of classes.
            block (callable, optional): Module specifying the layer which consists of the attention and mlp.
            norm_layer (callable, optional): Module specifying the normalization layer to use.
            patch_embed_kernel (tuple of ints): The kernel of the convolution that patchifies the input.
            patch_embed_stride (tuple of ints): The stride of the convolution that patchifies the input.
            patch_embed_padding (tuple of ints): The padding of the convolution that patchifies the input.
        r   z+The configuration parameter can't be empty.Ngư>)epsr	   )in_channelsout_channelskernel_sizer   r   c                     g | ]
\  }}||z  S r)   r)   r   rq   r   s      r+   r   z!MViT.__init__.<locals>.<listcomp>  s     uuuvdfnuuur*   r0   r   )r  r  r  r   r   )	r   r   r   r   r   r   r   r   r   c                     g | ]
\  }}||z  S r)   r)   r!  s      r+   r   z!MViT.__init__.<locals>.<listcomp>  s     ___vdfn___r*   rW   Tr   r   r   r   )#rI   rJ   r   r   r:   r   r   rL   r   r   r   	conv_projzipr   r   pos_encoding
ModuleListblocks	enumeraterK   r!   r   rE   rM   r   r   headmodulesr   r   r   weightr   	constant_
parameters)rO   r  r  r  r   r   r   r   r   r  r   r  r  r   r  r  r  total_stage_blocksr   stage_block_idr   sd_probmweightsrQ   s                           r+   rJ   zMViT.__init__  s_   R 	
 	D!!! //""JKKK=#E 4888J &q)8*%'
 
 
 vum=MP\=\^b^l^s9t9tuuu
 /$Q'6$Q-A7$Q-'	
 
 
 moo#,]#;#; 	` 	`NC+n<@RUX@XYGK)"/,C"/$3-*1)
 
 
   3<  1$$__ZQTQ]A^A^___
J}R0@AA	 MJw---ImB'7EE
 
	
  	= 	=A!RY'' =%%ahD%999a++ 30BG%%afc222Ar|,, =8'G%%ah4446%G%%afc222A122 = ||~~ = =GG))'t)<<<<	= 	=r*   r4   c                    t          |dd          d         }|                     |          }|                    d                              dd          }|                     |          }| j        j        f| j        j        z   }| j        D ]} |||          \  }}|                     |          }|d d df         }| 	                    |          }|S )Nr   r   r   r0   )
r>   r#  flattenr[   r%  r  r  r'  rE   r)  )rO   r4   rR   r  s       r+   rf   zMViT.forward"  s    q!Q"NN1IIaLL""1a(( a    .043D3QQ[ 	# 	#EU1c]]FAssIIaLL aaadGIIaLLr*   )	r  r   r   r  NNr  r  r  )r#   r$   r%   rj   r&   r   r   rh   r   r   r   rL   rg   rJ   rY   ri   rf   rk   rl   s   @r+   r   r     s        #&'*489=3<3<4=#v= v=CHov= v=  .	v=
 v= "&v= v= v= v= !v=  %v= v= bi01v= Xc29n56v= "#sC-0v=  "#sC-0!v=" #3S=1#v=$ 
%v= v= v= v= v= v=p %,        r*   r   r  r   r2  progresskwargsc                    |t          |dt          |j        d                              |j        d         d         |j        d         d         k    sJ t          |d|j        d                    t          |d|j        d                    |                    dd	          }|                    dd
          }t	          d||| |                    dd          |                    dd          |                    dd          |                    dd          |d|}|*|                    |                    |d                     |S )Nr  
categoriesmin_sizer   r0   r  r  min_temporal_size   r<     r   Fr   Tr   r   )r  r  r  r   r   r   r   r   )r5  
check_hashr)   )r   r   metapopr   load_state_dictget_state_dict)r  r   r2  r5  r6  r  r  models           r+   _mvitrD  9  sW    fmSl9S5T5TUUU|J'*gl:.Fq.IIIIIfngl:6NOOOfow|DW7XYYY::nj99LJJ33M 
!##jj%88 &

+Dd K Kjj%88

#4e<<3
 
 
 
E g44hSW4XXYYYLr*   c                   h    e Zd Z ed eedddd          ddedd	d
ddddiddd	          ZeZdS )r   z:https://download.pytorch.org/models/mvit_v1_b-dbeb1030.pthr;     ?rI  rI  ?rK  rK  	crop_sizeresize_sizemeanr   r=  zShttps://github.com/facebookresearch/pytorchvideo/blob/main/docs/source/model_zoo.mdThe weights were ported from the paper. The accuracies are estimated on video-level with parameters `frame_rate=7.5`, `clips_per_video=5`, and `clip_len=16`ip.Kinetics-400gJ+S@gh|?eW@zacc@1zacc@5guVQ@g rxa@	r9  r:  r8  recipe_docs
num_params_metrics_ops
_file_sizeurl
transformsr?  N	r#   r$   r%   r   r   r   r   KINETICS400_V1DEFAULTr)   r*   r+   r   r   Z  s        WH7 #%
 
 
 #!#1k[ ###! ! !#
 
  N: GGGr*   r   c                   h    e Zd Z ed eedddd          ddedd	d
ddddiddd	          ZeZdS )r   z:https://download.pytorch.org/models/mvit_v2_s-ae3be167.pthr;  rF  rH  rJ  rL  r=  zChttps://github.com/facebookresearch/SlowFast/blob/main/MODEL_ZOO.mdrP  irQ  g r0T@g(\W@rR  guVP@g?5^I|`@rS  rZ  Nr]  r)   r*   r+   r   r   {  s        WH7 #%
 
 
 #!#1[[ ###! ! !#
 
  N: GGGr*   r   
pretrained)r2  T)r2  r5  c                    t                               |           } g dg dg dg g dg g dg g g g g g g g g g g dg gg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dgg g dg g dg g g g g g g g g g g dg gg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dgd	}g }t          t          |d
                             D ]}|                    t          |d
         |         |d         |         |d         |         |d         |         |d         |         |d         |         |d         |         	                     t          ddd|dd|                    dd          | |d|S )a  
    Constructs a base MViTV1 architecture from
    `Multiscale Vision Transformers <https://arxiv.org/abs/2104.11227>`__.

    .. betastatus:: video module

    Args:
        weights (:class:`~torchvision.models.video.MViT_V1_B_Weights`, optional): The
            pretrained weights to use. See
            :class:`~torchvision.models.video.MViT_V1_B_Weights` below for
            more details, and possible values. By default, no pre-trained
            weights are used.
        progress (bool, optional): If True, displays a progress bar of the
            download to stderr. Default is True.
        **kwargs: parameters passed to the ``torchvision.models.video.MViT``
            base class. Please refer to the `source code
            <https://github.com/pytorch/vision/blob/main/torchvision/models/video/mvit.py>`_
            for more details about this class.

    .. autoclass:: torchvision.models.video.MViT_V1_B_Weights
        :members:
    r0   r   r   rT   rT   rT   rT   rT   rT   rT   rT   rT   rT   rT      rd  `      rg    rh  rh  rh  rh  rh  rh  rh  rh  rh  rh     ri  )rg  rg  rh  rh  rh  rh  rh  rh  rh  rh  rh  rh  rh  ri  ri  ri  r	   r	   r	   r0   r   r   r0   rd  rd  r0   rT   rT   r0   r0   r0   r   r   r   r   r    r!   r"   r   r   r   r   r    r!   r"   r;  r=  Fr   皙?)r  r  r  r   r   r   r2  r5  r)   )r   verifyranger   rK   r   rD  r@  r2  r5  r6  configr  is         r+   r   r     s   2  &&w//G FEEiiikkkB			2r2r2r2rSUWY[d[d[dfhiIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII!
$ B			2r2r2r2rSUWY[d[d[dfhiIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII!
1* *FX M3vk*++,, 
 
 -a0%&67: &'8 9! <
+A. -a0
+A. -a0  
	
 
	
 
	
 
	
  
# %$jj)@#FF
 
 
 
 
r*   c                 t   t                               |           } g dg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dgg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dgg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dgg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dg dgd	}g }t          t          |d
                             D ]}|                    t          |d
         |         |d         |         |d         |         |d         |         |d         |         |d         |         |d         |         	                     t          ddd|dddd|                    dd          | |d
|S )aC  Constructs a small MViTV2 architecture from
    `Multiscale Vision Transformers <https://arxiv.org/abs/2104.11227>`__ and
    `MViTv2: Improved Multiscale Vision Transformers for Classification
    and Detection <https://arxiv.org/abs/2112.01526>`__.

    .. betastatus:: video module

    Args:
        weights (:class:`~torchvision.models.video.MViT_V2_S_Weights`, optional): The
            pretrained weights to use. See
            :class:`~torchvision.models.video.MViT_V2_S_Weights` below for
            more details, and possible values. By default, no pre-trained
            weights are used.
        progress (bool, optional): If True, displays a progress bar of the
            download to stderr. Default is True.
        **kwargs: parameters passed to the ``torchvision.models.video.MViT``
            base class. Please refer to the `source code
            <https://github.com/pytorch/vision/blob/main/torchvision/models/video/mvit.py>`_
            for more details about this class.

    .. autoclass:: torchvision.models.video.MViT_V2_S_Weights
            :members:
    rc  )rf  rf  rg  rg  rh  rh  rh  rh  rh  rh  rh  rh  rh  rh  rh  ri  re  rj  rn  rk  rl  rm  ro  r   r   r   r   r    r!   r"   r;  r=  TFr   rp  )
r  r  r  r   r   r   r   r   r2  r5  r)   )r   rq  rr  r   rK   r   rD  r@  rs  s         r+   r   r     s+   4  &&w//G FEEhhhjjjIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII!
& IIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII!
& IIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII!
& IIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII!
uL LF\ M3vk*++,, 
 
 -a0%&67: &'8 9! <
+A. -a0
+A. -a0  
	
 
	
 
	
 
	
  # %$jj)@#FF    r*   );r   collections.abcr   dataclassesr   	functoolsr   typingr   r   r   rY   torch.fxtorch.nnrL   opsr
   r   transforms._presetsr   utilsr   _apir   r   r   _metar   _utilsr   r   __all__r   r&   r3   ri   rj   r>   rA   fxwraprg   rC   rv   r   rh   r   r   r   r   r   r(   r   rD  r   r   r^  r   r   r)   r*   r+   <module>r     s    $ $ $ $ $ $ ! ! ! ! ! !       * * * * * * * * * *         ' ' ' ' ' ' ' ' 6 6 6 6 6 6 ( ( ( ( ( ( 7 7 7 7 7 7 7 7 7 7 + + + + + + C C C C C C C C          Xc] s    %, C S U5<Y\K\E]     # 3 C TYT`     l    j   ) ) ) ) )29 ) ) )XEL S U\    9
,9|9 c39 c3	9
 |9 |9 |9 \9 9 9 9xU\ U\ TX     n    o   } } } } }") } } }@EJ EJ EJ EJ EJbi EJ EJ EJP       :M M M M M29 M M M`&  k" 	
  
   B       B       B ,0A0P!QRRR8<t ] ] ](#45 ] ]_b ]gk ] ] ] SR ]@ ,0A0P!QRRR8<t B B B(#45 B B_b Bgk B B B SR B B Br*   