CoolFace
Apppublic

fred-dev/comfy_ui_ali

sourceHugging Facemitupdated 1y agoView on Hugging Face
0likes
model_base.cpython-312.pyc319 linesDownload Raw Back to __pycache__
1�

2���gt����dZddlZddlZddlmZmZddlmZddlm	Z	ddl3mZddlm
Z
ddlmZddlZddlZddlZddlZddlZddlZddlZddlZddlZddlZddlZddlZddlZddlZddlZddl Zdd	l!m"Z"d4dl#m$Z$ddl%Zddl&Z&ddl'm(Z(e(rdd
l)m*Z*Gd�de"�Z+ddl,m-Z-m.Z.m/Z/m0Z0m1Z1m2Z2m3Z3d�Z4Gd�dejjjl�Z7dLd�Z8Gd�de7�Z9d�Z:Gd�de7�Z;Gd�de7�Z<Gd�de7�Z=Gd�de=�Z>Gd �d!e=�Z?Gd"�d#e7�Z@Gd$�d%e7�ZAGd&�d'�ZBGd(�d)eBe7�ZCGd*�d+eBe<�ZDGd,�d-e7�ZEGd.�d/e7�ZFGd0�d1e7�ZGGd2�d3e7�ZHGd4�d5e7�ZIGd6�d7e7�ZJGd8�d9e7�ZKGd:�d;e7�ZLGd<�d=e7�ZMGd>�d?e7�ZNGd@�dAe7�ZOGdB�dCeO�ZPGdD�dEeO�ZQGdF�dGe7�ZRGdH�dIe7�ZSGdJ�dKe7�ZTy)Ma�5    This file is part of ComfyUI.6    Copyright (C) 2024 Comfy7 8    This program is free software: you can redistribute it and/or modify9    it under the terms of the GNU General Public License as published by10    the Free Software Foundation, either version 3 of the License, or11    (at your option) any later version.12 13    This program is distributed in the hope that it will be useful,14    but WITHOUT ANY WARRANTY; without even the implied warranty of15    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the16    GNU General Public License for more details.17 18    You should have received a copy of the GNU General Public License19    along with this program.  If not, see <https://www.gnu.org/licenses/>.20�N)�	UNetModel�Timestep)�StageC)�StageB)�CLIPEmbeddingNoiseAugmentation)� ImageConcatWithNoiseAugmentation)�OpenAISignatureMMDITWrapper)�Enum�)�utils)�
TYPE_CHECKING)�ModelPatcherc�,�eZdZdZdZdZdZdZdZdZ	dZ21y	)22�	ModelTyper�������N)�__name__�23__module__�__qualname__�EPS�V_PREDICTION�V_PREDICTION_EDM�STABLE_CASCADE�EDM�FLOW�V_PREDICTION_CONTINUOUS�FLUX���[/Users/fredrodrigues/Library/CloudStorage/Dropbox/Code/comfyui_hardware/comfy/model_base.pyrr4s,��24�C��L����N�25�C��D����Dr$r)rrr�ModelSamplingDiscrete�ModelSamplingContinuousEDM�StableCascadeSampling�ModelSamplingContinuousVc���t}|tjk(rt}�n*|tjk(rt}�n|tjk(r
t}t26}n�|tjk(r5tjj}tjj}n�|tjk(r
t}t}n�|tjk(r
t}t27}ng|tjk(r
t}t}nG|tj k(r4tjj}tjj"}Gd�d|�}||�S)Nc��eZdZy)�%model_sampling.<locals>.ModelSamplingN)rrrr#r$r%�
ModelSamplingr,\s��r$r-)r&rrrrr'r �comfy�model_sampling�CONST�ModelSamplingDiscreteFlowrr(rr!r)r"�ModelSamplingFlux)�model_config�28model_type�s�cr-s     r%r/r/Bs���A��Y�]�]�"���	�y�-�-�	-���	�y�1�1�	1���&��	�y�~�~�	%�� � �&�&��� � �:�:��	�y�/�/�	/���!��	�y�}�}�	$���&��	�y�8�8�	8���$��	�y�~�~�	%�� � �&�&��� � �2�2��
��1�
���&�&r$c���eZdZejdef�fd�	Zdddifd�Zdddifd�Zd�Z	d�Z29d�Zd�Zd	�Z
dd30�Zd�Zd�Zdd
�Zd�Zd�Zd�Z�xZS)�	BaseModelNc���t�|��|j}|j|_||_|j31|_||_d|_|jdd��s(|j�d|jjdd�}tjj|jdd�|j32||j��}n|j}|di|�||d���|_tj j#�r?|jj%t&j(��t+j,d�t+j.d	j1|j3�|j33��||_t7||�|_|jd34d�|_|j8�d|_d|_t+j.d
j1|j<��t+j,dj1|j8��|j>|_y)N�disable_unet_model_creationF�fp8�dtype)�fp8_optimizations�35scaled_fp8)�device�36operations)�
memory_formatz,using channels last mode for diffusion modelz&model weight dtype {}, manual cast: {}�adm_in_channelsrr#z
model_type {}zadm {}) �super�__init__�unet_config�
latent_formatr3�manual_cast_dtyper?�current_patcher�get�custom_operations�
optimizationsr.�ops�pick_operationsr>�diffusion_model�model_management�force_channels_last�to�torch�
channels_last�logging�debug�info�format�	get_dtyper4r/�adm_channels�concat_keys�name�memory_usage_factor)	�selfr3r4r?�37unet_modelrEr;r@�	__class__s	        �r%rDzBaseModel.__init__cs����
����"�.�.��)�7�7���(���!-�!?�!?������/3������<�e�D��-�-�5�"�0�0�4�4�U�E�B��"�Y�Y�6�6�{���w�PT�7U�W[�Wm�Wm�BE�R^�Ri�Ri�6�j�38�)�;�;�39�#-�#b��#b�F�Wa�#b�D� ��%�%�9�9�;��$�$�'�'�e�6I�6I�'�J��
�
�L�M��L�L�A�H�H����IY�[_�[q�[q�r�s�$���,�\�:�F���'�O�O�,=�t�D������$� !�D��������_�+�+�J�O�O�<�=��
�
�h�o�o�d�&7�&7�8�9�#/�#C�#C�� r$c	��tjjj|j|tjjtjjj|��j||||||fi|��S�N)	r.�patcher_extension�WrapperExecutor�new_class_executor�_apply_model�get_all_wrappers�40WrappersMP�APPLY_MODEL�execute)r]�x�t�c_concat�c_crossattn�control�transformer_options�kwargss        r%�apply_modelzBaseModel.apply_model�s~���&�&�6�6�I�I������#�#�4�4�U�5L�5L�5W�5W�5c�5c�ex�y�41��'�!�Q��+�w�8K�	W�PV�	W�	Wr$c���|}|jj||�}	|�tj|	g|gzd��}	|}42|j	�}|j43�|j44}|	j
|�}	|jj|�j�}|45�|46j
|�}47i}|D]c}
||
}t|d�rK|jtjk7r.|jtjk7r|j
|�}|||
<�e|j|fd|i|��}|j|	|f|48||d�|��j�}|jj|||�S)Nr��dimr<rj)�contextrnro)r/�calculate_inputrR�catrXrGrQ�timestep�float�hasattrr<�int�long�process_timesteprN�calculate_denoised)r]rjrkrlrmrnrorp�sigma�xcrur<�extra_conds�o�extra�model_outputs                r%rezBaseModel._apply_model�sy����
�
 �
 �
0�
0���
:�������B�4�8�*�,�!�4�B������ ���!�!�-��*�*�E�
�U�U�5�\�����(�(��+�1�1�3�����j�j��'�G����A��1�I�E��u�g�&��;�;�%�)�)�+����u�z�z�0I�!�H�H�U�O�E�"�K��N��
"�D�!�!�!�8�q�8�K�8��+�t�+�+�B��M�7�G�i|�M�AL�M�S�S�U���"�"�5�5�e�\�1�M�Mr$c��|Srar#)r]rxrps   r%r}zBaseModel.process_timestep�s���r$c�.�|jjSra)rNr<)r]s r%rXzBaseModel.get_dtype�s���#�#�)�)�)r$c��yrar#�r]rps  r%�49encode_admzBaseModel.encode_adm�s��r$c��t|j�dkD�r�g}|jd|jdd��}|jdd�}|�|jdd�}n|j|�}|jdd�}|d}|jdd|jddk7r�tj||jd	|jd50dd�}|jd
k(r�|jd|jdkrPtjjj|ddddd|jd|jdz51fdd�}n|dd�dd�d|jd�f}tj||jd�}|��t|j�t|j�k(r|dd�dd�f}|jdz52}|jdt|j|d�z�}|jd53d|jd54dk7r3tj||jd	|jd55dd�}tj|j�|jd�}|jD�]T}|�u|dk(r!|j!|j#|��n�|dk(r!|j!|j#|��n�|dk(r�|j!d|j#|�z56�n�|dk(r.|j!tj$|�dd�dd�f�nX|dk(r!|j!|j'|��n2|dk(r-|j!tj(|�dd�dd�f�|dk(s��
|�"|j!|j#|����1|j!tj(|����Wtj*|d��}	|	Sy)Nr�concat_mask�denoise_mask�concat_latent_image�latent_image�noiser?r�����������bilinear�centerr������constantr)r�r�mask�masked_image�
mask_inverted��?�concat_imagers)�lenrZrI�process_latent_in�shaper�common_upscale�ndimrR�nn�57functional�pad�resize_to_batch_size�reshape�tuple�round�appendrQ�	ones_like�blank_inpaint_image_like�58zeros_likerw)59r]rp�cond_concatr�r�r�r?�num_dim�ck�datas60          r%�concat_condzBaseModel.concat_cond�s����t��� �1�$��K�!�:�:�m�V�Z�Z��PT�5U�V�L�"(�*�*�-B�D�"I��"�*�&,�j�j���&F�#�&*�&<�&<�=P�&Q�#��J�J�w��-�E��H�%�F�"�(�(���,����A�B��?�&+�&:�&:�;N�PU�P[�P[�\^�P_�af�al�al�mo�ap�r|�G�'H�#��:�:��?�*�0�0��4�u�{�{�2��F�.3�h�h�.A�.A�.E�.E�FY�\]�_`�bc�ef�hi�kp�kv�kv�wy�kz�~Q�~W�~W�XZ�~[�l[�\\�^h�jk�/l�+�.A�!�Q�HX����UW��HX�BX�.Y�+�"'�"<�"<�=P�RW�R]�R]�^_�R`�"a���'��|�)�)�*�c�%�+�+�.>�>�#/��2�A�2��#6�L��*�*�q�.��+�3�3�G�e�L�DV�DV�X_�W_�W`�Da�>b�4b�c���%�%�b�c�*�e�k�k�"�#�.>�>�#(�#7�#7��e�k�k�RT�o�W\�Wb�Wb�ce�Wf�hr�t|�#}�L�$�9�9�,�:L�:L�:N�PU�P[�P[�\]�P^�_���&�&���+��V�|�#�*�*�<�?�?�6�+B�C��~�-�#�*�*�+>�+A�+A�&�+I�J���.�#�*�*�3�����1H�+H�I��V�|�#�*�*�5�?�?�5�+A�!�R�a�R�%�+H�I��~�-�#�*�*�4�+H�+H��+O�P���.�#�*�*�5�+;�+;�E�+B�1�b�q�b�5�+I�J���'�*�6�#�*�*�+>�+A�+A�&�+I�J�#�*�*�5�+;�+;�E�+B�C�''�(�9�9�[�a�0�D��K�r$c�&�i}|jdi|��}|�"tjj|�|d<|jdi|��}|�"tjj|�|d<|j
dd�}|�"tjj|�|d<|j
dd�}|�"tjj|�|d<|j
dd�}|�"tjj|�|d<|S)	Nrl�y�61cross_attnrm�cross_attn_controlnet�crossattn_controlnet�noise_concatr#)r�r.�conds�CONDNoiseShaper��CONDRegularrI�
CONDCrossAttn)r]rp�outr��admr��cross_attn_cnetrls        r%r�zBaseModel.extra_conds�s�����&�d�&�&�0��0���"�#�k�k�8�8��E�C�62�O��d�o�o�'��'���?��{�{�.�.�s�3�C��H��Z�Z��d�3�63��!�!&���!:�!:�:�!F�C�
�� �*�*�%<�d�C���&�*/�+�+�*C�*C�O�*T�C�&�'��:�:�n�d�3����#�k�k�8�8��B�C�64�O��65r$c���i}t|j��}|D]4}|j|�s�|j|�||t	|�d<�6|j66j
|�}|jj|d��\}}t	|�dkDr$tjdj|��t	|�dkDr$tjdj|��~|S)NF)�strictrzunet missing: {}zunet unexpected: {})�list�keys�67startswith�popr�r3�process_unet_state_dictrN�load_state_dictrT�warningrW)r]�sd�unet_prefix�to_loadr��k�m�us        r%�load_model_weightszBaseModel.load_model_weightss������B�G�G�I����A��|�|�K�(�02���q�	���#�k�*�+�,�-���#�#�;�;�G�D���#�#�3�3�G�E�3�J���1��q�6�A�:��O�O�.�5�5�a�8�9��q�6�A�:��O�O�1�8�8��;�<���r$c�8�|jj|�Sra)rF�68process_in�r]�latents  r%r�zBaseModel.process_latent_ins���!�!�,�,�V�4�4r$c�8�|jj|�Sra)rF�process_outr�s  r%�process_latent_outzBaseModel.process_latent_outs���!�!�-�-�f�5�5r$c��g}|�*|j|jj|��|�*|j|jj|��|�*|j|jj	|��|j69j
�}|jj�.tjg|jj��|d<|jj|�}|jtjk(rtjg�|d<|D]}|j|��|S)N)r<r>�v_pred)r�r3�"process_clip_state_dict_for_saving�!process_vae_state_dict_for_saving�)process_clip_vision_state_dict_for_savingrN�70state_dictr>rR�tensor�"process_unet_state_dict_for_savingr4rr�update)r]�clip_state_dict�vae_state_dict�clip_vision_state_dict�	extra_sds�unet_state_dictr�s       r%�state_dict_for_savingzBaseModel.state_dict_for_savings���	��&����T�.�.�Q�Q�Ra�b�c��%����T�.�.�P�P�Q_�`�a�!�-����T�.�.�X�X�Yo�p�q��.�.�9�9�;�����'�'�3�,1�L�L��4�CT�CT�C_�C_�,`�O�L�)��+�+�N�N��_���?�?�i�4�4�4�(-���R�(8�O�H�%��B��"�"�2�&���r$c�&�d|_d�}||_y)N)r�r�c��tj|�}|dd�dfxxdzcc<|dd�dfxxdzcc<|dd�dfxxdzcc<|dd�dfxxdzcc<|S)	Nrg��HP�?rg�X���rg��e�c]�?rgD�l����?)rRr�)r��blank_images  r%r�z7BaseModel.set_inpaint.<locals>.blank_inpaint_image_like3sf���/�/�,�7�K���!����&����!����'����!����&����!����&���r$)rZr�)r]r�s  r%�set_inpaintzBaseModel.set_inpaint1s��3���	�)A��%r$c	��|jj|j|jdgdgt	|j�dz71zz�||�S�Nrr)r/�
noise_scalingr�r�r��r]rr�r�rps     r%�scale_latent_inpaintzBaseModel.scale_latent_inpaint=s^���"�"�0�0�������A��?O�ST�RU�Y\�]b�]h�]h�Yi�lm�Ym�Rn�?n�1o�qv�yE�F�	Fr$c���tjj�stjj�r{|j	�}|j72�|j73}|dt
j|dd�z}|tjj|�zdz|jzdzS|dt
j|dd�z}|dz|jzdzS)Nrrg{�G�z�?ig333333�?)74r.rO�xformers_enabled�!pytorch_attention_flash_attentionrXrG�math�prod�75dtype_sizer\)r]�input_shaper<�areas    r%�memory_requiredzBaseModel.memory_required@s����!�!�2�2�4��8N�8N�8p�8p�8r��N�N�$�E��%�%�1��.�.���q�>�D�I�I�k�!�"�o�$>�>�D��5�1�1�<�<�U�C�C�d�J�T�Me�Me�e�ju�v�v��q�>�D�I�I�k�!�"�o�$>�>�D��4�K�$�":�":�:�{�K�Kr$)��NNN)rrrrrrrDrqrer}rXr�r�r�r�r�r�r�r�r�r��
__classcell__�r_s@r%r8r8bsz���09�
�
�d�W`�D�B*.�4��ce�W�+/�D�$�df�N�:�*��7�r�0�"5�6��076A�F�Lr$r8c	���g}g}g}|D]�}|djD]�}	|d}77|d}t|jdz78|z�}||	j|�t	j79|g|��|��\}
}t	j|
|fd�|80z}|j|81�|j|�|j|�����t|�dkDr�t	j|�jd�}|}t|jdz82|z�}||dd�d|jj�ft	j83|g|����\}
}t	j|
|fd�}S)	N�clip_vision_output�strength�noise_augmentationr�r?��noise_level�seedr)r�)
�image_embedsr��max_noise_levelrQrRr�rwr�r��stack�sum�84time_embedrt)�unclip_conditioningr?�noise_augmentor�noise_augment_merger�85adm_inputs�weights�	noise_aug�unclip_cond�adm_cond�weight�
noise_augmentr��c_adm�noise_level_emb�adm_outs                r%�86unclip_admrNs����J��G��I�*��#�$8�9�F�F�H� ��,�F�'�(<�=�M���!@�!@�1�!D�
� U�V�K�%4�X�[�[��5H�V[�Vb�Vb�do�cp�y�WA�HL�&M�"�E�?��i�i��� 8�!�<�v�E�G��N�N�6�"����]�+����g�&�G�+��9�~����+�+�j�)�-�-�a�0��+�
��_�<�<�q�@�M�Q�R��!0���<[�_�=W�=W�=[�=[�<[�9[�1\�jo�jv�jv�yD�xE�NT�kU�"V�����)�)�U�O�4�a�8���Nr$c�>��eZdZejdf�fd�	Zd�Z�xZS)�87SD21UNCLIPNc�J��t�|�|||��tdi|��|_y)Nr�r#)rCrDrr)r]r3�noise_aug_configr4r?r_s     �r%rDzSD21UNCLIP.__init__gs(���
����z�&��A�=�Q�@P�Q��r$c88���|jdd�}|d}|�!tjd|jf�St	|||j89|jdd�|jdd�dz90�S)	Nrr?r�unclip_noise_augment_mergeg�������?rr�91)rIrR�zerosrYrr)r]rprr?s    r%r�zSD21UNCLIP.encode_admks���$�j�j�)>��E����!���&��;�;��4�#4�#4�5�6�6��1�6�4�;O�;O�QW�Q[�Q[�\x�z~�Q�BH�BL�BL�MS�UV�BW�Z\�B\�]�
]r$)rrrrrrDr�r�r�s@r%rrfs���BK�BX�BX�ae�R�]r$rc	��d|vr>t|jdd�|d||jdd�dz92��dd�dd�fS|dS)	Nrr?rrr)r��
pooled_output)rrI)�argsrs  r%�sdxl_pooledr ss|����$��$�(�(�#8�$�?��h��Q`�gk�go�go�pv�xy�gz�}�g�A�BC�DI�EI�DI�BI�J�	J��O�$�$r$c�>��eZdZejdf�fd�	Zd�Z�xZS)�SDXLRefinerNc�v��t�|�|||��td�|_t	diddd�dd���|_y�	Nr��i��squaredcos_cap_v2)�	timesteps�
beta_scheduler)�noise_schedule_config�timestep_dimr#�rCrDr�embedderrr�r]r3r4r?r_s    �r%rDzSDXLRefiner.__init__z�U���
����z�&��A� ��
��
�=� l�hl�@S�[T�fj�Ak� l��r$c�L�t||j�}|jdd�}|jdd�}|jdd�}|jdd�}|jdd�d	k(r|jd93d�}n|jd94d�}g}|j|j	tj|g���|j|j	tj|g���|j|j	tj|g���|j|j	tj|g���|j|j	tj|g���tjtj|��jd�
�j|jdd�}	tj|j|	j�|	fd�
�S)N�width��height�crop_wr�crop_h�prompt_typer��negative�aesthetic_scoreg@rrsr�r rrIr�r,rR�Tensor�flattenrw�	unsqueeze�repeatr�rQr?)95r]rp�clip_pooledr0r2r3r4r7r��flats96          r%r�zSDXLRefiner.encode_adms���!�&�$�*>�*>�?���97�98�7�C�(�����H�c�*�����H�a�(�����H�a�(���:�:�m�R�(�J�6�$�j�j�):�C�@�O�$�j�j�):�A�>�O����99�100�4�=�=����v�h�!7�8�9��101�102�4�=�=����u�g�!6�7�8��103�104�4�=�=����v�h�!7�8�9��105�106�4�=�=����v�h�!7�8�9��107�108�4�=�=�����.?�!@�A�B��}�}�U�Y�Y�s�^�,�6�6�1�6�=�D�D�[�EV�EV�WX�EY�[\�]���y�y�+�.�.����5�t�<�!�D�Dr$�rrrrrrDr�r�r�s@r%r"r"ys���09�
�
�d�l�109Er$r"c�>��eZdZejdf�fd�	Zd�Z�xZS)�SDXLNc�v��t�|�|||��td�|_t	diddd�dd���|_yr$r+r-s    �r%rDz
SDXL.__init__�r.r$c��t||j�}|jdd�}|jdd�}|jdd�}|jdd�}|jd|�}|jd|�}g}	|	j|j	tj|g���|	j|j	tj|g���|	j|j	tj|g���|	j|j	tj|g���|	j|j	tj|g���|	j|j	tj|g���tjtj|	��jd�	�j|jdd110�}111tj|j|112j�|113fd114�	�S)Nr0r1r2r3rr4�target_width�
target_heightrsrr8)r]rpr=r0r2r3r4rDrEr�r>s           r%r�zSDXL.encode_adm�s���!�&�$�*>�*>�?���115�116�7�C�(�����H�c�*�����H�a�(�����H�a�(���z�z�.�%�8���117�118�?�F�;�
����119�120�4�=�=����v�h�!7�8�9��121�122�4�=�=����u�g�!6�7�8��123�124�4�=�=����v�h�!7�8�9��125�126�4�=�=����v�h�!7�8�9��127�128�4�=�=����}�o�!>�?�@��129�130�4�=�=����|�n�!=�>�?��}�}�U�Y�Y�s�^�,�6�6�1�6�=�D�D�[�EV�EV�WX�EY�[\�]���y�y�+�.�.����5�t�<�!�D�Dr$r?r�s@r%rArA�s���09�
�
�d�l�131Er$rAc�D��eZdZejdf�fd�	Zd�Zd�Z�xZS)�SVD_img2vidNc�J��t�|�|||��td�|_y)Nr�r%)rCrDrr,r-s    �r%rDzSVD_img2vid.__init__�s"���
����z�&��A� ��
��
r$c�$�|jdd�dz132}|jdd�}|jdd�}g}|j|jtj|g���|j|jtj|g���|j|jtj|g���tj133tj|��jd��}|S)	N�fpsrr�motion_bucket_id��augmentation_levelrrs)rIr�r,rRr9r:rwr;)r]rp�fps_idrK�augmentationr�r>s       r%r�zSVD_img2vid.encode_adm�s������E�1�%��)��!�:�:�&8�#�>���z�z�"6��:�����134�135�4�=�=����v�h�!7�8�9��136�137�4�=�=����/?�.@�!A�B�C��138�139�4�=�=����|�n�!=�>�?��}�}�U�Y�Y�s�^�,�6�6�1�6�=���r$c�6�i}|jdi|��}|�"tjj|�|d<|j	dd�}|j	dd�}|�tj|�}|jdd|jddk7r3tj||jd|jddd�}tj||jd	�}tjj|�|d140<|j	dd�}|�"tjj|�|d<d
|vr%tjj|d
�|d<tjj|jd	�|d<|S)Nr�r�r�rr�r�r�r�rrlr�rm�time_conditioning�time_context�num_video_framesr#)r�r.r�r�rIrRr�r�rr�r�r�r��CONDConstant)r]rpr�r�r�r�r�s       r%r�zSVD_img2vid.extra_conds�sg�����d�o�o�'��'���?��{�{�.�.�s�3�C��H��z�z�"7��>���141�142�7�D�)���� �+�+�E�2�L����a�b�!�U�[�[���_�4� �/�/��e�k�k�"�o�u�{�{�[]��`j�lt�u�L��1�1�,����A��O���+�+�4�4�\�B��J���Z�Z��d�3�143��!�!&���!:�!:�:�!F�C�
���&�(�"'�+�+�";�";�F�CV�<W�"X�C���"'�+�+�":�":�5�;�;�q�>�"J�����144r$)	rrrrrrDr�r�r�r�s@r%rGrG�s���09�0J�0J�SW�&��r$rGc��eZdZd�Zy)�SV3D_uc	�,�|jdd�}g}|j|jtjtj145|g����tjtj|��jd��}|S)NrMrrs)rIr�r,rRr:r9rwr;)r]rprOr�r>s     r%r�zSV3D_u.encode_adm�sj���z�z�"6��:�����146�147�4�=�=����u�|�|�\�N�/K�!L�M�N��}�}�U�Y�Y�s�^�,�6�6�1�6�=���r$N)rrrr�r#r$r%rVrV�s��r$rVc�>��eZdZejdf�fd�	Zd�Z�xZS)�SV3D_pNc�J��t�|�|||��td�|_y)Nr��)rCrDr�embedder_512r-s    �r%rDzSV3D_p.__init__�s#���
����z�&��A�$�S�M��r$c�B��|jdd�}|jdd�}|jdd�}|jdd��g}|j|jtjtj148|g����|j|j
tjtjtjdtj149|g�z150�d����|j|j
tjtjtjtj151|g��d����tt�fd�|��}tj|d	�152�S)NrMr�	elevation�azimuthr��Zg�v@c�J��tj|�jd�S)Nr)rr�r�)�ar�s �r%�<lambda>z#SV3D_p.encode_adm.<locals>.<lambda>�s����!;�!;�A�u�{�{�1�~�!Nr$rrs)rIr�r,rRr:r9r\�deg2rad�fmodr��maprw)r]rprOr^r_r�r�s      @r%r�zSV3D_p.encode_adm�s"����z�z�"6��:���J�J�{�A�.�	��*�*�Y��*���153�154�7�D�)�����155�156�4�=�=����u�|�|�\�N�/K�!L�M�N��157�158�4�$�$�U�]�]�5�:�:�e�m�m�B�QV�Q]�Q]�_h�^i�Qj�Lj�>k�mr�3s�%t�u�v��159�160�4�$�$�U�]�]�5�:�:�e�m�m�E�L�L�Za�Yb�Lc�>d�fk�3l�%m�n�o��3�N�PS�T�U���y�y��!�$�$r$)rrrrrrDr�r�r�s@r%rYrY�s���09�0J�0J�SW�*�%r$rYc�B��eZdZejdddf�fd�	Zd�Z�xZS)�Stable_Zero123Nc�r��t�|�|||��tjjj|jd|jd|j�|��|_|jjj|�|jjj|�y)Nr�rr)r<r?)rCrDr.rL�manual_cast�Linearr�rX�
cc_projectionr�copy_�bias)r]r3r4r?�cc_projection_weight�cc_projection_biasr_s      �r%rDzStable_Zero123.__init__�s����
����z�&��A�"�Y�Y�2�2�9�9�:N�:T�:T�UV�:W�Ym�Ys�Ys�tu�Yv�C�M�M�O�X^�9�_������!�!�'�'�(<�=������%�%�&8�9r$c�`�i}|jdd�}|jdd�}|�tj|�}|jdd|jddk7r3t	j161||jd|jddd�}t	j||jd�}tjj|�|d	<|jd162d�}|�E|jddk7r|j|�}tjj|�|d<|S)
Nr�r�rr�r�r�r�rrlr�r1rm)rIrRr�r�rr�r�r.r�r�rlr�)r]rpr�r�r�r�s      r%r�zStable_Zero123.extra_condss�����z�z�"7��>���163�164�7�D�)���� �+�+�E�2�L����a�b�!�U�[�[���_�4� �/�/��e�k�k�"�o�u�{�{�[]��`j�lt�u�L��1�1�,����A��O���+�+�4�4�\�B��J���Z�Z��d�3�165��!�����#�s�*�!�/�/�166�;�167�!&���!:�!:�:�!F�C�
���168r$�rrrrrrDr�r�r�s@r%rhrh�s���09�
�
�d�ae�z~�:�r$rhc�>��eZdZejdf�fd�	Zd�Z�xZS)�
SD_X4UpscalerNc�T��t�|�|||��tddd�d��|_y)Nr�g-C��6?g{�G�z�?)�linear_start�169linear_endi^)r)r)rCrDrrr-s    �r%rDzSD_X4Upscaler.__init__s=���
����z�&��A�?�gm�~B�WC�UX� Y��r$c��i}|jdd�}|jdd�}|jdd�}|d}|ddz170}t|jj|z�}|�t	j171|�dd�dd�f}|jd	d|jd	dk7rBtj|j|�|jd172|jddd
�}t	j|g|��}|dkDr&|j|j|�||��\}}tj||jd�}tjj|�|d<tjj|�|d<|jdd�}	|	�"tjj!|	�|d<|S)Nr�r�r��r?rrrrr�r�r�r�r�rr�rlr�r�rm)rIr�rrrRr�r�rr�rQr�r�r.r�r�r�r�)173r]rpr��imager�rr?rr�r�s174          r%r�zSD_X4Upscaler.extra_condss������175�176�>�4�0���177�178�7�D�)���179�180�#7��=�
���!���f�~��"���T�1�1�A�A�]�R�S���=��$�$�U�+�A�b�q�b�D�1�E��;�;�q�r�?�e�k�k�!�"�o�-��(�(����&�)9�5�;�;�r�?�E�K�K�XZ�O�]g�iq�r�E��l�l�K�=��@���1��!%�!5�!5�e�h�h�v�6F�T_�fj�!5�!k��E�;��*�*�5�%�+�+�a�.�A���+�+�4�4�U�;��J���;�;�*�*�;�7��C���Z�Z��d�3�181��!�!&���!:�!:�:�!F�C�
���182r$�rrrrrrDr�r�r�s@r%rtrts���09�0F�0F�t�Y�r$rtc��eZdZd�Zy)�IP2Pc��|jdd�}|jdd�}|d}|�tj|�}|jdd|jddk7rBt	j183|j
|�|jd|jddd�}t	j||jd	�}|j|�S)184Nr�r�r?rr�r�r�r�r)	rIrRr�r�rr�rQr��process_ip2p_image_in)r]rprzr�r?s     r%r�zIP2P.concat_cond>s����185�186�0�$�7���187�188�7�D�)����!���=��$�$�U�+�E��;�;�q�r�?�e�k�k�!�"�o�-��(�(����&�)9�5�;�;�r�?�E�K�K�XZ�O�]g�iq�r�E��*�*�5�%�+�+�a�.�A���)�)�%�0�0r$N)rrrr�r#r$r%r}r}=s��1r$r}c�8��eZdZejdf�fd�	Z�xZS)�SD15_instructpix2pixNc�:��t�|�|||��d�|_y)Nr�c��|Srar#�rzs r%rcz/SD15_instructpix2pix.__init__.<locals>.<lambda>Ps��5r$)rCrDrr-s    �r%rDzSD15_instructpix2pix.__init__Ns���
����z�&��A�%8��"r$�rrrrrrDr�r�s@r%r�r�Ms���09�
�
�d�9�9r$r�c�8��eZdZejdf�fd�	Z�xZS)�SDXL_instructpix2pixNc�r��t�|�|||��|tjk(r	d�|_yd�|_y)Nr�c�\�tjj�j|�Sra)r.�latent_formatsrAr�r�s r%rcz/SDXL_instructpix2pix.__init__.<locals>.<lambda>Ws��u�7K�7K�7P�7P�7R�7]�7]�^c�7dr$c��|Srar#r�s r%rcz/SDXL_instructpix2pix.__init__.<locals>.<lambda>Ys��ur$)rCrDrrrr-s    �r%rDzSDXL_instructpix2pix.__init__Ts5���
����z�&��A���3�3�3�)d�D�&�)<�D�&r$r�r�s@r%r�r�Ss���09�
�
�d�=�=r$r�c�>��eZdZejdf�fd�	Zd�Z�xZS)�StableCascade_CNc���t�|�|||t��|jj	�jd�y�N�r?r^F)rCrDrrN�eval�requires_grad_r-s    �r%rDzStableCascade_C.__init__]�6���
����z�&�V��T����!�!�#�2�2�5�9r$c���i}|d}|�"tjj|�|d<d|vrYg}|dD]7}|d}|j|djjd�|z��9t
j|d��}nt
jd	�}tjj|�|d189<tjjt
jd��|d<tjjt
jd��|d
<|jdd�}|�"tjj|�|d<|S)Nr�clip_text_pooledrr�r�rrrs)rrr1�clip_img�r�sca�crpr��	clip_text)r.r�r�r�rr;rRrwrrIr�)	r]rpr�r��embedsrrr�r�s	         r%r�zStableCascade_C.extra_condsas2����!�/�2���'�&+�k�k�&=�&=�>N�&O�C�"�#� �F�*��F�%�&;�<��$�Z�0���
�
�k�*>�?�L�L�V�V�WX�Y�\b�b�c� =��y�y��Q�/�H��{�{�;�/�H��+�+�1�1�(�;��J���[�[�,�,�U�[�[��->�?��E�190��[�[�,�,�U�[�[��->�?��E�191��Z�Z��d�3�192��!�$�{�{�8�8��D�C����193r$�rrrrrrDr�r�r�s@r%r�r�\s���09�0H�0H�QU�:�r$r�c�>��eZdZejdf�fd�	Zd�Z�xZS)�StableCascade_BNc���t�|�|||t��|jj	�jd�yr�)rCrDrrNr�r�r-s    �r%rDzStableCascade_B.__init__zr�r$c194��i}|jdd�}|d}|�"tjj|�|d<|jdt	j195dd|jddzd	z|jd196dzd	zf|j|j|j���}tjj|�|d<tjjt	j197d
��|d<|S)Nr�r�clip�stable_cascade_priorr�rr�*r)r<�layoutr?�effnetr�r�)198rIr.r�r�rRrr�r<r�r?)r]rpr�r�r��priors      r%r�zStableCascade_B.extra_conds~s�����199�200�7�D�)��!�/�2���'��+�+�1�1�2B�C�C��K��201�202�1�5�;�;��2����TU��YZ�HZ�_a�Ga�di�do�do�pq�dr�uv�dv�{}�c}�?~�GL�GR�GR�[`�[g�[g�pu�p|�p|�4}�~�����/�/��6��H�
��[�[�,�,�U�[�[��->�?��E�203��204r$r�r�s@r%r�r�ys���09�0H�0H�QU�:�
r$r�c�H��eZdZejdf�fd�	Zd�Z�fd�Z�xZS)�SD3Nc�4��t�|�|||t��y�Nr�)rCrDr	r-s    �r%rDzSD3.__init__�s���
����z�&�Mh��ir$c��|dS�Nrr#r�s  r%r�zSD3.encode_adm��
���o�&�&r$c���t�|�di|��}|jdd�}|�"tjj|�|d<|S�Nr�rmr#�rCr�rIr.r�r��r]rpr�r�r_s    �r%r�zSD3.extra_conds��I����g�!�+�F�+���Z�Z��d�3�205��!�!&���!8�!8��!D�C�
���206r$)	rrrrr rDr�r�r�r�s@r%r�r��s!���09���t�j�'��r$r�c�B��eZdZejdf�fd�	Z�fd�Z�xZS)�AuraFlowNc���t�|�|||tjjj207j��yr�)rCrDr.�ldm�aura�mmdit�MMDiTr-s    �r%rDzAuraFlow.__init__�s-���
����z�&�U�Y�Y�^�^�Ma�Ma�Mg�Mg��hr$c���t�|�di|��}|jdd�}|�"tjj|�|d<|Sr�r�r�s    �r%r�zAuraFlow.extra_conds�r�r$�rrrrr rDr�r�r�s@r%r�r��s���09���t�i��r$r�c�J��eZdZejdf�fd�	Zd�Zd�fd�	Z�xZS)�StableAudio1Nc����t�|�|||tjjj208j��tjjjjddd��|_	tjjjjddd��|_209|jj|�|jj|�y)Nr�r1rr[)�min_val�max_val)rCrDr.r��audio�dit�AudioDiffusionTransformer�	embedders�NumberConditioner�seconds_start_embedder�seconds_total_embedderr�)r]r3�seconds_start_embedder_weights�seconds_total_embedder_weightsr4r?r_s      �r%rDzStableAudio1.__init__�s����
����z�&�U�Y�Y�_�_�M`�M`�Mz�Mz��{�&+�i�i�o�o�&?�&?�&Q�&Q�RU�_`�jm�&Q�&n��#�&+�i�i�o�o�&?�&?�&Q�&Q�RU�_`�jm�&Q�&n��#��#�#�3�3�4R�S��#�#�3�3�4R�Sr$c	��i}|jdd�}|d}|jdd�}|jdt|jddz��}|j|g�dj	|�}|j|g�dj	|�}t
j||gd��jd	�}	tjj|	�|d210<|jdd�}211|212��t
j|213j	|�|j|214jdddf�|j|215jdddf�gd��}216tjj|217�|d
<|S)Nr�r?�
seconds_startr�
seconds_totalr�gH�z��5@rs)rr��global_embedr�rrm)
rIr{r�r�rQr�rRrwr�r.r�r�r<)r]rpr�r�r?r�r��seconds_start_embed�seconds_total_embedr�r�s           r%r�zStableAudio1.extra_conds�s������218�219�7�D�)����!���220�221�?�A�6�
��222�223�?�C����B��%�8O�4P�Q�
�"�9�9�=�/�J�1�M�P�P�QW�X��"�9�9�=�/�J�1�M�P�P�QW�X���y�y�"5�7J�!K�QS�T�\�\�]d�e��#�k�k�5�5�l�C��N���Z�Z��d�3�224��!����J�M�M�&�$9�;N�;U�;U�Wa�Wg�Wg�hi�Wj�lm�op�Vq�;r�uH�uO�uO�Q[�Qa�Qa�bc�Qd�fg�ij�Pk�ul�$m�st�u�J�!&���!8�!8��!D�C�
���225r$c����t�	|�|||��}|jj�|jj�d�}|D]&}||}|D]}|||dj||�<��(|S)N)r�r�r�)z'conditioner.conditioners.seconds_start.z'conditioner.conditioners.seconds_total.z{}{})rCr�r�r�r�rW)226r]r�r�r�r��dr�r5�lr_s227         �r%r�z"StableAudio1.state_dict_for_saving�s����
�W�
*�?�[i�CY�
*�Z��8<�8S�8S�8^�8^�8`�NR�Ni�Ni�Nt�Nt�Nv�
w���A��!��A���*+�A�$��6�=�=��A�&�'����	r$r�)	rrrrr!rDr�r�r�r�s@r%r�r��s0���py�qR�qR�[_�T��*�r$r�c�B��eZdZejdf�fd�	Z�fd�Z�xZS)�228HunyuanDiTNc���t�|�|||tjjj229j��yr�)rCrDr.r��hydit�models�230HunYuanDiTr-s    �r%rDzHunyuanDiT.__init__�s-���
����z�&�U�Y�Y�_�_�Mc�Mc�Mn�Mn��or$c231����t�|�di|��}|jdd�}|�"tjj|�|d<|jdd�}|�"tjj|�|d<|jdd�}|�"tjj|�|d<|jdd�}|�"tjj|�|d<|jd	d232�}|jdd233�}|jd|�}	|jd
|�}234tjjt
j|||235|	ddgg��|d<|S)Nr�rm�attention_mask�text_embedding_mask�conditioning_mt5xl�encoder_hidden_states_t5�attention_mask_mt5xl�text_embedding_mask_t5r0r1r2rDrEr�image_meta_sizer#�rCr�rIr.r�r�rR�FloatTensor)r]rpr�r�r�r�r�r0r2rDrEr_s           �r%r�zHunyuanDiT.extra_conds�s`����g�!�+�F�+���Z�Z��d�3�236��!�!&���!8�!8��!D�C�
�����$4�d�;���%�).���)@�)@��)P�C�%�&�#�Z�Z�(<�d�C���)�.3�k�k�.E�.E�FX�.Y�C�*�+�%�z�z�*@�$�G���+�,1�K�K�,C�,C�DX�,Y�C�(�)��237�238�7�C�(�����H�c�*���z�z�.�%�8���239�240�?�F�;�
�!&���!8�!8��9J�9J�V�UZ�\i�kw�yz�|}�L~�K�:A�"B�����241r$r{r�s@r%r�r��s���09�0F�0F�t�p��r$r�c�B��eZdZejdf�fd�	Z�fd�Z�xZS)�PixArtNc���t�|�|||tjjj242j��yr�)rCrDr.r��pixart�pixartms�PixArtMSr-s    �r%rDzPixArt.__init__�s/���
����z�&�U�Y�Y�M]�M]�Mf�Mf�Mo�Mo��pr$c	����t�|�di|��}|jdd�}|�"tjj|�|d<|jdd�}|jdd�}|��|��tjjt
j||gg��|d<tjjt
j|jd||z�gg��|d<|S)	Nr�rmr0r2�c_size�aspect_ratio�c_arr#r�)r]rpr�r�r0r2r_s      �r%r�zPixArt.extra_conds�s�����g�!�+�F�+���Z�Z��d�3�243��!�!&���!8�!8��!D�C�
���244�245�7�D�)�����H�d�+�����!3�!�K�K�3�3�E�4E�4E��PU��FW�4X�Y�C��M��+�+�1�1�%�2C�2C�f�j�j�Q_�ag�hm�am�Fn�Eo�Dp�2q�r�C��K��246r$rrr�s@r%r�r��s���09�
�
�d�q�
�
r$r�c�N��eZdZejdf�fd�	Zd�Zd�Z�fd�Z�xZ	S)�FluxNc���t�|�|||tjjj247j��yr�)rCrDr.r��flux�modelr�r-s    �r%rDz
Flux.__init__s-���
����z�&�U�Y�Y�^�^�Ma�Ma�Mf�Mf��gr$c��	|jjjjd|jj|jjzz}|j248jd}||kry|jdd�}|jdd�}|d}|�tj|�}tj|j|�|jd|jdd	d249�}tj||jd�}|j|�}||dzkr|S|jd
|jdd��}|�tj|�dd�dd�f}tj |dd��}tj|j|�|jddz|jddzd	d250�}|j#|jd|jddzd|jddzd�j%ddddd�j'|jdd|jddz|jddz�}tj||jd�}tj(||fd��S#|j251jd}Y��oxYw)Nr�in_channels�out_channelsr�r�r?r�r�r�r�rrr�r�T�rt�keepdimrrrrs)rN�img_inrr��252patch_sizer3rErIrRr�rr�rQr�r�r��mean�view�permuter�rw)r]rp�num_channelsr�rzr�r?r�s        r%r�zFlux.concat_conds���	H��/�/�6�6�=�=�C�C�A�F�4�K_�K_�Kj�Kj�mq�nB�nB�nM�nM�LM�N�L�253�(�(�4�4�^�D���<�'���254�255�0�$�7���256�257�7�D�)����!���=��$�$�U�+�E��$�$�U�X�X�f�%5�u�{�{�2�����TV��Yc�em�n���*�*�5�%�+�+�a�.�A���&�&�u�-���<�!�+�+��L��z�z�-����N�D�)I�J���<��?�?�5�)�!�R�a�R�%�0�D��z�z�$�A�t�4���#�#�D�G�G�F�O�U�[�[��_�q�5H�%�+�+�VX�/�\]�J]�_i�ks�t���y�y����A���258�259�1�
��(:�A�t�z�z�!�}�PQ�?Q�ST�U�]�]�^_�ab�de�gh�jk�l�t�t�uy�u�u�AB�vC�EG�IM�IS�IS�TU�IV�Z[�I[�]a�]g�]g�hi�]j�no�]o�p���)�)�$����A��?���y�y�%���A�.�.��?	H��,�,�8�8��G�L�s�AJ.�.Kc��|dSr�r#r�s  r%r�zFlux.encode_adm*r�r$c���t�260|�d261i|��}|jdd�}|�"tjj|�|d<|jdd�}|��|dj}|d}tj|d|jjz�tj|d|jjz�}}tj||||f�}tjj|�|d<|jdd	�}	|	�6tjjtj|	g��|d<|S)Nr�rmr�r��attention_mask_img_shaperr�guidanceg@r#)rCr�rIr.r�r�r�r��ceilrNr�r�upscale_dit_maskrRr�)r]rpr�r�r�r��
mask_ref_size�h_tok�w_tokrr_s          �r%r�zFlux.extra_conds-sH����g�!�+�F�+���Z�Z��d�3�262��!�!&���!8�!8��!D�C�
�����$4�d�;���%��7�O�)�)�E�"�#=�>�M�#�i�i��a��4�3G�3G�3R�3R�(R�S�UY�U^�U^�_d�ef�_g�jn�j~�j~�kJ�kJ�`J�VK�E�U�"�3�3�N�M�TY�[`�Sa�b�N�$)�K�K�$;�$;�N�$K�C� �!��:�:�j�#�.����#�k�k�5�5�e�6G�6G��263�6S�T�C�264�O��265r$)266rrrrr"rDr�r�r�r�r�s@r%r�r�s'���09���t�h�#/�J'��r$r�c�B��eZdZejdf�fd�	Z�fd�Z�xZS)�267GenmoMochiNc���t�|�|||tjjj268jj��yr�)rCrDr.r��genmo�joint_model�asymm_models_joint�
AsymmDiTJointr-s    �r%rDzGenmoMochi.__init__Cs>���
����z�&�U�Y�Y�_�_�Mh�Mh�M{�M{�NJ�NJ��	Kr$c���t�|�di|��}|jdd�}|�otjj|�|d<tjj
tdtj|�j���|d<|jdd�}|�"tjj|�|d<|S�Nr�r�269num_tokensr�rmr#)rCr�rIr.r�r�rT�maxrRr�item�r]rpr�r�r�r_s     �r%r�zGenmoMochi.extra_condsFs�����g�!�+�F�+�����$4�d�;���%�$)�K�K�$;�$;�N�$K�C� �!� %��� 8� 8��Q��	�	�.�@Y�@^�@^�@`�9a� b�C����Z�Z��d�3�270��!�!&���!8�!8��!D�C�
���271r$r�r�s@r%rrBs���09���t�K�	�	r$rc�P��eZdZejdf�fd�	Z�fd�Zdd�Zd�Z�xZ	S)�LTXVNc���t�|�|||tjjj272j��yr�)rCrDr.r��273lightricksr��	LTXVModelr-s    �r%rDz
LTXV.__init__Rs/���
����z�&�U�Y�Y�Ma�Ma�Mg�Mg�Mq�Mq��rr$c�Z��t�|�d	i|��}|jdd�}|�"tjj|�|d<|jdd�}|�"tjj|�|d<tjj
|jdd��|d<|jd|jdd��}|�"tjj|�|d<|jdd�}|�"tjj|�|d<|S)274Nr�r�rm�275frame_rate�r�r��
keyframe_idxsr#�rCr�rIr.r�r�rT)r]rpr�r�r�r�rr_s       �r%r�zLTXV.extra_condsUs����g�!�+�F�+�����$4�d�;���%�$)�K�K�$;�$;�N�$K�C� �!��Z�Z��d�3�276��!�!&���!8�!8��!D�C�
��!�K�K�4�4�V�Z�Z��b�5Q�R��L���z�z�-����N�D�1Q�R���#�"'�+�+�"9�"9�,�"G�C����277�278�?�D�9�
��$�#(�;�;�#:�#:�=�#I�C�� ��279r$c	���|�|S|jjj||j|jdgdg|j280dz281zz�zdd�dd�f�dSr�)rN�282patchifier�patchifyr�r�r�)r]rxrjr�rps     r%r}zLTXV.process_timestepjs������O��#�#�.�.�7�7�,�(�-�-�Ya�Yg�Yg�hi�Yj�Xk�op�nq�vB�vG�vG�JK�vK�oL�YL�KM�:M�OP�RT�ST�RT�OT�9U�V�WX�Y�	Yr$c��|Srar#r�s     r%r�zLTXV.scale_latent_inpainto����r$ra)283rrrrr"rDr�r}r�r�r�s@r%rrQs"���09���t�s��*Y�284r$rc�N��eZdZejdf�fd�	Zd�Z�fd�Zd�Z�xZ	S)�HunyuanVideoNc���t�|�|||tjjj285j��yr�)rCrDr.r��
hunyuan_videor�r$r-s    �r%rDzHunyuanVideo.__init__ss/���
����z�&�U�Y�Y�Md�Md�Mj�Mj�Mw�Mw��xr$c��|dSr�r#r�s  r%r�zHunyuanVideo.encode_admvr�r$c�&��t�|�di|��}|jdd�}|�"tjj|�|d<|jdd�}|�"tjj|�|d<|jdd�}|�6tjjt
j|g��|d<|jdd�}|�6tjjt
j|g��|d<|S)Nr�r�rmrg@�guiding_frame_indexr#r�)r]rpr�r�r�rr)r_s       �r%r�zHunyuanVideo.extra_condsys�����g�!�+�F�+�����$4�d�;���%�$)�K�K�$;�$;�N�$K�C� �!��Z�Z��d�3�286��!�!&���!8�!8��!D�C�
���:�:�j�#�.����#�k�k�5�5�e�6G�6G��287�6S�T�C�288�O�$�j�j�)>��E���*�).���)@�)@��AR�AR�Tg�Sh�Ai�)j�C�%�&��289r$c��|Srar#)r]r�rps   r%r�z!HunyuanVideo.scale_latent_inpaint�r"r$)290rrrrr rDr�r�r�r�r�s@r%r$r$rs!���09���t�y�'��&r$r$c�B��eZdZejdf�fd�	Z�fd�Z�xZS)�HunyuanVideoI2VNc�8��t�|�|||��d|_y)Nr�)r�r��rCrDrZr-s    �r%rDzHunyuanVideoI2V.__init__�s���
����z�&��A�<��r$c�&��t�|�dd|i|��S�Nr�r#�rCr��r]r�rpr_s   �r%r�z$HunyuanVideoI2V.scale_latent_inpaint������w�+�P��P��P�Pr$�rrrrr rDr�r�r�s@r%r,r,�s���09���t�=�Q�Qr$r,c�B��eZdZejdf�fd�	Z�fd�Z�xZS)�HunyuanVideoSkyreelsI2VNc�8��t�|�|||��d|_y)Nr�)r�r.r-s    �r%rDz HunyuanVideoSkyreelsI2V.__init__�s���
����z�&��A�,��r$c�&��t�|�dd|i|��Sr0r1r2s   �r%r�z,HunyuanVideoSkyreelsI2V.scale_latent_inpaint�r3r$r4r�s@r%r6r6�s���09���t�-�Q�Qr$r6c�J��eZdZejddf�fd�	Z�fd�Zd�Z�xZS)�CosmosVideoFNc���t�|�|||tjjj291j��||_|jrd|_yy)Nr�)r�)	rCrDr.r��cosmosr��292GeneralDIT�image_to_videorZ�r]r3r4r>r?r_s     �r%rDzCosmosVideo.__init__�sM���
����z�&�U�Y�Y�M]�M]�Mc�Mc�Mn�Mn��o�,������1�D��r$c�b��t�|�di|��}|jdd�}|�"tjj|�|d<|jdd�}|�"tjj|�|d<tjj
|jdd��|d<|S)Nr�r�rmrrJr#rrs     �r%r�zCosmosVideo.extra_conds�s�����g�!�+�F�+�����$4�d�;���%�$)�K�K�$;�$;�N�$K�C� �!��Z�Z��d�3�293��!�!&���!8�!8��!D�C�
���[�[�-�-�f�j�j��t�.L�M��E�294��295r$c�j�|j|jdgdgt|j�dz296zz�}d}|dk7r||z}|jj	tj|g|j|j��|�}||dz|jjdzzdzzS)Nrr)r?r<rg�?)297r�r�r�r/rvrRr�r?r<�298sigma_data)r]rr�r�rp�sigma_noise_augmentations      r%r�z CosmosVideo.scale_latent_inpaint�s����
�
�u�{�{�1�~�.�!���E�K�K�8H�1�8L�1M�M�N��#$� �#�q�(�'�%�/�L��*�*�:�:�5�<�<�Ia�Hb�kw�k~�k~�GS�GY�GY�<Z�\h�i�����299�T�-@�-@�-K�-K�q�-P� P�UX�X�Y�Yr$)	rrrrrrDr�r�r�r�s@r%r:r:�s���09�
�
�e�\`�2�300�Zr$r:c�B��eZdZejdf�fd�	Z�fd�Z�xZS)�Lumina2Nc���t�|�|||tjjj301j��yr�)rCrDr.r��luminar��NextDiTr-s    �r%rDzLumina2.__init__�s/���
����z�&�U�Y�Y�M]�M]�Mc�Mc�Mk�Mk��lr$c����t�|�di|��}|jdd�}|��tj|�|j�k7r"tjj|�|d<tjjtdtj302|�j���|d<|jdd�}|�"tjj|�|d<|Sr)rCr�rIrR�numelrr.r�r�rTrrrs     �r%r�zLumina2.extra_conds�s�����g�!�+�F�+�����$4�d�;���%��{�{�>�*�n�.@�.@�.B�B�(-���(?�(?��(O��$�%� %��� 8� 8��Q��	�	�.�@Y�@^�@^�@`�9a� b�C����Z�Z��d�3�303��!�!&���!8�!8��!D�C�
���304r$r�r�s@r%rErE�s���09���t�m�305�306r$rEc�J��eZdZejddf�fd�	Zd�Z�fd�Z�xZS)�WAN21FNc���t�|�|||tjjj307j��||_yr�)rCrDr.r��wanr��WanModelr>r?s     �r%rDzWAN21.__init__�s5���
����z�&�U�Y�Y�]�]�M`�M`�Mi�Mi��j�,��r$c���|jdd�}|jjjjd|jdk(ry|jdd�}|d}|�tj|�}tj|j|�|jd|jddd�}|j|�}tj||jd	�}|js|S|jd308|jdd��}|� tj|�dd�dd�f}�nd
tj|dd��z309}tj|j|�|jd|jddd�}|jd|jdkrPt310jjj!|d	d	d	d	d	|jd|jdz311fdd	��}|j#ddddd�}tj||jd	�}tj$||fd��S)Nr�rr�r?r�r�r�r�rr�r�rr�Tr�r�r�)�mode�valuers)rIrN�patch_embeddingrr�rRr�rr�rQr�r�r>r�r�r�r�r<rw)r]rpr�rzr?r�s      r%r�zWAN21.concat_cond�s���312�313�7�D�)�����/�/�6�6�<�<�Q�?�5�;�;�q�>�Q���314�315�0�$�7����!���=��$�$�U�+�E��$�$�U�X�X�f�%5�u�{�{�2�����TV��Yc�em�n���&�&�u�-���*�*�5�%�+�+�a�.�A���"�"��L��z�z�-����N�D�)I�J���<��#�#�E�*�1�b�q�b�5�1�D�����D�a��>�>�D��'�'���������R��%�+�+�VX�/�[e�go�p�D��z�z�"�~����B��/��x�x�*�*�.�.�t�a��A�q�!�U�[�[�QS�_�W[�Wa�Wa�bd�We�Ee�5f�mw�@A�.�B���;�;�q�!�Q��1�-�D��-�-�d�E�K�K��N�C�D��y�y�$���A�.�.r$c���t�|�di|��}|jdd�}|�"tjj|�|d<|jdd�}|�,tjj|j�|d<|S)Nr�rmr��clip_fear#)rCr�rIr.r�r��penultimate_hidden_states)r]rpr�r�r�r_s     �r%r�zWAN21.extra_conds�s�����g�!�+�F�+���Z�Z��d�3�316��!�!&���!8�!8��!D�C�
��#�Z�Z�(<�d�C���)�#�k�k�5�5�6H�6b�6b�c�C�317�O��318r$)	rrrrr rDr�r�r�r�s@r%rLrL�s#���09���u�]a�-�/�>	�	r$rL)ryN)U�__doc__rRrT�.comfy.ldm.modules.diffusionmodules.openaimodelrr�comfy.ldm.cascade.stage_cr�comfy.ldm.cascade.stage_br�,comfy.ldm.modules.encoders.noise_aug_modulesr�,comfy.ldm.modules.diffusionmodules.upscalingr�(comfy.ldm.modules.diffusionmodules.mmditr	�.comfy.ldm.genmo.joint_model.asymm_models_jointr.�comfy.ldm.aura.mmdit�comfy.ldm.pixart.pixartms�comfy.ldm.hydit.models�comfy.ldm.audio.dit�comfy.ldm.audio.embedders�comfy.ldm.flux.model�comfy.ldm.lightricks.model�comfy.ldm.hunyuan_video.model�comfy.ldm.cosmos.model�comfy.ldm.lumina.model�comfy.ldm.wan.model�comfy.model_management�comfy.patcher_extension�comfy.conds�	comfy.ops�enumr319r�r�comfy.latent_formatsr��typingr
�comfy.model_patcherrr�comfy.model_samplingrrrr&r'r(r)r/r��Moduler8rrr r"rArGrVrYrhrtr}r�r�r�r�r�r�r�r�r�r�rrr$r,r6r:rErLr#r$r%�<module>rts���$
��N�,�,�W�Y�P�5�� ��� ��!�$������������ ��0�
��
�\�\�\�'�@iL������iL�X�0]��]�%�E�)�E�6E�9�E�4-�)�-�^�[��%�[�%�(�Y��<"�I�"�H
1�
1� 9�4��9�=�4��=��i��:�i��*�)��	�y�	�$�9�$�N���<�Y��&?�9�?�B
��
��9��B�9��:Q�l�Q�Q�l�Q�Z�)�Z�6�i�� -�I�-r$