o
    iN                     @   sN  d dl Z d dl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Zd dlZG dd de	jZG dd	 d	e	jZG d
d deZG dd de	jZG dd de	jZG dd de	jZG dd de	jZdZG dd de	jZG dd de	jZG dd de	jZG dd de	jZG dd deZdefd d!ZdS )"    N)override)ComfyExtensionio)HunyuanVideo15SRModel)LatentUpsamplerc                   @   0   e Zd Zedd ZedejfddZeZdS )CLIPTextEncodeHunyuanDiTc              
   C   sD   t jddt jdt jjddddt jjddddgt j gdS )	Nr   advanced/conditioningclipbertT	multilinedynamic_promptsmt5xlnode_idcategoryinputsoutputs)r   SchemaClipInputStringConditioningOutputcls r   8/mnt/c/Users/fbmor/ComfyUI/comfy_extras/nodes_hunyuan.pydefine_schema   s   
z&CLIPTextEncodeHunyuanDiT.define_schemareturnc                 C   s,   | |}| |d |d< t||S )Nr   )tokenizer   
NodeOutputencode_from_tokens_scheduled)r   r
   r   r   tokensr   r   r   execute   s   
z CLIPTextEncodeHunyuanDiT.executeN	__name__
__module____qualname__classmethodr   r   r"   r%   encoder   r   r   r   r      s    
r   c                   @   2   e Zd Zedd ZeddejfddZeZdS )	EmptyHunyuanLatentVideoc                 C   sp   t jdddt jjdddtjddt jjdd	dtjddt jjd
ddtjddt jjdddddgt j gdS )Nr-   zEmpty HunyuanVideo 1.0 Latentzlatent/videowidthP     defaultminmaxstepheight  length         
batch_size   r2   r3   r4   r   display_namer   r   r   r   r   Intr   nodesMAX_RESOLUTIONLatentr   r   r   r   r   r   '   s   z%EmptyHunyuanLatentVideo.define_schemar:   r    c                 C   B   t j|d|d d d |d |d gtj d}t|ddS )Nr0   r:   r;      devicesamplesdownscale_ratio_spacialtorchzeroscomfymodel_managementintermediate_devicer   r"   r   r.   r6   r8   r<   latentr   r   r   r%   8   s   2zEmptyHunyuanLatentVideo.executeNr:   	r'   r(   r)   r*   r   r   r"   r%   generater   r   r   r   r-   &   s    
r-   c                       s6   e Zd Ze fddZeddejfddZ  ZS )EmptyHunyuanVideo15Latentc                    s   t   }d|_d|_|S )NrX   zEmpty HunyuanVideo 1.5 Latent)superr   r   r@   )r   schema	__class__r   r   r   A   s   
z'EmptyHunyuanVideo15Latent.define_schemar:   r    c                 C   rF   )N    r:   r;   r0   rH   rJ   rM   rS   r   r   r   r%   H   s   2z!EmptyHunyuanVideo15Latent.executerU   )	r'   r(   r)   r*   r   r   r"   r%   __classcell__r   r   r[   r   rX   @   s
    rX   c                   @   .   e Zd Zedd ZeddejfddZdS )HunyuanVideo15ImageToVideoc                 C   s   t jddt jdt jdt jdt jjdddtjdd	t jjd
ddtjdd	t jjdddtjdd	t jjdddddt jjdddt j	jdddg	t jj
ddt jj
ddt jj
ddgdS )Nr`   conditioning/video_modelspositivenegativevaer.   r/   r0   r1   r6   r7   r8   !   r:   r;   r<   r=   r>   start_imageToptionalclip_vision_outputr@   rT   r   )r   r   r   r   VaerB   rC   rD   ImageClipVisionOutputr   rE   r   r   r   r   r   P   s$   


z(HunyuanVideo15ImageToVideo.define_schemaNr    c
                 C   s  t j|d|d d d |d |d gtj d}
|d urtj|d | dd||dddd}||d d d d d d d d	f }t j|
j	d
 d|
j	d |
j	d	 |
j	d ftj d}||d d d d d |j	d d d d d f< t j
dd|
j	d |j	d |j	d f|j|jd}d|d d d d d |j	d
 d d d f< t|||d}t|||d}|	d urt|d|	i}t|d|	i}i }|
|d< t|||S )Nr]   r:   r;   r0   rH   bilinearcenter   r      rI   dtype        concat_latent_imageconcat_maskri   rK   )rN   rO   rP   rQ   rR   utilscommon_upscalemovedimr+   shapeonesrI   ru   node_helpersconditioning_set_valuesr   r"   )r   rb   rc   rd   r.   r6   r8   r<   rf   ri   rT   encodedrx   mask
out_latentr   r   r   r%   g   s    2,&6,0,z"HunyuanVideo15ImageToVideo.execute)NNr'   r(   r)   r*   r   r   r"   r%   r   r   r   r   r`   O   s
    
r`   c                   @   r_   )HunyuanVideo15SuperResolutionc                 C   s   t jdt jdt jdt jjdddt jjdddt jjdddt jd	t jjd
ddddddgt jj	ddt jj	ddt jj	d	dgdS )Nr   rb   rc   rd   Trg   rf   ri   rT   noise_augmentationgffffff?rv         ?{Gz?r2   r3   r4   r5   advancedrj   r   r   r   )
r   r   r   r   rk   rl   rm   rE   Floatr   r   r   r   r   r      s   


z+HunyuanVideo15SuperResolution.define_schemaNr    c                 C   s  |d }|j d }	tj|j d |	d d |j d |j d |j d gtj d}
||
d d |	d d|	 d f< d|
d d d|	 d f< |d urtj|dd|j d d	 |j d d	 d
ddd}|	|d d d d d d d df }||
d d d |	d |j d d d d d f< d|
d d |	d df< t
||
|d}t
||
|d}|d urt
|d|i}t
|d|i}t|||S )NrK   r:   r   rr   rs   rn   rH   r0   ro   rp   rq   rx   r   ri   )r}   rN   rO   rP   rQ   rR   rz   r{   r|   r+   r   r   r   r"   )r   rb   rc   rT   r   rd   rf   ri   	in_latentin_channelscond_latentr   r   r   r   r%      s    
> 8&,z%HunyuanVideo15SuperResolution.execute)NNNr   r   r   r   r   r      s
    
r   c                   @   ,   e Zd Zedd ZedejfddZdS )LatentUpscaleModelLoaderc              	   C   s0   t jdddt jjdtddgt j gdS )Nr   zLoad Latent Upscale Modelloaders
model_namelatent_upscale_modelsoptionsr?   )r   r   Combor   folder_pathsget_filename_listLatentUpscaleModelr   r   r   r   r   r      s   z&LatentUpscaleModelLoader.define_schemar    c              	      sN  t d|}tjj|ddd\ }d v rB d jd  d jd  d jd td	d
   D dd}d}t||}|	  n`d v rdd  
 D   d jd  d jd t fddttdd
   D D d}d}t||}|	  n"d v rt|d }t|jtjjtjtjgdd}|  t|S )Nr   T)	safe_loadreturn_metadatazblocks.0.block.0.conv.weightzin_conv.conv.weightr:   zout_conv.conv.weightr   c                 S   $   g | ]}| d r|dr|qS )zblocks.z.block.0.conv.weight
startswithendswith.0kr   r   r   
<listcomp>      $ z4LatentUpscaleModelLoader.execute.<locals>.<listcomp>F)r   out_channelshidden_channels
num_blocksglobal_residual720pzup.0.block.0.conv1.conv.weightc                 S   s    i | ]\}}| d dd|qS )nin_shortcutznin_shortcut.convr:   )replace)r   keyvaluer   r   r   
<dictcomp>   s     z4LatentUpscaleModelLoader.execute.<locals>.<dictcomp>zconv_in.conv.weightzconv_out.conv.weightc                 3   s&    | ]} d | d j d V  qdS )up..block.0.conv1.conv.weightr   N)r}   )r   isdr   r   	<genexpr>   s   $ z3LatentUpscaleModelLoader.execute.<locals>.<genexpr>c                 S   r   )r   r   r   r   r   r   r   r      r   )
z_channelsr   block_out_channels1080pz%post_upsample_res_blocks.0.conv2.biasconfig)allowed_dtypes)ru   )r   get_full_path_or_raiserP   rz   load_torch_filer}   lenkeysr   load_sditemstuplerangejsonloadsr   from_configtorQ   	vae_dtyperN   bfloat16float32load_state_dictr   r"   )r   r   
model_pathmetadatar   
model_typemodelr   r   r   r%      s4   
*
$

z LatentUpscaleModelLoader.executeNr   r   r   r   r   r      s
    
r   c                   @   r   )$HunyuanVideo15LatentUpscaleWithModelc                 C   s~   t jdddt jdt jdt jjdg ddd	t jjd
dddddt jjddddddt jjdddgdgt j gdS )Nr   z*Hunyuan Video 15 Latent Upscale With ModelrT   r   rK   upscale_method)znearest-exactro   areabicubicbislerpro   )r   r2   r.   i   r   i @  rG   r1   r6   i  cropdisabledrp   r   r?   )r   r   r   r   rE   r   rB   r   r   r   r   r   r      s   

	z2HunyuanVideo15LatentUpscaleWithModel.define_schemar    c                 C   s   |dkr|dkrt |S |dkr,td|}tdt|d jd | |d jd  }n)|dkrKtd|}tdt|d jd | |d jd  }n
td|}td|}tj|d |d |d ||}||}t d|	 
 iS )Nr   @   rK   rn   rs   r0   )r   r"   r4   roundr}   rP   rz   r{   resample_latentcpufloat)r   r   rK   r   r.   r6   r   sr   r   r   r%      s   

,
,

 
z,HunyuanVideo15LatentUpscaleWithModel.executeNr   r   r   r   r   r      s
    
r   aZ  <|start_header_id|>system<|end_header_id|>

<image>
Describe the video by detailing the following aspects according to the reference image: 1. The main content and theme of the video.2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.4. background environment, light, style and atmosphere.5. camera angles, movements, and transitions used in the video:<|eot_id|>

<|start_header_id|>user<|end_header_id|>

{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>

c                   @   r   )#TextEncodeHunyuanVideo_ImageToVideoc                 C   sT   t jddt jdt jdt jjddddt jjdd	d
ddddgt j gdS )Nr   r	   r
   ri   promptTr   image_interleaverr   r:   i   zqHow much the image influences things vs the text prompt. Higher number means more influence from the text prompt.)r2   r3   r4   tooltipr   r   )	r   r   r   r   rm   r   rB   r   r   r   r   r   r   r     s$   

z1TextEncodeHunyuanVideo_ImageToVideo.define_schemar    c                 C   s$   |j |t|j|d}t||S )N)llama_templateimage_embedsr   )r!    PROMPT_TEMPLATE_ENCODE_VIDEO_I2Vmm_projectedr   r"   r#   )r   r
   ri   r   r   r$   r   r   r   r%   (  s   z+TextEncodeHunyuanVideo_ImageToVideo.executeNr&   r   r   r   r   r     s    
r   c                   @   s2   e Zd Zedd ZeddejfddZeZdS )HunyuanImageToVideoc                 C   s   t jddt jdt jdt jjdddtjddt jjd	d
dtjddt jjdddtjddt jjdddddt jjdg dddt j	jdddgt jj
ddt jj
ddgdS )Nr   ra   rb   rd   r.   r/   r0   r1   r6   r7   r8   5   r:   r;   r<   r=   r>   guidance_type)v1 (concat)v2 (replace)customT)r   r   rf   rg   rj   rT   r   )r   r   r   r   rk   rB   rC   rD   r   rl   r   rE   r   r   r   r   r   1  s    

z!HunyuanImageToVideo.define_schemaNr    c	                 C   sb  t j|d|d d d |d |d gtj d}	i }
|d urtj|d |d d d d d df dd||dd	dd}||}t j	dd|	j
d
 |j
d |j
d f|j|jd}d|d d d d d |j
d d d d f< |dkr|||d}n%|dkrddi}||	d d d d d |j
d
 f< ||
d< n|dkrd|i}t||}|	|
d< t||
S )Nr0   r:   r;   rG   rH   rq   rn   ro   rp   rr   rs   rt   rv   r   r   rw   r   guiding_frame_index
noise_maskr   
ref_latentrK   )rN   rO   rP   rQ   rR   rz   r{   r|   r+   r~   r}   rI   ru   r   r   r   r"   )r   rb   rd   r.   r6   r8   r<   r   rf   rT   r   rx   r   condr   r   r   r%   F  s$   2@
0, 
zHunyuanImageToVideo.executeNr&   r   r   r   r   r   0  s    
r   c                   @   r,   )	EmptyHunyuanImageLatentc                 C   sX   t jddt jjdddtjddt jjdddtjddt jjd	d
d
ddgt j gdS )Nr   rT   r.   i   r   r]   r1   r6   r<   r:   r=   r>   r   rA   r   r   r   r   r   d  s   z%EmptyHunyuanImageLatent.define_schemar:   r    c                 C   s2   t j|d|d |d gtj d}td|iS )Nr   r]   rH   rK   rM   )r   r.   r6   r<   rT   r   r   r   r%   s  s   $zEmptyHunyuanImageLatent.executeNrU   rV   r   r   r   r   r   c  s    
r   c                   @   r   )HunyuanRefinerLatentc                 C   sh   t jdt jdt jdt jdt jjddddd	d
dgt jjddt jjddt jjddgdS )Nr   rb   rc   rT   r   g?rv   r   r   Tr   rj   r   )r   r   r   r   rE   r   r   r   r   r   r   r   |  s   


z"HunyuanRefinerLatent.define_schemar    c                 C   sx   |d }t |||d}t |||d}i }tj|jd d|jd |jd |jd gtj d|d< t	|||S )	NrK   r   r   r]   r   rs   rn   rH   )
r   r   rN   rO   r}   rP   rQ   rR   r   r"   )r   rb   rc   rT   r   r   r   r   r   r%     s   :zHunyuanRefinerLatent.executeNr   r   r   r   r   r   {  s
    
r   c                   @   s(   e Zd Zedeeej  fddZdS )HunyuanExtensionr    c                    s   t ttttttttt	t
gS r   )r   r   r-   rX   r`   r   r   r   r   r   r   )selfr   r   r   get_node_list  s   zHunyuanExtension.get_node_listN)	r'   r(   r)   r   listtyper   	ComfyNoder   r   r   r   r   r     s    r   r    c                      s   t  S r   )r   r   r   r   r   comfy_entrypoint  s   r   )rC   r   rN   comfy.model_managementrP   typing_extensionsr   comfy_api.latestr   r   !comfy.ldm.hunyuan_video.upsamplerr   %comfy.ldm.lightricks.latent_upsamplerr   r   r   r   r   r-   rX   r`   r   r   r   r   r   r   r   r   r   r   r   r   r   r   <module>   s2    2,1( 3