
    +j                     f   S SK rS SKrS SKJr  S SKJr  S SKrS SK	rS SK
rS SKJr  S SKJrJr  \R                   " 5       SS j5       r " S S\R$                  5      r\R                   " 5       SS	 j5       r " S
 S\R$                  5      r " S S\R$                  5      r " S S\5      rS\4S jrg)    N)trange)override)to_d)ComfyExtensionioc	           
         Uc  0 OUnUc  [        [        U5      S-  S-   S5      nOUS-  n[        U[        U5      S-   5      n[        R                  " SXh5      SS  n	UR                  5       n
UR                  UR                  S   /5      n[        [        U5      S-
  US9 H  nU " XU   U-  40 UD6nUb  U" XX,   X,   US.5        UnU[        U	5      :  aE  [        R                  R                  U[        U
S   X   -  5      [        U
S   X   -  5      US	5      nX,S-      S:  d  M  XUS-      [        R                  " U5      -  -  nM     U$ )
N            ?r   disablexisigma	sigma_hatdenoiseddisabled)maxlenminnplinspacesizenew_onesshaper   comfyutilscommon_upscaleroundtorch
randn_like)modelr   sigmas
extra_argscallbackr   total_upscaleupscale_methodupscale_stepsupscales
orig_shapes_inr   r   s                 D/home/wildlama/comfy/ComfyUI/comfy_extras/nodes_advanced_samplers.pysample_lcm_upscaler0      sh   !)zJCK1,q0!4M3v;?;{{3=abAHJ::qwwqzl#DCK!OW51I,;
;1vyvy^fghs8}**1eJrNX[4P.QSXYcdfYgjrjuYuSv  yG  IS  TAa%=1A!1!1!!444A 6 H    c                   t    \ rS rSr/ SQr\S\R                  4S j5       r\S\R                  4S j5       r
Srg)SamplerLCMUpscale)   )bislerpznearest-exactbilinearareabicubicreturnc                 <   [         R                  " SS[         R                  R                  SSSSSSS	9[         R                  R                  S
SSSSSS	9[         R
                  R                  SU R                  S9/[         R                  R                  5       /S9$ )Nr3   model/sampling/samplersscale_ratior   g?g      4@{Gz?T)defaultr   r   stepadvancedscale_stepsr   i  r
   r*   )options)node_idcategoryinputsoutputs)	r   SchemaFloatInputIntComboUPSCALE_METHODSSamplerOutputclss    r/   define_schemaSamplerLCMUpscale.define_schema,   s    yy'.}csSWbfg]BBDq[_`/9L9LM
 ZZ&&()	
 		
r1   c                     US:  a  S n[         R                  R                  [        XUS.S9n[        R
                  " U5      $ )Nr   )r)   r+   r*   )extra_options)r   samplersKSAMPLERr0   r   
NodeOutput)rP   r<   rA   r*   samplers        r/   executeSamplerLCMUpscale.execute9   sI    ?K..))*<^i  \j  Mk)  l}}W%%r1    N)__name__
__module____qualname____firstlineno__rL   classmethodr   rG   rQ   rW   rY   __static_attributes__r[   r1   r/   r3   r3   )   sC    QO

bii 

 

 &"-- & &r1   r3   c           	        ^ Uc  0 OUnS/mU4S jnUR                  S0 5      R                  5       n[        R                  R	                  XvSS9US'   UR                  UR                  S   /5      n[        [        U5      S-
  US9 HL  n	X)   n
U " XU-  40 UD6n[        X-
  TS   -   X)   U5      nUb  U" XX)   XS.5        X)S-      U
-
  nXU-  -   nMN     U$ )	Nr   c                    > U S   TS'   U S   $ )Nuncond_denoisedr   r   r[   )argstemps    r/   post_cfg_function*sample_euler_pp.<locals>.post_cfg_functionF   s    ()QJr1   model_optionsT)disable_cfg1_optimizationr
   r   r   )
getcopyr   model_patcher#set_model_options_post_cfg_functionr   r   r   r   r   )r%   r   r&   r'   r(   r   rg   ri   r.   r   r   r   ddtrf   s                 @r/   sample_euler_pprq   A   s    !)zJ3D  NN?B7<<>M"'"5"5"Y"YZg  W["Y  #\J::qwwqzl#DCK!OW5I	,;
;Q'H=1vyyghE]Y&BJ 6 Hr1   c                   l    \ rS rSr\S\R                  4S j5       r\S\R                  4S j5       r	Sr
g)
SamplerLCMY   r9   c                 2   [         R                  " SSS[         R                  R                  SSSSSS	S
9[         R                  R                  SSSSSSS
9[         R                  R                  SSSSSSS
9/[         R                  R                  5       /S9$ )Nrs   r;   zdLCM sampler with tunable per-step noise. s_noise is a multiplier on the model's training noise scales_noiser   g        g      P@r=   zCPer-step noise multiplier at the first step (1.0 = match training).)r>   r   r   r?   tooltips_noise_endzYPer-step noise multiplier at the last step. Set equal to s_noise for a constant schedule.noise_clip_stdg      $@z.Clamp per-step noise to +/- N*std. 0 disables.)rC   rD   descriptionrE   rF   )r   rG   rH   rI   rM   rN   rO   s    r/   rQ   SamplerLCM.define_schemaZ   s    yy .y#3Dt'l  n}csSW (C  D/#4VZ'W  Y ZZ&&()
 	
r1   c                     [         R                  R                  S[        U5      [        U5      [        U5      S.5      n[        R
                  " U5      $ )Nlcm)rv   rx   ry   )r   rU   ksamplerfloatr   rW   )rP   rv   rx   ry   rX   s        r/   rY   SamplerLCM.executek   sF    ..)) >$[1"'"7
 }}W%%r1   r[   Nr\   r]   r^   r_   r`   r   rG   rQ   rW   rY   ra   r[   r1   r/   rs   rs   Y   s>    
bii 
 
  	&bmm 	& 	&r1   rs   c                   l    \ rS rSr\S\R                  4S j5       r\S\R                  4S j5       r	Sr
g)SamplerEulerCFGppx   r9   c                     [         R                  " SSS[         R                  R                  SSS/SS9/[         R                  R                  5       /SS	9$ )
Nr   zSamplerEulerCFG++experimentalversionregularalternativeT)rB   r@   )rC   display_namerD   rE   rF   is_experimental)r   rG   rK   rI   rM   rN   rO   s    r/   rQ   SamplerEulerCFGpp.define_schemay   sU    yy',#y9m2LW[\ ZZ&&() 	
 		
r1   c                     US:X  a$  [         R                  R                  [        5      nO[         R                  R	                  S5      n[
        R                  " U5      $ )Nr   euler_cfg_pp)r   rU   rV   rq   r~   r   rW   )rP   r   rX   s      r/   rY   SamplerEulerCFGpp.execute   sA    m#nn--o>Gnn--n=G}}W%%r1   r[   Nr   r[   r1   r/   r   r   x   s>    

bii 

 

 & & &r1   r   c                   L    \ rS rSr\S\\\R                        4S j5       r	Sr
g)AdvancedSamplersExtension   r9   c                 ,   #    [         [        [        /$ 7fN)r3   rs   r   )selfs    r/   get_node_list'AdvancedSamplersExtension.get_node_list   s      
 	
s   r[   N)r\   r]   r^   r_   r   listtyper   	ComfyNoder   ra   r[   r1   r/   r   r      s)    
T$r||*<%= 
 
r1   r   r9   c                     #    [        5       $ 7fr   )r   r[   r1   r/   comfy_entrypointr      s     $&&s   )NNNg       @r5   N)NNN)numpyr   r#   	tqdm.autor   typing_extensionsr   comfy.model_patcherr   comfy.samplerscomfy.utilscomfy.k_diffusion.samplingr   comfy_api.latestr   r   no_gradr0   r   r3   rq   rs   r   r   r   r[   r1   r/   <module>r      s       &    + /  6& &0  .& &>& &.
 
' 9 'r1   