
    +jo              	           S SK r S SKJr  S SKJr  S SKJr  S SKJrJ	r	  SS\
S\
S\
S	\
4S
 jjr " S S\	R                  5      r " S S\5      rS	\4S jrg)    N)override	rearrange)randint)ComfyExtensioniomax_optionsvalue	min_valuereturnc                ,   [        X5      n[        XS-   5       Vs/ s H  o0U-  S:X  d  M  UPM     nnUS U  Vs/ s H  o0U-  PM	     nn[        U5      S-
  S:  a)  [        S[        U5      S-
  SS9R	                  5       nXV   $ SnXV   $ s  snf s  snf )N   r   r   )lowhighsize)minrangelenr   item)r
   r   r	   idivisorsnsidxs          </home/wildlama/comfy/ComfyUI/comfy_extras/nodes_hypertile.pyrandom_divisorr   
   s    I%I !AI6I6a!)q.6HI&|4	541*4B	5
2w{Q!#b'A+D9>>@ 7N 7N J	5s   BBBc                   P    \ rS rSr\S 5       r\S\R                  4S j5       rSr	g)	HyperTile   c                    [         R                  " SS[         R                  R                  S5      [         R                  R                  SSSSSS	9[         R                  R                  S
SSSSS	9[         R                  R                  SSSSSS	9[         R
                  R                  SSSS9/[         R                  R                  5       /S9$ )Nr   zmodel/patch/unetmodel	tile_size   r   i   T)defaultr   maxadvanced	swap_size      	max_depthr   
   scale_depthF)r$   r&   )node_idcategoryinputsoutputs)r   SchemaModelInputIntBooleanOutput)clss    r   define_schemaHyperTile.define_schema   s    yy'w'[#1$QUV[!dS[!TR

   M !
 	
    r   c                    ^^^^	^
 [        SU5      S-  m	S m
U	UUUU
4S jnU
4S jnUR                  5       nUR                  U5        UR                  U5        U4$ )N       c           	      8  > U R                   S   nUS   n/ n[        TS-   5       H)  nUR                  US   SU-  -  US   SU-  -  -  5        M+     XF;   a  US   nUS   US   -  n	U R                  S5      n
[	        [
        R                  " X-  5      5      [	        [
        R                  " X-  5      5      pT(       a  SUR                  U5      -  OSn[        UTU-  T5      n[        UTU-  T5      nX-  S:  a  [        U SX-  X-  XS9n XX4mXU4$ XU4$ )Noriginal_shaper   r(   z$b (nh h nw w) c -> (b nh nw) (h w) c)hwnhnw)
shaper   appendr   roundmathsqrtindexr   r   )qkvextra_optionsmodel_chans
orig_shapeapply_tor   rF   aspect_ratiohwrB   rC   factorrD   rE   latent_tile_sizer*   r,   r'   temps                   r   hypertile_in'HyperTile.execute.<locals>.hypertile_in0   s:   ''"+K&'78JH9q=)B16!:z"~QRVWQW?X YZ * &%&67$Ry594VVAYTYYr'89:E$))BL]B^<_1?J!x~~k::PQ#A'7&'@)L#A'7&'@)L7Q;!!%KqwZ[ZafhpAA>DQw7Nr:   c                 R   > Tb"  Tu  p#pES m[        U SX#S9n [        U SXB-  XS-  S9n U $ )Nz(b nh nw) hw c -> b nh nw hw c)rD   rE   z"b nh nw (h w) c -> b (nh h nw w) c)rB   rC   r   )outrO   rD   rE   rB   rC   rW   s         r   hypertile_out(HyperTile.execute.<locals>.hypertile_outI   sB    #%E"T%IQWXYX_`Jr:   )r%   cloneset_model_attn1_patchset_model_attn1_output_patch)r7   r!   r"   r'   r*   r,   rX   r\   mrV   rW   s      ```   @@r   executeHyperTile.execute+   sY    r9-2	 	2	 KKM	-	&&}5ur:    N)
__name__
__module____qualname____firstlineno__classmethodr8   r   
NodeOutputrb   __static_attributes__rd   r:   r   r   r      s3    
 
  *R]] * *r:   r   c                   L    \ rS rSr\S\\\R                        4S j5       r	Sr
g)HyperTileExtensionY   r   c                    #    [         /$ 7fN)r   )selfs    r   get_node_list HyperTileExtension.get_node_listZ   s      
 	
s   
rd   N)re   rf   rg   rh   r   listtyper   	ComfyNoderr   rk   rd   r:   r   rm   rm   Y   s)    
T$r||*<%= 
 
r:   rm   c                     #    [        5       $ 7frp   )rm   rd   r:   r   comfy_entrypointrx   a   s     s   r   )rI   typing_extensionsr   einopsr   torchr   comfy_api.latestr   r   intr   rv   r   rm   rx   rd   r:   r   <module>r~      sa     &   /s # # 3 = =@
 
  2  r:   