
    +j                     >   S SK r S SKrS SKrS SKrS SKrS SKrS SKrS SKrS SK	r	S SK
r
S SKrS SKrS SKJr  S SKJr  S SKJrJrJr  S SKJr  S SKJr  \R4                  R6                  r\ R8                  r " S S\R:                  5      r " S	 S
\R:                  5      r " S S\R:                  5      r  " S S\R:                  5      r! " S S\R:                  5      r" " S S\R:                  5      r# " S S\R:                  5      r$ " S S\R:                  5      r% " S S\R:                  5      r& " S S\R:                  5      r' " S S\R:                  5      r( " S S\R:                  5      r) " S S \R:                  5      r* " S! S"\R:                  5      r+ " S# S$\R:                  5      r, " S% S&\R:                  5      r- " S' S(\R:                  5      r.S)\R^                  S*S*S+.S)\R^                  S,S,S+.S-\R`                  S.S/S+.S-\R`                  S0S1S+.S2\Rb                  S3S3S+.S2\Rb                  S4S4S+.S5.r2S6\
Rf                  S7\
Rf                  4S8 jr4S9r5S:r6S;r7S6\
Rf                  S7\
Rf                  4S< jr8S=r9S>\:S?\:S7\:4S@ jr;SA\<SB\<S7\:4SC jr=SD\:SE\>S-  SF\>S-  S7\:4SG jr?SHSISJ.r@SK\AS7\:4SL jrBSM\<SN\<SO\:S7\:4SP jrC S]SQ\:SE\>S-  SF\>S-  SR\<S-  S7\:4
SS jjrDST\
Rf                  SU\<SV\<SR\<S7\:4
SW jrE " SX SY\R:                  5      rF " SZ S[\5      rGS7\G4S\ jrHg)^    N)Fraction)PromptServer)ComfyExtensionIOUI)args)overridec                   T    \ rS rSr\S 5       r\S\R                  4S j5       r\r	Sr
g)	ImageCrop   c                     [         R                  " SS/SSSS[         R                  R                  S5      [         R                  R                  SS	S
[
        R                  S
S9[         R                  R                  SS	S
[
        R                  S
S9[         R                  R                  SSS[
        R                  S
S9[         R                  R                  SSS[
        R                  S
S9/[         R                  R                  5       /S9$ )Nr   trimzCrop Image (DEPRECATED)image/transformTImage Toolsimagewidth      defaultminmaxstepheightxr   y)node_idsearch_aliasesdisplay_namecategoryis_deprecatedessentials_categoryinputsoutputs)r   SchemaImageInputIntnodesMAX_RESOLUTIONOutputclss    9/home/wildlama/comfy/ComfyUI/comfy_extras/nodes_images.pydefine_schemaImageCrop.define_schema   s    yy"82& -w'Wcqe>R>RYZ[Xsu?S?SZ[\S!8L8LSTUS!8L8LSTU XX__&'
 	
    returnc                     [        XAR                  S   S-
  5      n[        XQR                  S   S-
  5      nX$-   nX5-   nUS S 2XW2XF2S S 24   n[        R                  " U5      $ )N   r   )r   shaper   
NodeOutput)	r-   r   r   r   r   r   to_xto_yimgs	            r.   executeImageCrop.execute/   sd    ;;q>A%&;;q>A%&yzAafafa'(}}S!!r1    N)__name__
__module____qualname____firstlineno__classmethodr/   r   r6   r:   crop__static_attributes__r<   r1   r.   r   r      s:    
 
$ "BMM " " Dr1   r   c                   P    \ rS rSr\S 5       r\S\R                  4S j5       rSr	g)ImageCropV2;   c                     [         R                  " S/ SQSSSSS[         R                  R                  S5      [         R                  R                  S	S
S9/[         R                  R                  5       /S9	$ )NrE   )rB   cutr   z
Crop Imager   z*Crop an image to the specified dimensions.r   Tr   crop_regionr   )	component)	r   r   r   r    descriptionr"   has_intermediate_outputr#   r$   )r   r%   r&   r'   BoundingBoxr+   r,   s    r.   r/   ImageCropV2.define_schema<   sg    yy!2%&F -$(w'$$]k$J XX__&'
 	
r1   r2   c                 ~   UR                  SS5      nUR                  SS5      nUR                  SS5      nUR                  SS5      n[        X1R                  S   S-
  5      n[        XAR                  S   S-
  5      nXS-   nXd-   nUS S 2XH2X72S S 24   n	[        R                  " U	[
        R                  " U	5      S	9$ )
Nr   r   r   r   r   r   r4   r   ui)getr   r5   r   r6   r   PreviewImage)
r-   r   rI   r   r   r   r   r7   r8   r9   s
             r.   r:   ImageCropV2.executeM   s    OOC#OOC#-3/;;q>A%&;;q>A%&yzAafafa'(}}SR__S%9::r1   r<   N
r=   r>   r?   r@   rA   r/   r   r6   r:   rC   r<   r1   r.   rE   rE   ;   s3    
 
  ;BMM ; ;r1   rE   c                   P    \ rS rSr\S 5       r\S\R                  4S j5       rSr	g)rM   \   c                    [         R                  " SSS[         R                  R                  SSS[        S9[         R                  R                  SSS[        S9[         R                  R                  SS	S
[        S9[         R                  R                  SS	S
[        S9/[         R
                  R                  5       /S9$ )NPrimitiveBoundingBoxzBounding Boxzutilities/primitiver   r   r   r   r   r   r   r   r   r   r   r   r    r#   r$   )r   r%   r(   r'   r*   rM   r+   r,   s    r.   r/   BoundingBox.define_schema]   s    yy*'*S!GS!GWcqnMXs~N	 ^^**,-
 	
r1   r2   c                 4    [         R                  " XX4S.5      $ )N)r   r   r   r   )r   r6   )r-   r   r   r   r   s        r.   r:   BoundingBox.executel   s    }}1uOPPr1   r<   NrU   r<   r1   r.   rM   rM   \   s6    
 
 QR]] Q Qr1   rM   c                   T    \ rS rSr\S 5       r\S\R                  4S j5       r\r	Sr
g)RepeatImageBatchq   c                     [         R                  " SSS/SS[         R                  R                  S5      [         R                  R                  SSSS	S
9/[         R                  R                  5       /S9$ )Nr`   zduplicate imagezclone imagezRepeat Image Batchimage/batchr   amountr      rZ   r   r   r   r    r#   r$   r   r%   r&   r'   r(   r+   r,   s    r.   r/   RepeatImageBatch.define_schemar   sd    yy&-}=-"w'XqaTB XX__&'

 
	
r1   r2   c                 X    UR                  USSS45      n[        R                  " U5      $ )Nr   )repeatr   r6   )r-   r   rd   ss       r.   r:   RepeatImageBatch.execute   s'    LL&!Aa)}}Qr1   r<   Nr=   r>   r?   r@   rA   r/   r   r6   r:   rj   rC   r<   r1   r.   r`   r`   q   s:    
 
  r}}     Fr1   r`   c                   T    \ rS rSr\S 5       r\S\R                  4S j5       r\r	Sr
g)ImageFromBatch   c                 <   [         R                  " S/ SQSS[         R                  R                  S5      [         R                  R                  SS[
        * [
        S9[         R                  R                  S	S
S
SS9/[         R                  R                  5       /S9$ )Nro   )zselect imagezpick from batchzextract imagezGet Image from Batchrc   r   batch_indexr   rZ   lengthr   re   rf   r   r%   r&   r'   r(   r*   r+   r,   s    r.   r/   ImageFromBatch.define_schema   sz    yy$O/"w']AN?P^_XqaTB
 XX__&'
 	
r1   r2   c                    UnUS:  a  X$R                   S   -  n[        S[        UR                   S   S-
  U5      5      n[        UR                   S   U-
  U5      nXBX#-    R                  5       n[        R
                  " U5      $ )Nr   r   )r5   r   r   cloner   r6   )r-   r   rr   rs   s_inrk   s         r.   r:   ImageFromBatch.execute   s|    ?::a=(K!SA!2K@ATZZ][0&9[1288:}}Qr1   r<   N)r=   r>   r?   r@   rA   r/   r   r6   r:   	frombatchrC   r<   r1   r.   ro   ro      s:    
 
  BMM     Ir1   ro   c                   T    \ rS rSr\S 5       r\S\R                  4S j5       r\r	Sr
g)ImageAddNoise   c                 .   [         R                  " SS/SS[         R                  R                  S5      [         R                  R                  SSSSS	S
S9[         R
                  R                  SSSSSS9/[         R                  R                  5       /S9$ )Nr|   z
film grainzAdd Noise to Imagezimage/filtersr   seedr   l    Tz,The random seed used for creating the noise.)r   r   r   control_after_generatetooltipstrength      ?              ?{Gz?r   rf   )r   r%   r&   r'   r(   Floatr+   r,   s    r.   r/   ImageAddNoise.define_schema   s    yy#(>-$w'*+/J   z3CStT XX__&'#
 	
r1   r2   c           
          [         R                  " U5      n[         R                  " X[         R                  " UR	                  5       USS9R                  U5      -  -   SSS9n[        R                  " U5      $ )Ncpu)	generatordevicer   r   )r   r   )torchmanual_seedcliprandnsizetor   r6   )r-   r   r   r   r   rk   s         r.   r:   ImageAddNoise.execute   sd    %%d+	JJ5;;uzz|yaf+g+j+jkp+q qqx{  BE  F}}Qr1   r<   Nrm   r<   r1   r.   r|   r|      s:    
 
*  r}}    
 Fr1   r|   c                   d    \ rS rSrSSSS.r\S 5       r\SS\R                  4S jj5       r	\	r
S	rg
)SaveAnimatedWEBP      r      )r   fastestslowestc                 B   [         R                  " SSS[         R                  R                  S5      [         R                  R                  SSS9[         R
                  R                  SS	S
SS
S9[         R                  R                  SSS9[         R                  R                  SSSSS9[         R                  R                  S[        U R                  R                  5       5      S9/[         R                  R                  [         R                  R                  /SS9$ )Nr   zSave Animated WEBPr   imagesfilename_prefixComfyUIr   fps      @r        @@r   losslessTqualityP   r   d   rZ   methodoptionsr   r   r    r#   hiddenis_output_node)r   r%   r&   r'   Stringr   Booleanr(   CombolistCOMPRESS_METHODSkeysHiddenpromptextra_pnginfor,   s    r.   r/   SaveAnimatedWEBP.define_schema   s    yy&-x(		 19EuctdS

  T :YsCxc6J6J6O6O6Q1RS II$$bii&=&=>
 	
r1   r2   c                     [         R                  " [        R                  R	                  UUU UUUU R
                  R                  U5      S9S9$ )N)r   r   r-   r   r   r   r   rP   )r   r6   r   ImageSaveHelperget_save_animated_webp_uir   rR   )r-   r   r   r   r   r   r   
num_framess           r.   r:   SaveAnimatedWEBP.execute   sR    }}!!;; /!++//7 < 

 
	
r1   r<   N)r   )r=   r>   r?   r@   r   rA   r/   r   r6   r:   save_imagesrC   r<   r1   r.   r   r      sJ    #$qA
 
$ 
_a_l_l 
 
 Kr1   r   c                   X    \ rS rSr\S 5       r\SS\R                  4S jj5       r\r	Sr
g)SaveAnimatedPNG   c                    [         R                  " SSS[         R                  R                  S5      [         R                  R                  SSS9[         R
                  R                  SS	S
SS
S9[         R                  R                  SSSSSS9/[         R                  R                  [         R                  R                  /SS9$ )Nr   zSave Animated PNGr   r   r   r   r   r   r   r   r   r   compress_levelr   r   	   T)r   r   r   advancedr   )
r   r%   r&   r'   r   r   r(   r   r   r   r,   s    r.   r/   SaveAnimatedPNG.define_schema   s    yy%,x(		 19EuctdS-qaQQUV	 II$$bii&=&=>
 	
r1   r2   c           
      h    [         R                  " [        R                  R	                  UUU UUS9S9$ )N)r   r   r-   r   r   rP   )r   r6   r   r   get_save_animated_png_ui)r-   r   r   r   r   s        r.   r:   SaveAnimatedPNG.execute   s;    }}!!:: /- ; 
 	
r1   r<   N)r   )r=   r>   r?   r@   rA   r/   r   r6   r:   r   rC   r<   r1   r.   r   r      s=    
 
 	
PRP]P] 	
 	
 Kr1   r   c                   ^    \ rS rSrSr\S 5       r\ SS\R                  4S jj5       r	\	r
Srg)	ImageStitchi  z8Upstreamed from https://github.com/kijai/ComfyUI-KJNodesc                    [         R                  " S/ SQSSS[         R                  R                  S5      [         R                  R                  S/ SQS	S
9[         R
                  R                  SSS9[         R                  R                  SSSSSSS9[         R                  R                  S/ SQSSS9[         R                  R                  SSS9/[         R                  R                  5       /S9$ )Nr   )zcombine imageszjoin imageszconcatenate imageszside by sidezStitch ImageszStitches image2 to image1 in the specified direction.
If image2 is not provided, returns image1 unchanged.
Optional spacing can be added between images.r   image1	direction)rightdownleftupr   )r   r   match_image_sizeTr   spacing_widthr      r4   )r   r   r   r   r   spacing_colorwhiteblackredgreenbluer   r   r   r   image2)optional)r   r   r   rK   r    r#   r$   )r   r%   r&   r'   r   r   r(   r+   r,   s    r.   r/   ImageStitch.define_schema  s    yy!b(< 'x({4S]de

  !3T B_aQDq[_`8bls  C  Dx$7 XX__&'!
 	
r1   Nr2   c           	         Uc  [         R                  " U5      $ UR                  S   UR                  S   :w  a  [        UR                  S   UR                  S   5      nUR                  S   U:  a;  [        R
                  " XSS  R                  XqR                  S   -
  SSS5      /5      nUR                  S   U:  a;  [        R
                  " XfSS  R                  XvR                  S   -
  SSS5      /5      nU(       a  UR                  SS u  pUR                  SS u  pX-  nUS;   a  U[        X-  5      pOU	[        X-  5      p[        R                  R                  UR                  SS5      XSS5      R                  SS5      nSS	S
SSS.nX   nU(       Gdv  UR                  SS u  pUR                  SS u  pS	n[        U[        5      (       d  UnUS;   a  X:w  a  [        X5      nX:  aA  X-
  nUS-  UUS-  -
  nn[        R                  R                  R!                  USSSSUU4SUS9nX:  aA  X-
  nUS-  UUS-  -
  nn[        R                  R                  R!                  USSSSUU4SUS9nOX:w  a  [        X5      nX:  a?  X-
  nUS-  UUS-  -
  nn[        R                  R                  R!                  USSUU4SUS9nX:  a?  X-
  nUS-  UUS-  -
  nn[        R                  R                  R!                  USSUU4SUS9nUR                  S   UR                  S   :w  a  [        UR                  S   UR                  S   5      nUR                  S   U:  aW  [        R
                  " U[        R"                  " / UR                  S S QUUR                  S   -
  P7SUR$                  06/SS9nUR                  S   U:  aW  [        R
                  " U[        R"                  " / UR                  S S QUUR                  S   -
  P7SUR$                  06/SS9nUS:  GaM  XDS-  -   nUS;   aE  UR                  S   [        UR                  S   UR                  S   5      UUR                  S   4nODUR                  S   U[        UR                  S   UR                  S   5      UR                  S   4n[        R&                  " US	UR$                  S9n[        U[        5      (       aJ  [)        U5       H"  u  nnUUR                  S   :  d  M  UUSU4'   M$     UR                  S   S:X  a  SUS'   O8UUSS [+        SUR                  S   5      24'   UR                  S   S:X  a  SUS'   US;   a  Xa/OX/nUS:  a  UR-                  SW5        US;   a  SOSn[         R                  " [        R
                  " UUS95      $ )Nr   r      )r   r   lanczosdisabledr   r   )r   r   r   )r   r   r   )r   r   r   r   r4   constant)modevaluer   dim)r   .r   ).r   )r   r   )r   r6   r5   r   r   catrj   intcomfyutilscommon_upscalemovedim
isinstancetuplenn
functionalpadonesr   full	enumerater   insert)r-   r   r   r   r   r   r   	max_batchh1w1h2w2aspect_ratiotarget_htarget_w	color_map	color_val	pad_valuepad_hpad_top
pad_bottompad_wpad_left	pad_rightmax_channelsspacing_shapespacingicr   
concat_dims                                  r.   r:   ImageStitch.execute$  s
    >==(( <<?fll1o-FLLOV\\!_=I||A*BC[//	LLO0KQPQSTUV ||A*BC[//	LLO0KQPQSTUV
 \\!A&FB\\!A&FB7L--%'R->)?(%'R->)?([[//r1%x9jgan 
 "$#
	 ,	  \\!A&FB\\!A&FBIi//%	--8"2{H} (.3qj%%1*:L!&!4!4!8!8!Q1gWaAbis  |E!8  "F} (.3qj%%1*:L!&!4!4!8!8!Q1gWaAbis  |E!8  "F 8"2{H} (.3qj%%1*:L)!&!4!4!8!8!QR[A\cmu~!8!} (.3qj%%1*:L)!&!4!4!8!8!QR[A\cmu~!8! <<v||B//v||B/b1ABL||B,.

 #\\#2.(6<<+;; $*== 
 ||B,.

 #\\#2.(6<<+;; $*== 
 1)Q->?M--LLOQa9!LL$	! LLO!Qa9LL$	! jjFMMJG)U++%i0DAq7==,,*+Q 1 ==$)&)GFO<E8s1gmmB&78889==$)&)GFO &/.%@&!vFV1MM!W%#'88Qa
}}UYYv:>??r1   r<   N)r=   r>   r?   r@   __doc__rA   r/   r   r6   r:   stitchrC   r<   r1   r.   r   r     sL    B
 
(  S@ 
S@ S@j Fr1   r   c                   T    \ rS rSr\S 5       r\S\R                  4S j5       r\r	Sr
g)ResizeAndPadImagei  c                    [         R                  " SS/SS[         R                  R                  S5      [         R                  R                  SSS[
        R                  SS	9[         R                  R                  S
SS[
        R                  SS	9[         R                  R                  SSS/SS9[         R                  R                  S/ SQSS9/[         R                  R                  5       /S9$ )Nr  zfit to sizezResize And Pad Imager   r   target_widthr   r   r   target_heightpadding_colorr   r   T)r   r   interpolation)areabicubicnearest-exactbilinearr   rf   )	r   r%   r&   r'   r(   r)   r*   r   r+   r,   s    r.   r/   ResizeAndPadImage.define_schema  s    yy')?/&w'^SaUEYEY`ab_cqeFZFZabc'8JUYZ8s  C  D XX__&'
 	
r1   r2   c                    UR                   u  pgpX(-  n
X7-  n[        X5      n[        X-  5      n[        X|-  5      nUR                  SSSS5      n[        R
                  R                  XXS5      nUS:X  a  SOSn[        R                  " XiX24UUR                  UR                  S	9nX>-
  S-  nX--
  S-  nUUS S 2S S 2UUU-   2UUU-   24'   UR                  SSSS5      n[        R                  " U5      $ )
Nr   r   r   r4   r   r   r   r   )dtyper   )r5   r   r   permuter   r   r   r   r   r   r   r   r6   )r-   r   r  r  r  r  
batch_sizeorig_height
orig_widthchannelsscale_wscale_hscale	new_width
new_heightimage_permutedresizedr  paddedy_offsetx_offsetoutputs                         r.   r:   ResizeAndPadImage.execute  s   8=5
+-G%
*+	,-
q!Q2++,,^
cmn(G3C	=?++<<	
 ".14 ,2V]q!Xh33Xh>R5RRS1a+}}V$$r1   r<   N)r=   r>   r?   r@   rA   r/   r   r6   r:   resize_and_padrC   r<   r1   r.   r  r    s=    
 
  %Z\ZgZg % %< Nr1   r  c                       \ rS rSr\S 5       r\SS\R                  R                  S\R                  4S jj5       r
\
rSrg)	SaveSVGNodei  c                    [         R                  " SSS/SSS[         R                  R                  S5      [         R                  R                  SS	S
S9/[         R
                  R                  [         R
                  R                  /SS9$ )Nr4  zexport vectorzsave vector graphicszSave SVGzSave SVG files on disk.r   svgr   svg/ComfyUIzThe prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes.r   r   T)r   r   r   rK   r    r#   r   r   )r   r%   SVGr'   r   r   r   r   r,   s    r.   r/   SaveSVGNode.define_schema  s    yy!+-CD#1U#		%) }    II$$bii&=&=>
 	
r1   r6  r2   c           	      F  ^ [         R                  " U[         R                  " 5       5      u  p4pVn/ n0 nU R                  R                  b  U R                  R                  US'   U R                  R
                  b%  UR                  U R                  R
                  5        U(       a  [        R                  " USS9OS n	[        UR                  5       GH$  u  pUR                  S[        U
5      5      nU SUS S3nUR                  S5        UR                  5       R                  S	5      nU	(       a0  S
U	 S3mU4S jn[         R"                  " SX[         R$                  S9n['        [(        R*                  R-                  X=5      S5       nUR/                  UR1                  S	5      5        S S S 5        UR3                  [4        R6                  " X[8        R:                  R<                  S95        US-  nGM'     [8        R>                  " SU0S9$ ! , (       d  f       Nj= f)Nr   r4   )indent%batch_num%_05z_.svgr   utf-8z3  <metadata>
                <![CDATA[
            z9
                ]]>
            </metadata>
            c                 2   > U R                  S5      S-   T-   $ )Nr   
)group)matchmetadata_elements    r.   replacement(SaveSVGNode.execute.<locals>.replacement(  s     ;;q>D03CCCr1   z(<svg[^>]*>))flagswbfilename	subfoldertyper   r   rP   ) folder_pathsget_save_image_pathget_output_directoryr   r   r   updatejsondumpsr   datareplacestrseekreaddecoderesubUNICODEopenospathjoinwriteencodeappendr   SavedResultr   
FolderTyper0  r6   )r-   r6  r   full_output_folderrK  counterrL  resultsmetadata_dictmetadata_jsonbatch_number	svg_bytesfilename_with_batch_numfilesvg_contentrF  svg_filerE  s                    @r.   r:   SaveSVGNode.execute  s   LXLlLlm|  K  `  `  b  McIg/(* ::(&)jj&7&7M(#::##/  !9!9: @M

=;RV (1':#L&.&6&6}c,FW&X#-.a|5AD NN1#..*11':K (O $ D
 !ff_kVXV`V`a bggll#5<dCx{11':; D NN2>>4SUS`S`SgSghiqLG= (;> }}7 344 DCs   !H
H 	r<   N)r7  )r=   r>   r?   r@   rA   r/   r   r9  Typer6   r:   save_svgrC   r<   r1   r.   r4  r4    sJ    
 
& .5"&&++ .5 .5 .5` Hr1   r4  c                   T    \ rS rSr\S 5       r\S\R                  4S j5       r\r	Sr
g)GetImageSizei:  c                 R   [         R                  " S/ SQSSS[         R                  R                  S5      /[         R                  R                  SS9[         R                  R                  SS9[         R                  R                  S	S9/[         R                  R                  /S
9$ )Nru  )
dimensions
resolutionz
image infozGet Image SizezGReturns width and height of the image, and passes it through unchanged.r   r   )r   r   r"  )r   r   r   rK   r    r#   r$   r   )r   r%   r&   r'   r(   r+   r   	unique_idr,   s    r.   r/   GetImageSize.define_schema;  s    yy"E)aw' 7384<8
 II''(
 	
r1   r2   c                 :   UR                   S   nUR                   S   nUR                   S   nU R                  R                  (       a=  [        R                  R                  SU SU SU 3U R                  R                  5        [        R                  " X2U5      $ )Nr   r4   r   zwidth: z
, height: z
 batch size: )r5   r   ry  r   instancesend_progress_textr   r6   )r-   r   r   r   r"  s        r.   r:   GetImageSize.executeN  s    QA[[^
 ::!!44wugZPVxWfgqfr5suxuu  vJ  vJ  K}}UJ77r1   r<   N)r=   r>   r?   r@   rA   r/   r   r6   r:   get_sizerC   r<   r1   r.   ru  ru  :  s:    
 
$ 	8r}} 	8 	8 Hr1   ru  c                   T    \ rS rSr\S 5       r\S\R                  4S j5       r\r	Sr
g)ImageRotatei]  c                     [         R                  " SSSS/SS[         R                  R                  S5      [         R                  R                  S/ S	QS
9/[         R                  R                  5       /S9$ )Nr  zRotate Imageturnzflip orientationr   r   r   rotation)nonez
90 degreesz180 degreesz270 degreesr   )r   r   r   r    r"   r#   r$   r   r%   r&   r'   r   r+   r,   s    r.   r/   ImageRotate.define_schema^  se    yy!'"$67& -w'z3gh XX__&'
 	
r1   r2   c                     SnUR                  S5      (       a  SnO1UR                  S5      (       a  SnOUR                  S5      (       a  Sn[        R                  " XSS/S9n[        R                  " U5      $ )	Nr   90r   180r4   270r   )kdims)
startswithr   rot90r   r6   )r-   r   r  	rotate_bys       r.   r:   ImageRotate.executem  sl    	t$$I  ''I  ''IEaV<}}U##r1   r<   N)r=   r>   r?   r@   rA   r/   r   r6   r:   rotaterC   r<   r1   r.   r  r  ]  s:    
 
 
$ 
$ 
$ Fr1   r  c                   T    \ rS rSr\S 5       r\S\R                  4S j5       r\r	Sr
g)	ImageFlipi}  c                     [         R                  " SSS/SS[         R                  R                  S5      [         R                  R                  SSS	/S
9/[         R                  R                  5       /S9$ )Nr  mirrorreflectz
Flip Imager   r   flip_methodzx-axis: verticallyzy-axis: horizontallyr   rf   r  r,   s    r.   r/   ImageFlip.define_schema~  sg    yy$i0%&w'}7KMc6de XX__&'

 
	
r1   r2   c                     UR                  S5      (       a  [        R                  " US/S9nO,UR                  S5      (       a  [        R                  " US/S9n[        R                  " U5      $ )Nr   r   )r  r   r4   )r  r   flipr   r6   )r-   r   r  s      r.   r:   ImageFlip.execute  sW    !!#&&JJuA3/E##C((JJuA3/E}}U##r1   r<   N)r=   r>   r?   r@   rA   r/   r   r6   r:   r  rC   r<   r1   r.   r  r  }  s:    
 
 $BMM $ $ Dr1   r  c                   T    \ rS rSr\S 5       r\S\R                  4S j5       r\r	Sr
g)ImageScaleToMaxDimensioni  c                 .   [         R                  " SSS[         R                  R                  S5      [         R                  R                  S/ SQS9[         R
                  R                  SS	S
[        SS9/[         R                  R                  5       /S9$ )Nr  zScale Image to Max Dimensionzimage/upscalingr   upscale_method)r  r   r  r  r  r  r   largest_sizer   r   r   r   r[   )r   r%   r&   r'   r   r(   r*   r+   r,   s    r.   r/   &ImageScaleToMaxDimension.define_schema  s|    yy.7&w'$c   ^Sa^Z[\ XX__&'
 	
r1   r2   c                 \   UR                   S   nUR                   S   nXE:  a  [        XT-  U-  5      nUnOXT:  a  [        XE-  U-  5      nUnOUnUnUR                  SS5      n[        R                  R                  XeXBS5      nUR                  SS5      n[        R                  " U5      $ )Nr   r4   r   r   )r5   roundr   r   r   r   r   r6   )r-   r   r  r  r   r   samplesrk   s           r.   r:    ImageScaleToMaxDimension.execute  s    QA>5>\9:E!F^FNl:;F E!F E--A&KK&&wvzZIIa}}Qr1   r<   N)r=   r>   r?   r@   rA   r/   r   r6   r:   upscalerC   r<   r1   r.   r  r    s:    
 
   R]]    & Gr1   r  c                   D    \ rS rSr\S 5       r\S 5       r\S 5       rSr	g)SplitImageToTileListi  c                 x   [         R                  " SS/ SQSS[         R                  R                  S5      [         R                  R                  SSS	[
        S
9[         R                  R                  SSS	[
        S
9[         R                  R                  SSSSS
9/[         R                  R                  SS9/S9$ )Nr  rc   zsplit imagez
tile imagezslice imagezSplit Image into List of TileszFSplits an image into a batched list of tiles with a specified overlap.r   
tile_widthr   @   rZ   tile_heightoverlap   r   re   Tis_output_list)r   r    r   r   rK   r#   r$   rt   r,   s    r.   r/   "SplitImageToTileList.define_schema  s    yy*"G9`w'\4R^T]DbnUYE	 t4
 	
r1   c                 l   / n[        [        US-  X$-
  5      5      n[        [        US-  X4-
  5      5      nSnX:  av  Sn	[        X-   U5      n
[        SX-
  5      nX:  a@  [        X-   U 5      n[        SX-
  5      nUR                  XX45        X:  a  OX-  n	X:  a  M@  X:  a   U$ X-  nX:  a  Mv  U$ )Ng      ?r   )r  r   r   rc  )r   r   r  r  r  coordsstride_xstride_yr   r   y_endy_startx_endx_starts                 r.   get_grid_coords$SplitImageToTileList.get_grid_coords  s    Z$.
0DEF[4/1FGHjA0E!U01G)ANE2a!34w>?> )   MA# j& r1   c                     UR                   u  pVpxU R                  XvX#U5      n	/ n
U	 H&  u  ppUS S 2X2X2S S 24   nU
R                  U5        M(     [        R                  " U
5      $ r  )r5   r  rc  r   r6   )r-   r   r  r  r  bhwr  r  output_listr  r  r  r  tiles                   r.   r:   SplitImageToTileList.execute  sm    [[
a$$Q:GL06,WuGM7=!;<Dt$ 17 }}[))r1   r<   N)
r=   r>   r?   r@   rA   r/   staticmethodr  r:   rC   r<   r1   r.   r  r    s<    
 
$  6 	* 	*r1   r  c                   4    \ rS rSr\S 5       r\S 5       rSrg)ImageMergeTileListi  c                 h   [         R                  " SSS/ SQS[         R                  R                  S5      [         R                  R                  SSS	S
S9[         R                  R                  SSS	S
S9[         R                  R                  SSSSS9/[         R                  R                  SS9/S9$ )Nr  zMerge List of Tiles to Imagerc   r  T
image_listfinal_widthr   r  i   rZ   final_heightr  r  r   re   Fr  )r   r   r    r   is_input_listr#   r$   rg   r,   s    r.   r/    ImageMergeTileList.define_schema   s    yy(7"G|,]DbeL^TruMYE	 u5
 	
r1   c                 J   US   nUS   nUS   nSnUS   n	U	R                   u  ppU	R                  nU	R                  n[        R	                  XVXU5      n[
        R                  " XX]4XS9n[
        R                  " XUS4XS9nUS:  Ga  [
        R                  " [        R                  [
        R                  " SSXUS9-  5      n[
        R                  " [        R                  [
        R                  " SSXUS9-  5      n[
        R                  " USS9n[
        R                  " USS9nUR                  S5      UR                  S5      -  R                  S5      R                  S5      n[
        R                  " U5      n[
        R                  " UUU5      nO[
        R                  " SXS4XS9n[!        U5       H  u  nu  nnnnU[#        U5      :  a    OUU   nUU-
  nUU-
  n[%        UUR                   S   5      n [%        UUR                   S   5      n!UU -   n"UU!-   n#US S 2S U 2S U!2S S 24   n$US S 2S U 2S U!2S S 24   n%US S 2UU"2UU#2S S 24==   U$U%-  -  ss'   US S 2UU"2UU#2S S 24==   U%-  ss'   M     SUUS:H  '   UU-  n&[&        R(                  " U&5      $ )	Nr   r   )r   r   r   gh㈵>r   r   r4   )r5   r   r   r  r  r   zerossinmathpilinspaceclamp	unsqueeze	ones_likelerpr   r   lenr   r   r6   )'r-   r  r  r  r  r  r  ovlpfeather_str
first_tiler  t_ht_wr  r   r   r  canvasweightsy_wx_w	sine_mask	flat_maskweight_maskr  r  r  r  r  r  region_hregion_wreal_hreal_wy_end_actualx_end_actual	tile_crop	mask_cropmerged_images'                                          r.   r:   ImageMergeTileList.execute  s   NOqz]
#))""  %55aCdKaA\&F++qQl6G!8))DGGennQ3UZ&[[\C))DGGennQ3UZ&[[\C++ct,C++ct,Cq)CMM!,<<GGJTTUWXI	2I**Y	;GK**a1%5fRK3<V3D/A/%C
O#a=DwHwH4::a=1F4::a=1F"V+L"V+LQ&!34I#Aww$:;I1gl*GL,@!CD	T]H]]DAw|+W\-A1DERE' 4E* !$1'}}\**r1   r<   N)r=   r>   r?   r@   rA   r/   r:   rC   r<   r1   r.   r  r    s(    
 
$ 5+ 5+r1   r  g     o@rgb24)r(  r   	frame_fmt
stream_fmtrgbag    @rgb48lergb48bergba64lergba64ber   	gbrpf32le
gbrapf32le))png8-bitF)r  r  T)r  16-bitF)r  r  T)exr32-bit floatF)r  r  Ttr2   c                     U R                   S   S:X  a1  U SSS24   U SSS24   p![        R                  " [        U5      U/SS9$ U S-  nU R	                  SS	9S
-   S-  S-  n[        R
                  " U S:*  X45      $ )zInverse sRGB EOTF (IEC 61966-2-1). Operates on RGB channels only;
alpha (if present as the 4th channel) is passed through unchanged.r   r   .Nr   r   gףp=
)@r   r  g)\(?gzG?g333333@g?ܵ?)r5   r   r   srgb_to_linearr  wherer  rgbalphalowhighs        r.   r  r  `  s     	wwr{asBQBwZ37Uyy.-u52>> e)CWWW%.36D;;qG|S//r1   g2M?ge8?g?c                 L   U R                   S   S:X  a1  U SSS24   U SSS24   p![        R                  " [        U5      U/SS9$ U S-  S-  n[        R                  " U R                  [        S	9[        -
  [        -  5      [        -   S
-  n[        R                  " U S:*  X45      $ )aZ  Inverse HLG OETF (BT.2100). Maps a non-linear HLG signal in [0, 1] to
*scene*-linear light in [0, 1]. Per BT.2100 Note 5a, this is the correct
transform when converting HLG to a linear scene-light representation
(rather than display-light, which would also involve the HLG OOTF).

Operates on RGB channels only; alpha is passed through unchanged.r   r   .Nr   r   r4   g      @r  g      (@r   )
r5   r   r   hlg_to_linearexpr  _HLG_C_HLG_A_HLG_Br  r  s        r.   r	  r	  s  s     	wwr{asBQBwZ37Uyy-,e4"==
 6S.CIIqww6w*V3v=>G4OD;;qCx++r1   s   PNG


chunk_typerT  c                     [         R                  " X-   5      S-  n[        R                  " S[	        U5      5      U -   U-   [        R                  " SU5      -   $ )zBBuild a single PNG chunk: length | type | data | CRC32(type+data).l    >I)zlibcrc32structpackr  )r  rT  crcs      r.   
_png_chunkr    sH    
**Z&
'*
4C;;tSY'*4t;fkk$PS>TTTr1   keywordtextc                 d    U R                  S5      S-   UR                  SSS9-   n[        SU5      $ )z1tEXt chunk: latin-1 keyword + NUL + latin-1 text.latin-1    rU  errorss   tEXt)rb  r  )r  r  payloads      r.   _png_text_chunkr    s5    nnY''1DKK	R[K4\\Ggw''r1   	png_bytesr   r   c           	         U R                  [        5      (       d  U $ / nUb/  UR                  [        S[        R
                  " U5      5      5        U(       aH  UR                  5        H4  u  pEUR                  [        U[        R
                  " U5      5      5        M6     U(       d  U $ [        R                  " SU SS 5      S   nSU-   S-   nU SU S	R                  U5      -   XS -   $ )
z?Insert ComfyUI prompt/workflow as tEXt chunks right after IHDR.Nr   r        r      r   r1   )
r  _PNG_SIGNATURErc  r  rR  rS  itemsr  unpackr`  )r   r   r   chunkskeyr   ihdr_lengthihdr_ends           r.   inject_png_metadatar,    s    //Foh

60BCD'--/JCMM/#tzz%/@AB 0 --i"o6q9K{"Q&HYh#((6"22Yy5IIIr1   )g{Gz?gQ?g333333?g333333?g333333?gQ?c]F?uV?)g-?g㥛 ?g(\?g/$?gS㥛?gZd;O?r-  r.  )Rec.709Rec.2020	primariesc                 0    [         R                  " S/U Q76 $ )zFSerialize 8 chromaticity floats into the EXR `chromaticities` payload.z<8f)r  r  )r1  s    r.   _pack_chromaticitiesr3    s    ;;u)y))r1   name	attr_typer   c                     U R                  S5      S-   UR                  S5      -   S-   [        R                  " S[        U5      5      -   U-   $ )zCSerialize one EXR header attribute: name\0 type\0 size:int32 value.r@  r  <i)rb  r  r  r  )r4  r5  r   s      r.   _exr_attributer8    sV     	Gw&


7
#	$&-	.
++dCJ
'	( 	r1   	exr_bytes
colorspacec                   ^ [        U 5      S:  d	  U SS S:w  a  U $ SnUb3  U[        SS[        R                  " U5      R	                  S5      5      -  nU(       aL  UR                  5        H8  u  pVU[        US[        R                  " U5      R	                  S5      5      -  nM:     Ub6  S	S	S
S.R                  US	5      nU[        SS[        [        U   5      5      -  nU(       d  U $ SnSn	Sn
U[        U 5      :  a  X   S:w  a  U R                  SU5      nXU R                  SSS9nU R                  SUS-   5      nXS-   U R                  SSS9n[        R                  " SXS-   US-    5      S   nUS-   nU UUU-    nUS:X  a  US:X  a  [        R                  " SU5      n	OUS:X  a  US:X  a  US   n
UU-   nU[        U 5      :  a
  X   S:w  a  M  U	c  U $ SSSSSSSSSSS.
R                  U
S5      nU	u  nnnnUU-
  S-   nUU-   S-
  U-  nUnUS-   nUUS-  -   n[        U5      m[        R                  " SU S3U UU 5      n[        R                  " SU S3/U4S jU 5       Q76 nU SU U-   U UU -   U-   U US -   $ ) u  Insert ComfyUI metadata and color-space info into an EXR header.

Color: EXR pixels are linear by convention. The standard way to describe
their RGB→XYZ relationship is the `chromaticities` attribute. We pick the
primaries that match what the user told us their input was:

  colorspace="sRGB" → Rec. 709 / sRGB primaries (D65)
  colorspace="HDR"  → Rec. 2020 / BT.2100 primaries (D65)

Pixels are always converted to linear scene light upstream (sRGB EOTF
inverse for sRGB; HLG OETF inverse for HDR), so the file content is
scene-linear in the indicated gamut. OpenEXR has no standard transfer-
function attribute (the OpenEXR TSC has discussed adding one but it
doesn't exist), so we don't invent one — `chromaticities` plus the EXR
linear-by-convention rule fully specifies the color.

Prompt/workflow: written as plain `string` attributes using the same keys
(`prompt`, `workflow`, ...) that Comfy uses for PNG tEXt chunks, so the
same readers can pull them out symmetrically.

Implementation note: the chunk-offset table that follows the header stores
*absolute* byte offsets into the file. Inserting N bytes into the header
means every offset must be incremented by N or the file becomes unreadable.
r"  Nr   s   v/1r1   r   stringr@  r/  r0  )sRGBlinearHDRchromaticitiesr   r  r  rU  r  r   r7     
dataWindowbox2iz<iiiicompressionr$         )
r   r   r4   r   r   rA  r      r"  r   <Qc              3   ,   >#    U  H	  oT-   v   M     g 7fr  r<   ).0odeltas     r.   	<genexpr>&inject_exr_metadata.<locals>.<genexpr>.  s     0PKqUKs   )r  r8  rR  rS  rb  r&  rR   r3  _CHROMATICITIESindexrY  r  r'  r  )r9  r   r   r:  new_blobr)  r   primaries_nameposdata_windowrD  name_end	attr_nametype_endr5  r   value_startscanlines_per_blockr>  y_miny_maxr   
num_chunks
header_endtable_startpixel_startold_offsets	new_tablerM  s                               @r.   inject_exr_metadatarc    s+   < 9~Yr].AAHN8Xtzz&7I7P7PQX7YZZ'--/JCsHdjj6G6N6Nw6WXXH 0
   
 #j)
$	 	
 	N !@A
 	

  CKK
I
9>Q#6??7C0(+229Y2O	??7HqL9l84;;Ii;X	}}T9\(Q,#GHKl+kD&89$g)= --7K-'I,F(KD  I
9>Q#6    
c+q  %AuaU]QF..27JJJJq.K
Q.KME--!J<q 19[3UVKa
|1-Q0PK0PQI 	+:
	
J{
+	, 	 KL
!		"r1   
img_tensorfile_format	bit_depthc                    U R                   u  pEnUS:H  n[        XU4   nUS   [        R                  :X  a_  US:X  a  [	        U 5      n OUS:X  a  [        U 5      n U R                  5       R                  5       R                  [        R                  5      n	OgXS   -  R                  SUS   5      n
U
R                  [        R                  5      R                  5       R                  5       R                  US   5      n	[        R                  R                  US5      nX[l        XKl        US   Ul        ['        S	S	5      Ul        [        R*                  R-                  XS
   S9nUS
   US   :w  a  UR/                  US   S9nSUl        UR(                  Ul        [3        UR5                  U5      5      [3        UR5                  S5      5      -   nSR7                  S U 5       5      $ )u  Encode a single HxWxC tensor to PNG or EXR bytes in memory.

For EXR the input is interpreted according to `colorspace` and converted
to scene-linear (EXR's convention) before writing:

  "sRGB"   → input is sRGB-encoded Rec. 709; apply inverse sRGB EOTF.
  "HDR"    → input is HLG-encoded Rec. 2020 (BT.2100); apply inverse HLG
             OETF to get scene-linear, per BT.2100 Note 5a.
  "linear" → input is already scene-linear (Rec. 709 primaries); write
             through unchanged. Use this for renderer/compositor output.

For PNG, colorspace selection does not modify pixels — PNG is delivered
sRGB-encoded and there is no PNG path for wide-gamut HDR in this node.
r   r   r=  r?  r(  r   r  r  r   r  )formatNr1   c              3   8   #    U  H  n[        U5      v   M     g 7fr  )bytes)rK  ps     r.   rN   _encode_image.<locals>.<genexpr>r  s     .gE!HHgs   )r5   _FORMAT_SPECSnpfloat32r  r	  r   numpyastyper  r   r   int32avCodecContextcreater   r   pix_fmtr   	time_base
VideoFramefrom_ndarrayreformatptsr   rb  r`  )rd  re  rf  r:  r   r   num_channels	has_alphaspecimg_npscaledcodecframepacketss                 r.   _encode_imager  =  s   ( #-"2"2F<!I+)<=DG}

"'
3J5 &z2J!'')00< G},33AtG}E5;;'++-335<<T']K
 OO"";4EKL&EMq!nEOMM&&v;6G&HEKD..d<&89EIooEO5<<&'$u||D/A*BBG88.g...r1   c                   X    \ rS rSr\S 5       r\S\S\S\R                  4S j5       r
Srg)	SaveImageAdvancediy  c                    [         R                  " S/ SQSSSS[         R                  R                  SSS	9[         R                  R                  S
SSS9[         R
                  R                  S[         R
                  R                  S[         R                  R                  SSS/SSS9[         R                  R                  SS/SSS9/5      [         R
                  R                  S[         R                  R                  SS/SSS9[         R                  R                  S/ SQSSSS9/5      /SS9/[         R                  R                  [         R                  R                  /SS9	$ )Nr  )savez
save imagezexport imagezoutput imagezwrite imagezSave Image (Advanced)z8Saves the input images to your ComfyUI output directory.r   Basicsr   zThe images to save.)r   r   r   zwThe prefix for the file to save. May include formatting tokens such as %date:yyyy-MM-dd% or %Empty Latent Image.width%.r8  rh  r  rf  r  r  Tr   input_color_spacer=  r  r  )r=  r?  r>  u  Colorspace of the input tensor. The EXR is always written as scene-linear in the matching gamut.
  'sRGB'   — input is sRGB-encoded Rec.709; the inverse sRGB EOTF is applied.
  'HDR'    — input is HLG-encoded Rec.2020 (BT.2100); the inverse HLG OETF is applied to get scene-linear light.
  'linear' — input is already scene-linear (Rec.709 primaries); written through unchanged. Use this for renderer/compositor output.)r   r   r   r   z+The file format in which to save the image.)r   r   )	r   r   r   rK   r    r"   r#   r   r   )r   r%   r&   r'   r   DynamicComboOptionr   r   r   r   r,   s    r.   r/   SaveImageAdvanced.define_schemaz  sn   yy'`0R (x1FG		%%S	    %%..uHHNN;(@S3:T + KHHNN+>39D + J7  ..uHHNN;@P3AD + RHHNN 3(A(.)-
%O + 7 > JC & "-\ II$$bii&=&=>m7
 7	
r1   r   rh  r2   c                 V   US   nUS   nUR                  SS5      n[        R                  " 5       n[        R                  " X'US   R                  S   US   R                  S   5      u  ppnU R
                  R                  nU R
                  R                  n[        R                  (       + n/ n[        U5       H  u  nn[        UXEU5      nU(       a&  US:X  a  [        UX5      nOUS:X  a  [        UXU5      nU	R                  S	[        U5      5      nU S
U
S SU 3n[!        ["        R$                  R'                  UU5      S5       nUR)                  U5        S S S 5        UR+                  UUSS.5        U
S-  n
M     [,        R.                  " SU0S9$ ! , (       d  f       NB= f)Nrh  rf  r  r=  r   r   r  r  r=  r>  r?  .rI  r0  rJ  r   rP   )rR   rN  rP  rO  r5   r   r   r   r   disable_metadatar   r  r,  rc  rU  rV  r]  r^  r_  r`  ra  rc  r   r6   )r-   r   r   rh  re  rf  r:  
output_dirrf  rK  rg  rL  r   r   write_metadatarh  rk  r   encodedr4  rn  fs                         r.   r:   SaveImageAdvanced.execute  s   X&;'	ZZ 3V<
!668
,,VAY__Q-?QRAS 	Jg/ ""

00!222#,V#4L%#E;:NG%'1'6QG E)1'6R\]G##M3|3DEDV1WRL+7Dbggll#5t<dCq  D NN9hWXqLG $5" }}7 344 DCs   F
F(	r<   N)r=   r>   r?   r@   rA   r/   rV  dictr   r6   r:   rC   r<   r1   r.   r  r  y  sB    8
 8
t "5c "54 "5BMM "5 "5r1   r  c                   L    \ rS rSr\S\\\R                        4S j5       r	Sr
g)ImagesExtensioni  r2   c                    #    [         [        [        [        [        [
        [        [        [        [        [        [        [        [        [        [        [         ["        /$ 7fr  )r   rE   rM   r`   ro   r|   r   r   r  r4  r   r  ru  r  r  r  r  r  )selfs    r.   get_node_listImagesExtension.get_node_list  sG      $ %
 	
s   AAr<   N)r=   r>   r?   r@   r	   r   rM  r   	ComfyNoder  rC   r<   r1   r.   r  r    s)    
T$r||*<%= 
 
r1   r  c                     #    [        5       $ 7fr  )r  r<   r1   r.   comfy_entrypointr    s     s   r  )Ir)   rN  rs  rR  r^  rZ  r  rp  rn  r  r   r  comfy.utilsr   	fractionsr   serverr   comfy_api.latestr   r   r   comfy.cli_argsr   typing_extensionsr	   r9  rr  r*   r  r   rE   rM   r`   ro   r|   r   r   r   r  r4  ru  r  r  r  r  r  uint8uint16ro  rm  Tensorr  r  r  r  r	  r%  rj  r  rV  r  r  r,  rP  r   r3  r8  rc  r  r  r  r  r<   r1   r.   <module>r     s     	  	 	         3 3  &ffkk%% @;",, ;BQ",, Q*r|| .R\\ :BLL @$r|| $Nbll Bm",, m`1 1hF",, FR 2<<  F",, @ 6'r|| 'T:*2<< :*zJ+ J+h )."((SZnuv(-"((SYntu(/"))S\nwx(/"))S]nxy.1BJJU`q|$}.1BJJUaq}$~
0ell 
0u|| 
0 
		,U\\ ,ell ,. &U5 U U% U(S ( ( (J5 J$+ JdUYk J^c J0 QP	*E *e *
  U u  "	pp4Kp $;p d
	p
 pn5/5/5/ 5/ 	5/
 5/x_5 _5D
n 
2 r1   