Compare commits
35
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e1d2cff331 | ||
|
|
08f892eace | ||
|
|
42f509b3e8 | ||
|
|
445e087a43 | ||
|
|
ce66eea33e | ||
|
|
8b2d38c2ec | ||
|
|
36fdaf500b | ||
|
|
35ec43af1f | ||
|
|
8761bf3c00 | ||
|
|
ee11be3ce2 | ||
|
|
ba2626426f | ||
|
|
688efe72d2 | ||
|
|
e87087e95a | ||
|
|
653a95bdd1 | ||
|
|
5862cc77e6 | ||
|
|
ca193cdfe8 | ||
|
|
6e9e6de056 | ||
|
|
b19dc8618f | ||
|
|
7875c1cf64 | ||
|
|
a877517b63 | ||
|
|
1c8909531c | ||
|
|
36516b4e71 | ||
|
|
e63938b095 | ||
|
|
d617d4bd31 | ||
|
|
fa234963e5 | ||
|
|
5f6848be70 | ||
|
|
bcfd123014 | ||
|
|
972af4120d | ||
|
|
466330e953 | ||
|
|
8ef49793d7 | ||
|
|
0992f16e0c | ||
|
|
864ffa6183 | ||
|
|
ddb96675f0 | ||
|
|
bdde649e25 | ||
|
|
e2a73515dd |
@@ -1,2 +1,10 @@
|
||||
**/*.png filter=lfs diff=lfs merge=lfs -text
|
||||
**/*.mp4 filter=lfs diff=lfs merge=lfs -text
|
||||
**/*.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
|
||||
example_workflows/low_level/end.jpg -filter -diff -merge -text
|
||||
example_workflows/low_level/fox.jpg -filter -diff -merge -text
|
||||
example_workflows/low_level/shrek2.jpg -filter -diff -merge -text
|
||||
example_workflows/low_level/start.jpg -filter -diff -merge -text
|
||||
example_workflows/tricks/shrek2.jpg -filter -diff -merge -text
|
||||
example_workflows/tricks/shrek3.jpg -filter -diff -merge -text
|
||||
|
||||
@@ -1,201 +1,381 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
LTX-2 Community License Agreement
|
||||
License date: January 5, 2026
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
By using or distributing any portion or element of LTX-2, you agree
|
||||
to be bound by this Agreement.
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
"Agreement" means the terms and conditions for the license, use,
|
||||
reproduction, and distribution of LTX-2 and the Complementary
|
||||
Materials, as specified in this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
"Control" means the direct or indirect ownership of more than
|
||||
fifty percent (50%) of the voting securities or other ownership
|
||||
interests, or the power to direct the management and policies of
|
||||
such Entity through voting rights, contract, or otherwise.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
"Data" means a collection of information and/or content extracted
|
||||
from the dataset used with LTX-2, including to train, pretrain,
|
||||
or otherwise evaluate LTX-2. The Data is not licensed under this
|
||||
Agreement.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
"Derivatives of LTX-2" means all modifications to LTX-2, works
|
||||
based on LTX-2, or any other model which is created or initialized
|
||||
by transfer of patterns of the weights, parameters, activations or
|
||||
output of LTX-2, to the other model, in order to cause the other
|
||||
model to perform similarly to LTX-2, including – but not limited
|
||||
to - distillation methods entailing the use of intermediate data
|
||||
representations or methods based on the generation of synthetic
|
||||
data by LTX-2 for training the other model. For clarity, Derivatives
|
||||
of LTX-2 include: (i) any fine-tuned or adapted weights, parameters,
|
||||
or checkpoints derived from LTX-2; (ii) derivative model architectures
|
||||
that incorporate or are based upon LTX-2's architecture; and
|
||||
(iii) any modified or extended versions of the Complementary
|
||||
Materials. All intellectual property rights in Derivatives of LTX-2
|
||||
shall be subject to the terms of this Agreement, and you may not
|
||||
claim exclusive ownership rights in any Derivatives of LTX-2 that
|
||||
would restrict the rights granted herein.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
"Entity" means any individual, corporation, partnership, limited
|
||||
liability company, or other legal entity. For purposes of this
|
||||
Agreement, an Entity shall be deemed to include, on an aggregative
|
||||
basis, all subsidiaries, affiliates, and other companies under
|
||||
common Control with such Entity. When determining whether an Entity
|
||||
meets any threshold under this Agreement (including revenue
|
||||
thresholds), all subsidiaries, affiliates, and companies under
|
||||
common Control shall be considered collectively.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
"Harm" includes but is not limited to physical, mental,
|
||||
psychological, financial and reputational damage, pain, or loss.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
"Licensor" or "Lightricks" means the owner that is granting the
|
||||
license under this Agreement. For the purposes of this Agreement,
|
||||
the Licensor is Lightricks Ltd.
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
"LTX-2" means the large language models, text/image/video/audio/3D
|
||||
generation models, and multimodal large language models and their
|
||||
software and algorithms, including trained model weights, parameters
|
||||
(including optimizer states), machine-learning model code,
|
||||
inference-enabling code, training-enabling code, fine-tuning
|
||||
enabling code, accompanying source code, scripts, documentation,
|
||||
tutorials, examples, and all other elements of the foregoing
|
||||
distributed and made publicly available by Lightricks (including,
|
||||
for example, at https://github.com/Lightricks/LTX-2) for the LTX-2
|
||||
model released on January 5, 2026. This license is applicable to
|
||||
all LTX-2 versions released since January 5, 2026, and all future
|
||||
releases of LTX-2 under this license.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
"Output" means the results of operating LTX-2 as embodied in
|
||||
informational content resulting therefrom.
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
"you" (or "your") means an individual or legal Entity licensing
|
||||
LTX-2 in accordance with this Agreement and/or making use of LTX-2
|
||||
for whichever purpose and in any field of use, including usage of
|
||||
LTX-2 in an end-use application - e.g. chatbot, translator, image
|
||||
generator.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
2. Grant of License. Subject to the terms and conditions of this
|
||||
Agreement, you are granted a non-exclusive, worldwide,
|
||||
non-transferable and royalty-free limited license under Licensor's
|
||||
intellectual property or other rights owned by Licensor embodied
|
||||
in LTX-2 to use, reproduce, prepare, distribute, publicly display,
|
||||
publicly perform, sublicense, copy, create derivative works of,
|
||||
and make modifications to LTX-2, for any purpose, subject to the
|
||||
restrictions set forth in Attachment A; provided however, that
|
||||
Entities with annual revenues of at least $10,000,000 (the
|
||||
"Commercial Entities") are required to obtain a paid commercial
|
||||
use license in order to use LTX-2 and Derivatives of LTX-2,
|
||||
subject to the terms and provisions of a different license (the
|
||||
"Commercial Use Agreement"), as will be provided by the Licensor.
|
||||
Commercial Entities interested in such a commercial license are
|
||||
required to contact Licensor. Any commercial use of LTX-2 or
|
||||
Derivatives of LTX-2 by the Commercial Entities not in accordance
|
||||
with this Agreement and/or the Commercial Use Agreement is strictly
|
||||
prohibited and shall be deemed a material breach of this Agreement.
|
||||
Such material breach will be subject, in addition to any license
|
||||
fees owed to Licensor for the period such Commercial Entity used
|
||||
LTX-2 (as will be determined by Licensor), to liquidated damages,
|
||||
which will be paid to Licensor immediately upon demand, in an
|
||||
amount equal to double the amount that would otherwise have been
|
||||
paid by you for the relevant period of time. Such amount reflects
|
||||
a reasonable estimation of the losses and administrative costs
|
||||
incurred due to such breach. You agree and understand that this
|
||||
remedy does not limit the Licensor's right to pursue other remedies
|
||||
available at law or equity.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
3. Distribution and Redistribution. You may host for third parties
|
||||
remote access purposes (e.g. software-as-a-service), reproduce
|
||||
and distribute copies of LTX-2 or Derivatives of LTX-2 thereof in
|
||||
any medium, with or without modifications, provided that you meet
|
||||
the following conditions:
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
(a) Use-based restrictions as referenced in paragraph 4 and all
|
||||
provisions of Attachment A MUST be included as an enforceable
|
||||
provision by you in any type of legal agreement (e.g. a
|
||||
license) governing the use and/or distribution of LTX-2 or
|
||||
Derivatives of LTX-2, and you shall give notice to subsequent
|
||||
users you distribute to, that LTX-2 or Derivatives of LTX-2
|
||||
are subject to paragraph 4 and Attachment A in their entirety,
|
||||
including all use restrictions and acceptable use policies;
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
(b) You must provide any third party recipients of LTX-2 or
|
||||
Derivatives of LTX-2 a copy of this Agreement, including all
|
||||
attachments and use policies. Any Derivative of LTX-2 (as
|
||||
defined in Section 1, including but not limited to fine-tuned
|
||||
weights, modified training code, models trained on Outputs, or
|
||||
any other derivative) must be distributed exclusively under
|
||||
the terms of this Agreement with a complete copy of this
|
||||
license included;
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
(c) You must cause any modified files to carry prominent notices
|
||||
stating that you changed the files;
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
(d) You must retain all copyright, patent, trademark, and
|
||||
attribution notices excluding those notices that do not
|
||||
pertain to any part of LTX-2, Derivatives of LTX-2.
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
You may add your own copyright statement to your modifications and
|
||||
may provide additional or different license terms and conditions -
|
||||
respecting paragraph 3(a) - for use, reproduction, or distribution
|
||||
of your modifications, or for any such Derivatives of LTX-2 as a
|
||||
whole, provided your use, reproduction, and distribution of LTX-2
|
||||
otherwise complies with the conditions stated in this Agreement,
|
||||
and you provide a complete copy of this Agreement with any such
|
||||
use, reproduction and distribution of LTX-2 and any Derivatives
|
||||
thereof.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
4. Use-based restrictions. The restrictions set forth in Attachment A
|
||||
are considered Use-based restrictions. Therefore, you cannot use
|
||||
LTX-2 and the Derivatives of LTX-2 in violation of the specified
|
||||
restricted uses. You may use LTX-2 subject to this Agreement,
|
||||
including only for lawful purposes and in accordance with the
|
||||
Agreement. "Use" may include creating any content with, fine-tuning,
|
||||
updating, running, training, evaluating and/or re-parametrizing
|
||||
LTX-2. You shall require all of your users who use LTX-2 or a
|
||||
Derivative of LTX-2 to comply with the terms of this paragraph 4.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
5. The Output You Generate. Except as set forth herein, Licensor
|
||||
claims no rights in the Output you generate using LTX-2. You are
|
||||
accountable for input you insert into LTX-2, the Output you
|
||||
generate and its subsequent uses. No use of the Output can
|
||||
contravene any provision as stated in the Agreement.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
6. Updates and Runtime Restrictions. To the maximum extent permitted
|
||||
by law, Licensor reserves the right to restrict (remotely or
|
||||
otherwise) usage of LTX-2 in violation of this Agreement, update
|
||||
LTX-2 through electronic means, or modify the Output of LTX-2
|
||||
based on updates. You shall undertake reasonable efforts to use
|
||||
the latest version of LTX-2. Any use of the non-current version
|
||||
of LTX-2 is done solely at your risk.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
7. Export Controls and Sanctions Compliance. You acknowledge that
|
||||
LTX-2, Derivatives of LTX-2 may be subject to export control laws
|
||||
and regulations, including but not limited to the U.S. Export
|
||||
Administration Regulations and sanctions programs administered by
|
||||
the Office of Foreign Assets Control (OFAC). You represent and
|
||||
warrant that you and any users of LTX-2 are not (i) located in,
|
||||
organized under the laws of, or ordinarily resident in any country
|
||||
or territory subject to comprehensive sanctions; (ii) identified
|
||||
on any U.S. government restricted party list, including the
|
||||
Specially Designated Nationals and Blocked Persons List; or
|
||||
(iii) otherwise prohibited from receiving LTX-2 under applicable
|
||||
law. You shall not export, re-export, or transfer LTX-2, directly
|
||||
or indirectly, in violation of any applicable export control or
|
||||
sanctions laws or regulations. You agree to comply with all
|
||||
applicable trade control laws and shall indemnify and hold
|
||||
Licensor harmless from any claims arising from your failure to
|
||||
comply with such laws.
|
||||
|
||||
8. Trademarks and related. Nothing in this Agreement permits you to
|
||||
make use of Licensor's trademarks, trade names, logos or to
|
||||
otherwise suggest endorsement or misrepresent the relationship
|
||||
between the parties; and any rights not expressly granted herein
|
||||
are reserved by the Licensor.
|
||||
|
||||
9. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides LTX-2 on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
implied, including, without limitation, any warranties or
|
||||
conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS
|
||||
FOR A PARTICULAR PURPOSE. You are solely responsible for
|
||||
determining the appropriateness of using or redistributing LTX-2
|
||||
and Derivatives of LTX-2 and assume any risks associated with
|
||||
your exercise of permissions under this Agreement.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
10. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall Licensor be liable
|
||||
to you for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as
|
||||
a result of this Agreement or out of the use or inability to use
|
||||
LTX-2 (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if Licensor has been
|
||||
advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
11. Accepting Warranty or Additional Liability. While redistributing
|
||||
LTX-2 and Derivatives of LTX-2, you may, provided you do not
|
||||
violate the terms of this Agreement, choose to offer and charge
|
||||
a fee for, acceptance of support, warranty, indemnity, or other
|
||||
liability obligations. However, in accepting such obligations,
|
||||
you may act only on your own behalf and on your sole
|
||||
responsibility, not on behalf of Licensor, and only if you agree
|
||||
to indemnify, defend, and hold Licensor harmless for any liability
|
||||
incurred by, or claims asserted against Licensor, by reason of
|
||||
your accepting any such warranty or additional liability.
|
||||
|
||||
12. Governing Law. This Agreement and all relations, disputes, claims
|
||||
and other matters arising hereunder (including non-contractual
|
||||
disputes or claims) will be governed exclusively by, and construed
|
||||
exclusively in accordance with, the laws of the State of New York.
|
||||
To the extent permitted by law, choice of laws rules and the
|
||||
United Nations Convention on Contracts for the International Sale
|
||||
of Goods will not apply. For the purposes of adjudicating any
|
||||
action or proceeding to enforce the terms of this Agreement, you
|
||||
hereby irrevocably consent to the exclusive jurisdiction of, and
|
||||
venue in, the federal and state courts located in the County of
|
||||
New York within the State of New York. The prevailing party in
|
||||
any claim or dispute between the parties under this Agreement
|
||||
will be entitled to reimbursement of its reasonable attorneys'
|
||||
fees and costs. You hereby waive the right to a trial by jury,
|
||||
to participate in a class or representative action (including in
|
||||
arbitration), or to combine individual proceedings in court or
|
||||
in arbitration without the consent of all parties.
|
||||
|
||||
13. Term and Termination. This Agreement is effective upon your
|
||||
acceptance and continues until terminated. Licensor may terminate
|
||||
this Agreement immediately upon written notice to you if you
|
||||
breach any provision of this Agreement, including but not limited
|
||||
to violations of the use restrictions in Attachment A or
|
||||
unauthorized commercial use. Upon termination: (a) all rights
|
||||
granted to you under this Agreement will immediately cease;
|
||||
(b) you must immediately cease all use of LTX-2 and Derivatives
|
||||
of LTX-2; (c) you must delete or destroy all copies of LTX-2
|
||||
and Derivatives of LTX-2 in your possession or control; and
|
||||
(d) you must notify any third parties to whom you distributed
|
||||
LTX-2 or Derivatives of LTX-2 of the termination. Sections 8-13,
|
||||
and Section 15 shall survive termination of this Agreement.
|
||||
Termination does not relieve you of any obligations incurred
|
||||
prior to termination, including payment obligations under
|
||||
Section 2. In addition, if You commence a lawsuit or other
|
||||
proceedings (including a cross-claim or counterclaim in a lawsuit)
|
||||
against Licensor or any person or entity alleging that LTX-2 or
|
||||
any Output, or any portion of any of the foregoing, infringe any
|
||||
intellectual property or other right owned or licensable by you,
|
||||
then all licenses granted to you under this Agreement shall
|
||||
terminate as of the date such lawsuit or other proceeding is filed.
|
||||
|
||||
14. Disputes and Arbitration. All disputes arising in connection with
|
||||
this Agreement shall be finally settled by arbitration under the
|
||||
Rules of Arbitration of the International Chamber of Commerce
|
||||
("ICC Rules"), by one (1) arbitrator appointed in accordance with
|
||||
the ICC Rules. The seat of arbitration shall be New York, NY, USA,
|
||||
and the proceedings shall be conducted in English. The arbitrator
|
||||
shall be empowered to grant any relief that a court could grant.
|
||||
Judgment on the arbitration award may be entered by any court
|
||||
having jurisdiction thereof. Each party waives its right to a
|
||||
trial by jury and to participate in any class or representative
|
||||
action.
|
||||
|
||||
15. If any provision of this Agreement is held to be
|
||||
invalid, illegal
|
||||
or unenforceable, the remaining provisions shall be unaffected
|
||||
thereby and remain valid as if such provision had not been set
|
||||
forth herein.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
ATTACHMENT A: Use Restrictions
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
When using the Outputs, LTX-2 and any Derivatives thereof, you
|
||||
will comply with the Acceptable Use Policy. In addition, you
|
||||
agree not to use the Outputs, LTX-2 or its Derivatives in any
|
||||
of the following ways:
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
1. In any way that violates any applicable national, federal,
|
||||
state, local or international law or regulation;
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
2. For the purpose of exploiting, Harming or attempting to
|
||||
exploit or Harm minors in any way;
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
3. To generate or disseminate false information and/or content
|
||||
with the purpose of Harming others;
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
4. To generate or disseminate personal identifiable information
|
||||
that can be used to Harm an individual;
|
||||
|
||||
5. To generate or disseminate information and/or content (e.g.
|
||||
images, code, posts, articles), and place the information
|
||||
and/or content in any context (e.g. bot generating tweets)
|
||||
without expressly and intelligibly disclaiming that the
|
||||
information and/or content is machine generated;
|
||||
|
||||
6. To defame, disparage or otherwise harass others;
|
||||
|
||||
7. To impersonate or attempt to impersonate (e.g. deepfakes)
|
||||
others without their consent;
|
||||
|
||||
8. For fully automated decision making that adversely impacts an
|
||||
individual's legal rights or otherwise creates or modifies a
|
||||
binding, enforceable obligation;
|
||||
|
||||
9. For any use intended to or which has the effect of
|
||||
discriminating against or Harming individuals or groups based
|
||||
on online or offline social behavior or known or predicted
|
||||
personal or personality characteristics;
|
||||
|
||||
10. To exploit any of the vulnerabilities of a specific group of
|
||||
persons based on their age, social, physical or mental
|
||||
characteristics, in order to materially distort the behavior
|
||||
of a person pertaining to that group in a manner that causes
|
||||
or is likely to cause that person or another person physical
|
||||
or psychological Harm;
|
||||
|
||||
11. For any use intended to or which has the effect of
|
||||
discriminating against individuals or groups based on legally
|
||||
protected characteristics or categories;
|
||||
|
||||
12. To provide medical advice and medical results interpretation;
|
||||
|
||||
13. To generate or disseminate information for the purpose to be
|
||||
used for administration of justice, law enforcement,
|
||||
immigration or asylum processes, such as predicting an
|
||||
individual will commit fraud/crime commitment (e.g. by text
|
||||
profiling, drawing causal relationships between assertions
|
||||
made in documents, indiscriminate and arbitrarily-targeted use);
|
||||
|
||||
14. To generate and/or disseminate malware (including – but not
|
||||
limited to – ransomware) or any other content to be used for
|
||||
the purpose of harming electronic systems;
|
||||
|
||||
15. To engage in, promote, incite, or facilitate discrimination
|
||||
or other unlawful or harmful conduct in the provision of
|
||||
employment, employment benefits, credit, housing, or other
|
||||
essential goods and services;
|
||||
|
||||
16. To engage in, promote, incite, or facilitate the harassment,
|
||||
abuse, threatening, or bullying of individuals or groups of
|
||||
individuals;
|
||||
|
||||
17. For military, warfare, nuclear industries or applications,
|
||||
weapons development, or any use in connection with activities
|
||||
that may cause death, personal injury, or severe physical or
|
||||
environmental damage;
|
||||
|
||||
18. For commercial use only: To train, improve, or fine-tune any
|
||||
other machine learning model, artificial intelligence system,
|
||||
or competing model, except for Derivatives of LTX-2 as
|
||||
expressly permitted under this Agreement;
|
||||
|
||||
19. To circumvent, disable, or interfere with any technical
|
||||
limitations, safety features, content filters, or use
|
||||
restrictions implemented in LTX-2 by Licensor;
|
||||
|
||||
20. To use LTX-2 or Derivatives of LTX-2 in any product, service,
|
||||
or application that directly competes with Licensor's
|
||||
commercial products or services, or is designed to replace or
|
||||
substitute Licensor's offerings in the market, without
|
||||
obtaining a separate commercial license from Licensor.
|
||||
|
||||
@@ -1,104 +1,97 @@
|
||||
# ComfyUI-LTXVideo
|
||||
|
||||
ComfyUI-LTXVideo is a collection of custom nodes for ComfyUI, designed to provide useful tools for working with the LTXV model.
|
||||
The model itself is supported in the core ComfyUI [code](https://github.com/comfyanonymous/ComfyUI/tree/master/comfy/ldm/lightricks).
|
||||
The main LTXVideo repository can be found [here](https://github.com/Lightricks/LTX-Video).
|
||||
[](https://github.com/Lightricks/LTX-2)
|
||||
[](https://ltx.io/model)
|
||||
[](https://huggingface.co/Lightricks/LTX-2)
|
||||
[](https://github.com/Lightricks/LTX-2/tree/main/packages/ltx-trainer)
|
||||
[](https://app.ltx.studio/ltx-2-playground/i2v)
|
||||
[](https://videos.ltx.io/LTX-2/grants/LTX_2_Technical_Report_compressed.pdf)
|
||||
[](https://discord.gg/ltxplatform)
|
||||
|
||||
# 5.03.2025 :star: LTXVideo 0.9.5 Release :star:
|
||||
|
||||
### LTXVideo 0.9.5 introduces:
|
||||
A collection of powerful custom nodes that extend ComfyUI's capabilities for the LTX-2 video generation model.
|
||||
|
||||
1. Improved quality with reduced artifacts.
|
||||
2. Support for higher resolution and longer sequences.
|
||||
3. Frame and sequence conditioning (beyond the first frame).
|
||||
4. Enhanced prompt understanding.
|
||||
5. Commercial license availability.
|
||||
LTX-2 is built into ComfyUI core ([see it here](https://github.com/comfyanonymous/ComfyUI/tree/master/comfy/ldm/lightricks)), making it readily accessible to all ComfyUI users. This repository hosts additional nodes and workflows to help you get the most out of LTX-2's advanced features.
|
||||
|
||||
### Technical Updates
|
||||
**To learn more about LTX-2** See the [main LTX-2 repository](https://github.com/Lightricks/LTX-2) for model details and additional resources.
|
||||
|
||||
Since LTXVideo is now fully supported in the ComfyUI core, we have removed the custom model implementation. Instead, we provide updated workflows to showcase the new features:
|
||||
|
||||
1. **Frame Conditioning** – Enables interpolation between given frames.
|
||||
2. **Sequence Conditioning** – Allows motion interpolation from a given frame sequence, enabling video extension from the beginning, end, or middle of the original video.
|
||||
3. **Prompt Enhancer** – A new node that helps generate prompts optimized for the best model performance.
|
||||
See the [Example Workflows](#example-workflows) section for more details.
|
||||
## Prerequisites
|
||||
Before you begin using an LTX-2 workflow in ComfyUI, make sure you have:
|
||||
|
||||
### LTXTricks Update
|
||||
* ComfyUI installed (Download here](https://www.comfy.org/download)
|
||||
* CUDA-compatible GPU with 32GB+ VRAM
|
||||
* 100GB+ free disk space for models and cache
|
||||
|
||||
The LTXTricks code has been integrated into this repository (in the `/tricks` folder) and will be maintained here. The original [repo](https://github.com/logtd/ComfyUI-LTXTricks) is no longer maintained, but all existing workflows should continue to function as expected.
|
||||
|
||||
## 22.12.2024
|
||||
## Quick Start 🚀
|
||||
|
||||
Fixed a bug which caused the model to produce artifacts on short negative prompts when using a native CLIP Loader node.
|
||||
We recommend using the LTX-2 workflows available in Comfy Manager.
|
||||
|
||||
## 19.12.2024 ⭐ Update ⭐
|
||||
1. Open ComfyUI
|
||||
2. Click the Manager button (or press Ctrl+M)
|
||||
3. Select Install Custom Nodes
|
||||
4. Search for “LTXVideo”
|
||||
5. Click Install
|
||||
6. Wait for installation to complete
|
||||
7. Restart ComfyUI
|
||||
|
||||
1. Improved model - removes "strobing texture" artifacts and generates better motion. Download from [here](https://huggingface.co/Lightricks/LTX-Video/resolve/main/ltx-video-2b-v0.9.1.safetensors).
|
||||
2. STG support
|
||||
3. Integrated image degradation system for improved motion generation.
|
||||
4. Additional initial latent optional input to chain latents for high res generation.
|
||||
5. Image captioning in image to video [flow](assets/ltxvideo-i2v.json).
|
||||
The nodes will appear in your node menu under the “LTXVideo” category. Required models will be downloaded on first use.
|
||||
|
||||
## Installation
|
||||
|
||||
Installation via [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) is preferred. Simply search for `ComfyUI-LTXVideo` in the list of nodes and follow installation instructions.
|
||||
## Example Workflows
|
||||
|
||||
### Manual installation
|
||||
The ComfyUI-LTXVideo installation includes several example workflows.
|
||||
You can see them all at:
|
||||
'''
|
||||
ComfyUI/custom_nodes/ComfyUI-LTXVideo/example_workflows/
|
||||
'''
|
||||
|
||||
1. Install ComfyUI
|
||||
2. Clone this repository to `custom-nodes` folder in your ComfyUI installation directory.
|
||||
3. Install the required packages:
|
||||
* [`Text to video full model`](./example_workflows/LTX-2_T2V_Full_wLora.json)
|
||||
* [`Text to video distilled model (Fast)`](./example_workflows/LTX-2_T2V_Distilled_wLora.json)
|
||||
* [`Image to video full model`](./example_workflows/LTX-2_I2V_Full_wLora.json)
|
||||
* [`Image to video distilled model (Fast)`](./example_workflows/LTX-2_I2V_Distilled_wLora.json)
|
||||
* [`Video to video detailer`](./example_workflows/LTX-2_V2V_Detailer.json)
|
||||
* [`IC-LoRA distilled model (depth + human pose + edges)`](./example_workflows/LTX-2_ICLoRA_All_Distilled.json)
|
||||
|
||||
```bash
|
||||
cd custom_nodes/ComfyUI-LTXVideo && pip install -r requirements.txt
|
||||
```
|
||||
## Required Models
|
||||
|
||||
For portable ComfyUI installations, run
|
||||
Download the following models:
|
||||
|
||||
```
|
||||
.\python_embeded\python.exe -m pip install -r .\ComfyUI\custom_nodes\ComfyUI-LTXVideo\requirements.txt
|
||||
```
|
||||
**LTX-2 Model Checkpoint** (choose and download one of the following)
|
||||
* [`ltx-2-19b-dev-fp8.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-dev-fp8.safetensors)
|
||||
* [`ltx-2-19b-distilled-fp8.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled-fp8.safetensors)
|
||||
* [`ltx-2-19b-dev.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-dev.safetensors)
|
||||
* [`ltx-2-19b-distilled.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled.safetensors)
|
||||
|
||||
### Models
|
||||
**Spatial Upsampler** - Required for current two-stage pipeline implementations in this repository
|
||||
* [`ltx-2-spatial-upsampler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-spatial-upsampler-x2-1.0.safetensors)
|
||||
|
||||
1. Download [ltx-video-2b-v0.9.1.safetensors](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltx-video-2b-v0.9.1.safetensors) from Hugging Face and place it under `models/checkpoints`.
|
||||
2. Install one of the t5 text encoders, for example [google_t5-v1_1-xxl_encoderonly](https://huggingface.co/mcmonkey/google_t5-v1_1-xxl_encoderonly/tree/main). You can install it using ComfyUI Model Manager.
|
||||
**Temporal Upsampler** - Required for current two-stage pipeline implementations in this repository
|
||||
* [`ltx-2-temporal-upsampler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-temporal-upsampler-x2-1.0.safetensors)
|
||||
|
||||
## Example workflows
|
||||
**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline and ICLoraPipeline)
|
||||
* [`ltx-2-19b-distilled-lora-384.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled-lora-384.safetensors)
|
||||
|
||||
Note that to run the example workflows, you need to have some additional custom nodes, like [ComfyUI-VideoHelperSuite](https://github.com/kosinkadink/ComfyUI-VideoHelperSuite) and others, installed. You can do it by pressing "Install Missing Custom Nodes" button in ComfyUI Manager.
|
||||
**Gemma Text Encoder** (download all assets from the repository)
|
||||
* [`Gemma 3`](https://huggingface.co/google/gemma-3-12b-it-qat-q4_0-unquantized/blob/main)
|
||||
|
||||
### Image-to-video
|
||||
**LoRAs**
|
||||
* [`ltx-2-19b-ic-lora-canny-control.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Canny-Control/blob/main/ltx-2-19b-ic-lora-canny-control.safetensors)
|
||||
* [`ltx-2-19b-ic-lora-depth-control.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Depth-Control/blob/main/ltx-2-19b-ic-lora-depth-control.safetensors)
|
||||
* [`ltx-2-19b-ic-lora-detailer.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Detailer/blob/main/ltx-2-19b-ic-lora-detailer.safetensors)
|
||||
* [`ltx-2-19b-ic-lora-pose-control.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Pose-Control/blob/main/ltx-2-19b-ic-lora-pose-control.safetensors)
|
||||
* [`ltx-2-19b-lora-camera-control-dolly-in.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In/blob/main/ltx-2-19b-lora-camera-control-dolly-in.safetensors)
|
||||
* [`ltx-2-19b-lora-camera-control-dolly-left.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Left/blob/main/ltx-2-19b-lora-camera-control-dolly-left.safetensors)
|
||||
* [`ltx-2-19b-lora-camera-control-dolly-out.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Out/blob/main/ltx-2-19b-lora-camera-control-dolly-out.safetensors)
|
||||
* [`ltx-2-19b-lora-camera-control-dolly-right.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Right/blob/main/ltx-2-19b-lora-camera-control-dolly-right.safetensors)
|
||||
* [`ltx-2-19b-lora-camera-control-jib-down.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down/blob/main/ltx-2-19b-lora-camera-control-jib-down.safetensors)
|
||||
* [`ltx-2-19b-lora-camera-control-jib-up.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up/blob/main/ltx-2-19b-lora-camera-control-jib-up.safetensors)
|
||||
* [`ltx-2-19b-lora-camera-control-static.safetensors`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static/blob/main/ltx-2-19b-lora-camera-control-static.safetensors)
|
||||
|
||||
[Download workflow](assets/ltxvideo-i2v.json)
|
||||

|
||||
|
||||
### Text-to-video
|
||||
## Advanced Techniques
|
||||
|
||||
[Download workflow](assets/ltxvideo-t2v.json)
|
||||

|
||||
|
||||
### Frame Interpolation
|
||||
|
||||
[Download workflow](assets/ltxvideo-frame-interpolation.json)
|
||||

|
||||
|
||||
### First Sequence Conditioning
|
||||
|
||||
[Download workflow](assets/ltxvideo-first-sequence-conditioning.json)
|
||||

|
||||
|
||||
### Last Sequence Conditioning
|
||||
|
||||
[Download workflow](assets/ltxvideo-last-sequence-conditioning.json)
|
||||

|
||||
|
||||
### Flow Edit
|
||||
|
||||
[Download workflow](tricks/assets/ltxvideo-flow-edit.json)
|
||||

|
||||
|
||||
### RF Edit
|
||||
|
||||
[Download workflow](tricks/assets/ltxvideo-rf-edit.json)
|
||||

|
||||
### Low VRAM
|
||||
* For systems with low VRAM you can use the model loader nodes from [low_vram_loaders.py](./low_vram_loaders.py). Those nodes ensure the correct order of execution and perform the model offloading such that generation fits in 32 GB VRAM.
|
||||
* For complete information about using LTX-2 models, workflows, and nodes in ComfyUI, please visit our [Open Source documentation](https://docs.ltx.video/open-source-model/integration-tools/comfy-ui).
|
||||
|
||||
+104
-6
@@ -1,15 +1,113 @@
|
||||
from .nodes_registry import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
from .prompt_enhancer_nodes import ( # noqa: F401
|
||||
LTXVPromptEnhancer,
|
||||
LTXVPromptEnhancerLoader,
|
||||
from .decoder_noise import DecoderNoise
|
||||
from .dynamic_conditioning import DynamicConditioning
|
||||
from .easy_samplers import (
|
||||
LinearOverlapLatentTransition,
|
||||
LTXVBaseSampler,
|
||||
LTXVExtendSampler,
|
||||
LTXVInContextSampler,
|
||||
)
|
||||
from .stg import LTXVApplySTG, STGGuiderNode # noqa: F401
|
||||
from .gemma_encoder import LTXVGemmaCLIPModelLoader, LTXVGemmaEnhancePrompt
|
||||
from .gemma_encoder_mgpu_loader import LTXVGemmaCLIPModelLoaderMGPU
|
||||
from .guide import LTXVAddGuideAdvanced
|
||||
from .guiders import GuiderParametersNode, MultimodalGuiderNode
|
||||
from .latent_norm import (
|
||||
LTXVAdainLatent,
|
||||
LTXVPerStepAdainPatcher,
|
||||
LTXVPerStepStatNormPatcher,
|
||||
LTXVStatNormLatent,
|
||||
)
|
||||
from .latents import (
|
||||
LTXVAddLatentGuide,
|
||||
LTXVImgToVideoConditionOnly,
|
||||
LTXVSelectLatents,
|
||||
LTXVSetVideoLatentNoiseMasks,
|
||||
)
|
||||
from .looping_sampler import LTXVLoopingSampler, MultiPromptProvider
|
||||
from .low_vram_loaders import (
|
||||
LowVRAMAudioVAELoader,
|
||||
LowVRAMCheckpointLoader,
|
||||
LowVRAMLatentUpscaleModelLoader,
|
||||
)
|
||||
from .masks import LTXVPreprocessMasks
|
||||
from .nodes_registry import NODE_CLASS_MAPPINGS as RUNTIME_NODE_CLASS_MAPPINGS
|
||||
from .nodes_registry import (
|
||||
NODE_DISPLAY_NAME_MAPPINGS as RUNTIME_NODE_DISPLAY_NAME_MAPPINGS,
|
||||
)
|
||||
from .nodes_registry import NODES_DISPLAY_NAME_PREFIX, camel_case_to_spaces
|
||||
from .prompt_enhancer_nodes import LTXVPromptEnhancer, LTXVPromptEnhancerLoader
|
||||
from .q8_nodes import LTXVQ8LoraModelLoader, LTXVQ8Patch
|
||||
from .stg import (
|
||||
LTXVApplySTG,
|
||||
STGAdvancedPresetsNode,
|
||||
STGGuiderAdvancedNode,
|
||||
STGGuiderNode,
|
||||
)
|
||||
from .tiled_sampler import LTXVTiledSampler
|
||||
from .tiled_vae_decode import LTXVTiledVAEDecode
|
||||
from .tricks import NODE_CLASS_MAPPINGS as TRICKS_NODE_CLASS_MAPPINGS
|
||||
from .tricks import NODE_DISPLAY_NAME_MAPPINGS as TRICKS_NODE_DISPLAY_NAME_MAPPINGS
|
||||
from .utiltily_nodes import ImageToCPU
|
||||
from .vae_patcher import LTXVPatcherVAE
|
||||
|
||||
# Static node mappings, required for ComfyUI-Manager mapping to work
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Set VAE Decoder Noise": DecoderNoise,
|
||||
"LTXVLinearOverlapLatentTransition": LinearOverlapLatentTransition,
|
||||
"LTXVAddGuideAdvanced": LTXVAddGuideAdvanced,
|
||||
"LTXVAddLatentGuide": LTXVAddLatentGuide,
|
||||
"LTXVAdainLatent": LTXVAdainLatent,
|
||||
"LTXVImgToVideoConditionOnly": LTXVImgToVideoConditionOnly,
|
||||
"LTXVPerStepAdainPatcher": LTXVPerStepAdainPatcher,
|
||||
"LTXVApplySTG": LTXVApplySTG,
|
||||
"LTXVBaseSampler": LTXVBaseSampler,
|
||||
"LTXVInContextSampler": LTXVInContextSampler,
|
||||
"LTXVExtendSampler": LTXVExtendSampler,
|
||||
"LTXVPreprocessMasks": LTXVPreprocessMasks,
|
||||
"LTXVPatcherVAE": LTXVPatcherVAE,
|
||||
"LTXVPromptEnhancer": LTXVPromptEnhancer,
|
||||
"LTXVPromptEnhancerLoader": LTXVPromptEnhancerLoader,
|
||||
"LTXQ8Patch": LTXVQ8Patch,
|
||||
"LTXVQ8LoraModelLoader": LTXVQ8LoraModelLoader,
|
||||
"LTXVSelectLatents": LTXVSelectLatents,
|
||||
"LTXVSetVideoLatentNoiseMasks": LTXVSetVideoLatentNoiseMasks,
|
||||
"LTXVTiledSampler": LTXVTiledSampler,
|
||||
"LTXVLoopingSampler": LTXVLoopingSampler,
|
||||
"LTXVTiledVAEDecode": LTXVTiledVAEDecode,
|
||||
"MultimodalGuider": MultimodalGuiderNode,
|
||||
"GuiderParameters": GuiderParametersNode,
|
||||
"STGAdvancedPresets": STGAdvancedPresetsNode,
|
||||
"STGGuiderAdvanced": STGGuiderAdvancedNode,
|
||||
"STGGuiderNode": STGGuiderNode,
|
||||
"LTXVMultiPromptProvider": MultiPromptProvider,
|
||||
"ImageToCPU": ImageToCPU,
|
||||
"LTXVStatNormLatent": LTXVStatNormLatent,
|
||||
"LTXVPerStepStatNormPatcher": LTXVPerStepStatNormPatcher,
|
||||
"LTXVGemmaCLIPModelLoader": LTXVGemmaCLIPModelLoader,
|
||||
"LTXVGemmaCLIPModelLoaderMGPU": LTXVGemmaCLIPModelLoaderMGPU,
|
||||
"LTXVGemmaEnhancePrompt": LTXVGemmaEnhancePrompt,
|
||||
"DynamicConditioning": DynamicConditioning,
|
||||
"LowVRAMCheckpointLoader": LowVRAMCheckpointLoader,
|
||||
"LowVRAMAudioVAELoader": LowVRAMAudioVAELoader,
|
||||
"LowVRAMLatentUpscaleModelLoader": LowVRAMLatentUpscaleModelLoader,
|
||||
}
|
||||
|
||||
# Consistent display names between static and dynamic node mappings in nodes_registry.py,
|
||||
# to prevent ComfyUI initializing them with default display names.
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
name: f"{NODES_DISPLAY_NAME_PREFIX} {camel_case_to_spaces(name)}"
|
||||
for name in NODE_CLASS_MAPPINGS.keys()
|
||||
}
|
||||
|
||||
# Merge the node mappings from tricks into the main mappings
|
||||
NODE_CLASS_MAPPINGS.update(TRICKS_NODE_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(TRICKS_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
|
||||
# Update with runtime mappings (these will override static mappings if there are any differences)
|
||||
NODE_CLASS_MAPPINGS.update(RUNTIME_NODE_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(RUNTIME_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
|
||||
# Export so that ComfyUI can pick them up.
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
__all__ = [
|
||||
"NODE_CLASS_MAPPINGS",
|
||||
"NODE_DISPLAY_NAME_MAPPINGS",
|
||||
]
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 474 KiB |
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:5b8ac2148bdf1092f4342eb4f08961458a9e22743e8896f5dbcf749166f4f526
|
||||
size 1582922
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:adab606ec23e98c5aadb7adbfa9d16cb1e9d460004bbe648cb03f5e89972229a
|
||||
size 1267760
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:97ea4a9a45004726ae84971bef44578270a7be986b1fce533cc8d3ec1a213fdd
|
||||
size 1390379
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3e29e3b630ab07c3404e715dd0f5d62ebdb4ef24af4279239e546f2a401824a3
|
||||
size 1167784
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:7ba0a0ab34867f35268823c6ef8ca4cd8067cd6484e0c6f5b1277c9778bb243e
|
||||
size 1379716
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:0823a77dec3324a8f696072a3988fe50431ce89482698e09047d32852f886cdc
|
||||
size 770420
|
||||
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:8681ff86edb43d0a12472175d5503a590910b24d5ec5d5d5b78eed079960eac0
|
||||
size 998115
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 349 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 570 KiB |
@@ -0,0 +1,57 @@
|
||||
from copy import copy
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="Set VAE Decoder Noise")
|
||||
class DecoderNoise:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE",),
|
||||
"timestep": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.05,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.001,
|
||||
"tooltip": "The timestep used for decoding the noise.",
|
||||
},
|
||||
),
|
||||
"scale": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.025,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.001,
|
||||
"tooltip": "The scale of the noise added to the decoder.",
|
||||
},
|
||||
),
|
||||
"seed": (
|
||||
"INT",
|
||||
{
|
||||
"default": 42,
|
||||
"min": 0,
|
||||
"max": 0xFFFFFFFFFFFFFFFF,
|
||||
"tooltip": "The random seed used for creating the noise.",
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
FUNCTION = "add_noise"
|
||||
RETURN_TYPES = ("VAE",)
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
def add_noise(self, vae, timestep, scale, seed):
|
||||
result = copy(vae)
|
||||
if hasattr(result, "first_stage_model"):
|
||||
result.first_stage_model.decode_timestep = timestep
|
||||
result.first_stage_model.decode_noise_scale = scale
|
||||
result._decode_timestep = timestep
|
||||
result.decode_noise_scale = scale
|
||||
result.seed = seed
|
||||
return (result,)
|
||||
@@ -0,0 +1,57 @@
|
||||
import torch
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="DynamicConditioning")
|
||||
class DynamicConditioning:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"power": ("FLOAT", {"default": 1.3, "min": 1, "max": 2, "step": 0.01}),
|
||||
"only_first_frame": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "apply"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
INIT = False
|
||||
|
||||
def apply(self, model, power, only_first_frame):
|
||||
self.only_first_frame = only_first_frame
|
||||
self.power = power
|
||||
model = model.clone()
|
||||
model.set_model_denoise_mask_function(self.forward)
|
||||
return (model,)
|
||||
|
||||
def find_step(self, sigma: torch.Tensor, step_sigmas: torch.Tensor):
|
||||
for i, step_sigma in enumerate(step_sigmas):
|
||||
if step_sigma <= sigma:
|
||||
return i
|
||||
return len(step_sigmas) - 1
|
||||
|
||||
def forward(
|
||||
self, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict
|
||||
):
|
||||
model = extra_options["model"]
|
||||
step_sigmas = extra_options["sigmas"]
|
||||
step = self.find_step(sigma, step_sigmas)
|
||||
# In order to apply power multiple times, this is the same as applying power number of times equal to step
|
||||
power = self.power**step
|
||||
denoise_mask = denoise_mask.clone()
|
||||
if self.only_first_frame:
|
||||
num_channels = model.model_patcher.model.diffusion_model.in_channels
|
||||
denoise_mask[:, :num_channels, :1] **= power
|
||||
else:
|
||||
denoise_mask **= power
|
||||
# make sure to update the denoise mask in the model, to get correct timestep values for all tokens
|
||||
for k in model.conds:
|
||||
if "positive" in k or "negative" in k:
|
||||
for cond in model.conds[k]:
|
||||
if "model_conds" in cond and "denoise_mask" in cond["model_conds"]:
|
||||
cond["model_conds"]["denoise_mask"].cond = denoise_mask
|
||||
# print(f"DynamicConditioning: power: {power}, step: {step}, sigma: {sigma}, step_sigmas: {step_sigmas}")
|
||||
return denoise_mask
|
||||
@@ -0,0 +1,909 @@
|
||||
import copy
|
||||
|
||||
import comfy
|
||||
import comfy_extras
|
||||
import nodes
|
||||
import torch
|
||||
from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced, SplitSigmas
|
||||
from comfy_extras.nodes_lt import EmptyLTXVLatentVideo, LTXVAddGuide, LTXVCropGuides
|
||||
|
||||
from .guide import blur_internal
|
||||
from .latent_norm import LTXVAdainLatent
|
||||
from .latents import LTXVAddLatentGuide, LTXVSelectLatents
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
def _get_raw_conds_from_guider(guider):
|
||||
if not hasattr(guider, "raw_conds"):
|
||||
if "negative" not in guider.original_conds:
|
||||
# for BasicGuider
|
||||
raise ValueError(
|
||||
"Guider does not have negative conds, cannot use it as a guider."
|
||||
)
|
||||
raw_pos = guider.original_conds["positive"]
|
||||
positive = [[raw_pos[0]["cross_attn"], copy.deepcopy(raw_pos[0])]]
|
||||
raw_neg = guider.original_conds["negative"]
|
||||
negative = [[raw_neg[0]["cross_attn"], copy.deepcopy(raw_neg[0])]]
|
||||
guider.raw_conds = (positive, negative)
|
||||
return guider.raw_conds
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVBaseSampler",
|
||||
)
|
||||
class LTXVBaseSampler:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", {"tooltip": "The diffusion model to use."}),
|
||||
"vae": ("VAE", {"tooltip": "The VAE to use."}),
|
||||
"width": (
|
||||
"INT",
|
||||
{
|
||||
"default": 768,
|
||||
"min": 64,
|
||||
"max": nodes.MAX_RESOLUTION,
|
||||
"step": 32,
|
||||
},
|
||||
),
|
||||
"height": (
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": nodes.MAX_RESOLUTION,
|
||||
"step": 32,
|
||||
},
|
||||
),
|
||||
"num_frames": (
|
||||
"INT",
|
||||
{"default": 97, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 8},
|
||||
),
|
||||
"guider": (
|
||||
"GUIDER",
|
||||
{"tooltip": "The guider to use, must be a STGGuiderAdvanced."},
|
||||
),
|
||||
"sampler": ("SAMPLER", {"tooltip": "The sampler to use."}),
|
||||
"sigmas": ("SIGMAS", {"tooltip": "The sigmas to use."}),
|
||||
"noise": ("NOISE", {"tooltip": "The noise to use for the sampling."}),
|
||||
},
|
||||
"optional": {
|
||||
"optional_cond_images": (
|
||||
"IMAGE",
|
||||
{"tooltip": "The images to use for conditioning the sampling."},
|
||||
),
|
||||
"optional_cond_indices": (
|
||||
"STRING",
|
||||
{
|
||||
"tooltip": "The indices of the images to use for conditioning the sampling."
|
||||
},
|
||||
),
|
||||
"strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.9,
|
||||
"min": 0,
|
||||
"max": 1,
|
||||
"tooltip": "The strength of the conditioning on the images.",
|
||||
},
|
||||
),
|
||||
"crop": (
|
||||
["center", "disabled"],
|
||||
{
|
||||
"default": "disabled",
|
||||
"tooltip": "The crop mode to use for the images.",
|
||||
},
|
||||
),
|
||||
"crf": (
|
||||
"INT",
|
||||
{
|
||||
"default": 35,
|
||||
"min": 0,
|
||||
"max": 100,
|
||||
"tooltip": "The CRF value to use for preprocessing the images.",
|
||||
},
|
||||
),
|
||||
"blur": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 10,
|
||||
"tooltip": "The blur value to use for preprocessing the images.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT", "CONDITIONING", "CONDITIONING")
|
||||
RETURN_NAMES = ("denoised", "positive", "negative")
|
||||
FUNCTION = "sample"
|
||||
CATEGORY = "sampling"
|
||||
|
||||
def sample(
|
||||
self,
|
||||
model,
|
||||
vae,
|
||||
width,
|
||||
height,
|
||||
num_frames,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
noise,
|
||||
optional_cond_images=None,
|
||||
optional_cond_indices=None,
|
||||
strength=0.9,
|
||||
crop="disabled",
|
||||
crf=35,
|
||||
blur=0,
|
||||
optional_negative_index_latents=None,
|
||||
optional_negative_index=-1,
|
||||
optional_negative_index_strength=1.0,
|
||||
optional_initialization_latents=None,
|
||||
guiding_start_step=0,
|
||||
guiding_end_step=1000,
|
||||
):
|
||||
guider = copy.copy(guider)
|
||||
guider.original_conds = copy.deepcopy(guider.original_conds)
|
||||
positive, negative = _get_raw_conds_from_guider(guider)
|
||||
|
||||
if optional_cond_images is not None:
|
||||
optional_cond_images = (
|
||||
comfy.utils.common_upscale(
|
||||
optional_cond_images.movedim(-1, 1),
|
||||
width,
|
||||
height,
|
||||
"bilinear",
|
||||
crop=crop,
|
||||
)
|
||||
.movedim(1, -1)
|
||||
.clamp(0, 1)
|
||||
)
|
||||
optional_cond_images = comfy_extras.nodes_lt.LTXVPreprocess.execute(
|
||||
optional_cond_images, crf
|
||||
)[0]
|
||||
for i in range(optional_cond_images.shape[0]):
|
||||
optional_cond_images[i] = blur_internal(
|
||||
optional_cond_images[i].unsqueeze(0), blur
|
||||
)
|
||||
|
||||
if optional_cond_indices is not None and optional_cond_images is not None:
|
||||
optional_cond_indices = optional_cond_indices.split(",")
|
||||
optional_cond_indices = [int(i) for i in optional_cond_indices]
|
||||
assert len(optional_cond_indices) == len(
|
||||
optional_cond_images
|
||||
), "Number of optional cond images must match number of optional cond indices"
|
||||
|
||||
if optional_initialization_latents is None:
|
||||
(latents,) = EmptyLTXVLatentVideo().execute(width, height, num_frames, 1)
|
||||
else:
|
||||
latents = optional_initialization_latents
|
||||
|
||||
if optional_cond_images is not None and 0 in optional_cond_indices:
|
||||
# apply classical i2v conditioning on the first frame
|
||||
idx_0 = optional_cond_indices.index(0)
|
||||
encode_pixels = optional_cond_images[idx_0 : idx_0 + 1, :, :, :3]
|
||||
t = vae.encode(encode_pixels)
|
||||
latents["samples"][:, :, : t.shape[2]] = t
|
||||
|
||||
if "noise_mask" not in latents:
|
||||
conditioning_latent_frames_mask = torch.ones(
|
||||
(1, 1, latents["samples"].shape[2], 1, 1),
|
||||
dtype=torch.float32,
|
||||
device=latents["samples"].device,
|
||||
)
|
||||
conditioning_latent_frames_mask[:, :, : t.shape[2]] = 1.0 - strength
|
||||
latents["noise_mask"] = conditioning_latent_frames_mask
|
||||
else:
|
||||
latents["noise_mask"][:, :, : t.shape[2]] = 1.0 - strength
|
||||
conditioning_latent_frames_mask = latents["noise_mask"]
|
||||
else:
|
||||
conditioning_latent_frames_mask = None
|
||||
|
||||
high_sigmas, rest_sigmas = SplitSigmas().get_sigmas(sigmas, guiding_start_step)
|
||||
middle_sigmas, low_sigmas = SplitSigmas().get_sigmas(
|
||||
rest_sigmas, guiding_end_step - guiding_start_step
|
||||
)
|
||||
|
||||
if len(high_sigmas) > 1:
|
||||
print("Denoising with no conditioning on sigmas: ", high_sigmas)
|
||||
(_, new_latents) = SamplerCustomAdvanced().sample(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=high_sigmas,
|
||||
latent_image=latents,
|
||||
)
|
||||
|
||||
if optional_cond_images is not None:
|
||||
# add conditioning on keyframes with index > 0
|
||||
for cond_image, cond_idx in zip(
|
||||
optional_cond_images, optional_cond_indices
|
||||
):
|
||||
if cond_idx == 0:
|
||||
# 0 is handled by classical i2v conditioning
|
||||
continue
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
latents,
|
||||
) = LTXVAddGuide.execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
latent=latents,
|
||||
image=cond_image.unsqueeze(0),
|
||||
frame_idx=cond_idx,
|
||||
strength=strength,
|
||||
)
|
||||
|
||||
if optional_negative_index_latents is not None:
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=latents,
|
||||
guiding_latent=optional_negative_index_latents,
|
||||
latent_idx=optional_negative_index,
|
||||
strength=optional_negative_index_strength,
|
||||
)
|
||||
|
||||
guider.set_conds(positive, negative)
|
||||
|
||||
# Denoise the latent video
|
||||
print("Denoising with conditioning on sigmas: ", middle_sigmas)
|
||||
(output_latents, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=middle_sigmas,
|
||||
latent_image=latents,
|
||||
)
|
||||
|
||||
# Clean up guides if image conditioning was used
|
||||
positive, negative, denoised_output_latents = LTXVCropGuides.execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=denoised_output_latents,
|
||||
)
|
||||
|
||||
denoised_output_latents["noise_mask"] = conditioning_latent_frames_mask
|
||||
|
||||
if len(low_sigmas) > 1:
|
||||
print(
|
||||
"Denoising with no conditioning but with classical i2v noise mask on sigmas: ",
|
||||
low_sigmas,
|
||||
)
|
||||
(_, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=low_sigmas,
|
||||
latent_image=denoised_output_latents,
|
||||
)
|
||||
|
||||
return (denoised_output_latents, positive, negative)
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVExtendSampler",
|
||||
)
|
||||
class LTXVExtendSampler:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", {"tooltip": "The diffusion model to use."}),
|
||||
"vae": ("VAE", {"tooltip": "The VAE to use."}),
|
||||
"latents": (
|
||||
"LATENT",
|
||||
{"tooltip": "The latents of the video to extend."},
|
||||
),
|
||||
"num_new_frames": (
|
||||
"INT",
|
||||
{
|
||||
"default": 80,
|
||||
"min": -1,
|
||||
"max": nodes.MAX_RESOLUTION,
|
||||
"step": 1,
|
||||
"tooltip": "If -1, the number of frames will be based on the number of frames in the optional_guiding_latents.",
|
||||
},
|
||||
),
|
||||
"frame_overlap": (
|
||||
"INT",
|
||||
{
|
||||
"default": 16,
|
||||
"min": 16,
|
||||
"max": 128,
|
||||
"step": 8,
|
||||
"tooltip": "The overlap region to use for conditioning the new frames on the end of the provided latents.",
|
||||
},
|
||||
),
|
||||
"guider": (
|
||||
"GUIDER",
|
||||
{"tooltip": "The guider to use, must be a STGGuiderAdvanced."},
|
||||
),
|
||||
"sampler": ("SAMPLER", {"tooltip": "The sampler to use."}),
|
||||
"sigmas": ("SIGMAS", {"tooltip": "The sigmas to use."}),
|
||||
"noise": ("NOISE", {"tooltip": "The noise to use for the sampling."}),
|
||||
"strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"tooltip": "The strength of the conditioning on the overlapping latents, when using optional_guiding_latents.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"optional_guiding_latents": (
|
||||
"LATENT",
|
||||
{"tooltip": "Optional latents to guide the sampling."},
|
||||
),
|
||||
"optional_cond_images": (
|
||||
"IMAGE",
|
||||
{"tooltip": "The images to use for conditioning the sampling."},
|
||||
),
|
||||
"optional_cond_indices": ("STRING",),
|
||||
"cond_image_strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"tooltip": "The strength of the conditioning on the images.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT", "CONDITIONING", "CONDITIONING")
|
||||
RETURN_NAMES = ("denoised_video", "positive", "negative")
|
||||
FUNCTION = "sample"
|
||||
CATEGORY = "sampling"
|
||||
|
||||
def sample(
|
||||
self,
|
||||
model,
|
||||
vae,
|
||||
latents,
|
||||
num_new_frames,
|
||||
frame_overlap,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
noise,
|
||||
strength=0.5,
|
||||
guiding_strength=1.0,
|
||||
cond_image_strength=1.0,
|
||||
optional_guiding_latents=None,
|
||||
optional_cond_images=None,
|
||||
optional_cond_indices=None,
|
||||
optional_reference_latents=None,
|
||||
optional_initialization_latents=None,
|
||||
adain_factor=0.0,
|
||||
optional_negative_index_latents=None,
|
||||
optional_negative_index=-1,
|
||||
optional_negative_index_strength=1.0,
|
||||
guiding_start_step=0,
|
||||
guiding_end_step=1000,
|
||||
normalize_per_frame=False,
|
||||
):
|
||||
guider = copy.copy(guider)
|
||||
guider.original_conds = copy.deepcopy(guider.original_conds)
|
||||
|
||||
if optional_cond_indices is not None and optional_cond_images is not None:
|
||||
optional_cond_indices = optional_cond_indices.split(",")
|
||||
optional_cond_indices = [int(i) for i in optional_cond_indices]
|
||||
assert len(optional_cond_indices) == len(
|
||||
optional_cond_images
|
||||
), "Number of optional cond images must match number of optional cond indices"
|
||||
|
||||
positive, negative = _get_raw_conds_from_guider(guider)
|
||||
|
||||
samples = latents["samples"]
|
||||
batch, channels, frames, height, width = samples.shape
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
overlap = frame_overlap // time_scale_factor
|
||||
|
||||
if num_new_frames == -1 and optional_guiding_latents is not None:
|
||||
num_new_frames = (
|
||||
optional_guiding_latents["samples"].shape[2] - overlap
|
||||
) * time_scale_factor
|
||||
|
||||
(last_overlap_latents,) = LTXVSelectLatents().select_latents(
|
||||
latents, -overlap, -1
|
||||
)
|
||||
|
||||
if optional_initialization_latents is None:
|
||||
new_latents = EmptyLTXVLatentVideo.execute(
|
||||
width=width * width_scale_factor,
|
||||
height=height * height_scale_factor,
|
||||
length=overlap * time_scale_factor + num_new_frames,
|
||||
batch_size=1,
|
||||
)[0]
|
||||
else:
|
||||
new_latents = optional_initialization_latents
|
||||
|
||||
last_overlap_latents["samples"] = last_overlap_latents["samples"].to(
|
||||
new_latents["samples"].device
|
||||
)
|
||||
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=new_latents,
|
||||
guiding_latent=last_overlap_latents,
|
||||
latent_idx=0,
|
||||
strength=strength,
|
||||
)
|
||||
|
||||
high_sigmas, rest_sigmas = SplitSigmas().get_sigmas(sigmas, guiding_start_step)
|
||||
middle_sigmas, low_sigmas = SplitSigmas().get_sigmas(
|
||||
rest_sigmas, guiding_end_step - guiding_start_step
|
||||
)
|
||||
|
||||
if optional_cond_images is not None:
|
||||
print("Adding conditioning on keyframes")
|
||||
for cond_image, cond_idx in zip(
|
||||
optional_cond_images, optional_cond_indices
|
||||
):
|
||||
if optional_guiding_latents is not None and cond_idx % 8 == 1:
|
||||
raise ValueError(
|
||||
f"Conditioning image index {cond_idx} (relative to the in the temporal chunk) is a multiple of 8 + 1,"
|
||||
"and guiding latents are used. Please provide other conditioning image indices"
|
||||
)
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddGuide.execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
latent=new_latents,
|
||||
image=cond_image.unsqueeze(0),
|
||||
frame_idx=cond_idx,
|
||||
strength=cond_image_strength,
|
||||
)
|
||||
|
||||
if len(high_sigmas) > 1:
|
||||
guider.set_conds(positive, negative)
|
||||
print("Denoising with overlap conditioning only on sigmas: ", high_sigmas)
|
||||
(_, new_latents) = SamplerCustomAdvanced().sample(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=high_sigmas,
|
||||
latent_image=new_latents,
|
||||
)
|
||||
|
||||
if optional_guiding_latents is not None:
|
||||
optional_guiding_latents = LTXVSelectLatents().select_latents(
|
||||
optional_guiding_latents, overlap, -1
|
||||
)[0]
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=new_latents,
|
||||
guiding_latent=optional_guiding_latents,
|
||||
latent_idx=last_overlap_latents["samples"].shape[2],
|
||||
strength=guiding_strength,
|
||||
)
|
||||
|
||||
if optional_negative_index_latents is not None:
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=new_latents,
|
||||
guiding_latent=optional_negative_index_latents,
|
||||
latent_idx=optional_negative_index,
|
||||
strength=optional_negative_index_strength,
|
||||
)
|
||||
|
||||
guider.set_conds(positive, negative)
|
||||
|
||||
# Denoise the latent video
|
||||
print("Denoising with full conditioning on sigmas: ", middle_sigmas)
|
||||
(output_latents, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=middle_sigmas,
|
||||
latent_image=new_latents,
|
||||
)
|
||||
|
||||
positive, negative, denoised_output_latents = LTXVCropGuides.execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=denoised_output_latents,
|
||||
)
|
||||
|
||||
if len(low_sigmas) > 1:
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
denoised_output_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=denoised_output_latents,
|
||||
guiding_latent=last_overlap_latents,
|
||||
latent_idx=0,
|
||||
strength=strength,
|
||||
)
|
||||
|
||||
if optional_cond_images is not None:
|
||||
print("Adding conditioning on keyframes")
|
||||
for cond_image, cond_idx in zip(
|
||||
optional_cond_images, optional_cond_indices
|
||||
):
|
||||
if optional_guiding_latents is not None and cond_idx % 8 == 1:
|
||||
raise ValueError(
|
||||
f"Conditioning image index {cond_idx} (relative to the in the temporal chunk) is a multiple of 8 + 1,"
|
||||
"and guiding latents are used. Please provide other conditioning image indices"
|
||||
)
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
denoised_output_latents,
|
||||
) = LTXVAddGuide.execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
latent=denoised_output_latents,
|
||||
image=cond_image.unsqueeze(0),
|
||||
frame_idx=cond_idx,
|
||||
strength=cond_image_strength,
|
||||
)
|
||||
|
||||
guider.set_conds(positive, negative)
|
||||
print(
|
||||
"Denoising with overlap + keyframes conditioning only on sigmas: ",
|
||||
low_sigmas,
|
||||
)
|
||||
(_, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=low_sigmas,
|
||||
latent_image=denoised_output_latents,
|
||||
)
|
||||
positive, negative, denoised_output_latents = LTXVCropGuides.execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=denoised_output_latents,
|
||||
)
|
||||
|
||||
if optional_reference_latents is not None:
|
||||
denoised_output_latents = LTXVAdainLatent().batch_normalize(
|
||||
latents=denoised_output_latents,
|
||||
reference=optional_reference_latents,
|
||||
factor=adain_factor,
|
||||
per_frame=normalize_per_frame,
|
||||
)[0]
|
||||
|
||||
# drop first output latent as it's a reinterpreted 8-frame latent understood as a 1-frame latent
|
||||
truncated_denoised_output_latents = LTXVSelectLatents().select_latents(
|
||||
denoised_output_latents, 1, -1
|
||||
)[0]
|
||||
|
||||
# Fuse new frames with old ones by calling LinearOverlapLatentTransition
|
||||
(latents,) = LinearOverlapLatentTransition().process(
|
||||
latents, truncated_denoised_output_latents, overlap - 1, axis=2
|
||||
)
|
||||
return (latents, positive, negative)
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVInContextSampler",
|
||||
)
|
||||
class LTXVInContextSampler:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE", {"tooltip": "The VAE to use."}),
|
||||
"guider": (
|
||||
"GUIDER",
|
||||
{"tooltip": "The guider to use, must be a STGGuiderAdvanced."},
|
||||
),
|
||||
"sampler": ("SAMPLER", {"tooltip": "The sampler to use."}),
|
||||
"sigmas": ("SIGMAS", {"tooltip": "The sigmas to use."}),
|
||||
"noise": ("NOISE", {"tooltip": "The noise to use for the sampling."}),
|
||||
"guiding_latents": (
|
||||
"LATENT",
|
||||
{
|
||||
"tooltip": "The latents to use for guiding the sampling, typically with an IC-LoRA."
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"optional_cond_images": (
|
||||
"IMAGE",
|
||||
{
|
||||
"tooltip": "The image to use for conditioning the sampling, if not provided, the sampling will be unconditioned (t2v setup). The image will be resized to the size of the first frame."
|
||||
},
|
||||
),
|
||||
"num_frames": (
|
||||
"INT",
|
||||
{
|
||||
"default": -1,
|
||||
"min": -1,
|
||||
"max": 1000,
|
||||
"step": 1,
|
||||
"tooltip": "If -1, the number of frames will be based on the number of frames in the guiding_latents.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT", "CONDITIONING", "CONDITIONING")
|
||||
RETURN_NAMES = ("denoised_video", "positive", "negative")
|
||||
FUNCTION = "sample"
|
||||
CATEGORY = "sampling"
|
||||
|
||||
def sample(
|
||||
self,
|
||||
vae,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
noise,
|
||||
guiding_latents,
|
||||
optional_cond_images=None,
|
||||
optional_cond_indices=None,
|
||||
num_frames=0,
|
||||
optional_initialization_latents=None,
|
||||
optional_negative_index_latents=None,
|
||||
optional_negative_index=-1,
|
||||
optional_negative_index_strength=1.0,
|
||||
cond_image_strength=1.0,
|
||||
guiding_strength=1.0,
|
||||
guiding_start_step=0,
|
||||
guiding_end_step=1000,
|
||||
):
|
||||
guider = copy.copy(guider)
|
||||
guider.original_conds = copy.deepcopy(guider.original_conds)
|
||||
if optional_cond_images is None:
|
||||
optional_cond_indices = None
|
||||
|
||||
if optional_cond_indices is not None and optional_cond_images is not None:
|
||||
optional_cond_indices = optional_cond_indices.split(",")
|
||||
optional_cond_indices = [int(i) for i in optional_cond_indices]
|
||||
assert len(optional_cond_indices) == len(
|
||||
optional_cond_images
|
||||
), "Number of optional cond images must match number of optional cond indices"
|
||||
|
||||
positive, negative = _get_raw_conds_from_guider(guider)
|
||||
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
|
||||
batch, channels, frames, height, width = guiding_latents["samples"].shape
|
||||
if num_frames != -1:
|
||||
frames = (num_frames - 1) // time_scale_factor + 1
|
||||
|
||||
if optional_initialization_latents is not None:
|
||||
new_latents = optional_initialization_latents
|
||||
else:
|
||||
new_latents = EmptyLTXVLatentVideo.execute(
|
||||
width=width * width_scale_factor,
|
||||
height=height * height_scale_factor,
|
||||
length=(frames - 1) * time_scale_factor + 1,
|
||||
batch_size=1,
|
||||
)[0]
|
||||
|
||||
high_sigmas, rest_sigmas = SplitSigmas().get_sigmas(sigmas, guiding_start_step)
|
||||
middle_sigmas, low_sigmas = SplitSigmas().get_sigmas(
|
||||
rest_sigmas, guiding_end_step - guiding_start_step
|
||||
)
|
||||
|
||||
if len(high_sigmas) > 1:
|
||||
print(
|
||||
"Denoising with keyframes only [if available] on sigmas: ",
|
||||
high_sigmas,
|
||||
)
|
||||
(_, new_latents) = SamplerCustomAdvanced().sample(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=high_sigmas,
|
||||
latent_image=new_latents,
|
||||
)
|
||||
|
||||
if optional_cond_indices is not None and 0 in optional_cond_indices:
|
||||
guiding_latents = LTXVSelectLatents().select_latents(
|
||||
guiding_latents, 1, -1
|
||||
)[0]
|
||||
skip_one_guiding_latent = True
|
||||
else:
|
||||
skip_one_guiding_latent = False
|
||||
|
||||
print("Adding conditioning on guiding latents")
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=new_latents,
|
||||
guiding_latent=guiding_latents,
|
||||
latent_idx=1 if skip_one_guiding_latent else 0,
|
||||
strength=guiding_strength,
|
||||
)
|
||||
if optional_cond_images is not None:
|
||||
print("Adding conditioning on keyframes", optional_cond_indices)
|
||||
for cond_image, cond_idx in zip(
|
||||
optional_cond_images, optional_cond_indices
|
||||
):
|
||||
if cond_idx % 8 == 1:
|
||||
raise ValueError(
|
||||
f"Conditioning image index {cond_idx} is a multiple of 8 + 1 and guiding latents are used. Please provide other cond image indices"
|
||||
)
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddGuide.execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
latent=new_latents,
|
||||
image=cond_image.unsqueeze(0),
|
||||
frame_idx=cond_idx,
|
||||
strength=cond_image_strength,
|
||||
)
|
||||
if optional_negative_index_latents is not None:
|
||||
(
|
||||
positive,
|
||||
negative,
|
||||
new_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
vae=vae,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=new_latents,
|
||||
guiding_latent=optional_negative_index_latents,
|
||||
latent_idx=optional_negative_index,
|
||||
strength=optional_negative_index_strength,
|
||||
)
|
||||
|
||||
guider.set_conds(positive, negative)
|
||||
|
||||
# Denoise the latent video
|
||||
print("Denoising with full conditioning on sigmas: ", middle_sigmas)
|
||||
(_, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=middle_sigmas,
|
||||
latent_image=new_latents,
|
||||
)
|
||||
|
||||
# Clean up guides if image conditioning was used
|
||||
positive, negative, denoised_output_latents = LTXVCropGuides.execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=denoised_output_latents,
|
||||
)
|
||||
|
||||
if len(low_sigmas) > 1:
|
||||
guider.set_conds(positive, negative)
|
||||
print(
|
||||
"Denoising with keyframes only [if available] conditioning on sigmas: ",
|
||||
low_sigmas,
|
||||
)
|
||||
(_, denoised_output_latents) = SamplerCustomAdvanced().sample(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=low_sigmas,
|
||||
latent_image=denoised_output_latents,
|
||||
)
|
||||
positive, negative, denoised_output_latents = LTXVCropGuides.execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
latent=denoised_output_latents,
|
||||
)
|
||||
|
||||
return (denoised_output_latents, positive, negative)
|
||||
|
||||
|
||||
@comfy_node(description="Linear transition with overlap")
|
||||
class LinearOverlapLatentTransition:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"samples1": ("LATENT",),
|
||||
"samples2": ("LATENT",),
|
||||
"overlap": ("INT", {"default": 1, "min": 1, "max": 256}),
|
||||
},
|
||||
"optional": {
|
||||
"axis": ("INT", {"default": 0}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "process"
|
||||
|
||||
CATEGORY = "Lightricks/latent"
|
||||
|
||||
def get_subbatch(self, samples):
|
||||
s = samples.copy()
|
||||
samples = s["samples"]
|
||||
return samples
|
||||
|
||||
def process(self, samples1, samples2, overlap, axis=0):
|
||||
samples1 = self.get_subbatch(samples1)
|
||||
samples2 = self.get_subbatch(samples2)
|
||||
|
||||
# Create transition coefficients
|
||||
alpha = torch.linspace(1, 0, overlap + 2)[1:-1].to(samples1.device)
|
||||
|
||||
# Create shape for broadcasting based on the axis
|
||||
shape = [1] * samples1.dim()
|
||||
shape[axis] = alpha.size(0)
|
||||
alpha = alpha.reshape(shape)
|
||||
|
||||
# Create slices for the overlap regions
|
||||
slice_all = [slice(None)] * samples1.dim()
|
||||
slice_overlap1 = slice_all.copy()
|
||||
slice_overlap1[axis] = slice(-overlap, None)
|
||||
slice_overlap2 = slice_all.copy()
|
||||
slice_overlap2[axis] = slice(0, overlap)
|
||||
slice_rest1 = slice_all.copy()
|
||||
slice_rest1[axis] = slice(None, -overlap)
|
||||
slice_rest2 = slice_all.copy()
|
||||
slice_rest2[axis] = slice(overlap, None)
|
||||
|
||||
# Combine samples
|
||||
parts = [
|
||||
samples1[tuple(slice_rest1)],
|
||||
alpha * samples1[tuple(slice_overlap1)]
|
||||
+ (1 - alpha) * samples2[tuple(slice_overlap2)],
|
||||
samples2[tuple(slice_rest2)],
|
||||
]
|
||||
|
||||
combined_samples = torch.cat(parts, dim=axis)
|
||||
combined_batch_index = torch.arange(0, combined_samples.shape[0]).to(
|
||||
dtype=torch.float32
|
||||
)
|
||||
|
||||
return (
|
||||
{
|
||||
"samples": combined_samples,
|
||||
"batch_index": combined_batch_index,
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,331 @@
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import comfy.ldm.common_dit
|
||||
import torch
|
||||
from comfy.ldm.lightricks.model import (
|
||||
CrossAttention,
|
||||
FeedForward,
|
||||
generate_freq_grid_np,
|
||||
interleaved_freqs_cis,
|
||||
split_freqs_cis,
|
||||
)
|
||||
from torch import nn
|
||||
|
||||
from .pos_embedding_exp_values import POS_EMBEDDING_EXP_VALUES
|
||||
|
||||
|
||||
class BasicTransformerBlock1D(nn.Module):
|
||||
r"""
|
||||
A basic Transformer block.
|
||||
|
||||
Parameters:
|
||||
|
||||
dim (`int`): The number of channels in the input and output.
|
||||
num_attention_heads (`int`): The number of heads to use for multi-head attention.
|
||||
attention_head_dim (`int`): The number of channels in each head.
|
||||
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
||||
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
|
||||
attention_bias (:
|
||||
obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
|
||||
upcast_attention (`bool`, *optional*):
|
||||
Whether to upcast the attention computation to float32. This is useful for mixed precision training.
|
||||
norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
|
||||
Whether to use learnable elementwise affine parameters for normalization.
|
||||
standardization_norm (`str`, *optional*, defaults to `"layer_norm"`): The type of pre-normalization to use. Can be `"layer_norm"` or `"rms_norm"`.
|
||||
norm_eps (`float`, *optional*, defaults to 1e-5): Epsilon value for normalization layers.
|
||||
qk_norm (`str`, *optional*, defaults to None):
|
||||
Set to 'layer_norm' or `rms_norm` to perform query and key normalization.
|
||||
final_dropout (`bool` *optional*, defaults to False):
|
||||
Whether to apply a final dropout after the last feed-forward layer.
|
||||
ff_inner_dim (`int`, *optional*): Dimension of the inner feed-forward layer. If not provided, defaults to `dim * 4`.
|
||||
ff_bias (`bool`, *optional*, defaults to `True`): Whether to use bias in the feed-forward layer.
|
||||
attention_out_bias (`bool`, *optional*, defaults to `True`): Whether to use bias in the attention output layer.
|
||||
use_rope (`bool`, *optional*, defaults to `False`): Whether to use Rotary Position Embeddings (RoPE).
|
||||
ffn_dim_mult (`int`, *optional*, defaults to 4): Multiplier for the inner dimension of the feed-forward layer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
n_heads,
|
||||
d_head,
|
||||
context_dim=None,
|
||||
attn_precision=None,
|
||||
dtype=None,
|
||||
device=None,
|
||||
operations=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Define 3 blocks. Each block has its own normalization layer.
|
||||
# 1. Self-Attn
|
||||
self.attn1 = CrossAttention(
|
||||
query_dim=dim,
|
||||
heads=n_heads,
|
||||
dim_head=d_head,
|
||||
context_dim=None,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
)
|
||||
|
||||
# 3. Feed-forward
|
||||
self.ff = FeedForward(
|
||||
dim,
|
||||
dim_out=dim,
|
||||
glu=True,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
)
|
||||
|
||||
def forward(self, hidden_states, attention_mask=None, pe=None) -> torch.FloatTensor:
|
||||
|
||||
# Notice that normalization is always applied before the real computation in the following blocks.
|
||||
|
||||
# 1. Normalization Before Self-Attention
|
||||
norm_hidden_states = comfy.ldm.common_dit.rms_norm(hidden_states)
|
||||
|
||||
norm_hidden_states = norm_hidden_states.squeeze(1)
|
||||
|
||||
# 2. Self-Attention
|
||||
attn_output = self.attn1(norm_hidden_states, mask=attention_mask, pe=pe)
|
||||
|
||||
hidden_states = attn_output + hidden_states
|
||||
if hidden_states.ndim == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
|
||||
# 3. Normalization before Feed-Forward
|
||||
norm_hidden_states = comfy.ldm.common_dit.rms_norm(hidden_states)
|
||||
|
||||
# 4. Feed-forward
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
|
||||
hidden_states = ff_output + hidden_states
|
||||
if hidden_states.ndim == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class Embeddings1DConnector(nn.Module):
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels=128,
|
||||
cross_attention_dim=2048,
|
||||
attention_head_dim=128,
|
||||
num_attention_heads=30,
|
||||
num_layers=2,
|
||||
positional_embedding_theta=10000.0,
|
||||
positional_embedding_max_pos=[1],
|
||||
causal_temporal_positioning=False,
|
||||
num_learnable_registers: Optional[int] = 128,
|
||||
dtype=None,
|
||||
device=None,
|
||||
operations=None,
|
||||
split_rope=False,
|
||||
double_precision_rope=False,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.dtype = dtype
|
||||
self.out_channels = in_channels
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.inner_dim = num_attention_heads * attention_head_dim
|
||||
self.causal_temporal_positioning = causal_temporal_positioning
|
||||
self.positional_embedding_theta = positional_embedding_theta
|
||||
self.positional_embedding_max_pos = positional_embedding_max_pos
|
||||
self.split_rope = split_rope
|
||||
self.double_precision_rope = double_precision_rope
|
||||
self.transformer_1d_blocks = nn.ModuleList(
|
||||
[
|
||||
BasicTransformerBlock1D(
|
||||
self.inner_dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
context_dim=cross_attention_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
self.num_learnable_registers = num_learnable_registers
|
||||
if self.num_learnable_registers:
|
||||
self.learnable_registers = nn.Parameter(
|
||||
torch.rand(
|
||||
self.num_learnable_registers, inner_dim, dtype=dtype, device=device
|
||||
)
|
||||
* 2.0
|
||||
- 1.0
|
||||
)
|
||||
|
||||
def get_fractional_positions(self, indices_grid):
|
||||
fractional_positions = torch.stack(
|
||||
[
|
||||
indices_grid[:, i] / self.positional_embedding_max_pos[i]
|
||||
for i in range(1)
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
return fractional_positions
|
||||
|
||||
def precompute_freqs(self, indices_grid, spacing):
|
||||
source_dtype = indices_grid.dtype
|
||||
dtype = (
|
||||
torch.float32
|
||||
if source_dtype in (torch.bfloat16, torch.float16)
|
||||
else source_dtype
|
||||
)
|
||||
|
||||
fractional_positions = self.get_fractional_positions(indices_grid)
|
||||
indices = (
|
||||
generate_freq_grid_np(
|
||||
self.positional_embedding_theta,
|
||||
indices_grid.shape[1],
|
||||
self.inner_dim,
|
||||
)
|
||||
if self.double_precision_rope
|
||||
else self.generate_freq_grid(spacing, dtype, fractional_positions.device)
|
||||
).to(device=fractional_positions.device)
|
||||
|
||||
if spacing == "exp_2":
|
||||
freqs = (
|
||||
(indices * fractional_positions.unsqueeze(-1))
|
||||
.transpose(-1, -2)
|
||||
.flatten(2)
|
||||
)
|
||||
else:
|
||||
freqs = (
|
||||
(indices * (fractional_positions.unsqueeze(-1) * 2 - 1))
|
||||
.transpose(-1, -2)
|
||||
.flatten(2)
|
||||
)
|
||||
return freqs
|
||||
|
||||
def generate_freq_grid(self, spacing, dtype, device):
|
||||
dim = self.inner_dim
|
||||
theta = self.positional_embedding_theta
|
||||
n_pos_dims = 1
|
||||
n_elem = 2 * n_pos_dims # 2 for cos and sin e.g. x 3 = 6
|
||||
start = 1
|
||||
end = theta
|
||||
|
||||
if spacing == "exp":
|
||||
# Workaround float precision issues when calculating positional embeddings exponents to ensure compatibility with JAX.
|
||||
indices = torch.tensor(POS_EMBEDDING_EXP_VALUES, dtype=dtype, device=device)
|
||||
elif spacing == "exp_2":
|
||||
indices = 1.0 / theta ** (torch.arange(0, dim, n_elem, device=device) / dim)
|
||||
indices = indices.to(dtype=dtype)
|
||||
elif spacing == "linear":
|
||||
indices = torch.linspace(
|
||||
start, end, dim // n_elem, device=device, dtype=dtype
|
||||
)
|
||||
elif spacing == "sqrt":
|
||||
indices = torch.linspace(
|
||||
start**2, end**2, dim // n_elem, device=device, dtype=dtype
|
||||
).sqrt()
|
||||
|
||||
indices = indices * math.pi / 2
|
||||
|
||||
return indices
|
||||
|
||||
def precompute_freqs_cis(self, indices_grid, spacing="exp"):
|
||||
dim = self.inner_dim
|
||||
n_elem = 2 # 2 because of cos and sin
|
||||
freqs = self.precompute_freqs(indices_grid, spacing)
|
||||
if self.split_rope:
|
||||
expected_freqs = dim // 2
|
||||
current_freqs = freqs.shape[-1]
|
||||
pad_size = expected_freqs - current_freqs
|
||||
cos_freq, sin_freq = split_freqs_cis(
|
||||
freqs, pad_size, self.num_attention_heads
|
||||
)
|
||||
else:
|
||||
cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
|
||||
return cos_freq.to(self.dtype), sin_freq.to(self.dtype), self.split_rope
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""
|
||||
The [`Transformer2DModel`] forward method.
|
||||
|
||||
Args:
|
||||
hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
|
||||
Input `hidden_states`.
|
||||
indices_grid (`torch.LongTensor` of shape `(batch size, 3, num latent pixels)`):
|
||||
attention_mask ( `torch.Tensor`, *optional*):
|
||||
An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
|
||||
is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
|
||||
negative values to the attention scores corresponding to "discard" tokens.
|
||||
Returns:
|
||||
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
|
||||
`tuple` where the first element is the sample tensor.
|
||||
"""
|
||||
# 1. Input
|
||||
|
||||
if self.num_learnable_registers:
|
||||
# replace all padded tokens with learnable registers
|
||||
assert (
|
||||
hidden_states.shape[1] % self.num_learnable_registers == 0
|
||||
), f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by num_learnable_registers {self.num_learnable_registers}."
|
||||
|
||||
num_registers_duplications = (
|
||||
hidden_states.shape[1] // self.num_learnable_registers
|
||||
)
|
||||
learnable_registers = torch.tile(
|
||||
self.learnable_registers, (num_registers_duplications, 1)
|
||||
)
|
||||
attention_mask_binary = (
|
||||
attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0
|
||||
).int()
|
||||
|
||||
non_zero_hidden_states = hidden_states[
|
||||
:, attention_mask_binary.squeeze().bool(), :
|
||||
]
|
||||
non_zero_nums = non_zero_hidden_states.shape[1]
|
||||
pad_length = hidden_states.shape[1] - non_zero_nums
|
||||
adjusted_hidden_states = torch.nn.functional.pad(
|
||||
non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0
|
||||
)
|
||||
flipped_mask = torch.flip(attention_mask_binary, dims=[1])
|
||||
hidden_states = (
|
||||
flipped_mask * adjusted_hidden_states
|
||||
+ (1 - flipped_mask) * learnable_registers
|
||||
)
|
||||
|
||||
attention_mask = torch.full_like(
|
||||
attention_mask,
|
||||
0.0,
|
||||
dtype=attention_mask.dtype,
|
||||
device=attention_mask.device,
|
||||
)
|
||||
|
||||
indices_grid = torch.arange(
|
||||
hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device
|
||||
)
|
||||
indices_grid = indices_grid[None, None, :]
|
||||
freqs_cis = self.precompute_freqs_cis(indices_grid)
|
||||
|
||||
# 2. Blocks
|
||||
for block_idx, block in enumerate(self.transformer_1d_blocks):
|
||||
hidden_states = block(
|
||||
hidden_states, attention_mask=attention_mask, pe=freqs_cis
|
||||
)
|
||||
|
||||
# 3. Output
|
||||
# if self.output_scale is not None:
|
||||
# hidden_states = hidden_states / self.output_scale
|
||||
|
||||
hidden_states = comfy.ldm.common_dit.rms_norm(hidden_states)
|
||||
|
||||
return hidden_states, attention_mask
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:6557896dab6ca788224de10cf66bb3b46fce990b921f4bb9783c2f43de77c209
|
||||
size 1761966
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:0f32314bb85cde1b2fcc978a71d0fd93242c1132c397989a3fa1f4486574a3ff
|
||||
size 1760295
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:4cfb677a93331ab9614c820ab4d7436eac0f54a4685a6bf8f9083000ecc66f6f
|
||||
size 6625560
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:64a49c9e6d5c8d9426afa1b3b3fdbac6cbf491e7e65cc21bdf2c677743a3982d
|
||||
size 1130439
|
||||
@@ -0,0 +1,714 @@
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.ops
|
||||
import comfy.sd
|
||||
import comfy.supported_models_base
|
||||
import folder_paths
|
||||
import safetensors
|
||||
import torch
|
||||
from comfy.ldm.lightricks.model import LTXFrequenciesPrecision, LTXRopeType
|
||||
from comfy.utils import load_torch_file
|
||||
from einops import rearrange
|
||||
from PIL import Image
|
||||
from transformers import (
|
||||
AutoImageProcessor,
|
||||
AutoTokenizer,
|
||||
Gemma3ForConditionalGeneration,
|
||||
Gemma3Processor,
|
||||
)
|
||||
|
||||
from .embeddings_connector import Embeddings1DConnector
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PREFIX_BASE = "model.diffusion_model."
|
||||
|
||||
|
||||
def _load_system_prompt(filename: str) -> str:
|
||||
"""Load system prompt from file at module level."""
|
||||
try:
|
||||
prompt_path = Path(__file__).parent / "system_prompts" / filename
|
||||
if prompt_path.exists():
|
||||
return prompt_path.read_text(encoding="utf-8").strip()
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not load {filename}: {e}")
|
||||
return ""
|
||||
|
||||
|
||||
DEFAULT_T2V_SYSTEM_PROMPT = _load_system_prompt("gemma_t2v_system_prompt.txt")
|
||||
DEFAULT_I2V_SYSTEM_PROMPT = _load_system_prompt("gemma_i2v_system_prompt.txt")
|
||||
|
||||
|
||||
def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
|
||||
"""Convert ComfyUI image tensor to PIL Image."""
|
||||
if tensor.dim() == 4:
|
||||
tensor = tensor[0]
|
||||
numpy_image = (tensor.cpu().numpy() * 255).astype("uint8")
|
||||
return Image.fromarray(numpy_image)
|
||||
|
||||
|
||||
def load_video_embeddings_connector(ltxv_path, dtype=torch.bfloat16):
|
||||
sd, metadata = load_torch_file(str(ltxv_path), return_metadata=True)
|
||||
config = json.loads(metadata.get("config", "{}"))
|
||||
transformer_config = config.get("transformer", {})
|
||||
rope_type = LTXRopeType.from_dict(transformer_config)
|
||||
frequencies_precision = LTXFrequenciesPrecision.from_dict(transformer_config)
|
||||
pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
|
||||
sd_keys = list(sd.keys())
|
||||
|
||||
video_only_connector_prefix = f"{PREFIX_BASE}embeddings_connector."
|
||||
av_connector_prefix = f"{PREFIX_BASE}video_embeddings_connector."
|
||||
prefix = (
|
||||
av_connector_prefix
|
||||
if f"{PREFIX_BASE}audio_adaln_single.linear.weight" in sd_keys
|
||||
else video_only_connector_prefix
|
||||
)
|
||||
return load_embeddings_connector(
|
||||
sd, prefix, dtype, rope_type, frequencies_precision, pe_max_pos
|
||||
)
|
||||
|
||||
|
||||
def load_audio_embeddings_connector(ltxv_path, dtype=torch.bfloat16):
|
||||
sd, metadata = load_torch_file(str(ltxv_path), return_metadata=True)
|
||||
config = json.loads(metadata.get("config", "{}"))
|
||||
transformer_config = config.get("transformer", {})
|
||||
rope_type = LTXRopeType.from_dict(transformer_config)
|
||||
frequencies_precision = LTXFrequenciesPrecision.from_dict(transformer_config)
|
||||
pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
|
||||
return load_embeddings_connector(
|
||||
sd,
|
||||
f"{PREFIX_BASE}audio_embeddings_connector.",
|
||||
dtype,
|
||||
rope_type,
|
||||
frequencies_precision,
|
||||
pe_max_pos,
|
||||
)
|
||||
|
||||
|
||||
def load_embeddings_connector(
|
||||
sd,
|
||||
connector_prefix,
|
||||
dtype=torch.bfloat16,
|
||||
rope_type=LTXRopeType.INTERLEAVED,
|
||||
frequencies_precision=LTXFrequenciesPrecision.FLOAT32,
|
||||
pe_max_pos=None,
|
||||
):
|
||||
sd_connector = {
|
||||
k[len(connector_prefix) :]: v
|
||||
for k, v in sd.items()
|
||||
if k.startswith(connector_prefix)
|
||||
}
|
||||
|
||||
if len(sd_connector) == 0:
|
||||
return None
|
||||
|
||||
operations = comfy.ops.pick_operations(dtype, dtype, disable_fast_fp8=True)
|
||||
connector = Embeddings1DConnector(
|
||||
dtype=dtype,
|
||||
operations=operations,
|
||||
positional_embedding_max_pos=pe_max_pos if pe_max_pos is not None else [1],
|
||||
split_rope=rope_type == LTXRopeType.SPLIT,
|
||||
double_precision_rope=frequencies_precision == LTXFrequenciesPrecision.FLOAT64,
|
||||
)
|
||||
connector.load_state_dict(sd_connector)
|
||||
return connector
|
||||
|
||||
|
||||
def load_proj_matrix_from_ltxv(ltxv_path: Path, prefix=""):
|
||||
with safetensors.safe_open(ltxv_path, framework="pt", device="cpu") as f:
|
||||
keys = filter(lambda key: key.startswith(prefix), f.keys())
|
||||
sd = {k.removeprefix(prefix): f.get_tensor(k) for k in keys if k in f.keys()}
|
||||
if not sd:
|
||||
return None
|
||||
model = GemmaFeaturesExtractorProjLinear()
|
||||
model.load_state_dict(sd)
|
||||
return model
|
||||
|
||||
|
||||
def load_proj_matrix_from_checkpoint(checkpoint_path: Path):
|
||||
"""
|
||||
Load model weights from a checkpoint file.
|
||||
:param checkpoint_path: Path to the checkpoint file.
|
||||
"""
|
||||
model = GemmaFeaturesExtractorProjLinear()
|
||||
loaded_state_dict = load_torch_file(str(checkpoint_path), return_metadata=False)
|
||||
if "aggregate_embed.weight" not in loaded_state_dict:
|
||||
raise ValueError(
|
||||
f"Checkpoint {checkpoint_path} does not contain 'aggregate_embed.weight'."
|
||||
)
|
||||
model.load_state_dict(loaded_state_dict)
|
||||
return model
|
||||
|
||||
|
||||
class LTXVGemmaTokenizer:
|
||||
def __init__(self, tokenizer_path: str, max_length: int = 1024):
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
tokenizer_path, local_files_only=True, model_max_length=max_length
|
||||
)
|
||||
# Gemma expects left padding for chat-style prompts; for plain text it doesn't matter much.
|
||||
self.tokenizer.padding_side = "left"
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
self.max_length = max_length
|
||||
|
||||
def tokenize_with_weights(self, text: str, return_word_ids: bool = False):
|
||||
text = text.strip()
|
||||
encoded = self.tokenizer(
|
||||
text,
|
||||
padding="max_length",
|
||||
max_length=self.max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
input_ids = encoded.input_ids
|
||||
attention_mask = encoded.attention_mask
|
||||
tuples = [
|
||||
(token_id, attn, i)
|
||||
for i, (token_id, attn) in enumerate(zip(input_ids[0], attention_mask[0]))
|
||||
]
|
||||
out = {"gemma": tuples}
|
||||
|
||||
if not return_word_ids:
|
||||
out = {k: [(t, w) for t, w, _ in v] for k, v in out.items()}
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class GemmaFeaturesExtractorProjLinear(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.aggregate_embed = torch.nn.Linear(3840 * 49, 3840, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
return self.aggregate_embed(x)
|
||||
|
||||
|
||||
class LTXVGemmaTextEncoderModel(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
model: Gemma3ForConditionalGeneration,
|
||||
feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
|
||||
embeddings_connector: Embeddings1DConnector,
|
||||
audio_embeddings_connector: Embeddings1DConnector | None = None,
|
||||
processor: Gemma3Processor | None = None,
|
||||
dtype=torch.bfloat16,
|
||||
device="cpu",
|
||||
):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.processor = processor
|
||||
self.feature_extractor_linear = feature_extractor_linear.to(dtype=dtype)
|
||||
self.embeddings_connector = embeddings_connector.to(dtype=dtype)
|
||||
self.audio_embeddings_connector = (
|
||||
audio_embeddings_connector.to(dtype=dtype)
|
||||
if audio_embeddings_connector is not None
|
||||
else None
|
||||
)
|
||||
self.dtypes = set([dtype])
|
||||
# Cache an estimate of memory required to load/keep the model on device
|
||||
# weights size + small overhead
|
||||
self._model_memory_required = (
|
||||
comfy.model_management.module_size(self.model) + 256 * 1024 * 1024
|
||||
)
|
||||
|
||||
def set_clip_options(self, options):
|
||||
pass
|
||||
|
||||
def reset_clip_options(self):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def norm_and_concat_padded_batch(
|
||||
encoded_text: torch.Tensor,
|
||||
sequence_lengths: torch.Tensor,
|
||||
padding_side: str = "right",
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Normalize a 4D tensor [B, T, D, L] per sample and per layer, using sequence_lengths to mask.
|
||||
Returns [B, T, D * L] tensor with original padding preserved.
|
||||
|
||||
Args:
|
||||
encoded_text: 4D tensor [B, T, D, L]
|
||||
sequence_lengths: 1D tensor [B] with actual sequence lengths
|
||||
padding_side: "left" or "right" to indicate which side has padding
|
||||
"""
|
||||
B, T, D, L = encoded_text.shape
|
||||
device = encoded_text.device
|
||||
|
||||
# Build mask: [B, T, 1, 1]
|
||||
token_indices = torch.arange(T, device=device)[None, :] # [1, T]
|
||||
|
||||
if padding_side == "right":
|
||||
# For right padding, valid tokens are from 0 to sequence_length-1
|
||||
mask = token_indices < sequence_lengths[:, None] # [B, T]
|
||||
elif padding_side == "left":
|
||||
# For left padding, valid tokens are from (T - sequence_length) to T-1
|
||||
start_indices = T - sequence_lengths[:, None] # [B, 1]
|
||||
mask = token_indices >= start_indices # [B, T]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"padding_side must be 'left' or 'right', got {padding_side}"
|
||||
)
|
||||
|
||||
mask = rearrange(mask, "b t -> b t 1 1")
|
||||
|
||||
# Compute masked mean: [B, 1, 1, L]
|
||||
masked = encoded_text.masked_fill(~mask, 0.0)
|
||||
# denom = mask.sum(dim=(1, 2), keepdim=True) # [B, 1, 1, 1]
|
||||
denom = (sequence_lengths * D).view(B, 1, 1, 1)
|
||||
mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + 1e-6)
|
||||
|
||||
# Compute masked min/max: [B, 1, 1, L]
|
||||
x_min = encoded_text.masked_fill(~mask, float("inf")).amin(
|
||||
dim=(1, 2), keepdim=True
|
||||
)
|
||||
x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(
|
||||
dim=(1, 2), keepdim=True
|
||||
)
|
||||
range_ = x_max - x_min
|
||||
|
||||
# Normalize only the valid tokens
|
||||
normed = 8 * (encoded_text - mean) / (range_ + 1e-6)
|
||||
|
||||
# concat to be [Batch, T, D * L] - this preserves the original structure
|
||||
normed = normed.reshape(B, T, -1) # [B, T, D * L]
|
||||
|
||||
# Apply mask to preserve original padding (set padded positions to 0)
|
||||
mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, D * L)
|
||||
normed = normed.masked_fill(~mask_flattened, 0.0)
|
||||
|
||||
return normed
|
||||
|
||||
def forward(self, input_ids, attention_mask, padding_side="right"):
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
output_hidden_states=True,
|
||||
)
|
||||
encoded_text_features = torch.stack(outputs.hidden_states, dim=-1)
|
||||
encoded_text_features_dtype = encoded_text_features.dtype
|
||||
|
||||
sequence_lengths = attention_mask.sum(dim=-1)
|
||||
normed_concated_encoded_text_features = (
|
||||
LTXVGemmaTextEncoderModel.norm_and_concat_padded_batch(
|
||||
encoded_text_features, sequence_lengths, padding_side=padding_side
|
||||
)
|
||||
)
|
||||
|
||||
projected = self.feature_extractor_linear(
|
||||
normed_concated_encoded_text_features.to(encoded_text_features_dtype)
|
||||
)
|
||||
|
||||
return projected
|
||||
|
||||
def encode_token_weights(self, token_weight_pairs):
|
||||
token_pairs = token_weight_pairs["gemma"]
|
||||
input_ids = torch.tensor(
|
||||
[[t[0] for t in token_pairs]], device=self.model.device
|
||||
)
|
||||
attention_mask = torch.tensor(
|
||||
[[w[1] for w in token_pairs]], device=self.model.device
|
||||
)
|
||||
|
||||
self.to(self.model.device)
|
||||
|
||||
encoded_input = self(input_ids, attention_mask, padding_side="left")
|
||||
# convert attention mask to format embeddings connector expects.
|
||||
connector_attention_mask = (attention_mask - 1).to(encoded_input.dtype).reshape(
|
||||
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
|
||||
) * torch.finfo(encoded_input.dtype).max
|
||||
encoded, encoded_connector_attention_mask = self.embeddings_connector(
|
||||
encoded_input,
|
||||
connector_attention_mask,
|
||||
)
|
||||
|
||||
# restore the mask values to int64
|
||||
attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)
|
||||
|
||||
attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
|
||||
encoded = encoded * attention_mask
|
||||
|
||||
if self.audio_embeddings_connector is not None:
|
||||
encoded_for_audio, _ = self.audio_embeddings_connector(
|
||||
encoded_input, connector_attention_mask
|
||||
)
|
||||
encoded = torch.concat(
|
||||
[encoded, encoded_for_audio], dim=len(encoded.shape) - 1
|
||||
)
|
||||
|
||||
return encoded, None, {"attention_mask": attention_mask.squeeze(-1)}
|
||||
|
||||
def load_sd(self, sd):
|
||||
return self.model.load_state_dict(sd, strict=False)
|
||||
|
||||
def memory_required(self, input_shape):
|
||||
# Return a conservative estimate in bytesed(input_shape)
|
||||
return self._model_memory_required
|
||||
|
||||
@staticmethod
|
||||
def _cat_with_padding(
|
||||
tensor: torch.Tensor,
|
||||
padding_length: int,
|
||||
value: int | float,
|
||||
) -> torch.Tensor:
|
||||
"""Concatenate a tensor with a padding tensor of the given value."""
|
||||
return torch.cat(
|
||||
[
|
||||
tensor,
|
||||
torch.full(
|
||||
(1, padding_length),
|
||||
value,
|
||||
dtype=tensor.dtype,
|
||||
device=tensor.device,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
def _pad_inputs_for_attention_alignment(self, model_inputs, alignment: int = 8):
|
||||
"""Pad sequence length to multiple of alignment for Flash Attention compatibility.
|
||||
|
||||
Flash Attention within SDPA requires sequence lengths aligned to 8 bytes.
|
||||
This pads input_ids, attention_mask, and token_type_ids (if present) to prevent
|
||||
'p.attn_bias_ptr is not correctly aligned' errors.
|
||||
"""
|
||||
seq_len = model_inputs.input_ids.shape[1]
|
||||
padded_len = ((seq_len + alignment - 1) // alignment) * alignment
|
||||
padding_length = padded_len - seq_len
|
||||
|
||||
if padding_length > 0:
|
||||
pad_token_id = (
|
||||
self.processor.tokenizer.pad_token_id
|
||||
if self.processor.tokenizer.pad_token_id is not None
|
||||
else 0
|
||||
)
|
||||
|
||||
model_inputs["input_ids"] = self._cat_with_padding(
|
||||
model_inputs.input_ids, padding_length, pad_token_id
|
||||
)
|
||||
|
||||
model_inputs["attention_mask"] = self._cat_with_padding(
|
||||
model_inputs.attention_mask, padding_length, 0
|
||||
)
|
||||
|
||||
if (
|
||||
"token_type_ids" in model_inputs
|
||||
and model_inputs["token_type_ids"] is not None
|
||||
):
|
||||
model_inputs["token_type_ids"] = self._cat_with_padding(
|
||||
model_inputs["token_type_ids"], padding_length, 0
|
||||
)
|
||||
|
||||
return model_inputs
|
||||
|
||||
def _enhance(
|
||||
self,
|
||||
messages: list,
|
||||
image: Optional[Image.Image] = None,
|
||||
max_new_tokens: int = 512,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Common enhancement logic for both T2V and I2V modes."""
|
||||
if self.processor is None:
|
||||
raise ValueError("Processor not loaded - enhancement not available")
|
||||
|
||||
text = self.processor.tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
model_inputs = self.processor(
|
||||
text=text,
|
||||
images=image,
|
||||
return_tensors="pt",
|
||||
).to(self.model.device)
|
||||
model_inputs = self._pad_inputs_for_attention_alignment(model_inputs)
|
||||
|
||||
with torch.inference_mode(), torch.random.fork_rng(devices=[self.model.device]):
|
||||
torch.manual_seed(seed)
|
||||
outputs = self.model.generate(
|
||||
**model_inputs,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
)
|
||||
generated_ids = outputs[0][len(model_inputs.input_ids[0]) :]
|
||||
enhanced_prompt = self.processor.tokenizer.decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)
|
||||
|
||||
return enhanced_prompt
|
||||
|
||||
def enhance_t2v(
|
||||
self,
|
||||
prompt: str,
|
||||
system_prompt: str,
|
||||
max_new_tokens: int = 512,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Enhance a text prompt for T2V generation."""
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": f"User Raw Input Prompt: {prompt}."},
|
||||
]
|
||||
return self._enhance(messages, max_new_tokens=max_new_tokens, seed=seed)
|
||||
|
||||
def enhance_i2v(
|
||||
self,
|
||||
prompt: str,
|
||||
image: Image.Image,
|
||||
system_prompt: str,
|
||||
max_new_tokens: int = 512,
|
||||
seed: int = 42,
|
||||
) -> str:
|
||||
"""Enhance a text prompt for I2V generation using a reference image."""
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image"},
|
||||
{"type": "text", "text": f"User Raw Input Prompt: {prompt}."},
|
||||
],
|
||||
},
|
||||
]
|
||||
return self._enhance(
|
||||
messages, image=image, max_new_tokens=max_new_tokens, seed=seed
|
||||
)
|
||||
|
||||
|
||||
def ltxv_gemma_tokenizer(tokenizer_path, max_length=256):
|
||||
class _LTXVGemmaTokenizer(LTXVGemmaTokenizer):
|
||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
||||
super().__init__(tokenizer_path, max_length=max_length)
|
||||
|
||||
return _LTXVGemmaTokenizer
|
||||
|
||||
|
||||
def ltxv_gemma_clip(encoder_path, ltxv_path, processor=None, dtype=None):
|
||||
class _LTXVGemmaTextEncoderModel(LTXVGemmaTextEncoderModel):
|
||||
def __init__(self, device="cpu", dtype=dtype, model_options={}):
|
||||
dtype = torch.bfloat16 # TODO: make this configurable
|
||||
gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
|
||||
encoder_path,
|
||||
local_files_only=True,
|
||||
torch_dtype=dtype,
|
||||
)
|
||||
|
||||
feature_extractor_linear = load_proj_matrix_from_ltxv(
|
||||
ltxv_path,
|
||||
"text_embedding_projection.",
|
||||
)
|
||||
if feature_extractor_linear is None:
|
||||
feature_extractor_linear = load_proj_matrix_from_checkpoint(
|
||||
encoder_path / "proj_linear.safetensors"
|
||||
)
|
||||
|
||||
embeddings_connector = load_video_embeddings_connector(ltxv_path)
|
||||
audio_embeddings_connector = load_audio_embeddings_connector(ltxv_path)
|
||||
super().__init__(
|
||||
model=gemma_model,
|
||||
feature_extractor_linear=feature_extractor_linear,
|
||||
embeddings_connector=embeddings_connector,
|
||||
audio_embeddings_connector=audio_embeddings_connector,
|
||||
processor=processor,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
return _LTXVGemmaTextEncoderModel
|
||||
|
||||
|
||||
def find_matching_dir(root_path: str, pattern: str) -> str:
|
||||
"""
|
||||
Recursively search for files matching a glob pattern and return the parent directory of the first match.
|
||||
"""
|
||||
|
||||
matches = list(Path(root_path).rglob(pattern))
|
||||
if not matches:
|
||||
raise FileNotFoundError(
|
||||
f"No files matching pattern '{pattern}' found under {root_path}"
|
||||
)
|
||||
return str(matches[0].parent)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVGemmaCLIPModelLoader", description="Gemma 3 Model Loader")
|
||||
class LTXVGemmaCLIPModelLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"gemma_path": (
|
||||
folder_paths.get_filename_list("text_encoders"),
|
||||
{"tooltip": "The name of the text encoder model to load."},
|
||||
),
|
||||
"ltxv_path": (
|
||||
folder_paths.get_filename_list("checkpoints"),
|
||||
{"tooltip": "The name of the ltxv model to load."},
|
||||
),
|
||||
"max_length": (
|
||||
"INT",
|
||||
{"default": 1024, "min": 16, "max": 131072, "step": 8},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
RETURN_NAMES = ("clip",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Gemma CLIP Loader"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def load_model(self, gemma_path: str, ltxv_path: str, max_length: int):
|
||||
path = Path(folder_paths.get_full_path("text_encoders", gemma_path))
|
||||
model_root = path.parents[1]
|
||||
tokenizer_path = Path(find_matching_dir(model_root, "tokenizer.model"))
|
||||
gemma_model_path = Path(find_matching_dir(model_root, "model*.safetensors"))
|
||||
processor_path = Path(find_matching_dir(model_root, "preprocessor_config.json"))
|
||||
tokenizer_class = ltxv_gemma_tokenizer(tokenizer_path, max_length=max_length)
|
||||
|
||||
processor = None
|
||||
try:
|
||||
image_processor = AutoImageProcessor.from_pretrained(
|
||||
str(processor_path),
|
||||
local_files_only=True,
|
||||
)
|
||||
processor = Gemma3Processor(
|
||||
image_processor=image_processor,
|
||||
tokenizer=tokenizer_class().tokenizer,
|
||||
)
|
||||
logger.info(f"Loaded processor from {model_root} - enhancement enabled")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not load processor from {model_root}: {e}")
|
||||
|
||||
clip_dtype = torch.bfloat16
|
||||
ltxv_full_path = folder_paths.get_full_path("checkpoints", ltxv_path)
|
||||
clip_target = comfy.supported_models_base.ClipTarget(
|
||||
tokenizer=tokenizer_class,
|
||||
clip=ltxv_gemma_clip(
|
||||
gemma_model_path, ltxv_full_path, processor=processor, dtype=clip_dtype
|
||||
),
|
||||
)
|
||||
|
||||
return (comfy.sd.CLIP(clip_target),)
|
||||
|
||||
|
||||
_UNICODE_REPLACEMENTS = str.maketrans(
|
||||
"\u2018\u2019\u201c\u201d\u2014\u2013\u00a0\u2032\u2212", "''\"\"-- '-"
|
||||
)
|
||||
|
||||
|
||||
def clean_response(text):
|
||||
text = text.translate(_UNICODE_REPLACEMENTS)
|
||||
|
||||
# Remove leading non-letter characters
|
||||
for i, char in enumerate(text):
|
||||
if char.isalpha():
|
||||
return text[i:]
|
||||
return text
|
||||
|
||||
|
||||
@comfy_node(name="LTXVGemmaEnhancePrompt", description="Gemma 3 Prompt Enhancer")
|
||||
class LTXVGemmaEnhancePrompt:
|
||||
"""Enhance prompts using Gemma 3 model. Supports T2V and I2V modes."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"clip": ("CLIP",),
|
||||
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
"system_prompt": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": DEFAULT_T2V_SYSTEM_PROMPT,
|
||||
},
|
||||
),
|
||||
"max_tokens": (
|
||||
"INT",
|
||||
{"default": 512, "min": 32, "max": 1024, "step": 16},
|
||||
),
|
||||
"bypass_i2v": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",),
|
||||
"seed": (
|
||||
"INT",
|
||||
{"default": 42, "min": 0, "max": 0xFFFFFFFF},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("enhanced_prompt",)
|
||||
FUNCTION = "enhance"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Gemma Enhance Prompt"
|
||||
OUTPUT_NODE = True
|
||||
DESCRIPTION = (
|
||||
"Enhance text prompts using Gemma 3 VLLM for improved video generation."
|
||||
)
|
||||
|
||||
def enhance(
|
||||
self,
|
||||
clip,
|
||||
prompt: str,
|
||||
system_prompt: str,
|
||||
max_tokens: int,
|
||||
bypass_i2v: bool,
|
||||
image: Optional[torch.Tensor] = None,
|
||||
seed: int = 42,
|
||||
):
|
||||
if not isinstance(seed, int):
|
||||
seed = 42
|
||||
|
||||
clip.load_model()
|
||||
encoder = clip.cond_stage_model
|
||||
|
||||
if encoder.processor is None:
|
||||
raise ValueError(
|
||||
"Processor not loaded - enhancement not available. "
|
||||
"Ensure your model directory has chat_template.json, processor_config.json, "
|
||||
"and preprocessor_config.json files."
|
||||
)
|
||||
|
||||
# Determine mode: use I2V if image is provided and not bypassed
|
||||
use_i2v = image is not None and not bypass_i2v
|
||||
|
||||
# Auto-select the appropriate system prompt if user is using default T2V prompt
|
||||
if use_i2v and system_prompt.strip() == DEFAULT_T2V_SYSTEM_PROMPT.strip():
|
||||
system_prompt = DEFAULT_I2V_SYSTEM_PROMPT
|
||||
logger.info("Auto-selected I2V system prompt for image-to-video mode")
|
||||
|
||||
if not system_prompt or not system_prompt.strip():
|
||||
raise ValueError(
|
||||
"system_prompt is required and cannot be empty or whitespace-only"
|
||||
)
|
||||
|
||||
if use_i2v:
|
||||
pil_image = tensor_to_pil(image)
|
||||
enhanced_prompt = encoder.enhance_i2v(
|
||||
prompt=prompt,
|
||||
image=pil_image,
|
||||
system_prompt=system_prompt,
|
||||
max_new_tokens=max_tokens,
|
||||
seed=seed,
|
||||
)
|
||||
else:
|
||||
enhanced_prompt = encoder.enhance_t2v(
|
||||
prompt=prompt,
|
||||
system_prompt=system_prompt,
|
||||
max_new_tokens=max_tokens,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
enhanced_prompt = clean_response(enhanced_prompt)
|
||||
|
||||
return (enhanced_prompt,)
|
||||
@@ -0,0 +1,297 @@
|
||||
import json
|
||||
import logging
|
||||
import weakref
|
||||
from pathlib import Path
|
||||
from types import MethodType
|
||||
|
||||
import accelerate
|
||||
import comfy.sd
|
||||
import comfy.supported_models_base
|
||||
import folder_paths
|
||||
import safetensors
|
||||
import torch
|
||||
from comfy.ldm.lightricks.model import LTXFrequenciesPrecision, LTXRopeType
|
||||
from comfy.utils import load_torch_file
|
||||
from transformers import (
|
||||
AutoImageProcessor,
|
||||
Gemma3ForConditionalGeneration,
|
||||
Gemma3Processor,
|
||||
)
|
||||
|
||||
from .embeddings_connector import Embeddings1DConnector
|
||||
from .gemma_encoder import (
|
||||
GemmaFeaturesExtractorProjLinear,
|
||||
LTXVGemmaTextEncoderModel,
|
||||
find_matching_dir,
|
||||
ltxv_gemma_tokenizer,
|
||||
)
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PREFIX_BASE = "model.diffusion_model."
|
||||
|
||||
|
||||
def _is_device_like(x):
|
||||
return isinstance(x, (str, torch.device, int)) or isinstance(
|
||||
x, torch.Tensor
|
||||
) # to(tensor) → matches tensor's device/dtype
|
||||
|
||||
|
||||
def _is_dtype_like(x):
|
||||
# torch.float16 / torch.bfloat16 / torch.float32, etc.
|
||||
return isinstance(x, torch.dtype)
|
||||
|
||||
|
||||
def protect_module_transfers(module: torch.nn.Module, warn=False):
|
||||
"""
|
||||
Make .to/.cuda/.cpu NO-OPs for device moves across an entire module tree,
|
||||
while still allowing dtype/memory_format conversions.
|
||||
"""
|
||||
|
||||
def wrap_one(m: torch.nn.Module):
|
||||
if hasattr(m, "__to_patched"): # idempotent
|
||||
return
|
||||
|
||||
orig_to = m.__class__.to
|
||||
proxy = weakref.proxy(m)
|
||||
|
||||
def safe_to(self, *args, **kwargs):
|
||||
# Block explicit/implicit device moves
|
||||
device_in_kwargs = "device" in kwargs
|
||||
device_like_positional = len(args) > 0 and _is_device_like(args[0])
|
||||
|
||||
# Allow pure dtype/memory_format conversions:
|
||||
# .to(torch.float16)
|
||||
# .to(dtype=torch.float16)
|
||||
# .to(memory_format=...)
|
||||
|
||||
if device_in_kwargs or device_like_positional:
|
||||
if warn:
|
||||
print("[protect_module_transfers] Ignoring device move .to(...).")
|
||||
return self # NO-OP for device moves
|
||||
|
||||
# Otherwise (dtype / memory_format / non_blocking without device) → delegate
|
||||
return orig_to(self, *args, **kwargs)
|
||||
|
||||
# Bind methods
|
||||
m.to = MethodType(safe_to, proxy)
|
||||
m.cuda = MethodType(lambda self, *a, **k: self, proxy) # NO-OP
|
||||
m.cpu = MethodType(lambda self, *a, **k: self, proxy) # NO-OP
|
||||
m.__to_patched = True
|
||||
|
||||
# Recurse over the whole tree
|
||||
for sub in module.modules(): # includes module itself and children
|
||||
wrap_one(sub)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
def load_proj_matrix_from_ltxv(ltxv_path: Path, prefix=""):
|
||||
with safetensors.safe_open(ltxv_path, framework="pt", device="cpu") as f:
|
||||
keys = filter(lambda key: key.startswith(prefix), f.keys())
|
||||
sd = {k.removeprefix(prefix): f.get_tensor(k) for k in keys if k in f.keys()}
|
||||
|
||||
if not sd:
|
||||
return None
|
||||
|
||||
with torch.device("meta"):
|
||||
model = GemmaFeaturesExtractorProjLinear()
|
||||
|
||||
for name, tensor in sd.items():
|
||||
accelerate.utils.set_module_tensor_to_device(model, name, "cpu", value=tensor)
|
||||
|
||||
device_map = accelerate.infer_auto_device_map(
|
||||
model,
|
||||
max_memory=None,
|
||||
)
|
||||
|
||||
model = accelerate.dispatch_model(
|
||||
model,
|
||||
device_map=device_map,
|
||||
)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def load_proj_matrix_from_checkpoint(checkpoint_path: Path):
|
||||
"""
|
||||
Load model weights from a checkpoint file.
|
||||
:param checkpoint_path: Path to the checkpoint file.
|
||||
"""
|
||||
with torch.device("meta"):
|
||||
model = GemmaFeaturesExtractorProjLinear()
|
||||
model = accelerate.load_checkpoint_and_dispatch(
|
||||
model, str(checkpoint_path.resolve()), device_map="auto"
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def ltxv_gemma_clip(encoder_path, ltxv_path, processor=None, dtype=None):
|
||||
class _LTXVGemmaTextEncoderModel(LTXVGemmaTextEncoderModel):
|
||||
def __init__(self, device="cpu", dtype=dtype, model_options={}):
|
||||
dtype = torch.bfloat16 # TODO: make this configurable
|
||||
gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
|
||||
encoder_path,
|
||||
local_files_only=True,
|
||||
torch_dtype=dtype,
|
||||
device_map="auto",
|
||||
)
|
||||
|
||||
feature_extractor_linear = load_proj_matrix_from_ltxv(
|
||||
ltxv_path, "text_embedding_projection."
|
||||
)
|
||||
if feature_extractor_linear is None:
|
||||
feature_extractor_linear = load_proj_matrix_from_checkpoint(
|
||||
encoder_path / "proj_linear.safetensors"
|
||||
)
|
||||
|
||||
embeddings_connector = load_video_embeddings_connector(ltxv_path)
|
||||
audio_embeddings_connector = load_audio_embeddings_connector(ltxv_path)
|
||||
super().__init__(
|
||||
model=gemma_model,
|
||||
feature_extractor_linear=feature_extractor_linear,
|
||||
embeddings_connector=embeddings_connector,
|
||||
audio_embeddings_connector=audio_embeddings_connector,
|
||||
processor=processor,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Due to memory management issue on mgpu we need to disable the offloading of sharded gemma model.
|
||||
protect_module_transfers(self)
|
||||
|
||||
return _LTXVGemmaTextEncoderModel
|
||||
|
||||
|
||||
def load_video_embeddings_connector(ltxv_path, dtype=torch.bfloat16):
|
||||
sd, metadata = load_torch_file(str(ltxv_path), return_metadata=True)
|
||||
config = json.loads(metadata.get("config", "{}"))
|
||||
transformer_config = config.get("transformer", {})
|
||||
rope_type = LTXRopeType.from_dict(transformer_config)
|
||||
frequencies_precision = LTXFrequenciesPrecision.from_dict(transformer_config)
|
||||
pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
|
||||
sd_keys = list(sd.keys())
|
||||
|
||||
video_only_connector_prefix = f"{PREFIX_BASE}embeddings_connector."
|
||||
av_connector_prefix = f"{PREFIX_BASE}video_embeddings_connector."
|
||||
prefix = (
|
||||
av_connector_prefix
|
||||
if f"{PREFIX_BASE}audio_adaln_single.linear.weight" in sd_keys
|
||||
else video_only_connector_prefix
|
||||
)
|
||||
return load_embeddings_connector(
|
||||
sd, prefix, dtype, rope_type, frequencies_precision, pe_max_pos
|
||||
)
|
||||
|
||||
|
||||
def load_audio_embeddings_connector(ltxv_path, dtype=torch.bfloat16):
|
||||
sd, metadata = load_torch_file(str(ltxv_path), return_metadata=True)
|
||||
config = json.loads(metadata.get("config", "{}"))
|
||||
transformer_config = config.get("transformer", {})
|
||||
rope_type = LTXRopeType.from_dict(transformer_config)
|
||||
frequencies_precision = LTXFrequenciesPrecision.from_dict(transformer_config)
|
||||
pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
|
||||
return load_embeddings_connector(
|
||||
sd,
|
||||
f"{PREFIX_BASE}audio_embeddings_connector.",
|
||||
dtype,
|
||||
rope_type,
|
||||
frequencies_precision,
|
||||
pe_max_pos,
|
||||
)
|
||||
|
||||
|
||||
def load_embeddings_connector(
|
||||
sd,
|
||||
connector_prefix,
|
||||
dtype=torch.bfloat16,
|
||||
rope_type=LTXRopeType.INTERLEAVED,
|
||||
frequencies_precision=LTXFrequenciesPrecision.FLOAT32,
|
||||
pe_max_pos=None,
|
||||
):
|
||||
sd_connector = {
|
||||
k[len(connector_prefix) :]: v
|
||||
for k, v in sd.items()
|
||||
if k.startswith(connector_prefix)
|
||||
}
|
||||
|
||||
if len(sd_connector) == 0:
|
||||
return None
|
||||
|
||||
operations = comfy.ops.pick_operations(dtype, dtype, disable_fast_fp8=True)
|
||||
connector = Embeddings1DConnector(
|
||||
dtype=dtype,
|
||||
operations=operations,
|
||||
positional_embedding_max_pos=pe_max_pos if pe_max_pos is not None else [1],
|
||||
split_rope=rope_type == LTXRopeType.SPLIT,
|
||||
double_precision_rope=frequencies_precision == LTXFrequenciesPrecision.FLOAT64,
|
||||
)
|
||||
connector.load_state_dict(sd_connector)
|
||||
device_map = accelerate.infer_auto_device_map(connector)
|
||||
connector = accelerate.dispatch_model(connector, device_map=device_map)
|
||||
return connector
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVGemmaCLIPModelLoaderMGPU", description="Gemma 3 Model Loader on MGPU"
|
||||
)
|
||||
class LTXVGemmaCLIPModelLoaderMGPU:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"gemma_path": (
|
||||
folder_paths.get_filename_list("text_encoders"),
|
||||
{"tooltip": "The name of the text encoder model to load."},
|
||||
),
|
||||
"ltxv_path": (
|
||||
folder_paths.get_filename_list("checkpoints"),
|
||||
{"tooltip": "The name of the ltxv model to load."},
|
||||
),
|
||||
"max_length": (
|
||||
"INT",
|
||||
{"default": 1024, "min": 16, "max": 131072, "step": 8},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
RETURN_NAMES = ("clip",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Gemma CLIP Loader"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def load_model(self, gemma_path: str, ltxv_path: str, max_length: int):
|
||||
path = Path(folder_paths.get_full_path("text_encoders", gemma_path))
|
||||
model_root = path.parents[1]
|
||||
tokenizer_path = Path(find_matching_dir(model_root, "tokenizer.model"))
|
||||
gemma_model_path = Path(find_matching_dir(model_root, "model*.safetensors"))
|
||||
|
||||
tokenizer_class = ltxv_gemma_tokenizer(tokenizer_path, max_length=max_length)
|
||||
|
||||
processor = None
|
||||
try:
|
||||
image_processor = AutoImageProcessor.from_pretrained(
|
||||
str(model_root),
|
||||
local_files_only=True,
|
||||
)
|
||||
processor = Gemma3Processor(
|
||||
image_processor=image_processor,
|
||||
tokenizer=tokenizer_class().tokenizer,
|
||||
)
|
||||
logger.info(f"Loaded processor from {model_root} - enhancement enabled")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not load processor from {model_root}: {e}")
|
||||
|
||||
clip_dtype = torch.bfloat16
|
||||
ltxv_full_path = folder_paths.get_full_path("checkpoints", ltxv_path)
|
||||
clip_target = comfy.supported_models_base.ClipTarget(
|
||||
tokenizer=tokenizer_class,
|
||||
clip=ltxv_gemma_clip(
|
||||
gemma_model_path, ltxv_full_path, processor=processor, dtype=clip_dtype
|
||||
),
|
||||
)
|
||||
|
||||
return (comfy.sd.CLIP(clip_target),)
|
||||
@@ -0,0 +1,254 @@
|
||||
import comfy
|
||||
import comfy_extras.nodes_lt as nodes_lt
|
||||
import comfy_extras.nodes_post_processing as post_processing
|
||||
import nodes
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
def blur_internal(image, blur_radius):
|
||||
if blur_radius > 0:
|
||||
# https://docs.opencv.org/2.4/modules/imgproc/doc/filtering.html#getgaussiankernel
|
||||
# sigma = 0.3 * blur_radius + 0.5 is what is recommended in the OpenCV doc for the
|
||||
# relationship between sigma and kernel size 2*blur_radius + 1, however we want somewhat weaker
|
||||
# blurring, so we use 0.3 * blur_radius instead, reducing the sigma value by 0.5
|
||||
sigma = 0.3 * blur_radius
|
||||
image = post_processing.Blur.execute(image, blur_radius, sigma)[0]
|
||||
return image
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddGuideAdvanced")
|
||||
class LTXVAddGuideAdvanced:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"vae": ("VAE",),
|
||||
"latent": ("LATENT",),
|
||||
"image": ("IMAGE",),
|
||||
"frame_idx": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": -9999,
|
||||
"max": 9999,
|
||||
"tooltip": "Frame index to start the conditioning at. For single-frame images or "
|
||||
"videos with 1-8 frames, any frame_idx value is acceptable. For videos with 9+ "
|
||||
"frames, frame_idx must be divisible by 8, otherwise it will be rounded down to "
|
||||
"the nearest multiple of 8. Negative values are counted from the end of the video.",
|
||||
},
|
||||
),
|
||||
"strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"tooltip": "Strength of the conditioning. Higher values will make the conditioning more exact.",
|
||||
},
|
||||
),
|
||||
"crf": (
|
||||
"INT",
|
||||
{
|
||||
"default": 29,
|
||||
"min": 0,
|
||||
"max": 51,
|
||||
"step": 1,
|
||||
"tooltip": "CRF value for the video. Higher values mean more motion, lower values mean higher quality.",
|
||||
},
|
||||
),
|
||||
"blur_radius": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 7,
|
||||
"step": 1,
|
||||
"tooltip": "Blur kernel radius size. Higher values mean more motion, lower values mean higher quality.",
|
||||
},
|
||||
),
|
||||
"interpolation": (
|
||||
[
|
||||
"lanczos",
|
||||
"bislerp",
|
||||
"nearest",
|
||||
"bilinear",
|
||||
"bicubic",
|
||||
"area",
|
||||
"nearest-exact",
|
||||
],
|
||||
{"default": "lanczos"},
|
||||
),
|
||||
"crop": (["center", "disabled"], {"default": "disabled"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
||||
RETURN_NAMES = ("positive", "negative", "latent")
|
||||
|
||||
CATEGORY = "conditioning/video_models"
|
||||
FUNCTION = "generate"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Adds a conditioning frame or a video at a specific frame index. "
|
||||
"This node is used to add a keyframe or a video segment which should appear in the "
|
||||
"generated video at a specified index. It resizes the image to the correct size and "
|
||||
"applies preprocessing to it."
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
latent,
|
||||
image,
|
||||
frame_idx,
|
||||
strength,
|
||||
crf,
|
||||
blur_radius,
|
||||
interpolation,
|
||||
crop,
|
||||
):
|
||||
_, width_scale_factor, height_scale_factor = vae.downscale_index_formula
|
||||
width, height = (
|
||||
latent["samples"].shape[4] * width_scale_factor,
|
||||
latent["samples"].shape[3] * height_scale_factor,
|
||||
)
|
||||
image = (
|
||||
comfy.utils.common_upscale(
|
||||
image.movedim(-1, 1), width, height, interpolation, crop=crop
|
||||
)
|
||||
.movedim(1, -1)
|
||||
.clamp(0, 1)
|
||||
)
|
||||
image = nodes_lt.LTXVPreprocess().execute(image, crf)[0]
|
||||
image = blur_internal(image, blur_radius)
|
||||
return nodes_lt.LTXVAddGuide().execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
latent=latent,
|
||||
image=image,
|
||||
frame_idx=frame_idx,
|
||||
strength=strength,
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVImgToVideoAdvanced")
|
||||
class LTXVImgToVideoAdvanced:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"vae": ("VAE",),
|
||||
"image": ("IMAGE",),
|
||||
"width": (
|
||||
"INT",
|
||||
{
|
||||
"default": 768,
|
||||
"min": 64,
|
||||
"max": nodes.MAX_RESOLUTION,
|
||||
"step": 32,
|
||||
},
|
||||
),
|
||||
"height": (
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": nodes.MAX_RESOLUTION,
|
||||
"step": 32,
|
||||
},
|
||||
),
|
||||
"length": (
|
||||
"INT",
|
||||
{"default": 97, "min": 9, "max": nodes.MAX_RESOLUTION, "step": 8},
|
||||
),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
||||
"crf": (
|
||||
"INT",
|
||||
{
|
||||
"default": 29,
|
||||
"min": 0,
|
||||
"max": 51,
|
||||
"step": 1,
|
||||
"tooltip": "CRF value for the video. Higher values mean more motion, lower values mean higher quality.",
|
||||
},
|
||||
),
|
||||
"blur_radius": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 7,
|
||||
"step": 1,
|
||||
"tooltip": "Blur kernel radius size. Higher values mean more motion, lower values mean higher quality.",
|
||||
},
|
||||
),
|
||||
"interpolation": (
|
||||
[
|
||||
"lanczos",
|
||||
"bislerp",
|
||||
"nearest",
|
||||
"bilinear",
|
||||
"bicubic",
|
||||
"area",
|
||||
"nearest-exact",
|
||||
],
|
||||
{"default": "lanczos"},
|
||||
),
|
||||
"crop": (["center", "disabled"], {"default": "disabled"}),
|
||||
"strength": ("FLOAT", {"default": 0.9, "min": 0, "max": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
||||
RETURN_NAMES = ("positive", "negative", "latent")
|
||||
|
||||
CATEGORY = "conditioning/video_models"
|
||||
FUNCTION = "generate"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Adds a conditioning frame or a video at index 0. "
|
||||
"This node is used to add a keyframe or a video segment which should appear in the "
|
||||
"generated video at index 0. It resizes the image to the correct size "
|
||||
"and applies preprocessing to it."
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
image,
|
||||
width,
|
||||
height,
|
||||
length,
|
||||
batch_size,
|
||||
crf,
|
||||
blur_radius,
|
||||
interpolation,
|
||||
crop,
|
||||
strength,
|
||||
):
|
||||
image = comfy.utils.common_upscale(
|
||||
image.movedim(-1, 1), width, height, interpolation, crop=crop
|
||||
).movedim(1, -1)
|
||||
image = nodes_lt.LTXVPreprocess().execute(image, crf)[0]
|
||||
image = blur_internal(image, blur_radius)
|
||||
return nodes_lt.LTXVImgToVideo().execute(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
image=image,
|
||||
width=width,
|
||||
height=height,
|
||||
length=length,
|
||||
batch_size=batch_size,
|
||||
strength=strength,
|
||||
)
|
||||
@@ -0,0 +1,2 @@
|
||||
from .multimodal_guider import MultimodalGuiderNode # noqa: F401
|
||||
from .parameters import GuiderParametersNode # noqa: F401
|
||||
@@ -0,0 +1,323 @@
|
||||
from typing import List
|
||||
|
||||
import comfy.samplers
|
||||
import torch
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
from comfy.patcher_extension import CallbacksMP
|
||||
|
||||
from ..nodes_registry import comfy_node
|
||||
from ..stg import STGBlockWrapper, STGFlag
|
||||
from .parameters import GuiderParameters, Modality
|
||||
|
||||
|
||||
class MultimodalGuider(comfy.samplers.CFGGuider):
|
||||
def __init__(
|
||||
self,
|
||||
model: ModelPatcher,
|
||||
parameters: dict[str, GuiderParameters],
|
||||
skip_blocks: List[int],
|
||||
):
|
||||
model = model.clone()
|
||||
self.current_step = 0
|
||||
self.last_denoised_v = None
|
||||
self.last_denoised_a = None
|
||||
|
||||
model.add_callback_with_key(
|
||||
CallbacksMP.ON_PRE_RUN, "mm_guider_on_pre_run", self.reset_current_step
|
||||
)
|
||||
super().__init__(model)
|
||||
|
||||
self.stg_flag = STGFlag(
|
||||
do_skip=False,
|
||||
skip_layers=skip_blocks,
|
||||
)
|
||||
|
||||
self.patch_model(model, self.stg_flag)
|
||||
self.parameters = parameters
|
||||
|
||||
def reset_current_step(self, model_patcher=None):
|
||||
self.current_step = 0
|
||||
self.last_denoised_v = None
|
||||
self.last_denoised_a = None
|
||||
|
||||
@classmethod
|
||||
def patch_model(cls, model: ModelPatcher, stg_flag: STGFlag):
|
||||
transformer_blocks = cls.get_transformer_blocks(model)
|
||||
|
||||
for i, block in enumerate(transformer_blocks):
|
||||
model.set_model_patch_replace(
|
||||
STGBlockWrapper(block, stg_flag, i), "dit", "double_block", i
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_transformer_blocks(model: ModelPatcher):
|
||||
diffusion_model = model.get_model_object("diffusion_model")
|
||||
key = "diffusion_model.transformer_blocks"
|
||||
if diffusion_model.__class__.__name__ == "LTXVTransformer3D":
|
||||
key = "diffusion_model.transformer.transformer_blocks"
|
||||
return model.get_model_object(key)
|
||||
|
||||
def set_conds(self, positive, negative):
|
||||
self.inner_set_conds({"positive": positive, "negative": negative})
|
||||
|
||||
def calc_stg_indexes(
|
||||
self,
|
||||
run_vx: bool,
|
||||
run_ax: bool,
|
||||
audio_length: int,
|
||||
audio_ptb: bool,
|
||||
video_ptb: bool,
|
||||
):
|
||||
stg_indexes = set()
|
||||
num_self_attns = 0
|
||||
if run_vx:
|
||||
num_self_attns += 1
|
||||
if run_ax and audio_length > 0:
|
||||
num_self_attns += 1
|
||||
|
||||
video_attn_idx = 0
|
||||
audio_attn_idx = 0 if num_self_attns == 1 else 2
|
||||
|
||||
if video_ptb:
|
||||
stg_indexes.add(video_attn_idx)
|
||||
if audio_ptb:
|
||||
stg_indexes.add(audio_attn_idx)
|
||||
|
||||
return list(stg_indexes)
|
||||
|
||||
def predict_noise(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
model_options: dict = {},
|
||||
seed=None,
|
||||
):
|
||||
x_shape = x.shape
|
||||
# in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
|
||||
# but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
|
||||
current_step = self.current_step
|
||||
self.current_step = current_step + 1
|
||||
|
||||
positive_cond = self.conds.get("positive", None)
|
||||
negative_cond = self.conds.get("negative", None)
|
||||
|
||||
audio_params: GuiderParameters = self.parameters.get(
|
||||
Modality.AUDIO.value, GuiderParameters()
|
||||
)
|
||||
video_params: GuiderParameters = self.parameters.get(
|
||||
Modality.VIDEO.value, GuiderParameters()
|
||||
)
|
||||
|
||||
audio_length = model_options["transformer_options"].get("audio_length", 0)
|
||||
|
||||
run_vx = not video_params.do_skip(current_step)
|
||||
run_ax = not audio_params.do_skip(current_step)
|
||||
run_a2v = video_params.do_cross_attn(current_step)
|
||||
run_v2a = audio_params.do_cross_attn(current_step)
|
||||
|
||||
vx, ax = self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
x, audio_length
|
||||
)
|
||||
if not run_vx:
|
||||
vx = self.last_denoised_v
|
||||
if not run_ax:
|
||||
ax = self.last_denoised_a
|
||||
x = self.inner_model.diffusion_model.recombine_audio_and_video_latents(
|
||||
vx, ax, x_shape
|
||||
)
|
||||
if not run_vx and not run_ax:
|
||||
return x
|
||||
|
||||
try:
|
||||
model_options["transformer_options"]["run_vx"] = run_vx
|
||||
model_options["transformer_options"]["run_ax"] = run_ax
|
||||
model_options["transformer_options"]["a2v_cross_attn"] = run_a2v
|
||||
model_options["transformer_options"]["v2a_cross_attn"] = run_v2a
|
||||
noise_pred_pos = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[positive_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
finally:
|
||||
del model_options["transformer_options"]["run_vx"]
|
||||
del model_options["transformer_options"]["run_ax"]
|
||||
del model_options["transformer_options"]["a2v_cross_attn"]
|
||||
del model_options["transformer_options"]["v2a_cross_attn"]
|
||||
v_noise_pred_pos, a_noise_pred_pos = (
|
||||
self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
noise_pred_pos, audio_length
|
||||
)
|
||||
)
|
||||
|
||||
a_noise_pred_neg, v_noise_pred_neg = 0, 0
|
||||
a_noise_pred_perturbed, v_noise_pred_perturbed = 0, 0
|
||||
a_noise_pred_modality, v_noise_pred_modality = 0, 0
|
||||
|
||||
if any(params.do_uncond() for params in [audio_params, video_params]):
|
||||
try:
|
||||
model_options["transformer_options"]["run_vx"] = run_vx
|
||||
model_options["transformer_options"]["run_ax"] = run_ax
|
||||
model_options["transformer_options"]["a2v_cross_attn"] = run_a2v
|
||||
model_options["transformer_options"]["v2a_cross_attn"] = run_v2a
|
||||
noise_pred_neg = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[negative_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
v_noise_pred_neg, a_noise_pred_neg = (
|
||||
self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
noise_pred_neg, audio_length
|
||||
)
|
||||
)
|
||||
finally:
|
||||
del model_options["transformer_options"]["run_vx"]
|
||||
del model_options["transformer_options"]["run_ax"]
|
||||
del model_options["transformer_options"]["a2v_cross_attn"]
|
||||
del model_options["transformer_options"]["v2a_cross_attn"]
|
||||
|
||||
if any(params.do_perturbed() for params in [audio_params, video_params]):
|
||||
try:
|
||||
stg_indexes = self.calc_stg_indexes(
|
||||
run_vx,
|
||||
run_ax,
|
||||
audio_length,
|
||||
audio_params.perturb_attn,
|
||||
video_params.perturb_attn,
|
||||
)
|
||||
model_options["transformer_options"]["run_vx"] = run_vx
|
||||
model_options["transformer_options"]["run_ax"] = run_ax
|
||||
model_options["transformer_options"]["a2v_cross_attn"] = run_a2v
|
||||
model_options["transformer_options"]["v2a_cross_attn"] = run_v2a
|
||||
model_options["transformer_options"]["ptb_index"] = 0
|
||||
model_options["transformer_options"]["stg_indexes"] = stg_indexes
|
||||
self.stg_flag.do_skip = True
|
||||
noise_pred_perturbed = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[positive_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
v_noise_pred_perturbed, a_noise_pred_perturbed = (
|
||||
self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
noise_pred_perturbed, audio_length
|
||||
)
|
||||
)
|
||||
finally:
|
||||
self.stg_flag.do_skip = False
|
||||
del model_options["transformer_options"]["ptb_index"]
|
||||
del model_options["transformer_options"]["run_vx"]
|
||||
del model_options["transformer_options"]["run_ax"]
|
||||
del model_options["transformer_options"]["a2v_cross_attn"]
|
||||
del model_options["transformer_options"]["v2a_cross_attn"]
|
||||
del model_options["transformer_options"]["stg_indexes"]
|
||||
|
||||
if any(params.do_modality() for params in [audio_params, video_params]):
|
||||
try:
|
||||
model_options["transformer_options"]["run_vx"] = run_vx
|
||||
model_options["transformer_options"]["run_ax"] = run_ax
|
||||
model_options["transformer_options"]["a2v_cross_attn"] = False
|
||||
model_options["transformer_options"]["v2a_cross_attn"] = False
|
||||
noise_pred_modality = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[positive_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
v_noise_pred_modality, a_noise_pred_modality = (
|
||||
self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
noise_pred_modality, audio_length
|
||||
)
|
||||
)
|
||||
finally:
|
||||
del model_options["transformer_options"]["a2v_cross_attn"]
|
||||
del model_options["transformer_options"]["v2a_cross_attn"]
|
||||
del model_options["transformer_options"]["run_vx"]
|
||||
del model_options["transformer_options"]["run_ax"]
|
||||
|
||||
if run_vx:
|
||||
vx = video_params.calculate(
|
||||
v_noise_pred_pos,
|
||||
v_noise_pred_neg,
|
||||
v_noise_pred_perturbed,
|
||||
v_noise_pred_modality,
|
||||
)
|
||||
else:
|
||||
vx = self.last_denoised_v
|
||||
|
||||
if run_ax:
|
||||
ax = audio_params.calculate(
|
||||
a_noise_pred_pos,
|
||||
a_noise_pred_neg,
|
||||
a_noise_pred_perturbed,
|
||||
a_noise_pred_modality,
|
||||
)
|
||||
else:
|
||||
ax = self.last_denoised_a
|
||||
|
||||
x = self.inner_model.diffusion_model.recombine_audio_and_video_latents(
|
||||
vx, ax, x_shape
|
||||
)
|
||||
|
||||
# normally this would be done in cfg_function, but we skipped
|
||||
# that for efficiency: we can compute the noise predictions in
|
||||
# a single call to calc_cond_batch() (rather than two)
|
||||
# so we replicate the hook here
|
||||
for fn in model_options.get("sampler_post_cfg_function", []):
|
||||
args = {
|
||||
"denoised": x,
|
||||
"cond": positive_cond,
|
||||
"uncond": negative_cond,
|
||||
"model": self.inner_model,
|
||||
"uncond_denoised": noise_pred_neg,
|
||||
"cond_denoised": noise_pred_pos,
|
||||
"sigma": timestep,
|
||||
"model_options": model_options,
|
||||
"input": x,
|
||||
# not in the original call in samplers.py:cfg_function, but made available for future hooks
|
||||
"perturbed_cond": positive_cond,
|
||||
"perturbed_cond_denoised": noise_pred_perturbed,
|
||||
}
|
||||
x = fn(args)
|
||||
|
||||
self.last_denoised_v, self.last_denoised_a = (
|
||||
self.inner_model.diffusion_model.separate_audio_and_video_latents(
|
||||
x, audio_length
|
||||
)
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
@comfy_node(name="MultimodalGuider")
|
||||
class MultimodalGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"parameters": ("GUIDER_PARAMETERS",),
|
||||
"skip_blocks": (
|
||||
"STRING",
|
||||
{"default": "", "multiline": True},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
|
||||
FUNCTION = "get_guider"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
def get_guider(self, model, positive, negative, parameters, skip_blocks):
|
||||
skip_blocks = [int(n.strip()) for n in skip_blocks.split(",") if n.strip()]
|
||||
guider = MultimodalGuider(model, parameters, skip_blocks)
|
||||
guider.set_conds(positive, negative)
|
||||
guider.raw_conds = (positive, negative)
|
||||
return (guider,)
|
||||
@@ -0,0 +1,164 @@
|
||||
import math
|
||||
from enum import Enum
|
||||
|
||||
from ..nodes_registry import comfy_node
|
||||
|
||||
|
||||
class Modality(Enum):
|
||||
VIDEO = "VIDEO"
|
||||
AUDIO = "AUDIO"
|
||||
|
||||
|
||||
class GuiderParameters:
|
||||
def __init__(
|
||||
self,
|
||||
cfg_scale: float = 1.0,
|
||||
stg_scale: float = 0.0,
|
||||
perturb_attn: bool = True,
|
||||
rescale_scale: float = 0.0,
|
||||
modality_scale: float = 1.0,
|
||||
skip_step: int = 0,
|
||||
cross_attn: bool = True,
|
||||
):
|
||||
self.cfg_scale = cfg_scale
|
||||
self.stg_scale = stg_scale
|
||||
self.perturb_attn = perturb_attn
|
||||
self.rescale_scale = rescale_scale
|
||||
self.modality_scale = modality_scale
|
||||
self.skip_step = skip_step
|
||||
self.cross_attn = cross_attn
|
||||
|
||||
def __str__(self):
|
||||
return f"cfg_scale: {self.cfg_scale}, stg_scale: {self.stg_scale}, rescale_scale: {self.rescale_scale}, modality_scale: {self.modality_scale}"
|
||||
|
||||
def __repr__(self):
|
||||
return f"cfg_scale: {self.cfg_scale}, stg_scale: {self.stg_scale}, rescale_scale: {self.rescale_scale}, modality_scale: {self.modality_scale}"
|
||||
|
||||
def calculate(
|
||||
self, noise_pred_pos, noise_pred_neg, noise_pred_pertubed, noise_pred_modality
|
||||
):
|
||||
noise_pred = (
|
||||
noise_pred_pos
|
||||
+ (self.cfg_scale - 1) * (noise_pred_pos - noise_pred_neg)
|
||||
+ self.stg_scale * (noise_pred_pos - noise_pred_pertubed)
|
||||
+ (self.modality_scale - 1) * (noise_pred_pos - noise_pred_modality)
|
||||
)
|
||||
|
||||
if self.rescale_scale != 0:
|
||||
factor = noise_pred_pos.std() / noise_pred.std()
|
||||
factor = self.rescale_scale * factor + (1 - self.rescale_scale)
|
||||
noise_pred = noise_pred * factor
|
||||
|
||||
return noise_pred
|
||||
|
||||
def do_uncond(self):
|
||||
return not math.isclose(self.cfg_scale, 1.0)
|
||||
|
||||
def do_perturbed(self):
|
||||
return not math.isclose(self.stg_scale, 0.0)
|
||||
|
||||
def do_modality(self):
|
||||
return not math.isclose(self.modality_scale, 1.0)
|
||||
|
||||
def do_skip(self, step: int) -> bool:
|
||||
if self.skip_step == 0:
|
||||
return False
|
||||
|
||||
return step % (self.skip_step + 1) != 0
|
||||
|
||||
def do_cross_attn(self, step: int) -> bool:
|
||||
return self.cross_attn and not self.do_skip(step)
|
||||
|
||||
|
||||
@comfy_node(name="GuiderParameters")
|
||||
class GuiderParametersNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"modality": (
|
||||
[m.value for m in Modality],
|
||||
{
|
||||
"default": Modality.VIDEO.value,
|
||||
},
|
||||
),
|
||||
"cfg": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"stg": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"perturb_attn": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
"rescale": (
|
||||
"FLOAT",
|
||||
{"default": 0.7, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"modality_scale": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"skip_step": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0, "max": 100, "step": 1},
|
||||
),
|
||||
"cross_attn": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"parameters": (
|
||||
"GUIDER_PARAMETERS",
|
||||
{"default": None},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER_PARAMETERS",)
|
||||
|
||||
FUNCTION = "get_parameters"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
def get_parameters(
|
||||
self,
|
||||
modality,
|
||||
cfg,
|
||||
stg,
|
||||
perturb_attn,
|
||||
rescale,
|
||||
modality_scale,
|
||||
skip_step,
|
||||
cross_attn,
|
||||
parameters=None,
|
||||
):
|
||||
parameters = parameters.copy() if parameters is not None else {}
|
||||
|
||||
if modality in parameters:
|
||||
raise ValueError(f"Modality {modality} already exists in parameters")
|
||||
|
||||
parameters.update(
|
||||
{
|
||||
modality: GuiderParameters(
|
||||
cfg,
|
||||
stg,
|
||||
perturb_attn,
|
||||
rescale,
|
||||
modality_scale,
|
||||
skip_step,
|
||||
cross_attn,
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
return (parameters,)
|
||||
+331
@@ -0,0 +1,331 @@
|
||||
import copy
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAdainLatent")
|
||||
class LTXVAdainLatent:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latents": ("LATENT",),
|
||||
"reference": ("LATENT",),
|
||||
"factor": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"per_frame": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "batch_normalize"
|
||||
|
||||
CATEGORY = "Lightricks/latents"
|
||||
|
||||
def batch_normalize(self, latents, reference, factor, per_frame=False):
|
||||
latents_copy = copy.deepcopy(latents)
|
||||
t = latents_copy["samples"] # B x C x F x H x W
|
||||
|
||||
if per_frame:
|
||||
if reference["samples"].size(2) == 1:
|
||||
print("Reference has only one frame, using it for all frames")
|
||||
reference["samples"] = reference["samples"].repeat(
|
||||
1, 1, t.size(2), 1, 1
|
||||
)
|
||||
|
||||
elif t.size(2) > reference["samples"].size(2):
|
||||
raise ValueError("Latents have more frames than reference")
|
||||
|
||||
for i in range(t.size(0)): # batch
|
||||
for c in range(t.size(1)):
|
||||
if not per_frame:
|
||||
r_sd, r_mean = torch.std_mean(
|
||||
reference["samples"][i, c], dim=None
|
||||
) # index by original dim order
|
||||
i_sd, i_mean = torch.std_mean(t[i, c], dim=None)
|
||||
|
||||
t[i, c] = ((t[i, c] - i_mean) / i_sd) * r_sd + r_mean
|
||||
else:
|
||||
for f in range(t.size(2)):
|
||||
r_sd, r_mean = torch.std_mean(
|
||||
reference["samples"][i, c, f], dim=None
|
||||
) # index by original dim order
|
||||
i_sd, i_mean = torch.std_mean(t[i, c, f], dim=None)
|
||||
t[i, c, f] = ((t[i, c, f] - i_mean) / i_sd) * r_sd + r_mean
|
||||
|
||||
latents_copy["samples"] = torch.lerp(latents["samples"], t, factor)
|
||||
return (latents_copy,)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVStatNormLatent")
|
||||
class LTXVStatNormLatent:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latents": ("LATENT",),
|
||||
"target_mean": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"target_std": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.01,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"percentile": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 95.0,
|
||||
"min": 50.0,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"round": 0.1,
|
||||
"tooltip": "Percentile of distribution to use for statistics calculation",
|
||||
},
|
||||
),
|
||||
"factor": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"clip_outliers": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "statistical_normalize"
|
||||
|
||||
CATEGORY = "Lightricks/latents"
|
||||
|
||||
def statistical_normalize(
|
||||
self, latents, target_mean, target_std, percentile, factor, clip_outliers
|
||||
):
|
||||
latents_copy = copy.deepcopy(latents)
|
||||
t = latents_copy["samples"] # B x C x F x H x W
|
||||
|
||||
# For 95% of distribution, we want to exclude 2.5% from each tail
|
||||
lower_percentile = (100 - percentile) / 2
|
||||
upper_percentile = 100 - lower_percentile
|
||||
|
||||
for i in range(t.size(0)): # batch
|
||||
for c in range(t.size(1)): # channel
|
||||
channel_data = t[i, c]
|
||||
original_shape = channel_data.shape
|
||||
channel_flat = channel_data.flatten()
|
||||
|
||||
# Calculate percentiles
|
||||
lower_bound = torch.quantile(channel_flat, lower_percentile / 100)
|
||||
upper_bound = torch.quantile(channel_flat, upper_percentile / 100)
|
||||
|
||||
# Create mask for values within the percentile range
|
||||
mask = (channel_flat >= lower_bound) & (channel_flat <= upper_bound)
|
||||
|
||||
# Calculate mean and std only on the masked values
|
||||
if mask.sum() > 0:
|
||||
filtered_data = channel_flat[mask]
|
||||
current_mean = filtered_data.mean()
|
||||
current_std = filtered_data.std()
|
||||
|
||||
# Avoid division by zero
|
||||
if current_std > 1e-8:
|
||||
# Normalize all values
|
||||
normalized_flat = (
|
||||
(channel_flat - current_mean) / current_std
|
||||
) * target_std + target_mean
|
||||
|
||||
if clip_outliers:
|
||||
# Calculate the normalized bounds
|
||||
normalized_lower = (
|
||||
(lower_bound - current_mean) / current_std
|
||||
) * target_std + target_mean
|
||||
normalized_upper = (
|
||||
(upper_bound - current_mean) / current_std
|
||||
) * target_std + target_mean
|
||||
|
||||
# Clip outliers to the normalized bounds
|
||||
normalized_flat = torch.where(
|
||||
channel_flat < lower_bound,
|
||||
normalized_lower,
|
||||
normalized_flat,
|
||||
)
|
||||
normalized_flat = torch.where(
|
||||
channel_flat > upper_bound,
|
||||
normalized_upper,
|
||||
normalized_flat,
|
||||
)
|
||||
|
||||
# Reshape back to original shape
|
||||
t[i, c] = normalized_flat.reshape(original_shape)
|
||||
else:
|
||||
# If std is too small, just shift by mean
|
||||
t[i, c] = channel_data - current_mean + target_mean
|
||||
|
||||
latents_copy["samples"] = torch.lerp(latents["samples"], t, factor)
|
||||
return (latents_copy,)
|
||||
|
||||
|
||||
class PerStepNormPatcher:
|
||||
"""
|
||||
Base class for per-step normalization nodes.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def required(s):
|
||||
return {
|
||||
"model": ("MODEL",),
|
||||
"factors": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "0.9, 0.75, 0.0",
|
||||
"tooltip": "Comma-separated list of factors, each factor will be used for one step.",
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def optional(s):
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def patch_smooth_norm(
|
||||
model: ModelPatcher, factors: str, normalize_fn: Callable
|
||||
) -> ModelPatcher:
|
||||
model = model.clone()
|
||||
factors = [float(x) for x in factors.split(",")]
|
||||
step = 0
|
||||
|
||||
def norm_fn(args):
|
||||
nonlocal step
|
||||
latent = args["denoised"]
|
||||
factor = factors[min(step, len(factors) - 1)]
|
||||
step += 1
|
||||
out_latent = normalize_fn(latent={"samples": latent}, factor=factor)
|
||||
return out_latent[0]["samples"]
|
||||
|
||||
model.set_model_sampler_post_cfg_function(norm_fn)
|
||||
return (model,)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(s, *args, **kwargs):
|
||||
return float("NaN")
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPerStepAdainPatcher")
|
||||
class LTXVPerStepAdainPatcher(PerStepNormPatcher):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": super().required()
|
||||
| {
|
||||
"reference": ("LATENT",),
|
||||
},
|
||||
"optional": super().optional()
|
||||
| {
|
||||
"per_frame": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch_model"
|
||||
|
||||
CATEGORY = "Lightricks/latents"
|
||||
|
||||
def patch_model(self, model, factors, reference, per_frame=False):
|
||||
def cfg_adain(latent, factor):
|
||||
return LTXVAdainLatent().batch_normalize(
|
||||
latent, reference, factor, per_frame
|
||||
)
|
||||
|
||||
return self.patch_smooth_norm(model, factors, cfg_adain)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPerStepStatNormPatcher")
|
||||
class LTXVPerStepStatNormPatcher(PerStepNormPatcher):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": super().required()
|
||||
| {
|
||||
"target_mean": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": -10.0,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"target_std": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.01,
|
||||
"max": 10.0,
|
||||
"step": 0.01,
|
||||
"round": 0.01,
|
||||
},
|
||||
),
|
||||
"percentile": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 95.0,
|
||||
"min": 50.0,
|
||||
"max": 100.0,
|
||||
"step": 0.1,
|
||||
"round": 0.1,
|
||||
"tooltip": "Percentile of distribution to use for statistics calculation",
|
||||
},
|
||||
),
|
||||
"clip_outliers": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch_model"
|
||||
|
||||
CATEGORY = "Lightricks/latents"
|
||||
|
||||
def patch_model(
|
||||
self, model, factors, target_mean, target_std, percentile, clip_outliers
|
||||
):
|
||||
def cfg_stat_norm(latent, factor):
|
||||
return LTXVStatNormLatent().statistical_normalize(
|
||||
latent, target_mean, target_std, percentile, factor, clip_outliers
|
||||
)
|
||||
|
||||
return self.patch_smooth_norm(model, factors, cfg_stat_norm)
|
||||
+830
@@ -0,0 +1,830 @@
|
||||
from typing import Optional
|
||||
|
||||
import comfy.utils
|
||||
import comfy_extras.nodes_lt as nodes_lt
|
||||
import numpy as np
|
||||
import torch
|
||||
from comfy.ldm.lightricks.vae.audio_vae import LATENT_DOWNSAMPLE_FACTOR
|
||||
from comfy.nested_tensor import NestedTensor
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVSelectLatents")
|
||||
class LTXVSelectLatents:
|
||||
"""
|
||||
Selects a range of frames from a video latent.
|
||||
|
||||
Features:
|
||||
- Supports positive and negative indexing
|
||||
- Preserves batch processing capabilities
|
||||
- Handles noise masks if present
|
||||
- Maintains 5D tensor format
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"samples": ("LATENT",),
|
||||
"start_index": (
|
||||
"INT",
|
||||
{"default": 0, "min": -9999, "max": 9999, "step": 1},
|
||||
),
|
||||
"end_index": (
|
||||
"INT",
|
||||
{"default": -1, "min": -9999, "max": 9999, "step": 1},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "select_latents"
|
||||
CATEGORY = "latent/video"
|
||||
DESCRIPTION = (
|
||||
"Selects a range of frames from the video latent. "
|
||||
"start_index and end_index define a closed interval (inclusive of both endpoints)."
|
||||
)
|
||||
|
||||
def select_latents(self, samples: dict, start_index: int, end_index: int) -> tuple:
|
||||
"""
|
||||
Selects a range of frames from the video latent.
|
||||
|
||||
Args:
|
||||
samples (dict): Video latent dictionary
|
||||
start_index (int): Starting frame index (supports negative indexing)
|
||||
end_index (int): Ending frame index (supports negative indexing)
|
||||
|
||||
Returns:
|
||||
tuple: Contains modified latent dictionary with selected frames
|
||||
|
||||
Raises:
|
||||
ValueError: If indices are invalid
|
||||
"""
|
||||
try:
|
||||
s = samples.copy()
|
||||
video_latent = s["samples"]
|
||||
batch, channels, frames, height, width = video_latent.shape
|
||||
|
||||
# Handle negative indices
|
||||
start_idx = frames + start_index if start_index < 0 else start_index
|
||||
end_idx = frames + end_index if end_index < 0 else end_index
|
||||
|
||||
# Validate and clamp indices
|
||||
start_idx = max(0, min(start_idx, frames - 1))
|
||||
end_idx = max(0, min(end_idx, frames - 1))
|
||||
if start_idx > end_idx:
|
||||
start_idx = min(start_idx, end_idx)
|
||||
|
||||
# Select frames while maintaining 5D format
|
||||
s["samples"] = video_latent[:, :, start_idx : end_idx + 1, :, :]
|
||||
|
||||
# Handle noise mask if present
|
||||
if "noise_mask" in s and s["noise_mask"] is not None:
|
||||
s["noise_mask"] = s["noise_mask"][:, :, start_idx : end_idx + 1, :, :]
|
||||
|
||||
return (s,)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[LTXVSelectLatents] Error: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddLatents")
|
||||
class LTXVAddLatents:
|
||||
"""
|
||||
Concatenates two video latents along the frames dimension.
|
||||
|
||||
Features:
|
||||
- Validates dimension compatibility
|
||||
- Handles device placement
|
||||
- Preserves noise masks with proper handling
|
||||
- Supports batch processing
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latents1": ("LATENT",),
|
||||
"latents2": ("LATENT",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "add_latents"
|
||||
CATEGORY = "latent/video"
|
||||
DESCRIPTION = (
|
||||
"Concatenates two video latents along the frames dimension. "
|
||||
"latents1 and latents2 must have the same dimensions except for the frames dimension."
|
||||
)
|
||||
|
||||
def add_latents(
|
||||
self, latents1: torch.Tensor, latents2: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Concatenates two video latents along the frames dimension.
|
||||
|
||||
Args:
|
||||
latents1 (dict): First video latent dictionary
|
||||
latents2 (dict): Second video latent dictionary
|
||||
|
||||
Returns:
|
||||
tuple: Contains concatenated latent dictionary
|
||||
|
||||
Raises:
|
||||
ValueError: If latent dimensions don't match
|
||||
RuntimeError: If tensor operations fail
|
||||
"""
|
||||
try:
|
||||
s = latents1.copy()
|
||||
video_latent1 = latents1["samples"]
|
||||
video_latent2 = latents2["samples"]
|
||||
|
||||
# Ensure tensors are on the same device
|
||||
target_device = video_latent1.device
|
||||
video_latent2 = video_latent2.to(target_device)
|
||||
|
||||
# Validate dimensions
|
||||
self._validate_dimensions(video_latent1, video_latent2)
|
||||
|
||||
# Concatenate along frames dimension
|
||||
s["samples"] = torch.cat([video_latent1, video_latent2], dim=2)
|
||||
|
||||
# Handle noise masks
|
||||
s["noise_mask"] = self._merge_noise_masks(
|
||||
latents1, latents2, video_latent1.shape[2], video_latent2.shape[2]
|
||||
)
|
||||
|
||||
return (s,)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[LTXVAddLatents] Error: {str(e)}")
|
||||
raise
|
||||
|
||||
def _validate_dimensions(self, latent1: torch.Tensor, latent2: torch.Tensor):
|
||||
"""Validates that latent dimensions match except for frames."""
|
||||
# video latents
|
||||
if latent1.ndim == 5 and latent2.ndim == 5:
|
||||
b1, c1, f1, h1, w1 = latent1.shape
|
||||
b2, c2, f2, h2, w2 = latent2.shape
|
||||
|
||||
if not (b1 == b2 and c1 == c2 and h1 == h2 and w1 == w2):
|
||||
raise ValueError(
|
||||
f"Latent dimensions must match (except frames dimension).\n"
|
||||
f"Got shapes {latent1.shape} and {latent2.shape}"
|
||||
)
|
||||
# audio latents
|
||||
elif latent1.ndim == 4 and latent2.ndim == 4:
|
||||
b1, c1, f1, s1 = latent1.shape
|
||||
b2, c2, f2, s2 = latent2.shape
|
||||
|
||||
if not (b1 == b2 and c1 == c2 and s1 == s2):
|
||||
raise ValueError(
|
||||
f"Latent dimensions must match (except frames dimension).\n"
|
||||
f"Got shapes {latent1.shape} and {latent2.shape}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Latent dimensions must be 4 (audio) or 5 (video) for both inputs.\n"
|
||||
f"Got shapes {latent1.shape} and {latent2.shape}"
|
||||
)
|
||||
|
||||
def _merge_noise_masks(
|
||||
self, latents1: torch.Tensor, latents2: torch.Tensor, frames1: int, frames2: int
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""Merges noise masks from both latents with proper handling."""
|
||||
if "noise_mask" in latents1 and "noise_mask" in latents2:
|
||||
return torch.cat([latents1["noise_mask"], latents2["noise_mask"]], dim=2)
|
||||
elif "noise_mask" in latents1 and latents1["noise_mask"] is not None:
|
||||
zeros = torch.zeros_like(latents1["noise_mask"][:, :, :frames2, :, :])
|
||||
return torch.cat([latents1["noise_mask"], zeros], dim=2)
|
||||
elif "noise_mask" in latents2 and latents2["noise_mask"] is not None:
|
||||
zeros = torch.zeros_like(latents2["noise_mask"][:, :, :frames1, :, :])
|
||||
return torch.cat([zeros, latents2["noise_mask"]], dim=2)
|
||||
return None
|
||||
|
||||
|
||||
@comfy_node(name="LTXVSetVideoLatentNoiseMasks")
|
||||
class LTXVSetVideoLatentNoiseMasks:
|
||||
"""
|
||||
Applies multiple masks to a video latent.
|
||||
|
||||
Features:
|
||||
- Supports multiple input mask formats (2D, 3D, 4D)
|
||||
- Automatically handles fewer masks than frames by reusing the last mask
|
||||
- Resizes masks to match latent dimensions
|
||||
- Preserves batch processing capabilities
|
||||
|
||||
Input Formats:
|
||||
- 2D mask: Single mask [H, W]
|
||||
- 3D mask: Multiple masks [M, H, W]
|
||||
- 4D mask: Multiple masks with channels [M, C, H, W]
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"samples": ("LATENT",),
|
||||
"masks": ("MASK",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "set_mask"
|
||||
CATEGORY = "latent/video"
|
||||
DESCRIPTION = (
|
||||
"Applies multiple masks to a video latent. "
|
||||
"masks can be 2D, 3D, or 4D tensors. "
|
||||
"If there are fewer masks than frames, the last mask will be reused."
|
||||
)
|
||||
|
||||
def set_mask(self, samples: dict, masks: torch.Tensor) -> tuple:
|
||||
"""
|
||||
Applies masks to video latent frames.
|
||||
|
||||
Args:
|
||||
samples (dict): Video latent dictionary containing 'samples' tensor
|
||||
masks (torch.Tensor): Mask tensor in various possible formats
|
||||
- 2D: [H, W] single mask
|
||||
- 3D: [M, H, W] multiple masks
|
||||
- 4D: [M, C, H, W] multiple masks with channels
|
||||
|
||||
Returns:
|
||||
tuple: Contains modified latent dictionary with applied masks
|
||||
|
||||
Raises:
|
||||
ValueError: If mask dimensions are unsupported
|
||||
RuntimeError: If tensor operations fail
|
||||
"""
|
||||
try:
|
||||
s = samples.copy()
|
||||
video_latent = s["samples"]
|
||||
batch_size, channels, num_frames, height, width = video_latent.shape
|
||||
|
||||
# Initialize noise_mask if not present or resize if dimensions don't match
|
||||
if "noise_mask" not in s:
|
||||
s["noise_mask"] = torch.zeros(
|
||||
(batch_size, 1, num_frames, height, width),
|
||||
dtype=video_latent.dtype,
|
||||
device=video_latent.device,
|
||||
)
|
||||
else:
|
||||
existing_shape = s["noise_mask"].shape
|
||||
# Check if noise_mask has the right number of frames
|
||||
if existing_shape[2] != num_frames:
|
||||
s["noise_mask"] = torch.zeros(
|
||||
(batch_size, 1, num_frames, height, width),
|
||||
dtype=video_latent.dtype,
|
||||
device=video_latent.device,
|
||||
)
|
||||
|
||||
# Process masks
|
||||
masks_reshaped = self._reshape_masks(masks)
|
||||
M = masks_reshaped.shape[0]
|
||||
resized_masks = self._resize_masks(masks_reshaped, height, width)
|
||||
|
||||
# Apply masks efficiently
|
||||
self._apply_masks(s["noise_mask"], resized_masks, num_frames, M)
|
||||
return (s,)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[LTXVSetVideoLatentNoiseMasks] Error: {str(e)}")
|
||||
raise
|
||||
|
||||
def _reshape_masks(self, masks: torch.Tensor) -> torch.Tensor:
|
||||
"""Reshapes input masks to consistent 4D format."""
|
||||
original_shape = tuple(masks.shape)
|
||||
ndims = masks.ndim
|
||||
|
||||
if ndims == 2:
|
||||
return masks.unsqueeze(0).unsqueeze(0)
|
||||
elif ndims == 3:
|
||||
return masks.reshape(masks.shape[0], 1, masks.shape[1], masks.shape[2])
|
||||
elif ndims == 4:
|
||||
return masks.reshape(masks.shape[0], 1, masks.shape[2], masks.shape[3])
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported 'masks' dimension: {original_shape}. "
|
||||
"Must be 2D (H,W), 3D (M,H,W), or 4D (M,C,H,W)."
|
||||
)
|
||||
|
||||
def _resize_masks(
|
||||
self, masks: torch.Tensor, height: int, width: int
|
||||
) -> torch.Tensor:
|
||||
"""Resizes all masks to match latent dimensions."""
|
||||
return torch.nn.functional.interpolate(
|
||||
masks, size=(height, width), mode="bilinear", align_corners=False
|
||||
)
|
||||
|
||||
def _apply_masks(
|
||||
self,
|
||||
noise_mask: torch.Tensor,
|
||||
resized_masks: torch.Tensor,
|
||||
num_frames: int,
|
||||
M: int,
|
||||
) -> None:
|
||||
"""Applies resized masks to all frames."""
|
||||
for f in range(num_frames):
|
||||
mask_idx = min(f, M - 1) # Reuse last mask if we run out
|
||||
noise_mask[:, :, f] = resized_masks[mask_idx]
|
||||
|
||||
|
||||
@comfy_node(name="LTXVDilateLatent")
|
||||
class LTXVDilateLatent:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latent": ("LATENT",),
|
||||
"horizontal_scale": (
|
||||
"INT",
|
||||
{"default": 1, "min": 1, "max": 100, "step": 1},
|
||||
),
|
||||
"vertical_scale": (
|
||||
"INT",
|
||||
{"default": 1, "min": 1, "max": 100, "step": 1},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "dilate_latent"
|
||||
CATEGORY = "latent/video"
|
||||
DESCRIPTION = "Dilates a latent by a grid size."
|
||||
|
||||
def dilate_latent(
|
||||
self, latent: dict, horizontal_scale: int, vertical_scale: int
|
||||
) -> tuple:
|
||||
if horizontal_scale == 1 and vertical_scale == 1:
|
||||
return (latent,)
|
||||
|
||||
samples = latent["samples"]
|
||||
mask = latent.get("noise_mask", None)
|
||||
dilated_shape = samples.shape[:3] + (
|
||||
samples.shape[3] * vertical_scale,
|
||||
samples.shape[4] * horizontal_scale,
|
||||
)
|
||||
|
||||
dilated_samples = torch.zeros(
|
||||
dilated_shape,
|
||||
device=samples.device,
|
||||
dtype=samples.dtype,
|
||||
requires_grad=False,
|
||||
)
|
||||
dilated_samples[..., ::vertical_scale, ::horizontal_scale] = samples
|
||||
|
||||
dilated_mask_shape = (
|
||||
dilated_samples.shape[0],
|
||||
1,
|
||||
dilated_samples.shape[2],
|
||||
dilated_samples.shape[3],
|
||||
dilated_samples.shape[4],
|
||||
)
|
||||
dilated_mask = torch.full(
|
||||
dilated_mask_shape,
|
||||
-1.0,
|
||||
device=samples.device,
|
||||
dtype=samples.dtype,
|
||||
requires_grad=False,
|
||||
)
|
||||
dilated_mask[..., ::vertical_scale, ::horizontal_scale] = (
|
||||
mask if mask is not None else 1.0
|
||||
)
|
||||
latent = {"samples": dilated_samples, "noise_mask": dilated_mask}
|
||||
|
||||
return (latent,)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVAddLatentGuide")
|
||||
class LTXVAddLatentGuide:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"latent": ("LATENT",),
|
||||
"guiding_latent": ("LATENT",),
|
||||
"latent_idx": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": -9999,
|
||||
"max": 9999,
|
||||
"step": 1,
|
||||
"tooltip": "Latent index to start the conditioning at. Can be negative to"
|
||||
"indicate that the conditioning is on the frames before the latent.",
|
||||
},
|
||||
),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
||||
RETURN_NAMES = ("positive", "negative", "latent")
|
||||
|
||||
CATEGORY = "ltxtricks"
|
||||
FUNCTION = "generate"
|
||||
|
||||
DESCRIPTION = "Adds a keyframe or a video segment at a specific frame index."
|
||||
|
||||
def generate(
|
||||
self, vae, positive, negative, latent, guiding_latent, latent_idx, strength
|
||||
):
|
||||
noise_mask = nodes_lt.get_noise_mask(latent)
|
||||
latent = latent["samples"]
|
||||
guide = guiding_latent["samples"]
|
||||
|
||||
assert (
|
||||
latent.shape[4] % guide.shape[4] == 0
|
||||
and latent.shape[3] % guide.shape[3] == 0
|
||||
), "The ratio of the height and width of the latents and optional_guiding_latents must be an integer"
|
||||
|
||||
guiding_latent = LTXVDilateLatent().dilate_latent(
|
||||
guiding_latent,
|
||||
horizontal_scale=latent.shape[4] // guide.shape[4],
|
||||
vertical_scale=latent.shape[3] // guide.shape[3],
|
||||
)[0]
|
||||
|
||||
guide = guiding_latent["samples"]
|
||||
guide_mask = guiding_latent.get("noise_mask", None)
|
||||
|
||||
scale_factors = vae.downscale_index_formula
|
||||
|
||||
if latent_idx <= 0:
|
||||
frame_idx = latent_idx * scale_factors[0]
|
||||
else:
|
||||
frame_idx = 1 + (latent_idx - 1) * scale_factors[0]
|
||||
|
||||
positive, negative, latent, noise_mask = nodes_lt.LTXVAddGuide.append_keyframe(
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
frame_idx=frame_idx,
|
||||
latent_image=latent,
|
||||
noise_mask=noise_mask,
|
||||
guiding_latent=guide,
|
||||
strength=strength,
|
||||
scale_factors=scale_factors,
|
||||
guide_mask=guide_mask,
|
||||
)
|
||||
|
||||
return (
|
||||
positive,
|
||||
negative,
|
||||
{"samples": latent, "noise_mask": noise_mask},
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVImgToVideoConditionOnly")
|
||||
class LTXVImgToVideoConditionOnly:
|
||||
"""
|
||||
Applies image conditioning to the first frames of a video latent.
|
||||
|
||||
Features:
|
||||
- Takes existing latent and applies image conditioning
|
||||
- Automatically resizes image to match latent dimensions
|
||||
- Creates noise mask for strength control
|
||||
- Supports bypass mode
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE",),
|
||||
"image": ("IMAGE",),
|
||||
"latent": ("LATENT",),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
|
||||
},
|
||||
"optional": {
|
||||
"bypass": (
|
||||
"BOOLEAN",
|
||||
{"default": False, "tooltip": "Bypass the conditioning."},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
RETURN_NAMES = ("latent",)
|
||||
CATEGORY = "conditioning/video_models"
|
||||
FUNCTION = "generate"
|
||||
DESCRIPTION = (
|
||||
"Applies image conditioning to the first frames of an existing latent. "
|
||||
"Creates a noise mask to control conditioning strength."
|
||||
)
|
||||
|
||||
def encode_image(self, image, shape, vae):
|
||||
|
||||
time_scale_factor, height_scale_factor, width_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
batch, channels, frames, height, width = shape
|
||||
width *= width_scale_factor
|
||||
height *= height_scale_factor
|
||||
|
||||
# Resize image to match latent dimensions
|
||||
if image.shape[1] != height or image.shape[2] != width:
|
||||
pixels = comfy.utils.common_upscale(
|
||||
image.movedim(-1, 1), width, height, "bilinear", "center"
|
||||
).movedim(1, -1)
|
||||
else:
|
||||
pixels = image
|
||||
|
||||
# Encode image (only RGB channels)
|
||||
encode_pixels = pixels[:, :, :, :3]
|
||||
t = vae.encode(encode_pixels)
|
||||
return t
|
||||
|
||||
def generate(self, image, vae, latent, strength, bypass=False):
|
||||
if bypass:
|
||||
return (latent,)
|
||||
|
||||
samples = latent["samples"]
|
||||
t = self.encode_image(image, samples.shape, vae)
|
||||
# Apply encoded image to first frames
|
||||
samples[:, :, : t.shape[2]] = t
|
||||
|
||||
# Create noise mask for conditioning
|
||||
conditioning_latent_frames_mask = torch.ones(
|
||||
(1, 1, samples.shape[2], 1, 1),
|
||||
dtype=torch.float32,
|
||||
device=samples.device,
|
||||
)
|
||||
conditioning_latent_frames_mask[:, :, : t.shape[2]] = 1.0 - strength
|
||||
|
||||
return ({"samples": samples, "noise_mask": conditioning_latent_frames_mask},)
|
||||
|
||||
|
||||
def get_video_latent_blend_coefficients(
|
||||
video_frame_index_start,
|
||||
video_frame_index_end,
|
||||
video_frame_count,
|
||||
slope_len=3,
|
||||
):
|
||||
"""
|
||||
Returns a blend coefficient list shaped such that:
|
||||
- It is 0.0 outside the range [video_frame_index_start, video_frame_index_end]
|
||||
- It ramps up from 0.0 to 1.0 starting at video_frame_index_start-slope_len to video_frame_index_start
|
||||
- It stays at 1.0 during [video_frame_index_start, video_frame_index_end]
|
||||
- It ramps down from 1.0 to 0.0 starting at video_frame_index_end to video_frame_index_end+slope_len
|
||||
- The 'slope_len' parameter controls the sharpness (in frames) of the slope.
|
||||
"""
|
||||
|
||||
coeffs = [0.0] * video_frame_count
|
||||
|
||||
# Clamp arguments to safe range
|
||||
video_frame_index_start = max(
|
||||
0, min(video_frame_count - 1, video_frame_index_start)
|
||||
)
|
||||
video_frame_index_end = max(
|
||||
video_frame_index_start, min(video_frame_count - 1, video_frame_index_end)
|
||||
)
|
||||
slope_len = max(1, slope_len)
|
||||
|
||||
# Ramp up
|
||||
rampl_start = max(0, video_frame_index_start - slope_len)
|
||||
for i in range(rampl_start, video_frame_index_start):
|
||||
# Linear ramp, but could be smoother (sigmoid/cosine)
|
||||
coeffs[i] = (i - rampl_start + 1) / slope_len
|
||||
|
||||
# Plateau
|
||||
for i in range(video_frame_index_start, video_frame_index_end + 1):
|
||||
coeffs[i] = 1.0
|
||||
|
||||
# Ramp down
|
||||
rampr_end = min(video_frame_count, video_frame_index_end + slope_len + 1)
|
||||
for i in range(video_frame_index_end + 1, rampr_end):
|
||||
coeffs[i] = 1.0 - ((i - (video_frame_index_end + 1) + 1) / slope_len)
|
||||
coeffs[i] = max(0.0, coeffs[i])
|
||||
|
||||
import numpy as np
|
||||
|
||||
num_coeffs = len(coeffs)
|
||||
pixel_frame_length = (num_coeffs - 1) * 8 + 1
|
||||
|
||||
xp = np.array([0] + list(range(1, pixel_frame_length, 8)))
|
||||
fp = np.array(coeffs)
|
||||
|
||||
# Calculate the upsampled coefficients using np.interp
|
||||
pixel_frame_positions = np.arange(pixel_frame_length)
|
||||
pixel_frame_coefficients = np.interp(pixel_frame_positions, xp, fp).tolist()
|
||||
|
||||
return coeffs, pixel_frame_coefficients
|
||||
|
||||
|
||||
@comfy_node(description="LTXV Set Audio Video Mask By Time")
|
||||
class LTXVSetAudioVideoMaskByTime:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"av_latent": ("LATENT",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"model": ("MODEL",),
|
||||
"vae": ("VAE",),
|
||||
"audio_vae": ("VAE",),
|
||||
"start_time": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 2000.0}),
|
||||
"end_time": ("FLOAT", {"default": 10.0, "min": 0.0, "max": 2000.0}),
|
||||
"video_fps": ("FLOAT", {"default": 24.0, "min": 0.0, "max": 500.0}),
|
||||
"mask_video": ("BOOLEAN", {"default": True}),
|
||||
"mask_audio": ("BOOLEAN", {"default": True}),
|
||||
"mask_init_value_video": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "min": 0.0, "max": 1.0},
|
||||
),
|
||||
"mask_init_value_audio": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "min": 0.0, "max": 1.0},
|
||||
),
|
||||
"slope_len": (
|
||||
"INT",
|
||||
{"default": 3, "min": 1, "max": 100, "step": 1},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"spatial_mask": (
|
||||
"MASK",
|
||||
{"default": None, "tooltip": "Spatial mask."},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT", "FLOAT", "FLOAT")
|
||||
RETURN_NAMES = (
|
||||
"positive",
|
||||
"negative",
|
||||
"av_latent",
|
||||
"video_latent_blend_coefficients",
|
||||
"video_pixel_blend_coefficients",
|
||||
)
|
||||
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "utility"
|
||||
DESCRIPTION = "Sets the audio and video mask by time."
|
||||
|
||||
def run(
|
||||
self,
|
||||
av_latent,
|
||||
positive,
|
||||
negative,
|
||||
model,
|
||||
vae,
|
||||
audio_vae,
|
||||
start_time,
|
||||
end_time,
|
||||
video_fps,
|
||||
mask_video,
|
||||
mask_audio,
|
||||
mask_init_value_video,
|
||||
mask_init_value_audio,
|
||||
slope_len,
|
||||
spatial_mask=None,
|
||||
):
|
||||
from comfy.ldm.lightricks.av_model import LTXAVModel
|
||||
|
||||
if model.model.diffusion_model.__class__.__name__ != "LTXAVModel":
|
||||
raise ValueError()
|
||||
|
||||
ltxav: LTXAVModel = model.model.diffusion_model
|
||||
|
||||
# Extract configuration from the audio VAE
|
||||
sampling_rate = audio_vae.autoencoder.sampling_rate
|
||||
mel_hop_length = audio_vae.autoencoder.mel_hop_length
|
||||
audio_latents_per_second = (
|
||||
sampling_rate / mel_hop_length / LATENT_DOWNSAMPLE_FACTOR
|
||||
)
|
||||
|
||||
time_scale_factor = vae.downscale_index_formula[0]
|
||||
video_latents_per_second = video_fps / time_scale_factor
|
||||
|
||||
if not isinstance(av_latent["samples"], NestedTensor):
|
||||
raise ValueError("av_latent must be a NestedTensor")
|
||||
|
||||
video_samples, audio_samples = ltxav.separate_audio_and_video_latents(
|
||||
av_latent["samples"].tensors,
|
||||
None,
|
||||
)
|
||||
video_mask = torch.full(
|
||||
(
|
||||
video_samples.shape[0],
|
||||
video_samples.shape[1],
|
||||
video_samples.shape[2],
|
||||
video_samples.shape[3],
|
||||
video_samples.shape[4],
|
||||
),
|
||||
fill_value=mask_init_value_video,
|
||||
)
|
||||
audio_mask = torch.full(
|
||||
(
|
||||
audio_samples.shape[0],
|
||||
audio_samples.shape[1],
|
||||
audio_samples.shape[2],
|
||||
audio_samples.shape[3],
|
||||
),
|
||||
fill_value=mask_init_value_audio,
|
||||
)
|
||||
if spatial_mask is not None:
|
||||
if spatial_mask.ndim == 3:
|
||||
spatial_mask = spatial_mask.unsqueeze(0)
|
||||
if spatial_mask.ndim == 2:
|
||||
spatial_mask = spatial_mask.unsqueeze(0).unsqueeze(0)
|
||||
spatial_mask = torch.nn.functional.interpolate(
|
||||
spatial_mask,
|
||||
size=(video_samples.shape[3], video_samples.shape[4]),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
video_latent_frame_count = video_samples.shape[2]
|
||||
audio_latent_frame_count = audio_samples.shape[2]
|
||||
video_pixel_frame_count = (video_latent_frame_count - 1) * time_scale_factor + 1
|
||||
xp = np.array(
|
||||
[0]
|
||||
+ list(
|
||||
range(1, video_pixel_frame_count + time_scale_factor, time_scale_factor)
|
||||
)
|
||||
)
|
||||
video_pixel_frame_start_raw = int(round(start_time * video_fps))
|
||||
# video_frame_index_start = index of the value in xp rounding up
|
||||
video_latent_frame_index_start = np.searchsorted(
|
||||
xp, video_pixel_frame_start_raw, side="left"
|
||||
)
|
||||
video_pixel_frame_end_raw = int(round(end_time * video_fps))
|
||||
# video_frame_index_end = index of the value in xp rounding down
|
||||
video_latent_frame_index_end = (
|
||||
np.searchsorted(xp, video_pixel_frame_end_raw, side="right") - 1
|
||||
)
|
||||
audio_latent_frame_index_start = int(
|
||||
round(start_time * audio_latents_per_second)
|
||||
)
|
||||
audio_latent_frame_index_end = (
|
||||
int(round(end_time * audio_latents_per_second)) + 1
|
||||
)
|
||||
# clamping
|
||||
video_latent_frame_index_start = max(0, video_latent_frame_index_start)
|
||||
video_latent_frame_index_end = min(
|
||||
video_latent_frame_index_end, video_latent_frame_count
|
||||
)
|
||||
audio_latent_frame_index_start = max(0, audio_latent_frame_index_start)
|
||||
audio_latent_frame_index_end = min(
|
||||
audio_latent_frame_index_end, audio_latent_frame_count
|
||||
)
|
||||
print(
|
||||
"noise mask start and end indices: video (%d %d), audio (%d %d), video fps: %f, video_latents_per_second: %f, audio_latents_per_second: %f, "
|
||||
"video_latent_frame_count: %d, video_pixel_frame_count: %d, video_pixel_frame_start_raw: %d, video_pixel_frame_end_raw: %d, start_time: %f, end_time: %f"
|
||||
% (
|
||||
video_latent_frame_index_start,
|
||||
video_latent_frame_index_end,
|
||||
audio_latent_frame_index_start,
|
||||
audio_latent_frame_index_end,
|
||||
video_fps,
|
||||
video_latents_per_second,
|
||||
audio_latents_per_second,
|
||||
video_latent_frame_count,
|
||||
video_pixel_frame_count,
|
||||
video_pixel_frame_start_raw,
|
||||
video_pixel_frame_end_raw,
|
||||
start_time,
|
||||
end_time,
|
||||
)
|
||||
)
|
||||
if mask_video:
|
||||
if spatial_mask is not None:
|
||||
video_mask[
|
||||
:,
|
||||
:,
|
||||
video_latent_frame_index_start:video_latent_frame_index_end,
|
||||
:,
|
||||
:,
|
||||
] = spatial_mask
|
||||
else:
|
||||
video_mask[
|
||||
:, :, video_latent_frame_index_start:video_latent_frame_index_end
|
||||
] = 1.0
|
||||
if mask_audio:
|
||||
audio_mask[
|
||||
:, :, audio_latent_frame_index_start:audio_latent_frame_index_end
|
||||
] = 1.0
|
||||
|
||||
av_latent["noise_mask"] = NestedTensor(
|
||||
ltxav.recombine_audio_and_video_latents(video_mask, audio_mask)
|
||||
)
|
||||
|
||||
video_latent_blend_coefficients, video_pixel_blend_coefficients = (
|
||||
get_video_latent_blend_coefficients(
|
||||
video_latent_frame_index_start,
|
||||
video_latent_frame_index_end,
|
||||
video_latent_frame_count,
|
||||
slope_len=slope_len,
|
||||
)
|
||||
)
|
||||
return (
|
||||
positive,
|
||||
negative,
|
||||
av_latent,
|
||||
video_latent_blend_coefficients,
|
||||
video_pixel_blend_coefficients,
|
||||
)
|
||||
@@ -0,0 +1,293 @@
|
||||
# LTXVLoopingSampler Documentation
|
||||
|
||||
## Overview
|
||||
|
||||
The **LTXVLoopingSampler** is a unified ComfyUI node designed for generating long, high-resolution videos through versatile guidance and conditioningi techniques, as well as temporal and spatial tiling. It overcomes memory and computational limitations by breaking video generation into manageable chunks, enabling the creation of videos that would otherwise be impossible to generate in a single model run.
|
||||
It provides a multitude of features (autoregressive long vidoes, keyframes, guidance using IC-LoRA video modalities, normalization, 'long memory' with conditioning on negative positional encodings, and more ) all in one place.
|
||||
|
||||
The **MultiPromptProvider** utility node for providing dynamic prompts per temporal tile.
|
||||
|
||||
---
|
||||
|
||||
## LTXVLoopingSampler
|
||||
|
||||
### Core Concept
|
||||
|
||||
The LTXVLoopingSampler addresses two fundamental challenges in video generation:
|
||||
|
||||
1. **Temporal Length**: Generating very long videos by dividing them into overlapping temporal segments
|
||||
2. **Spatial Resolution**: Creating high-resolution frames by splitting them into overlapping spatial regions
|
||||
|
||||
Each "tile" is processed independently and then seamlessly blended with its neighbors to create the final coherent video output.
|
||||
|
||||
### Key Features
|
||||
|
||||
- **🎬 Temporal Tiling**: Generates long videos by processing overlapping time segments
|
||||
- **🖼️ Spatial Tiling**: Creates high-resolution output by dividing frames into spatial regions
|
||||
- **🎯 Multiple Conditioning Modes**: Supports image conditioning, guiding latents, and negative index conditioning
|
||||
- **⚡ Memory Efficient**: Processes one spatial tile at a time to minimize memory usage
|
||||
- **🔄 Seamless Blending**: Uses weighted blending for smooth transitions between tiles
|
||||
- **🎲 Advanced Seeding**: Configurable per-tile seeding for reproducible results
|
||||
|
||||
### Interface Reference
|
||||
|
||||
#### Required Parameters
|
||||
|
||||
| Parameter | Type | Default | Range | Description |
|
||||
|-----------|------|---------|-------|-------------|
|
||||
| **model** | MODEL | - | - | The diffusion model for video generation |
|
||||
| **vae** | VAE | - | - | VAE for encoding/decoding between pixel and latent space |
|
||||
| **noise** | NOISE | - | - | Noise generator for the sampling process |
|
||||
| **sampler** | SAMPLER | - | - | Sampling algorithm (e.g., Euler, DPM++) |
|
||||
| **sigmas** | SIGMAS | - | - | Noise schedule defining the denoising steps |
|
||||
| **guider** | GUIDER | - | - | Conditioning guider (must be STGGuiderAdvanced to work with LTXVLoopingSampler) |
|
||||
| **latents** | LATENT | - | - | Input latent tensor defining video dimensions and length. Non empty latents can be passed here to enable a partial denoising use case. |
|
||||
|
||||
#### Temporal Configuration
|
||||
|
||||
| Parameter | Type | Default | Range | Description |
|
||||
|-----------|------|---------|-------|-------------|
|
||||
| **temporal_tile_size** | INT | 80 | 24-1000 (step: 8) | Frames per temporal tile (in pixel space). The first tile will be one frame longer than this number, and all new tiles will add `temporal_tile_size-temporal_overlap` new frames each. |
|
||||
| **temporal_overlap** | INT | 24 | 16-80 (step: 8) | Overlapping frames between consecutive tiles |
|
||||
| **temporal_overlap_cond_strength** | FLOAT | 0.5 | 0.0-1.0 | How strongly the final frames of the previous tile condition the new tile generation |
|
||||
|
||||
#### Spatial Configuration
|
||||
|
||||
| Parameter | Type | Default | Range | Description |
|
||||
|-----------|------|---------|-------|-------------|
|
||||
| **horizontal_tiles** | INT | 1 | 1-6 | Number of horizontal spatial divisions |
|
||||
| **vertical_tiles** | INT | 1 | 1-6 | Number of vertical spatial divisions |
|
||||
| **spatial_overlap** | INT | 1 | 1-8 | Overlapping latent 'pixels' between spatial tiles (latent space) |
|
||||
|
||||
#### Conditioning Controls
|
||||
|
||||
| Parameter | Type | Default | Range | Description |
|
||||
|-----------|------|---------|-------|-------------|
|
||||
| **guiding_strength** | FLOAT | 1.0 | 0.0-1.0 | Influence of guiding latents on generation, if provided. |
|
||||
| **cond_image_strength** | FLOAT | 1.0 | 0.0-1.0 | Influence of conditioning images on first image in classical i2v setup, and of all keyframes. |
|
||||
|
||||
#### Optional Advanced Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| **optional_cond_images** | IMAGE | None | Images for I2V conditioning (resized to match output using center crop operation) |
|
||||
| **optional_guiding_latents** | LATENT | None | Latents for guided generation (usually used with IC-LoRA loaded) |
|
||||
| **adain_factor** | FLOAT | 0.0 (0.0-1.0) | AdaIn normalization strength to prevent oversaturation. AdaIn is applied on each new temporal tile, with respect to the latent statistics of the first temporal tile. |
|
||||
| **optional_positive_conditionings** | CONDITIONING | None | Per-tile prompts (from MultiPromptProvider) |
|
||||
| **optional_negative_index_latents** | LATENT | None | Context latents for long-term coherence |
|
||||
| **guiding_start_step** | INT | 0 (0-1000) | Step to begin applying guiding latents, if provided. |
|
||||
| **guiding_end_step** | INT | 1000 (0-1000) | Step to stop applying guiding latents, if provided. |
|
||||
| **optional_cond_image_indices** | STRING | "0" | Keyframe indices for conditioning (comma-separated list of indices) |
|
||||
| **optional_normalizing_latents** | LATENT | None | Reference latents for output normalization. If provided, used in the AdaIn operation, each spatio-temporal tile with its counterpart, instead of using the first temporal tile as a normalization refernece. |
|
||||
|
||||
### Output
|
||||
|
||||
| Output | Type | Description |
|
||||
|--------|------|-------------|
|
||||
| **denoised_output** | LATENT | Complete video latents after tiling and blending |
|
||||
|
||||
### Generation Modes
|
||||
|
||||
#### 🎨 Text-to-Video (T2V)
|
||||
Pure text-driven video generation starting from noise.
|
||||
|
||||
**Setup:**
|
||||
- Don't connect `optional_cond_images`
|
||||
- Provide text conditioning through the guider
|
||||
- Video generates from pure noise
|
||||
|
||||
#### 🖼️ Image-to-Video (I2V)
|
||||
Video generation conditioned on input images.
|
||||
|
||||
**Setup:**
|
||||
- Connect a single image to `optional_cond_images` and set `optional_cond_image_indices` to "0".
|
||||
- Images are automatically resized to match output dimensions
|
||||
|
||||
#### 🗝️ Keyframe-Based Generation
|
||||
|
||||
Keyframe-based generation allows you to guide the video at specific frames using conditioning images (keyframes). This is useful for tasks like animating a sequence between important visual states or ensuring the video matches certain reference images at chosen points.
|
||||
|
||||
- Use `optional_cond_images` to provide one or more keyframe images.
|
||||
- Use `optional_cond_image_indices` to specify which frames should be conditioned on these images. This should be a comma-separated list of frame indices (e.g., `"0, 32, 64"`).
|
||||
- Each image in `optional_cond_images` will be applied to the corresponding frame index in `optional_cond_image_indices`.
|
||||
- The node will automatically resize and center-crop the images to match the output dimensions.
|
||||
|
||||
**Example:**
|
||||
- To condition the first and last frames of a 64-frame video, set:
|
||||
- `optional_cond_images`: [start_image, end_image]
|
||||
- `optional_cond_image_indices`: `"0, 63"`
|
||||
|
||||
This will ensure the generated video starts and ends with the provided keyframes, smoothly interpolating between them.
|
||||
|
||||
**Tip:** You can combine keyframe conditioning with text prompts and guiding latents for even more control over the video generation process.
|
||||
|
||||
#### 🎯 Guided Generation
|
||||
Advanced control using guiding latents (often with IC-LoRA).
|
||||
|
||||
**Setup:**
|
||||
- Connect guiding latents to `optional_guiding_latents`
|
||||
- Adjust `guiding_strength` (0.0-1.0)
|
||||
- Set `guiding_start_step` and `guiding_end_step` for control over the effect of the guidance during the denoising steps.
|
||||
|
||||
#### 📏 High-Resolution Generation
|
||||
Create videos larger than base model resolution using spatial tiling.
|
||||
|
||||
**Configuration Example:**
|
||||
```
|
||||
horizontal_tiles: 2
|
||||
vertical_tiles: 2
|
||||
spatial_overlap: 2
|
||||
```
|
||||
This creates a 2x2 grid of spatial tiles with 2-pixel overlap for seamless blending.
|
||||
|
||||
#### ⏱️ Long Video Generation
|
||||
Generate extended sequences using temporal tiling.
|
||||
|
||||
**How It Works:**
|
||||
- Each temporal tile is conditioned on the **end frames** of the previous tile
|
||||
- This creates seamless motion continuity across tile boundaries
|
||||
- The `temporal_overlap` defines how many frames from the previous tile are used for conditioning
|
||||
|
||||
**Best Practices:**
|
||||
- Use `temporal_overlap` = 1/3 of `temporal_tile_size`
|
||||
- Higher `temporal_overlap_cond_strength` = stronger temporal consistency
|
||||
- Set `adain_factor` = 0.1-0.3 for long sequences to prevent oversaturation
|
||||
- Consider using different prompts per tile via MultiPromptProvider
|
||||
|
||||
#### 🧩 Negative Index Latents for Enhanced Coherence
|
||||
|
||||
You can provide `optional_negative_index_latents` to further improve long video consistency. This feature allows each temporal chunk to attend to an additional latent (or set of latents), typically derived from another image or video, to guide the generation process and maintain global coherence.
|
||||
|
||||
- **How it works:**
|
||||
Each temporal chunk is conditioned on the provided negative index latent(s) using *negative positional embeddings*. This mechanism helps the model reference global context or style cues throughout the video, reducing drift and improving overall consistency, especially in long sequences.
|
||||
|
||||
- **Usage:**
|
||||
- Connect your reference latent(s) to the `optional_negative_index_latents` input.
|
||||
|
||||
- **Typical applications:**
|
||||
- Maintaining a consistent subject or style across long videos
|
||||
- Referencing a global template or anchor image
|
||||
- Enforcing scene or character coherence
|
||||
|
||||
**Tip:** You can combine negative index latents with keyframes, guiding latents, and text prompts for maximum control.
|
||||
|
||||
### Technical Implementation
|
||||
|
||||
#### Temporal Tiling Algorithm
|
||||
|
||||
The node processes video in overlapping temporal chunks, with each tile conditioned on the preceding tile's output:
|
||||
|
||||
1. **First Tile (0 → temporal_tile_size)**
|
||||
- Uses `LTXVInContextSampler` (with guidance) or `LTXVBaseSampler` (without)
|
||||
- Establishes the foundation for the video sequence
|
||||
- Generated from pure noise (T2V) or conditioning images (I2V / keyframes)
|
||||
|
||||
2. **Subsequent Tiles**
|
||||
- Uses `LTXVExtendSampler` to extend from previous tile's results
|
||||
- **Conditioning Mechanism**: Each new tile is conditioned on the **final frames** (overlap region) of the preceding tile
|
||||
- The overlap region from the previous tile provides temporal continuity and motion consistency
|
||||
- `temporal_overlap_cond_strength` controls how strongly the previous tile's end frames influence the new generation
|
||||
- Each tile advances by `(temporal_tile_size - temporal_overlap)` frames, creating seamless temporal progression
|
||||
- Guiding latents, if provided, are used outside of the overlap region.
|
||||
- Conditioning images (`optional_cond_images`), if provided, are used to condition the generation as well.
|
||||
|
||||
#### Spatial Tiling Algorithm
|
||||
|
||||
For high-resolution generation, frames are divided spatially:
|
||||
|
||||
1. **Tile Calculation**: Frame divided into overlapping regions
|
||||
2. **Independent Processing**: Each spatial region processed separately
|
||||
3. **Weight Generation**: Linear blending weights created for overlap regions
|
||||
4. **Final Compositing**: Weighted accumulation produces seamless result
|
||||
|
||||
#### Seed Management
|
||||
|
||||
Each tile receives a unique, deterministic seed:
|
||||
```
|
||||
tile_seed = base_seed + temporal_offset + spatial_offset + custom_offset (hidedn interface)
|
||||
```
|
||||
|
||||
This ensures:
|
||||
- Reproducible results across runs
|
||||
- Consistent randomness per tile
|
||||
- Ability to re-generate specific tiles
|
||||
|
||||
#### Memory Optimization Strategies
|
||||
|
||||
- **Sequential Spatial Processing**: Only one spatial tile in memory at a time
|
||||
- **Temporal Chunking**: Long videos processed in manageable segments
|
||||
- **Weighted Accumulation**: Final output built incrementally
|
||||
- **Garbage Collection**: Intermediate results freed after use
|
||||
|
||||
### Configuration Guidelines
|
||||
|
||||
#### Temporal Parameters
|
||||
- **Short Videos (< 5 seconds)**: `temporal_tile_size: 40-80`
|
||||
- **Medium Videos (5-15 seconds)**: `temporal_tile_size: 80-120`
|
||||
- **Long Videos (> 15 seconds)**: `temporal_tile_size: 120-200`
|
||||
- **Overlap Rule**: Use 25-30% of tile size for smooth transitions
|
||||
|
||||
#### Spatial Parameters
|
||||
- **Standard Resolution**: `horizontal_tiles: 1, vertical_tiles: 1`
|
||||
- **2K Generation**: `horizontal_tiles: 2, vertical_tiles: 1-2`
|
||||
- **4K Generation**: `horizontal_tiles: 2-3, vertical_tiles: 2-3`
|
||||
- **Overlap Minimum**: At least 1 pixel, increase for better blending
|
||||
|
||||
#### Quality Settings
|
||||
- **High Quality**: Lower `adain_factor` (0.0-0.1), higher overlap
|
||||
- **Fast Generation**: Higher `adain_factor` (0.2-0.5), lower overlap
|
||||
- **Temporal Consistency**: Higher `temporal_overlap_cond_strength`
|
||||
|
||||
### Common Use Cases
|
||||
|
||||
#### Creating Cinematic Sequences
|
||||
```
|
||||
temporal_tile_size: 80
|
||||
temporal_overlap: 24
|
||||
horizontal_tiles: 2
|
||||
vertical_tiles: 1
|
||||
adain_factor: 0.1
|
||||
```
|
||||
|
||||
#### Long-Form Content Generation
|
||||
```
|
||||
temporal_tile_size: 120
|
||||
temporal_overlap: 40
|
||||
adain_factor: 0.2
|
||||
+ MultiPromptProvider for narrative progression
|
||||
```
|
||||
|
||||
#### High-Resolution Showcase Videos
|
||||
```
|
||||
horizontal_tiles: 3
|
||||
vertical_tiles: 2
|
||||
spatial_overlap: 3
|
||||
temporal_tile_size: 60
|
||||
```
|
||||
|
||||
### Troubleshooting
|
||||
|
||||
#### Common Issues
|
||||
- **Visible Seams**: Increase spatial/temporal overlap
|
||||
- **Oversaturation**: Increase `adain_factor`
|
||||
- **Memory Errors**: Reduce tile sizes or use fewer spatial tiles
|
||||
- **Inconsistent Motion**: Increase `temporal_overlap` and `temporal_overlap_cond_strength`
|
||||
- **Temporal Discontinuities**: The end frames of each tile condition the next tile - ensure sufficient overlap
|
||||
- **Color Shifts**: Use `optional_normalizing_latents` for reference and/or increase `adain_factor`
|
||||
|
||||
#### Performance Tips
|
||||
- Process spatial tiles sequentially to minimize memory usage
|
||||
- Use appropriate tile sizes for your hardware capabilities
|
||||
- Consider temporal overlap vs. generation speed trade-offs
|
||||
|
||||
---
|
||||
|
||||
## MultiPromptProvider
|
||||
|
||||
A companion utility node that enables dynamic prompt changes across temporal tiles.
|
||||
- Connect multiple pipe-separated prompts (`"prompt1|prompt2|prompt3"`) to create evolving narratives throughout video generation.
|
||||
- Each prompt is used in one temporal_tile in `LTXVLoopingSampler`.
|
||||
- If not enough prompts are provided, the last one is repeated.
|
||||
- If too many prompts provided, the last ones are ignored.
|
||||
|
||||
---
|
||||
@@ -0,0 +1,974 @@
|
||||
import copy
|
||||
from dataclasses import dataclass
|
||||
|
||||
import comfy
|
||||
import torch
|
||||
|
||||
from .easy_samplers import LTXVBaseSampler, LTXVExtendSampler, LTXVInContextSampler
|
||||
from .latents import LTXVDilateLatent, LTXVSelectLatents
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@dataclass
|
||||
class TileConfig:
|
||||
"""Configuration for spatial tile processing."""
|
||||
|
||||
tile_latents: dict
|
||||
tile_guiding_latents: dict
|
||||
tile_negative_index_latents: dict
|
||||
tile_keyframes: torch.Tensor
|
||||
keyframe_per_tile_indices: list
|
||||
tile_normalizing_latents: dict
|
||||
tile_height: int
|
||||
tile_width: int
|
||||
v: int
|
||||
h: int
|
||||
vertical_tiles: int
|
||||
horizontal_tiles: int
|
||||
first_seed: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class SamplingConfig:
|
||||
"""Configuration for sampling parameters."""
|
||||
|
||||
temporal_tile_size: int
|
||||
temporal_overlap: int
|
||||
temporal_overlap_cond_strength: float
|
||||
cond_image_strength: float
|
||||
guiding_strength: float
|
||||
adain_factor: float
|
||||
optional_negative_index: int
|
||||
optional_negative_index_strength: float
|
||||
optional_positive_conditionings: list
|
||||
time_scale_factor: int
|
||||
width_scale_factor: int
|
||||
height_scale_factor: int
|
||||
per_tile_seed_offsets: list
|
||||
guiding_start_step: int
|
||||
guiding_end_step: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelConfig:
|
||||
"""Configuration for model components."""
|
||||
|
||||
model: object
|
||||
vae: object
|
||||
noise: object
|
||||
sampler: object
|
||||
sigmas: object
|
||||
guider: object
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVLoopingSampler",
|
||||
)
|
||||
class LTXVLoopingSampler:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", {"tooltip": "The diffusion model to use."}),
|
||||
"vae": ("VAE", {"tooltip": "The VAE to use."}),
|
||||
"noise": ("NOISE", {"tooltip": "The noise to use."}),
|
||||
"sampler": ("SAMPLER", {"tooltip": "The sampler to use."}),
|
||||
"sigmas": ("SIGMAS", {"tooltip": "The sigmas to use."}),
|
||||
"guider": (
|
||||
"GUIDER",
|
||||
{"tooltip": "The guider to use, must be a STGGuiderAdvanced."},
|
||||
),
|
||||
"latents": (
|
||||
"LATENT",
|
||||
{
|
||||
"tooltip": "The latents to use for creating the long video, they can be guiding latents or empty latents when no guidance is used."
|
||||
},
|
||||
),
|
||||
"temporal_tile_size": (
|
||||
"INT",
|
||||
{
|
||||
"default": 80,
|
||||
"min": 24,
|
||||
"max": 1000,
|
||||
"step": 8,
|
||||
"tooltip": "The size of the temporal tile to use for the sampling, in pixel frames, in addition to the overlapping region.",
|
||||
},
|
||||
),
|
||||
"temporal_overlap": (
|
||||
"INT",
|
||||
{
|
||||
"default": 24,
|
||||
"min": 16,
|
||||
"max": 80,
|
||||
"step": 8,
|
||||
"tooltip": "The overlap between the temporal tiles, in pixel frames.",
|
||||
},
|
||||
),
|
||||
"guiding_strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "The strength of the conditioning on guiding latents, when optional_guiding_latents are provided.",
|
||||
},
|
||||
),
|
||||
"temporal_overlap_cond_strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "The strength of the conditioning on the latents from the previous temporal tile.",
|
||||
},
|
||||
),
|
||||
"cond_image_strength": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "The strength of the conditioning on the optional_cond_images.",
|
||||
},
|
||||
),
|
||||
"horizontal_tiles": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 6,
|
||||
"tooltip": "Number of horizontal spatial tiles.",
|
||||
},
|
||||
),
|
||||
"vertical_tiles": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 6,
|
||||
"tooltip": "Number of vertical spatial tiles.",
|
||||
},
|
||||
),
|
||||
"spatial_overlap": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 8,
|
||||
"tooltip": "Overlap between spatial tiles.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"optional_cond_images": (
|
||||
"IMAGE",
|
||||
{
|
||||
"tooltip": "The image to use for conditioning the first frame in the video (i2v setup). If not provided, the first frame will be unconditioned (t2v setup). The image will be resized to the size of the first frame."
|
||||
},
|
||||
),
|
||||
"optional_guiding_latents": (
|
||||
"LATENT",
|
||||
{
|
||||
"tooltip": "The latents to use for guiding the sampling, typically with an IC-LoRA."
|
||||
},
|
||||
),
|
||||
"adain_factor": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "The strength of the AdaIn operation used to fix the statistics of each new generated temporal tile, to prevent accumulated oversaturation.",
|
||||
},
|
||||
),
|
||||
"optional_positive_conditionings": (
|
||||
"CONDITIONING",
|
||||
{
|
||||
"tooltip": "Optional way to provide changing positive prompts, one per temporal tile, using the MultiPromptProvider node."
|
||||
},
|
||||
),
|
||||
"optional_negative_index_latents": (
|
||||
"LATENT",
|
||||
{
|
||||
"tooltip": "Special optional latents to condition on a negative index before each new temporal tile as a way to provide long term context during video generation."
|
||||
},
|
||||
),
|
||||
"guiding_start_step": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 1000,
|
||||
"tooltip": "The step at which the guiding latents start to be used.",
|
||||
},
|
||||
),
|
||||
"guiding_end_step": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1000,
|
||||
"min": 0,
|
||||
"max": 1000,
|
||||
"tooltip": "The step at which the guiding latents stop to be used.",
|
||||
},
|
||||
),
|
||||
"optional_cond_image_indices": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "0",
|
||||
"tooltip": "The indices of the keyframes to use for the sampling, separated by commas. The indices are counted from the end of the video.",
|
||||
},
|
||||
),
|
||||
"optional_normalizing_latents": (
|
||||
"LATENT",
|
||||
{
|
||||
"tooltip": "The latents to use for normalizing the output latents, they will be used to normalize the output latents to the same statistics as the input latents."
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
RETURN_NAMES = ("denoised_output",)
|
||||
|
||||
FUNCTION = "sample"
|
||||
CATEGORY = "sampling"
|
||||
|
||||
def _extract_latent_spatial_tile(self, latent_dict, v_start, v_end, h_start, h_end):
|
||||
"""Extract spatial tile from a latent dictionary."""
|
||||
if latent_dict is None:
|
||||
return None
|
||||
tile_samples = latent_dict["samples"][:, :, :, v_start:v_end, h_start:h_end]
|
||||
if "noise_mask" in latent_dict and latent_dict["noise_mask"] is not None:
|
||||
tile_masks = latent_dict["noise_mask"][
|
||||
:, :, :, v_start:v_end, h_start:h_end
|
||||
]
|
||||
return {"samples": tile_samples, "noise_mask": tile_masks}
|
||||
else:
|
||||
return {"samples": tile_samples}
|
||||
|
||||
def _extract_spatial_tile(
|
||||
self,
|
||||
latents,
|
||||
optional_guiding_latents,
|
||||
optional_negative_index_latents,
|
||||
optional_normalizing_latents,
|
||||
optional_keyframes,
|
||||
v_start,
|
||||
v_end,
|
||||
h_start,
|
||||
h_end,
|
||||
height_scale_factor,
|
||||
width_scale_factor,
|
||||
):
|
||||
"""Extract spatial tiles from all inputs for a given spatial region."""
|
||||
# Extract spatial tile from latents
|
||||
tile_latents = self._extract_latent_spatial_tile(
|
||||
latents, v_start, v_end, h_start, h_end
|
||||
)
|
||||
|
||||
# Extract spatial tile from guiding latents if provided
|
||||
tile_guiding_latents = self._extract_latent_spatial_tile(
|
||||
optional_guiding_latents, v_start, v_end, h_start, h_end
|
||||
)
|
||||
|
||||
# Extract spatial tile from negative index latents if provided
|
||||
tile_negative_index_latents = self._extract_latent_spatial_tile(
|
||||
optional_negative_index_latents, v_start, v_end, h_start, h_end
|
||||
)
|
||||
|
||||
# Extract spatial tile from normalizing latents if provided
|
||||
tile_normalizing_latents = self._extract_latent_spatial_tile(
|
||||
optional_normalizing_latents, v_start, v_end, h_start, h_end
|
||||
)
|
||||
|
||||
if optional_keyframes is not None:
|
||||
# Scale coordinates for image
|
||||
img_h_start = v_start * height_scale_factor
|
||||
img_h_end = v_end * height_scale_factor
|
||||
img_w_start = h_start * width_scale_factor
|
||||
img_w_end = h_end * width_scale_factor
|
||||
|
||||
tile_keyframes = optional_keyframes[
|
||||
:, img_h_start:img_h_end, img_w_start:img_w_end, :
|
||||
]
|
||||
else:
|
||||
tile_keyframes = None
|
||||
|
||||
return (
|
||||
tile_latents,
|
||||
tile_guiding_latents,
|
||||
tile_negative_index_latents,
|
||||
tile_keyframes,
|
||||
tile_normalizing_latents,
|
||||
)
|
||||
|
||||
def _process_temporal_chunks(
|
||||
self,
|
||||
tile_config: TileConfig,
|
||||
sampling_config: SamplingConfig,
|
||||
model_config: ModelConfig,
|
||||
):
|
||||
"""Process all temporal chunks for a single spatial tile."""
|
||||
chunk_index = 0
|
||||
tile_out_latents = None
|
||||
first_tile_out_latents = None
|
||||
|
||||
for i_temporal_tile, (start_index, end_index) in enumerate(
|
||||
zip(
|
||||
range(
|
||||
0,
|
||||
tile_config.tile_latents["samples"].shape[2]
|
||||
+ sampling_config.temporal_tile_size
|
||||
- sampling_config.temporal_overlap,
|
||||
sampling_config.temporal_tile_size
|
||||
- sampling_config.temporal_overlap,
|
||||
),
|
||||
range(
|
||||
sampling_config.temporal_tile_size,
|
||||
tile_config.tile_latents["samples"].shape[2]
|
||||
+ sampling_config.temporal_tile_size
|
||||
- sampling_config.temporal_overlap,
|
||||
sampling_config.temporal_tile_size
|
||||
- sampling_config.temporal_overlap,
|
||||
),
|
||||
)
|
||||
):
|
||||
if tile_config.tile_guiding_latents is not None:
|
||||
guiding_latent_chunk = LTXVSelectLatents().select_latents(
|
||||
tile_config.tile_guiding_latents,
|
||||
start_index,
|
||||
min(
|
||||
end_index - 1,
|
||||
tile_config.tile_guiding_latents["samples"].shape[2] - 1,
|
||||
),
|
||||
)[0]
|
||||
else:
|
||||
guiding_latent_chunk = None
|
||||
|
||||
latent_chunk = LTXVSelectLatents().select_latents(
|
||||
tile_config.tile_latents,
|
||||
start_index,
|
||||
min(end_index - 1, tile_config.tile_latents["samples"].shape[2] - 1),
|
||||
)[0]
|
||||
|
||||
print(
|
||||
"Processing temporal chunk at index",
|
||||
start_index,
|
||||
"to",
|
||||
min(end_index - 1, tile_config.tile_latents["samples"].shape[2] - 1),
|
||||
)
|
||||
|
||||
seed_offset = self._get_per_tile_value(
|
||||
sampling_config.per_tile_seed_offsets, i_temporal_tile
|
||||
)
|
||||
|
||||
if tile_config.tile_normalizing_latents is not None:
|
||||
normalizing_latent_chunk = LTXVSelectLatents().select_latents(
|
||||
tile_config.tile_normalizing_latents,
|
||||
start_index,
|
||||
min(
|
||||
end_index - 1,
|
||||
tile_config.tile_normalizing_latents["samples"].shape[2] - 1,
|
||||
),
|
||||
)[0]
|
||||
normalize_per_frame = True
|
||||
print(
|
||||
"Normalizing latents provided, normalizing per frame and channel with factor",
|
||||
sampling_config.adain_factor,
|
||||
)
|
||||
else:
|
||||
normalizing_latent_chunk = first_tile_out_latents
|
||||
normalize_per_frame = False
|
||||
print(
|
||||
"No normalizing latents provided, normalizing per channel using first chunk with factor",
|
||||
sampling_config.adain_factor,
|
||||
)
|
||||
|
||||
model_config.noise.seed = self._calculate_tile_seed(
|
||||
tile_config.first_seed,
|
||||
start_index,
|
||||
tile_config.vertical_tiles,
|
||||
tile_config.horizontal_tiles,
|
||||
tile_config.v,
|
||||
tile_config.h,
|
||||
seed_offset,
|
||||
)
|
||||
|
||||
# Handle optional positive conditionings
|
||||
new_guider = self._prepare_guider_for_chunk(
|
||||
model_config.guider,
|
||||
sampling_config.optional_positive_conditionings,
|
||||
chunk_index,
|
||||
)
|
||||
|
||||
this_chunk_keyframe_indices = [
|
||||
in_tile_index
|
||||
for (tile_index, in_tile_index) in tile_config.keyframe_per_tile_indices
|
||||
if tile_index == i_temporal_tile
|
||||
]
|
||||
if this_chunk_keyframe_indices and tile_config.tile_keyframes is not None:
|
||||
this_chunk_keyframes = torch.cat(
|
||||
[
|
||||
tile_config.tile_keyframes[i_keyframe].unsqueeze(0)
|
||||
for (i_keyframe, (tile_index, in_tile_index)) in enumerate(
|
||||
tile_config.keyframe_per_tile_indices
|
||||
)
|
||||
if tile_index == i_temporal_tile
|
||||
]
|
||||
)
|
||||
print(
|
||||
f"Chunk {i_temporal_tile} keyframe indices: {this_chunk_keyframe_indices}"
|
||||
)
|
||||
else:
|
||||
this_chunk_keyframes = None
|
||||
print(f"Chunk {i_temporal_tile} has no keyframes")
|
||||
this_chunk_keyframe_indices = ",".join(
|
||||
[str(i) for i in this_chunk_keyframe_indices]
|
||||
)
|
||||
if start_index == 0:
|
||||
if tile_config.tile_guiding_latents is not None:
|
||||
tile_out_latents = LTXVInContextSampler().sample(
|
||||
vae=model_config.vae,
|
||||
guider=new_guider,
|
||||
sampler=model_config.sampler,
|
||||
sigmas=model_config.sigmas,
|
||||
noise=model_config.noise,
|
||||
guiding_latents=guiding_latent_chunk,
|
||||
optional_cond_images=this_chunk_keyframes,
|
||||
optional_cond_indices=this_chunk_keyframe_indices,
|
||||
num_frames=-1,
|
||||
optional_negative_index_latents=tile_config.tile_negative_index_latents,
|
||||
optional_negative_index=sampling_config.optional_negative_index,
|
||||
optional_negative_index_strength=sampling_config.optional_negative_index_strength,
|
||||
optional_initialization_latents=latent_chunk,
|
||||
cond_image_strength=sampling_config.cond_image_strength,
|
||||
guiding_strength=sampling_config.guiding_strength,
|
||||
guiding_start_step=sampling_config.guiding_start_step,
|
||||
guiding_end_step=sampling_config.guiding_end_step,
|
||||
)[0]
|
||||
else:
|
||||
tile_out_latents = LTXVBaseSampler().sample(
|
||||
model=model_config.model,
|
||||
vae=model_config.vae,
|
||||
noise=model_config.noise,
|
||||
sampler=model_config.sampler,
|
||||
sigmas=model_config.sigmas,
|
||||
guider=new_guider,
|
||||
num_frames=(
|
||||
min(
|
||||
sampling_config.temporal_tile_size,
|
||||
tile_config.tile_latents["samples"].shape[2],
|
||||
)
|
||||
- 1
|
||||
)
|
||||
* sampling_config.time_scale_factor
|
||||
+ 1,
|
||||
width=tile_config.tile_width
|
||||
* sampling_config.width_scale_factor,
|
||||
height=tile_config.tile_height
|
||||
* sampling_config.height_scale_factor,
|
||||
optional_cond_images=this_chunk_keyframes,
|
||||
optional_cond_indices=this_chunk_keyframe_indices,
|
||||
crop="center",
|
||||
crf=30,
|
||||
strength=sampling_config.cond_image_strength,
|
||||
optional_negative_index_latents=tile_config.tile_negative_index_latents,
|
||||
optional_negative_index=sampling_config.optional_negative_index,
|
||||
optional_negative_index_strength=sampling_config.optional_negative_index_strength,
|
||||
optional_initialization_latents=latent_chunk,
|
||||
guiding_start_step=sampling_config.guiding_start_step,
|
||||
guiding_end_step=sampling_config.guiding_end_step,
|
||||
)[0]
|
||||
first_tile_out_latents = copy.deepcopy(tile_out_latents)
|
||||
else:
|
||||
tile_out_latents = LTXVExtendSampler().sample(
|
||||
model=model_config.model,
|
||||
vae=model_config.vae,
|
||||
sampler=model_config.sampler,
|
||||
sigmas=model_config.sigmas,
|
||||
noise=model_config.noise,
|
||||
latents=tile_out_latents,
|
||||
num_new_frames=(
|
||||
latent_chunk["samples"].shape[2]
|
||||
- sampling_config.temporal_overlap
|
||||
)
|
||||
* sampling_config.time_scale_factor,
|
||||
frame_overlap=sampling_config.temporal_overlap
|
||||
* sampling_config.time_scale_factor,
|
||||
guider=new_guider,
|
||||
strength=sampling_config.temporal_overlap_cond_strength,
|
||||
guiding_strength=sampling_config.guiding_strength,
|
||||
cond_image_strength=sampling_config.cond_image_strength,
|
||||
optional_guiding_latents=guiding_latent_chunk,
|
||||
optional_cond_images=this_chunk_keyframes,
|
||||
optional_cond_indices=this_chunk_keyframe_indices,
|
||||
optional_reference_latents=normalizing_latent_chunk,
|
||||
normalize_per_frame=normalize_per_frame,
|
||||
adain_factor=sampling_config.adain_factor,
|
||||
optional_negative_index_latents=tile_config.tile_negative_index_latents,
|
||||
optional_negative_index=sampling_config.optional_negative_index,
|
||||
optional_negative_index_strength=sampling_config.optional_negative_index_strength,
|
||||
optional_initialization_latents=latent_chunk,
|
||||
guiding_start_step=sampling_config.guiding_start_step,
|
||||
guiding_end_step=sampling_config.guiding_end_step,
|
||||
)[0]
|
||||
|
||||
chunk_index += 1
|
||||
|
||||
return tile_out_latents
|
||||
|
||||
def _create_spatial_weights(
|
||||
self,
|
||||
tile_shape,
|
||||
v,
|
||||
h,
|
||||
horizontal_tiles,
|
||||
vertical_tiles,
|
||||
spatial_overlap,
|
||||
device,
|
||||
dtype,
|
||||
):
|
||||
"""Create blending weights for spatial tiles."""
|
||||
tile_weights = torch.ones(tile_shape, device=device, dtype=dtype)
|
||||
|
||||
# Apply horizontal blending weights
|
||||
if h > 0: # Left overlap
|
||||
h_blend = torch.linspace(0, 1, spatial_overlap, device=device, dtype=dtype)
|
||||
tile_weights[:, :, :, :, :spatial_overlap] *= h_blend.view(1, 1, 1, 1, -1)
|
||||
if h < horizontal_tiles - 1: # Right overlap
|
||||
h_blend = torch.linspace(1, 0, spatial_overlap, device=device, dtype=dtype)
|
||||
tile_weights[:, :, :, :, -spatial_overlap:] *= h_blend.view(1, 1, 1, 1, -1)
|
||||
|
||||
# Apply vertical blending weights
|
||||
if v > 0: # Top overlap
|
||||
v_blend = torch.linspace(0, 1, spatial_overlap, device=device, dtype=dtype)
|
||||
tile_weights[:, :, :, :spatial_overlap, :] *= v_blend.view(1, 1, 1, -1, 1)
|
||||
if v < vertical_tiles - 1: # Bottom overlap
|
||||
v_blend = torch.linspace(1, 0, spatial_overlap, device=device, dtype=dtype)
|
||||
tile_weights[:, :, :, -spatial_overlap:, :] *= v_blend.view(1, 1, 1, -1, 1)
|
||||
|
||||
return tile_weights
|
||||
|
||||
def _calculate_tile_seed(
|
||||
self,
|
||||
first_seed,
|
||||
start_index,
|
||||
vertical_tiles,
|
||||
horizontal_tiles,
|
||||
v,
|
||||
h,
|
||||
seed_offset,
|
||||
):
|
||||
"""Calculate the seed value for a specific temporal and spatial tile."""
|
||||
return (
|
||||
first_seed
|
||||
+ start_index * (vertical_tiles * horizontal_tiles)
|
||||
+ v * horizontal_tiles
|
||||
+ h
|
||||
+ seed_offset
|
||||
)
|
||||
|
||||
def _get_per_tile_value(self, value_list, tile_index):
|
||||
"""Get a value from a per-tile configuration list, using the last value if the list is shorter."""
|
||||
return value_list[min(tile_index, len(value_list) - 1)]
|
||||
|
||||
def _parse_comma_separated_string(
|
||||
self, config_string, default_value, converter_func, optional_total_size=None
|
||||
):
|
||||
"""Parse a comma-separated configuration string into a list with type conversion."""
|
||||
if config_string == "":
|
||||
config_string = default_value
|
||||
values = [converter_func(item.strip()) for item in config_string.split(",")]
|
||||
|
||||
def handle_negative_index(value):
|
||||
if value < 0:
|
||||
return value + optional_total_size
|
||||
return value
|
||||
|
||||
if optional_total_size is not None:
|
||||
values = [handle_negative_index(value) for value in values]
|
||||
return values
|
||||
|
||||
def _prepare_guider_for_chunk(
|
||||
self, guider, optional_positive_conditionings, chunk_index
|
||||
):
|
||||
"""Prepare the guider for a specific chunk, handling optional positive conditionings."""
|
||||
if optional_positive_conditionings is not None:
|
||||
new_guider = copy.copy(guider)
|
||||
positive, negative = guider.raw_conds
|
||||
# Use the conditioning at chunk_index, or the last one if we've run out
|
||||
conditioning_index = min(
|
||||
chunk_index, len(optional_positive_conditionings) - 1
|
||||
)
|
||||
new_guider.set_conds(
|
||||
optional_positive_conditionings[conditioning_index],
|
||||
negative,
|
||||
)
|
||||
new_guider.raw_conds = (
|
||||
optional_positive_conditionings[conditioning_index],
|
||||
negative,
|
||||
)
|
||||
return new_guider
|
||||
else:
|
||||
return guider
|
||||
|
||||
def _calculate_keyframe_per_tile_indices(
|
||||
self, keyframe_indices, temporal_tile_size, temporal_overlap, num_frames
|
||||
):
|
||||
"""
|
||||
Calculate which temporal tile each keyframe falls into.
|
||||
|
||||
Returns a list of tuples (temporal_tile_index, in_tile_index) for each keyframe.
|
||||
- temporal_tile_index: which temporal tile the keyframe falls in
|
||||
- in_tile_index: index within that specific temporal tile
|
||||
|
||||
First tile: frames [0, temporal_tile_size - 8] (size = temporal_tile_size - 8 + 1)
|
||||
Subsequent tiles follow the pattern from _process_temporal_chunks:
|
||||
- Tile n starts at: n * (temporal_tile_size - temporal_overlap)
|
||||
- Tile n ends at: temporal_tile_size + n * (temporal_tile_size - temporal_overlap) - 1
|
||||
|
||||
For subsequent tiles, keyframes with in_tile_index < temporal_overlap
|
||||
don't count as falling in that tile.
|
||||
"""
|
||||
result = []
|
||||
|
||||
for keyframe_index in keyframe_indices:
|
||||
if keyframe_index >= num_frames:
|
||||
print(
|
||||
f"Keyframe index {keyframe_index} is greater than num_frames {num_frames}, skipping"
|
||||
)
|
||||
continue
|
||||
# First tile (tile 0): covers frames [0, temporal_tile_size - 8]
|
||||
if keyframe_index < temporal_tile_size - 7:
|
||||
result.append((0, keyframe_index))
|
||||
continue
|
||||
|
||||
# Find which subsequent tile this keyframe could fall into
|
||||
# Tile n starts at: n * (temporal_tile_size - temporal_overlap) - 7
|
||||
# Tile n ends at: temporal_tile_size + n * (temporal_tile_size - temporal_overlap) - 1 - 7
|
||||
|
||||
tile_step = temporal_tile_size - temporal_overlap
|
||||
tile_index = 1
|
||||
|
||||
while True:
|
||||
tile_start = tile_index * tile_step - 7
|
||||
tile_end = temporal_tile_size + tile_index * tile_step - 1 - 7
|
||||
print(
|
||||
f"Tile {tile_index} starts at {tile_start} and ends at {tile_end}"
|
||||
)
|
||||
|
||||
# Check if keyframe falls within this tile's range
|
||||
if keyframe_index <= tile_end:
|
||||
in_tile_index = (
|
||||
keyframe_index - tile_start - 7
|
||||
) # this extra -7 is needed because the first latent is re-interpreted inside the temporal tile as 1 pixel frame
|
||||
|
||||
# For tiles > 0, if in_tile_index < temporal_overlap,
|
||||
# the keyframe doesn't "fall" in this tile - assign to previous tile
|
||||
if in_tile_index < temporal_overlap:
|
||||
tile_index -= 1
|
||||
if tile_index == 0:
|
||||
# Previous tile is the first tile
|
||||
in_tile_index = keyframe_index
|
||||
else:
|
||||
# Previous tile starts at (tile_index) * tile_step
|
||||
prev_start = tile_start - tile_step
|
||||
in_tile_index = (
|
||||
keyframe_index - prev_start - 7
|
||||
) # same reason as above
|
||||
|
||||
result.append((tile_index, in_tile_index))
|
||||
break
|
||||
|
||||
tile_index += 1
|
||||
|
||||
return result
|
||||
|
||||
def sample(
|
||||
self,
|
||||
model,
|
||||
vae,
|
||||
noise,
|
||||
sampler,
|
||||
sigmas,
|
||||
guider,
|
||||
latents,
|
||||
guiding_strength,
|
||||
adain_factor,
|
||||
temporal_tile_size,
|
||||
temporal_overlap,
|
||||
temporal_overlap_cond_strength,
|
||||
horizontal_tiles,
|
||||
vertical_tiles,
|
||||
spatial_overlap,
|
||||
optional_cond_images=None,
|
||||
cond_image_strength=1.0,
|
||||
optional_guiding_latents=None,
|
||||
optional_negative_index_latents=None,
|
||||
optional_negative_index_strength=1.0, # hidden interface
|
||||
optional_positive_conditionings=None,
|
||||
guiding_start_step=0,
|
||||
guiding_end_step=1000,
|
||||
optional_cond_image_indices="0",
|
||||
optional_normalizing_latents=None,
|
||||
per_tile_seed_offsets="0", # hidden interface
|
||||
):
|
||||
# Get dimensions and prepare for spatial tiling
|
||||
samples = latents["samples"]
|
||||
if (
|
||||
isinstance(samples, comfy.nested_tensor.NestedTensor)
|
||||
and len(samples.tensors) == 2
|
||||
):
|
||||
raise ValueError(
|
||||
"LoopingSampler currently does not support Audio Visual latents. please only use video latents."
|
||||
)
|
||||
batch, channels, frames, height, width = samples.shape
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
temporal_tile_size = temporal_tile_size // time_scale_factor
|
||||
temporal_overlap = temporal_overlap // time_scale_factor
|
||||
first_seed = noise.seed
|
||||
|
||||
per_tile_seed_offsets = self._parse_comma_separated_string(
|
||||
per_tile_seed_offsets, "0", int
|
||||
)
|
||||
|
||||
keyframe_indices = self._parse_comma_separated_string(
|
||||
optional_cond_image_indices,
|
||||
"0",
|
||||
int,
|
||||
optional_total_size=frames * time_scale_factor - 7,
|
||||
)
|
||||
keyframe_per_tile_indices = self._calculate_keyframe_per_tile_indices(
|
||||
keyframe_indices,
|
||||
temporal_tile_size * time_scale_factor,
|
||||
temporal_overlap * time_scale_factor,
|
||||
frames * time_scale_factor - 7,
|
||||
)
|
||||
print(f"Keyframe per tile indices: {keyframe_per_tile_indices}")
|
||||
if optional_cond_images is not None:
|
||||
optional_keyframes = (
|
||||
comfy.utils.common_upscale(
|
||||
optional_cond_images.movedim(-1, 1),
|
||||
width * width_scale_factor,
|
||||
height * height_scale_factor,
|
||||
"bilinear",
|
||||
crop="center",
|
||||
)
|
||||
.movedim(1, -1)
|
||||
.clamp(0, 1)
|
||||
)
|
||||
else:
|
||||
optional_keyframes = None
|
||||
|
||||
if optional_guiding_latents is not None:
|
||||
guide = optional_guiding_latents["samples"]
|
||||
assert (
|
||||
samples.shape[2] == guide.shape[2]
|
||||
), "The number of frames in the latents and optional_guiding_latents must be the same"
|
||||
assert (
|
||||
samples.shape[3] % guide.shape[3] == 0
|
||||
and samples.shape[4] % guide.shape[4] == 0
|
||||
), "The ratio of the height and width of the latents and optional_guiding_latents must be an integer"
|
||||
grid_size_h = samples.shape[3] // guide.shape[3]
|
||||
grid_size_w = samples.shape[4] // guide.shape[4]
|
||||
optional_guiding_latents = LTXVDilateLatent().dilate_latent(
|
||||
optional_guiding_latents, grid_size_w, grid_size_h
|
||||
)[0]
|
||||
|
||||
# Calculate tile sizes with overlap
|
||||
base_tile_height = (
|
||||
height + (vertical_tiles - 1) * spatial_overlap
|
||||
) // vertical_tiles
|
||||
base_tile_width = (
|
||||
width + (horizontal_tiles - 1) * spatial_overlap
|
||||
) // horizontal_tiles
|
||||
|
||||
# Output tensors will be initialized after first tile to get correct temporal dimension
|
||||
final_output = None
|
||||
weights = None
|
||||
|
||||
# Process each spatial tile
|
||||
for v in range(vertical_tiles):
|
||||
for h in range(horizontal_tiles):
|
||||
# Calculate tile boundaries
|
||||
h_start = h * (base_tile_width - spatial_overlap)
|
||||
v_start = v * (base_tile_height - spatial_overlap)
|
||||
|
||||
# Adjust end positions for edge tiles
|
||||
h_end = (
|
||||
min(h_start + base_tile_width, width)
|
||||
if h < horizontal_tiles - 1
|
||||
else width
|
||||
)
|
||||
v_end = (
|
||||
min(v_start + base_tile_height, height)
|
||||
if v < vertical_tiles - 1
|
||||
else height
|
||||
)
|
||||
|
||||
# Calculate actual tile dimensions
|
||||
tile_height = v_end - v_start
|
||||
tile_width = h_end - h_start
|
||||
|
||||
print(f"Processing spatial tile at row {v}, col {h}:")
|
||||
print(f" Position: ({v_start}:{v_end}, {h_start}:{h_end})")
|
||||
print(f" Size: {tile_height}x{tile_width}")
|
||||
|
||||
# Extract spatial tiles from all inputs
|
||||
(
|
||||
tile_latents,
|
||||
tile_guiding_latents,
|
||||
tile_negative_index_latents,
|
||||
tile_keyframes,
|
||||
tile_normalizing_latents,
|
||||
) = self._extract_spatial_tile(
|
||||
latents,
|
||||
optional_guiding_latents,
|
||||
optional_negative_index_latents,
|
||||
optional_normalizing_latents,
|
||||
optional_keyframes,
|
||||
v_start,
|
||||
v_end,
|
||||
h_start,
|
||||
h_end,
|
||||
height_scale_factor,
|
||||
width_scale_factor,
|
||||
)
|
||||
|
||||
# Process all temporal chunks for this spatial tile
|
||||
tile_config = TileConfig(
|
||||
tile_latents=tile_latents,
|
||||
tile_guiding_latents=tile_guiding_latents,
|
||||
tile_negative_index_latents=tile_negative_index_latents,
|
||||
tile_keyframes=tile_keyframes,
|
||||
keyframe_per_tile_indices=keyframe_per_tile_indices,
|
||||
tile_normalizing_latents=tile_normalizing_latents,
|
||||
tile_height=tile_height,
|
||||
tile_width=tile_width,
|
||||
v=v,
|
||||
h=h,
|
||||
vertical_tiles=vertical_tiles,
|
||||
horizontal_tiles=horizontal_tiles,
|
||||
first_seed=first_seed,
|
||||
)
|
||||
|
||||
sampling_config = SamplingConfig(
|
||||
temporal_tile_size=temporal_tile_size,
|
||||
temporal_overlap=temporal_overlap,
|
||||
temporal_overlap_cond_strength=temporal_overlap_cond_strength,
|
||||
cond_image_strength=cond_image_strength,
|
||||
guiding_strength=guiding_strength,
|
||||
adain_factor=adain_factor,
|
||||
optional_negative_index=(
|
||||
-1 * tile_negative_index_latents["samples"].shape[2]
|
||||
if tile_negative_index_latents is not None
|
||||
else -1
|
||||
),
|
||||
optional_negative_index_strength=optional_negative_index_strength,
|
||||
optional_positive_conditionings=optional_positive_conditionings,
|
||||
time_scale_factor=time_scale_factor,
|
||||
width_scale_factor=width_scale_factor,
|
||||
height_scale_factor=height_scale_factor,
|
||||
per_tile_seed_offsets=per_tile_seed_offsets,
|
||||
guiding_start_step=guiding_start_step,
|
||||
guiding_end_step=guiding_end_step,
|
||||
)
|
||||
|
||||
model_config = ModelConfig(
|
||||
model=model,
|
||||
vae=vae,
|
||||
noise=noise,
|
||||
sampler=sampler,
|
||||
sigmas=sigmas,
|
||||
guider=guider,
|
||||
)
|
||||
|
||||
tile_out_latents = self._process_temporal_chunks(
|
||||
tile_config,
|
||||
sampling_config,
|
||||
model_config,
|
||||
)
|
||||
|
||||
# Initialize output tensors on first tile (to get correct temporal dimension)
|
||||
if final_output is None:
|
||||
out_temporal = tile_out_latents["samples"].shape[2]
|
||||
final_output = torch.zeros(
|
||||
batch,
|
||||
channels,
|
||||
out_temporal,
|
||||
height,
|
||||
width,
|
||||
device=samples.device,
|
||||
dtype=samples.dtype,
|
||||
)
|
||||
weights = torch.zeros_like(final_output)
|
||||
|
||||
# Move tile samples to output device and create weight mask
|
||||
tile_samples = tile_out_latents["samples"].to(final_output.device)
|
||||
tile_weights = self._create_spatial_weights(
|
||||
tile_samples.shape,
|
||||
v,
|
||||
h,
|
||||
horizontal_tiles,
|
||||
vertical_tiles,
|
||||
spatial_overlap,
|
||||
final_output.device,
|
||||
final_output.dtype,
|
||||
)
|
||||
|
||||
# Add weighted tile to final output
|
||||
final_output[:, :, :, v_start:v_end, h_start:h_end] += (
|
||||
tile_samples * tile_weights
|
||||
)
|
||||
weights[:, :, :, v_start:v_end, h_start:h_end] += tile_weights
|
||||
|
||||
# Normalize by weights
|
||||
final_output = final_output / (weights + 1e-8)
|
||||
out_latents = {"samples": final_output}
|
||||
|
||||
noise.seed = first_seed
|
||||
return (out_latents,)
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="MultiPromptProvider",
|
||||
)
|
||||
class MultiPromptProvider:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"prompts": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"dynamicPrompts": True,
|
||||
"tooltip": "Prompts to encode, one per line. Each prompt will be encoded separately. Each prompt will be used in one temporal_tile in LTXVLoopingSampler.",
|
||||
},
|
||||
),
|
||||
"clip": ("CLIP", {"tooltip": "CLIP model to encode the prompts."}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING",)
|
||||
RETURN_NAMES = ("conditionings",)
|
||||
|
||||
FUNCTION = "get_prompt_list"
|
||||
CATEGORY = "prompt"
|
||||
|
||||
def get_prompt_list(self, prompts, clip):
|
||||
prompt_list = prompts.split("|")
|
||||
prompt_list = [prompt.strip() for prompt in prompt_list]
|
||||
encoded_prompt_list = [
|
||||
clip.encode_from_tokens_scheduled(clip.tokenize(prompt))
|
||||
for prompt in prompt_list
|
||||
]
|
||||
return (encoded_prompt_list,)
|
||||
@@ -0,0 +1,127 @@
|
||||
"""
|
||||
Low VRAM loader nodes for sequential model loading.
|
||||
|
||||
These nodes add a 'dependencies' input to control execution order,
|
||||
ensuring models load one after another instead of simultaneously.
|
||||
This reduces peak VRAM usage in memory-constrained environments.
|
||||
"""
|
||||
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
import nodes
|
||||
from comfy.ldm.lightricks.vae.audio_vae import AudioVAE
|
||||
from comfy_api.latest import io
|
||||
from comfy_extras.nodes_hunyuan import LatentUpscaleModelLoader
|
||||
|
||||
|
||||
class AnyType(str):
|
||||
"""Custom type that matches any ComfyUI type for dependency inputs."""
|
||||
|
||||
def __ne__(self, __value: object) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
any_type = AnyType("*")
|
||||
|
||||
|
||||
class LowVRAMCheckpointLoader(nodes.CheckpointLoaderSimple):
|
||||
"""Load Checkpoint with dependency input for sequential loading in low VRAM scenarios."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s) -> dict:
|
||||
base_types = super().INPUT_TYPES()
|
||||
base_types["optional"] = {
|
||||
"dependencies": (
|
||||
any_type,
|
||||
{
|
||||
"tooltip": "Connect any output from a previous loader to ensure sequential loading."
|
||||
},
|
||||
),
|
||||
}
|
||||
return base_types
|
||||
|
||||
CATEGORY = "LTXV/loaders"
|
||||
DESCRIPTION = (
|
||||
"Loads a diffusion model checkpoint with dependency support. "
|
||||
"Connect 'dependencies' to a previous loader's output to ensure sequential loading and reduce peak VRAM usage."
|
||||
)
|
||||
FUNCTION = "load_checkpoint_sequentially"
|
||||
|
||||
def load_checkpoint_sequentially(
|
||||
self,
|
||||
ckpt_name: str,
|
||||
dependencies: object = None, # noqa: ARG002 - required for ComfyUI optional inputs
|
||||
) -> tuple:
|
||||
return super().load_checkpoint(ckpt_name)
|
||||
|
||||
|
||||
class LowVRAMAudioVAELoader:
|
||||
"""Load LTXV Audio VAE with dependency input for sequential loading in low VRAM scenarios."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s) -> dict:
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (
|
||||
folder_paths.get_filename_list("checkpoints"),
|
||||
{"tooltip": "Audio VAE checkpoint to load."},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"dependencies": (
|
||||
any_type,
|
||||
{
|
||||
"tooltip": "Connect any output from a previous loader to ensure sequential loading."
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VAE",)
|
||||
RETURN_NAMES = ("audio_vae",)
|
||||
CATEGORY = "LTXV/loaders"
|
||||
DESCRIPTION = (
|
||||
"Loads an LTXV Audio VAE checkpoint with dependency support. "
|
||||
"Connect 'dependencies' to a previous loader's output to ensure sequential loading and reduce peak VRAM usage."
|
||||
)
|
||||
FUNCTION = "load_audio_vae_sequentially"
|
||||
|
||||
def load_audio_vae_sequentially(
|
||||
self,
|
||||
ckpt_name: str,
|
||||
dependencies: object = None, # noqa: ARG002 - required for ComfyUI optional inputs
|
||||
) -> tuple:
|
||||
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
|
||||
sd, metadata = comfy.utils.load_torch_file(ckpt_path, return_metadata=True)
|
||||
audio_vae = AudioVAE(sd, metadata)
|
||||
return (audio_vae,)
|
||||
|
||||
|
||||
class LowVRAMLatentUpscaleModelLoader(LatentUpscaleModelLoader):
|
||||
"""Load latent upsampler with dependency input for sequential loading in low VRAM scenarios."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LowVRAMLatentUpscaleModelLoader",
|
||||
display_name="Low VRAMLoad Latent Upscale Model",
|
||||
category="LTXV/loaders",
|
||||
inputs=[
|
||||
io.Combo.Input(
|
||||
"model_name",
|
||||
options=folder_paths.get_filename_list("latent_upscale_models"),
|
||||
),
|
||||
io.AnyType.Input(
|
||||
"dependencies",
|
||||
optional=True,
|
||||
tooltip="Connect any output from a previous loader to ensure sequential loading.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.LatentUpscaleModel.Output(),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_name, dependencies=None) -> io.NodeOutput:
|
||||
return LatentUpscaleModelLoader.execute(model_name=model_name)
|
||||
@@ -0,0 +1,212 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from comfy.model_management import is_device_cpu
|
||||
from comfy_extras.nodes_mask import GrowMask
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPreprocessMasks")
|
||||
class LTXVPreprocessMasks:
|
||||
"""
|
||||
Preprocesses masks for use with the LTXVideo model's latent masking.
|
||||
|
||||
This node handles temporal mask processing by:
|
||||
1. Validating mask dimensions against VAE downscaling factors
|
||||
2. Optionally inverting masks
|
||||
3. Handling the first frame mask separately (can be ignored)
|
||||
4. Combining temporal masks using max pooling
|
||||
5. Growing/shrinking masks with morphological operations
|
||||
6. Clamping mask values to ensure proper opacity
|
||||
|
||||
The output is a set of masks ready for latent-space masking operations.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: "LTXVPreprocessMasks") -> dict[str, Any]: # noqa: N802
|
||||
return {
|
||||
"required": {
|
||||
"masks": ("MASK",),
|
||||
"vae": ("VAE",),
|
||||
"invert_input_masks": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "Invert the input masks before processing."
|
||||
"Useful for masking vs unmasking operations.",
|
||||
},
|
||||
),
|
||||
"ignore_first_mask": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Zero out the first mask, typically used when"
|
||||
"it corresponds to the conditioning frame.",
|
||||
},
|
||||
),
|
||||
"pooling_method": (
|
||||
["max", "mean", "min"],
|
||||
{
|
||||
"default": "max",
|
||||
"tooltip": "Method to combine temporal masks. Max preserves strongest values, "
|
||||
"mean averages them, min takes weakest values.",
|
||||
},
|
||||
),
|
||||
"grow_mask": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": -16384,
|
||||
"max": 16384,
|
||||
"step": 1,
|
||||
"tooltip": "Pixels to grow (positive) or shrink (negative) the mask."
|
||||
"Uses morphological operations.",
|
||||
},
|
||||
),
|
||||
"tapered_corners": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "Use tapered corners in morphological operations for smoother mask edges.",
|
||||
},
|
||||
),
|
||||
"clamp_min": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Minimum value to clamp masks to, ensuring minimum opacity.",
|
||||
},
|
||||
),
|
||||
"clamp_max": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Maximum value to clamp masks to.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "preprocess_masks"
|
||||
CATEGORY = "Lightricks/mask_operations"
|
||||
DESCRIPTION = (
|
||||
"Preprocess masks to be used for masking latents in the LTXVideo model."
|
||||
)
|
||||
|
||||
def preprocess_masks(
|
||||
self,
|
||||
masks: torch.Tensor,
|
||||
vae, # noqa: ANN001
|
||||
pooling_method: str,
|
||||
grow_mask: int,
|
||||
tapered_corners: bool,
|
||||
clamp_min: float,
|
||||
clamp_max: float,
|
||||
ignore_first_mask: bool,
|
||||
invert_input_masks: bool,
|
||||
) -> tuple[torch.Tensor]:
|
||||
"""
|
||||
Preprocess masks for temporal latent masking.
|
||||
|
||||
Args:
|
||||
masks: Input masks tensor of shape (batch_size, H, W) where batch_size = 1 + (num_frames - 1)
|
||||
The first mask is for the conditioning frame, remaining masks are for generated frames
|
||||
vae: VAE model object containing downscale_index_formula attribute
|
||||
pooling_method: Method to combine temporal masks ('max', 'mean', 'min')
|
||||
grow_mask: Pixels to grow (positive) or shrink (negative) the mask
|
||||
tapered_corners: Whether to use tapered corners in morphological operations
|
||||
clamp_min: Minimum value to clamp masks to (ensures minimum opacity)
|
||||
clamp_max: Maximum value to clamp masks to
|
||||
ignore_first_mask: Whether to zero out the first mask (often used for conditioning)
|
||||
invert_input_masks: Whether to invert all input masks (1 - mask)
|
||||
|
||||
Returns:
|
||||
Tuple containing the processed masks tensor
|
||||
|
||||
Raises:
|
||||
ValueError: If mask dimensions are not compatible with VAE downscaling factors
|
||||
"""
|
||||
|
||||
if masks.ndim != 3:
|
||||
raise ValueError("Masks must be of shape (batch_size, H, W).")
|
||||
|
||||
batch_size, H, W = masks.shape
|
||||
|
||||
# Get VAE downscaling factors for temporal and spatial dimensions
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
|
||||
# Validate that mask dimensions are compatible with VAE downscaling
|
||||
# The batch size minus 1 (for the first frame) must be divisible by time_scale_factor
|
||||
if (batch_size - 1) % time_scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Masks batch size must have a multiple of {time_scale_factor} masks + 1."
|
||||
)
|
||||
|
||||
if H % height_scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Masks height must have a multiple of {height_scale_factor}."
|
||||
)
|
||||
|
||||
if W % width_scale_factor != 0:
|
||||
raise ValueError(
|
||||
f"Masks width must have a multiple of {width_scale_factor}."
|
||||
)
|
||||
|
||||
# Optionally invert all masks (useful for masking vs unmasking operations)
|
||||
if invert_input_masks:
|
||||
masks = 1.0 - masks
|
||||
|
||||
# Handle first mask separately - often used as conditioning frame
|
||||
if ignore_first_mask:
|
||||
first_mask = masks[0:1, :, :] * 0 # Zero out the first mask
|
||||
else:
|
||||
first_mask = masks[0:1, :, :]
|
||||
|
||||
# Process remaining masks by grouping them according to time_scale_factor
|
||||
# and taking the maximum value across each temporal group
|
||||
remaining = masks[1:].view(-1, time_scale_factor, H, W)
|
||||
|
||||
# Apply selected pooling method
|
||||
pooling_methods = {
|
||||
"max": lambda x: x.max(dim=1)[0],
|
||||
"mean": lambda x: x.mean(dim=1),
|
||||
"min": lambda x: x.min(dim=1)[0],
|
||||
}
|
||||
|
||||
combined = pooling_methods[pooling_method](remaining)
|
||||
|
||||
# Apply morphological operations to grow or shrink masks
|
||||
if grow_mask != 0:
|
||||
grow_mask_node = GrowMask()
|
||||
|
||||
# GrowMask node is not supported on GPU, so we need to move the combined tensor to CPU
|
||||
# if it is not already on CPU
|
||||
input_device = combined.device
|
||||
if not is_device_cpu(input_device):
|
||||
combined = combined.cpu()
|
||||
|
||||
(combined,) = grow_mask_node.expand_mask(
|
||||
mask=combined, expand=grow_mask, tapered_corners=tapered_corners
|
||||
)
|
||||
|
||||
# Move the combined tensor back to the original device
|
||||
if not is_device_cpu(input_device):
|
||||
combined = combined.to(input_device)
|
||||
|
||||
# Combine first mask with processed remaining masks and clamp values
|
||||
# Clamping ensures masks have proper opacity values
|
||||
processed_masks = torch.cat(
|
||||
[first_mask, torch.clamp(combined, clamp_min, clamp_max)], dim=0
|
||||
)
|
||||
|
||||
return (processed_masks,)
|
||||
+2
-2
@@ -99,7 +99,7 @@ def _format_description(
|
||||
|
||||
# If description is not provided, auto-generate one based on the class name
|
||||
if description is None:
|
||||
description = _camel_case_to_spaces(class_name)
|
||||
description = camel_case_to_spaces(class_name)
|
||||
|
||||
# Strip the prefix if it's already there
|
||||
prefix_len = len(NODES_DISPLAY_NAME_PREFIX)
|
||||
@@ -118,7 +118,7 @@ def _format_description(
|
||||
return description
|
||||
|
||||
|
||||
def _camel_case_to_spaces(text: str) -> str:
|
||||
def camel_case_to_spaces(text: str) -> str:
|
||||
# Add space before each capital letter except the first one
|
||||
spaced_text = re.sub(r"(?<=[a-z])(?=[A-Z])", " ", text)
|
||||
# Handle sequences of uppercase letters followed by a lowercase letter
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,55 @@
|
||||
[
|
||||
{
|
||||
"name": "Custom"
|
||||
},
|
||||
{
|
||||
"name": "13b Dynamic",
|
||||
"skip_steps_sigma_threshold": 0.997,
|
||||
"cfg_star_rescale": true,
|
||||
"sigmas": [1.0, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180],
|
||||
"cfg_values": [1, 6, 8, 6, 1, 1],
|
||||
"stg_scale_values": [0, 4, 4, 4, 2, 1],
|
||||
"stg_rescale_values": [1, 0.5, 0.5, 1, 1, 1],
|
||||
"stg_layers_indices": [[11, 25, 35, 39], [22, 35, 39], [28], [28], [28], [28]]
|
||||
},
|
||||
{
|
||||
"name": "13b Balanced",
|
||||
"skip_steps_sigma_threshold": 0.998,
|
||||
"cfg_star_rescale": true,
|
||||
"sigmas": [1.0, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180],
|
||||
"cfg_values": [1, 6, 8, 6, 1, 1],
|
||||
"stg_scale_values": [0, 4, 4, 4, 2, 1],
|
||||
"stg_rescale_values": [1, 0.5, 0.5, 1, 1, 1],
|
||||
"stg_layers_indices": [[12], [12], [5], [5], [28], [29]]
|
||||
},
|
||||
{
|
||||
"name": "13b Upscale",
|
||||
"skip_steps_sigma_threshold": 0.997,
|
||||
"cfg_star_rescale": true,
|
||||
"sigmas": [1.0, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180],
|
||||
"cfg_values": [1, 1, 1, 1, 1, 1],
|
||||
"stg_scale_values": [1, 1, 1, 1, 1, 1],
|
||||
"stg_rescale_values": [1, 1, 1, 1, 1, 1],
|
||||
"stg_layers_indices": [[42], [42], [42], [42], [42], [42]]
|
||||
},
|
||||
{
|
||||
"name": "13b Distilled",
|
||||
"skip_steps_sigma_threshold": 0.997,
|
||||
"cfg_star_rescale": true,
|
||||
"sigmas": [1.0],
|
||||
"cfg_values": [1],
|
||||
"stg_scale_values": [0],
|
||||
"stg_rescale_values": [1],
|
||||
"stg_layers_indices": [[25]]
|
||||
},
|
||||
{
|
||||
"name": "2b",
|
||||
"skip_steps_sigma_threshold": 0.997,
|
||||
"cfg_star_rescale": true,
|
||||
"sigmas": [1.0, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180],
|
||||
"cfg_values": [4, 4, 4, 4, 1, 1],
|
||||
"stg_scale_values": [2, 2, 2, 2, 1, 0],
|
||||
"stg_rescale_values": [1, 1, 1, 1, 1, 1],
|
||||
"stg_layers_indices": [[14], [14], [14], [14], [14], [14]]
|
||||
}
|
||||
]
|
||||
@@ -90,6 +90,7 @@ class LTXVPromptEnhancerLoader:
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Prompt Enhancer (Down)Loader"
|
||||
OUTPUT_NODE = False
|
||||
DESCRIPTION = "Downloads and initializes LLM and image captioning models from Hugging Face to enhance text prompts for image generation."
|
||||
|
||||
def model_path_download_if_needed(self, model_name):
|
||||
model_directory = os.path.join(folder_paths.models_dir, MODELS_PATH_KEY)
|
||||
@@ -180,6 +181,10 @@ class LTXVPromptEnhancer:
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Prompt Enhancer"
|
||||
OUTPUT_NODE = False
|
||||
DESCRIPTION = (
|
||||
"Enhances text prompts for image generation using LLMs. "
|
||||
"Optionally incorporates reference images to create more contextually relevant descriptions."
|
||||
)
|
||||
|
||||
def enhance(
|
||||
self,
|
||||
|
||||
+29
-2
@@ -9,14 +9,24 @@ requires-python = ">=3.10"
|
||||
readme = "README.md"
|
||||
license = { file = "LICENSE" }
|
||||
dependencies = [
|
||||
"diffusers",
|
||||
"huggingface_hub>=0.25.2",
|
||||
"transformers[timm]>=4.45.0"
|
||||
"transformers[timm]>=4.50.0",
|
||||
"einops",
|
||||
"ninja~=1.11.1.4",
|
||||
"accelerate"
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
internal = [
|
||||
"ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1",
|
||||
"av>=10.0.0"
|
||||
"ltx-media-io>=0.1.26",
|
||||
"av>=10.0.0",
|
||||
"q8-kernels~=0.1.4",
|
||||
"ring_buffer==0.0.9",
|
||||
"ltx-multigpu-inference~=0.0.13",
|
||||
"bitsandbytes",
|
||||
"ultralytics"
|
||||
]
|
||||
dev = [
|
||||
"pre-commit>=4.0.1",
|
||||
@@ -24,3 +34,20 @@ dev = [
|
||||
"websocket-client==1.6.1",
|
||||
"scikit-image==0.24.0"
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1",
|
||||
ring-buffer = {index = "ltx"}
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "ltx"
|
||||
url = "https://oauth2accesstoken@us-python.pkg.dev/ltx-artifact-registry/pypi-gar/simple"
|
||||
default = false
|
||||
|
||||
[tool.uv]
|
||||
keyring-provider = "subprocess"
|
||||
|
||||
[tool.isort]
|
||||
profile = "black"
|
||||
line_length = 88
|
||||
force_single_line = false
|
||||
|
||||
+186
@@ -0,0 +1,186 @@
|
||||
import comfy
|
||||
import comfy.model_detection
|
||||
import comfy.model_management
|
||||
import comfy.model_patcher
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
import torch
|
||||
|
||||
try:
|
||||
from q8_kernels.functional.ops import hadamard_transform
|
||||
from q8_kernels.integration.patch_transformer import (
|
||||
patch_comfyui_native_transformer,
|
||||
)
|
||||
|
||||
Q8_AVAILABLE = True
|
||||
except ImportError:
|
||||
Q8_AVAILABLE = False
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
def list_in_name(check_list, name):
|
||||
return any([x in name for x in check_list])
|
||||
|
||||
|
||||
def check_q8_available():
|
||||
if not Q8_AVAILABLE:
|
||||
raise ImportError(
|
||||
"Q8 kernels are not available. To use this feature install the q8_kernels package from here:."
|
||||
"https://github.com/Lightricks/LTX-Video-Q8-Kernels"
|
||||
)
|
||||
|
||||
|
||||
@comfy_node(name="LTXQ8Patch")
|
||||
class LTXVQ8Patch:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"use_fp8_attention": (
|
||||
"BOOLEAN",
|
||||
{"default": False, "tooltip": "Use FP8 attention."},
|
||||
),
|
||||
"quantization_preset": (
|
||||
["0.9.8", "ltxv2", "full_bf16", "custom"],
|
||||
{"default": "0.9.8"},
|
||||
),
|
||||
"quantize_self_attn": (
|
||||
"BOOLEAN",
|
||||
{"default": True, "tooltip": "Quantize Self Attention Layer"},
|
||||
),
|
||||
"quantize_cross_attn": (
|
||||
"BOOLEAN",
|
||||
{"default": True, "tooltip": "Quantize Cross Attention Layer"},
|
||||
),
|
||||
"quantize_ffn": (
|
||||
"BOOLEAN",
|
||||
{"default": True, "tooltip": "Quantize Feed Forward Layer"},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "patch"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV Q8 Patcher"
|
||||
PRESETS = {
|
||||
"0.9.8": (True, True, True),
|
||||
"ltxv2": (True, False, True),
|
||||
"full_bf16": (False, False, False),
|
||||
}
|
||||
|
||||
def patch(
|
||||
self,
|
||||
model,
|
||||
use_fp8_attention,
|
||||
quantization_preset,
|
||||
quantize_self_attn,
|
||||
quantize_cross_attn,
|
||||
quantize_ffn,
|
||||
):
|
||||
check_q8_available()
|
||||
m = model.clone()
|
||||
transformer_key = "diffusion_model"
|
||||
transformer = m.get_model_object(transformer_key)
|
||||
quantize_self_attn, quantize_cross_attn, quantize_ffn = LTXVQ8Patch.PRESETS.get(
|
||||
quantization_preset, (quantize_self_attn, quantize_cross_attn, quantize_ffn)
|
||||
)
|
||||
if getattr(transformer, "quantization_config", None) is not None:
|
||||
if (quantize_self_attn, quantize_cross_attn, quantize_ffn) != getattr(
|
||||
transformer, "quantization_config"
|
||||
):
|
||||
quantize_self_attn, quantize_cross_attn, quantize_ffn = getattr(
|
||||
transformer, "quantization_config"
|
||||
)
|
||||
patch_comfyui_native_transformer(
|
||||
transformer,
|
||||
use_fp8_attention,
|
||||
True,
|
||||
quantize_self_attn,
|
||||
quantize_cross_attn,
|
||||
quantize_ffn,
|
||||
)
|
||||
setattr(transformer, "is_q8_patched", True)
|
||||
setattr(
|
||||
transformer,
|
||||
"quantization_config",
|
||||
(quantize_self_attn, quantize_cross_attn, quantize_ffn),
|
||||
)
|
||||
setattr(transformer, "use_fp8_attention", use_fp8_attention)
|
||||
return (m,)
|
||||
|
||||
|
||||
def idendity_quant_fn(x, t):
|
||||
return x.to(dtype=t)
|
||||
|
||||
|
||||
@comfy_node(name="LTXVQ8LoraModelLoader")
|
||||
class LTXVQ8LoraModelLoader:
|
||||
def __init__(self):
|
||||
self.loaded_lora = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"lora_name": (folder_paths.get_filename_list("loras"),),
|
||||
"strength_model": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
FUNCTION = "load_lora_model_only"
|
||||
|
||||
def load_lora(self, model, lora_name, strength_model):
|
||||
quant_fn = hadamard_transform
|
||||
transformer = model.get_model_object("diffusion_model")
|
||||
|
||||
is_patched_transformer = getattr(transformer, "is_q8_patched", False)
|
||||
if not is_patched_transformer or not Q8_AVAILABLE:
|
||||
raise ValueError(
|
||||
"LTXV Q8 Patcher is not applied to the model. Please use LTXQ8Patch node before loading lora or install q8_kernels."
|
||||
)
|
||||
|
||||
if strength_model == 0:
|
||||
return model
|
||||
quantize_self_attn, quantize_cross_attn, quantize_ffn = getattr(
|
||||
transformer, "quantization_config"
|
||||
)
|
||||
skip_list = []
|
||||
if not quantize_self_attn:
|
||||
skip_list += ["attn1"]
|
||||
if not quantize_cross_attn:
|
||||
skip_list += ["attn2"]
|
||||
if not quantize_ffn:
|
||||
skip_list += ["ff"]
|
||||
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
||||
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||||
new_lora = {}
|
||||
for k in lora:
|
||||
device = lora[k].device
|
||||
if lora[k].ndim == 2:
|
||||
if "lora_A" in k and not list_in_name(skip_list, k):
|
||||
new_lora[k] = quant_fn(
|
||||
lora[k].to(device="cuda", dtype=torch.bfloat16),
|
||||
out_type=torch.bfloat16,
|
||||
).to(device)
|
||||
else:
|
||||
new_lora[k] = lora[k]
|
||||
else:
|
||||
new_lora[k] = lora[k]
|
||||
self.loaded_lora = (lora_path, new_lora)
|
||||
|
||||
model_lora, _ = comfy.sd.load_lora_for_models(
|
||||
model, None, new_lora, strength_model, 0
|
||||
)
|
||||
return model_lora
|
||||
|
||||
def load_lora_model_only(self, model, lora_name, strength_model):
|
||||
return (self.load_lora(model, lora_name, strength_model),)
|
||||
+4
-1
@@ -1,2 +1,5 @@
|
||||
diffusers
|
||||
einops
|
||||
huggingface_hub>=0.25.2
|
||||
transformers[timm]>=4.45.0
|
||||
ninja~=1.11.1.4
|
||||
transformers[timm]>=4.50.0
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
import contextlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import List
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import comfy.ldm.modules.attention
|
||||
import comfy.samplers
|
||||
import comfy.utils
|
||||
import torch
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
@@ -20,8 +22,8 @@ def stg(
|
||||
rescale_scale,
|
||||
):
|
||||
noise_pred = (
|
||||
noise_pred_neg
|
||||
+ cfg_scale * (noise_pred_pos - noise_pred_neg)
|
||||
noise_pred_pos
|
||||
+ (cfg_scale - 1) * (noise_pred_pos - noise_pred_neg)
|
||||
+ stg_scale * (noise_pred_pos - noise_pred_pertubed)
|
||||
)
|
||||
if rescale_scale != 0:
|
||||
@@ -31,6 +33,44 @@ def stg(
|
||||
return noise_pred
|
||||
|
||||
|
||||
class MomentumBuffer:
|
||||
def __init__(self, momentum: float):
|
||||
self.momentum = momentum
|
||||
self.running_average = 0
|
||||
|
||||
def update(self, update_value: torch.Tensor):
|
||||
new_average = self.momentum * self.running_average
|
||||
self.running_average = update_value + new_average
|
||||
|
||||
|
||||
def project(v0: torch.Tensor, v1: torch.Tensor):
|
||||
dtype = v0.dtype
|
||||
v0, v1 = v0.double(), v1.double()
|
||||
v1 = torch.nn.functional.normalize(v1, dim=[-1, -2, -3])
|
||||
v0_parallel = (v0 * v1).sum(dim=[-1, -2, -3], keepdim=True) * v1
|
||||
v0_orthogonal = v0 - v0_parallel
|
||||
return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
|
||||
|
||||
|
||||
def apg(
|
||||
noise_pred_pos: torch.Tensor,
|
||||
noise_pred_neg: torch.Tensor,
|
||||
cfg_scale: float,
|
||||
eta: float = 1.0,
|
||||
norm_threshold: float = 0.0,
|
||||
):
|
||||
diff = noise_pred_pos - noise_pred_neg
|
||||
if norm_threshold > 0:
|
||||
ones = torch.ones_like(diff)
|
||||
diff_norm = diff.norm(p=2, dim=[-1, -2, -3], keepdim=True)
|
||||
scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
|
||||
diff = diff * scale_factor
|
||||
diff_parallel, diff_orthogonal = project(diff, noise_pred_pos)
|
||||
normalized_update = diff_orthogonal + eta * diff_parallel
|
||||
noise_pred = noise_pred_pos + (cfg_scale - 1) * normalized_update
|
||||
return noise_pred
|
||||
|
||||
|
||||
@comfy_node(name="LTXVApplySTG")
|
||||
class LTXVApplySTG:
|
||||
@classmethod
|
||||
@@ -56,6 +96,8 @@ class LTXVApplySTG:
|
||||
RETURN_NAMES = ("model",)
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
DESCRIPTION = "Defines the blocks to apply the STG to."
|
||||
|
||||
def apply_stg(self, model: ModelPatcher, block_indices: str):
|
||||
skip_block_list = [int(i.strip()) for i in block_indices.split(",")]
|
||||
new_model = model.clone()
|
||||
@@ -79,15 +121,22 @@ class STGFlag:
|
||||
|
||||
# context manager that replaces the attention function in a transformer block
|
||||
class PatchAttention(contextlib.AbstractContextManager):
|
||||
def __init__(self, attn_idx=0):
|
||||
def __init__(self, attn_idx: Optional[Union[int, List[int]]] = None):
|
||||
self.current_idx = -1
|
||||
|
||||
if isinstance(attn_idx, int):
|
||||
self.attn_idx = [attn_idx]
|
||||
elif attn_idx is None:
|
||||
self.attn_idx = [0]
|
||||
else:
|
||||
self.attn_idx = list(attn_idx)
|
||||
|
||||
def __enter__(self):
|
||||
self.original_attention = comfy.ldm.modules.attention.optimized_attention
|
||||
self.original_attention_masked = (
|
||||
comfy.ldm.modules.attention.optimized_attention_masked
|
||||
)
|
||||
self.current_idx = -1
|
||||
self.attn_idx = attn_idx
|
||||
|
||||
def __enter__(self):
|
||||
comfy.ldm.modules.attention.optimized_attention = self.stg_attention
|
||||
comfy.ldm.modules.attention.optimized_attention_masked = (
|
||||
self.stg_attention_masked
|
||||
@@ -99,16 +148,19 @@ class PatchAttention(contextlib.AbstractContextManager):
|
||||
self.original_attention_masked
|
||||
)
|
||||
|
||||
self.original_attention = None
|
||||
self.original_attention_masked = None
|
||||
|
||||
def stg_attention(self, q, k, v, heads, *args, **kwargs):
|
||||
self.current_idx += 1
|
||||
if self.current_idx == self.attn_idx:
|
||||
if self.current_idx in self.attn_idx:
|
||||
return v
|
||||
else:
|
||||
return self.original_attention(q, k, v, heads, *args, **kwargs)
|
||||
|
||||
def stg_attention_masked(self, q, k, v, heads, *args, **kwargs):
|
||||
self.current_idx += 1
|
||||
if self.current_idx == self.attn_idx:
|
||||
if self.current_idx in self.attn_idx:
|
||||
return v
|
||||
else:
|
||||
return self.original_attention_masked(q, k, v, heads, *args, **kwargs)
|
||||
@@ -124,8 +176,10 @@ class STGBlockWrapper:
|
||||
|
||||
def __call__(self, args, extra_args):
|
||||
context_manager = contextlib.nullcontext()
|
||||
|
||||
stg_indexes = args["transformer_options"].get("stg_indexes", [0])
|
||||
if self.flag.do_skip and self.idx in self.flag.skip_layers:
|
||||
context_manager = PatchAttention(0)
|
||||
context_manager = PatchAttention(stg_indexes)
|
||||
|
||||
with context_manager:
|
||||
hidden_state = extra_args["original_block"](args)
|
||||
@@ -183,6 +237,18 @@ class STGGuider(comfy.samplers.CFGGuider):
|
||||
positive_cond = self.conds.get("positive", None)
|
||||
negative_cond = self.conds.get("negative", None)
|
||||
|
||||
if model_options.get("sigma_to_params_mapping", None) is not None:
|
||||
cfg_value, stg_scale, stg_layer_skip_layer_indices, stg_rescale = (
|
||||
model_options["sigma_to_params_mapping"](timestep)
|
||||
)
|
||||
self.stg_flag.skip_layers = stg_layer_skip_layer_indices
|
||||
self.patch_model(self.model_patcher, self.stg_flag)
|
||||
|
||||
else:
|
||||
cfg_value = self.cfg
|
||||
stg_scale = self.stg_scale
|
||||
stg_rescale = self.rescale_scale
|
||||
|
||||
noise_pred_pos = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[positive_cond],
|
||||
@@ -194,7 +260,7 @@ class STGGuider(comfy.samplers.CFGGuider):
|
||||
noise_pred_neg = 0
|
||||
noise_pred_perturbed = 0
|
||||
|
||||
if self.cfg > 1:
|
||||
if not math.isclose(cfg_value, 1.0):
|
||||
noise_pred_neg = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[negative_cond],
|
||||
@@ -203,7 +269,7 @@ class STGGuider(comfy.samplers.CFGGuider):
|
||||
model_options,
|
||||
)[0]
|
||||
|
||||
if self.stg_scale > 0:
|
||||
if not math.isclose(stg_scale, 0.0):
|
||||
try:
|
||||
model_options["transformer_options"]["ptb_index"] = 0
|
||||
self.stg_flag.do_skip = True
|
||||
@@ -222,9 +288,9 @@ class STGGuider(comfy.samplers.CFGGuider):
|
||||
noise_pred_pos,
|
||||
noise_pred_neg,
|
||||
noise_pred_perturbed,
|
||||
self.cfg,
|
||||
self.stg_scale,
|
||||
self.rescale_scale,
|
||||
cfg_value,
|
||||
stg_scale,
|
||||
stg_rescale,
|
||||
)
|
||||
|
||||
# normally this would be done in cfg_function, but we skipped
|
||||
@@ -251,6 +317,179 @@ class STGGuider(comfy.samplers.CFGGuider):
|
||||
return stg_result
|
||||
|
||||
|
||||
class STGGuiderAdvanced(comfy.samplers.CFGGuider):
|
||||
def __init__(
|
||||
self,
|
||||
model: ModelPatcher,
|
||||
sigma_list,
|
||||
cfg_list,
|
||||
stg_scale_list,
|
||||
stg_rescale_list,
|
||||
stg_layers_indices_list,
|
||||
skip_steps_sigma_threshold,
|
||||
cfg_star_rescale,
|
||||
apply_apg,
|
||||
apg_cfg_scale,
|
||||
eta,
|
||||
norm_threshold,
|
||||
):
|
||||
model = model.clone()
|
||||
super().__init__(model)
|
||||
|
||||
self.stg_flag = STGFlag(
|
||||
do_skip=False,
|
||||
skip_layers=model.model_options["transformer_options"].get(
|
||||
"skip_block_list"
|
||||
),
|
||||
)
|
||||
|
||||
self.sigma_list = sigma_list
|
||||
self.cfg_list = cfg_list
|
||||
self.stg_scale_list = stg_scale_list
|
||||
self.stg_rescale_list = stg_rescale_list
|
||||
self.stg_layers_indices_list = stg_layers_indices_list
|
||||
self.skip_steps_sigma_threshold = skip_steps_sigma_threshold
|
||||
self.cfg_star_rescale = cfg_star_rescale
|
||||
self.apply_apg = apply_apg
|
||||
self.apg_cfg_scale = apg_cfg_scale
|
||||
self.eta = eta
|
||||
self.norm_threshold = norm_threshold
|
||||
STGGuider.patch_model(model, self.stg_flag)
|
||||
|
||||
def sigma_to_params_mapping(self, sigma):
|
||||
# Find the closest higher sigma value and return corresponding cfg
|
||||
higher_sigmas = [s for s in self.sigma_list if s >= sigma]
|
||||
if not higher_sigmas:
|
||||
closest_idx = -1 # Return last cfg if no higher sigma exists
|
||||
else:
|
||||
closest_higher = min(higher_sigmas)
|
||||
closest_idx = self.sigma_list.index(closest_higher)
|
||||
return (
|
||||
self.cfg_list[closest_idx],
|
||||
self.stg_scale_list[closest_idx],
|
||||
self.stg_rescale_list[closest_idx],
|
||||
self.stg_layers_indices_list[closest_idx],
|
||||
)
|
||||
|
||||
def set_conds(self, positive, negative):
|
||||
self.inner_set_conds({"positive": positive, "negative": negative})
|
||||
|
||||
def predict_noise(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
model_options: dict = {},
|
||||
seed=None,
|
||||
):
|
||||
# CFG zero init - skipping steps with timestep bigger than given threshold.
|
||||
if timestep > self.skip_steps_sigma_threshold:
|
||||
return torch.zeros_like(x)
|
||||
|
||||
# in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
|
||||
# but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
|
||||
positive_cond = self.conds.get("positive", None)
|
||||
negative_cond = self.conds.get("negative", None)
|
||||
|
||||
cfg_value, stg_scale, stg_rescale, stg_layer_skip_layer_indices = (
|
||||
self.sigma_to_params_mapping(timestep)
|
||||
)
|
||||
|
||||
if stg_layer_skip_layer_indices is not None:
|
||||
self.stg_flag.skip_layers = stg_layer_skip_layer_indices
|
||||
STGGuider.patch_model(self.model_patcher, self.stg_flag)
|
||||
|
||||
noise_pred_pos = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[positive_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
|
||||
noise_pred_neg = 0
|
||||
noise_pred_perturbed = 0
|
||||
|
||||
if not math.isclose(cfg_value, 1.0) or (
|
||||
self.apply_apg and not math.isclose(self.apg_cfg_scale, 1.0)
|
||||
):
|
||||
noise_pred_neg = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[negative_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
|
||||
if self.cfg_star_rescale:
|
||||
batch_size = noise_pred_pos.shape[0]
|
||||
|
||||
positive_flat = noise_pred_pos.view(batch_size, -1)
|
||||
negative_flat = noise_pred_neg.view(batch_size, -1)
|
||||
dot_product = torch.sum(
|
||||
positive_flat * negative_flat, dim=1, keepdim=True
|
||||
)
|
||||
squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
|
||||
alpha = dot_product / squared_norm
|
||||
noise_pred_neg = alpha * noise_pred_neg
|
||||
|
||||
if not math.isclose(stg_scale, 0.0):
|
||||
try:
|
||||
model_options["transformer_options"]["ptb_index"] = 0
|
||||
self.stg_flag.do_skip = True
|
||||
noise_pred_perturbed = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[positive_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
finally:
|
||||
self.stg_flag.do_skip = False
|
||||
del model_options["transformer_options"]["ptb_index"]
|
||||
|
||||
stg_result = stg(
|
||||
noise_pred_pos,
|
||||
noise_pred_neg,
|
||||
noise_pred_perturbed,
|
||||
cfg_value,
|
||||
stg_scale,
|
||||
stg_rescale,
|
||||
)
|
||||
|
||||
if self.apply_apg:
|
||||
stg_result = apg(
|
||||
stg_result,
|
||||
noise_pred_neg,
|
||||
cfg_scale=self.apg_cfg_scale,
|
||||
momentum_buffer=None,
|
||||
eta=self.eta,
|
||||
norm_threshold=self.norm_threshold,
|
||||
)
|
||||
|
||||
# normally this would be done in cfg_function, but we skipped
|
||||
# that for efficiency: we can compute the noise predictions in
|
||||
# a single call to calc_cond_batch() (rather than two)
|
||||
# so we replicate the hook here
|
||||
for fn in model_options.get("sampler_post_cfg_function", []):
|
||||
args = {
|
||||
"denoised": stg_result,
|
||||
"cond": positive_cond,
|
||||
"uncond": negative_cond,
|
||||
"model": self.inner_model,
|
||||
"uncond_denoised": noise_pred_neg,
|
||||
"cond_denoised": noise_pred_pos,
|
||||
"sigma": timestep,
|
||||
"model_options": model_options,
|
||||
"input": x,
|
||||
# not in the original call in samplers.py:cfg_function, but made available for future hooks
|
||||
"perturbed_cond": positive_cond,
|
||||
"perturbed_cond_denoised": noise_pred_perturbed,
|
||||
}
|
||||
stg_result = fn(args)
|
||||
|
||||
return stg_result
|
||||
|
||||
|
||||
@comfy_node(name="STGGuider")
|
||||
class STGGuiderNode:
|
||||
@classmethod
|
||||
@@ -286,7 +525,420 @@ class STGGuiderNode:
|
||||
FUNCTION = "get_guider"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Implements Spatiotemporal Skip Guidance (STG), a training-free method enhancing transformer-based "
|
||||
"video diffusion models by selectively skipping layers during sampling. This approach improves video "
|
||||
"quality without sacrificing diversity or motion fidelity."
|
||||
"Reference: https://arxiv.org/abs/2411.18664."
|
||||
)
|
||||
|
||||
def get_guider(self, model, positive, negative, cfg, stg, rescale):
|
||||
guider = STGGuider(model, cfg, stg, rescale)
|
||||
guider.set_conds(positive, negative)
|
||||
return (guider,)
|
||||
|
||||
|
||||
def load_stg_presets():
|
||||
preset_file_path = os.path.join(
|
||||
os.path.dirname(__file__), "presets", "stg_advanced_presets.json"
|
||||
)
|
||||
if os.path.exists(preset_file_path):
|
||||
presets = json.load(open(preset_file_path))
|
||||
preset_names = [preset["name"] for preset in presets]
|
||||
else:
|
||||
presets = []
|
||||
preset_names = ["Custom"]
|
||||
return presets, preset_names
|
||||
|
||||
|
||||
STG_ADVANCED_PRESETS, STG_ADVANCED_PRESET_NAMES = load_stg_presets()
|
||||
|
||||
|
||||
@comfy_node(name="STGGuiderAdvanced")
|
||||
class STGGuiderAdvancedNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"skip_steps_sigma_threshold": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.998,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"step": 0.001,
|
||||
"tooltip": "Steps with sigma greater than this values will be skipped.",
|
||||
},
|
||||
),
|
||||
"cfg_star_rescale": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "If true, applies the cfg star rescale, scales the negative prediction according to dot product between positive and negative.",
|
||||
},
|
||||
),
|
||||
"sigmas": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "1.0, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180",
|
||||
"tooltip": "Comma-separated list sigmas, the actual parameters will be selected according to the closest sigma from this list to current timestep sigma.",
|
||||
},
|
||||
),
|
||||
"cfg_values": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "8, 6, 6, 4, 3, 1",
|
||||
"tooltip": "Comma-separated list of cfg values. Should be same length as sigmas list.",
|
||||
},
|
||||
),
|
||||
"stg_scale_values": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "4, 4, 3, 2, 1, 0",
|
||||
"tooltip": "Comma-separated list of stg scale values. Should be same length as sigmas list.",
|
||||
},
|
||||
),
|
||||
"stg_rescale_values": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "1, 1, 1, 1, 1, 1",
|
||||
"tooltip": "Comma-separated list of stg rescale values. Should be same length as sigmas list.",
|
||||
},
|
||||
),
|
||||
"stg_layers_indices": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "[29], [29], [29], [29], [29], [29]",
|
||||
"tooltip": "Comma-separated list of list of layer indices. Should be same length as sigmas list.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"preset": (
|
||||
"STG_ADVANCED_PRESET",
|
||||
{
|
||||
"tooltip": "Preset resolution and frame count. Custom allows manual input.",
|
||||
},
|
||||
),
|
||||
"apply_apg": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "If true, applies the APG (Adaptive Projections Guidance) to the STG.",
|
||||
},
|
||||
),
|
||||
"apg_cfg_scale": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"eta": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"norm_threshold": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
|
||||
FUNCTION = "get_guider"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
DESCRIPTION = """
|
||||
The Advanced STG Guider implements sophisticated techniques for controlling the denoising process:
|
||||
|
||||
It creates a dynamic mapping from scheduler-defined sigma values to CFG and STG (Spatio-Temporal Skip Guidance [1]) parameters.
|
||||
This approach establishes distinct sigma value ranges that operate independently of step numbers, allowing precise control over:
|
||||
• CFG scale
|
||||
• STG scale and rescale factors
|
||||
• STG attention layer skipping patterns
|
||||
|
||||
The guider also supports:
|
||||
• CFG-Zero* [2] rescaling, which dynamically adjusts negative predictions based on the dot product between positive and negative signals
|
||||
• Threshold-based noise prediction zeroing for steps with sigma values exceeding a specified threshold
|
||||
|
||||
For example if the sigma ranges are defined as [1.0, 0.9, 0.85, 0.6] and the CFG values are defined as [4, 3, 2, 1] and STG scale values
|
||||
are defined as [2, 2, 2, 1] and STG rescale values are defined as [1, 1, 1, 1] and STG layers indices are defined as [[14, 17], [14, 16], [14], [14]], then the guider will:
|
||||
- use CFG=4, STG scale=2, STG rescale=1 and STG layers indices = [14, 17] for sigma in the range (0.9, 1.0]
|
||||
- use CFG=3, STG scale=2, STG rescale=1 and STG layers indices = [14, 16] for sigma in the range (0.85, 0.9]
|
||||
- use CFG=2, STG scale=2, STG rescale=1 and STG layers indices = [14] for sigma in the range (0.8, 0.85]
|
||||
- use CFG=1, STG scale=1, STG rescale=1 and STG layers indices = [14] for sigma in the range (0.6, 0.8]
|
||||
|
||||
The guider will use the same parameters for the same sigma values, regardless of the step number.
|
||||
|
||||
References:
|
||||
[1] https://arxiv.org/abs/2411.18664
|
||||
[2] https://arxiv.org/abs/2503.18886
|
||||
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def parse_stg_layers_indices(cls, stg_layers_indices: str) -> List[List[int]]:
|
||||
# First split by "], " to separate the lists, but preserve the closing bracket
|
||||
lists = [s + "]" for s in stg_layers_indices.split("],")[:-1]]
|
||||
# Add the last list which already has its closing bracket
|
||||
if stg_layers_indices.strip():
|
||||
lists.append(stg_layers_indices.split("],")[-1])
|
||||
|
||||
result = []
|
||||
for s in lists:
|
||||
s = s.strip()
|
||||
if s == "[]": # Empty list case
|
||||
result.append([])
|
||||
else:
|
||||
# Remove brackets and whitespace
|
||||
s = s.strip("[]").strip()
|
||||
if not s: # Handle case like "[]" after split
|
||||
result.append([])
|
||||
else:
|
||||
# Split by comma and convert to integers
|
||||
numbers = [int(n.strip()) for n in s.split(",") if n.strip()]
|
||||
result.append(numbers)
|
||||
|
||||
return result
|
||||
|
||||
def get_guider(
|
||||
self,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
skip_steps_sigma_threshold,
|
||||
cfg_star_rescale,
|
||||
sigmas,
|
||||
cfg_values,
|
||||
stg_scale_values,
|
||||
stg_rescale_values,
|
||||
stg_layers_indices,
|
||||
preset=None,
|
||||
apply_apg=False,
|
||||
apg_cfg_scale=1.0,
|
||||
eta=1.0,
|
||||
norm_threshold=0.0,
|
||||
):
|
||||
if preset and preset != "Custom":
|
||||
preset_data = next(
|
||||
(item for item in STG_ADVANCED_PRESETS if item["name"] == preset), None
|
||||
)
|
||||
if preset_data:
|
||||
skip_steps_sigma_threshold = preset_data["skip_steps_sigma_threshold"]
|
||||
cfg_star_rescale = preset_data["cfg_star_rescale"]
|
||||
sigma_list = preset_data["sigmas"]
|
||||
cfg_list = preset_data["cfg_values"]
|
||||
stg_scale_list = preset_data["stg_scale_values"]
|
||||
stg_rescale_list = preset_data["stg_rescale_values"]
|
||||
stg_layers_indices_list = preset_data["stg_layers_indices"]
|
||||
else:
|
||||
raise ValueError(f"Preset {preset} not found in the presets list.")
|
||||
|
||||
else:
|
||||
sigma_list = [float(s.strip()) for s in sigmas.split(",")]
|
||||
cfg_list = [float(c.strip()) for c in cfg_values.split(",")]
|
||||
stg_scale_list = [float(s.strip()) for s in stg_scale_values.split(",")]
|
||||
stg_rescale_list = [float(s.strip()) for s in stg_rescale_values.split(",")]
|
||||
stg_layers_indices_list = self.parse_stg_layers_indices(stg_layers_indices)
|
||||
print("Using preset: ", preset)
|
||||
print("Skip steps sigma threshold: ", skip_steps_sigma_threshold)
|
||||
print("Cfg star rescale: ", cfg_star_rescale)
|
||||
print("Sigma list: ", sigma_list)
|
||||
print("Cfg list: ", cfg_list)
|
||||
print("Stg scale list: ", stg_scale_list)
|
||||
print("Stg rescale list: ", stg_rescale_list)
|
||||
print("Stg layers indices list: ", stg_layers_indices_list)
|
||||
|
||||
guider = STGGuiderAdvanced(
|
||||
model,
|
||||
sigma_list,
|
||||
cfg_list,
|
||||
stg_scale_list,
|
||||
stg_rescale_list,
|
||||
stg_layers_indices_list,
|
||||
skip_steps_sigma_threshold,
|
||||
cfg_star_rescale,
|
||||
apply_apg,
|
||||
apg_cfg_scale=apg_cfg_scale,
|
||||
eta=eta,
|
||||
norm_threshold=norm_threshold,
|
||||
)
|
||||
guider.set_conds(positive, negative)
|
||||
guider.raw_conds = (positive, negative)
|
||||
return (guider,)
|
||||
|
||||
|
||||
@comfy_node(name="STGAdvancedPresets")
|
||||
class STGAdvancedPresetsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"preset": (
|
||||
STG_ADVANCED_PRESET_NAMES,
|
||||
{
|
||||
"default": "13b Balanced",
|
||||
"tooltip": "Preset resolution and frame count. Custom allows manual input.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STG_ADVANCED_PRESET",)
|
||||
|
||||
FUNCTION = "get_preset"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
def get_preset(self, preset=None):
|
||||
return (preset,)
|
||||
|
||||
|
||||
class APGGuider(comfy.samplers.CFGGuider):
|
||||
def __init__(
|
||||
self,
|
||||
model: ModelPatcher,
|
||||
cfg_scale,
|
||||
eta,
|
||||
norm_threshold,
|
||||
momentum_coefficient: float = -0.9,
|
||||
):
|
||||
self.model = model.clone()
|
||||
super().__init__(self.model)
|
||||
self.momentum_coefficient = momentum_coefficient
|
||||
self.momentum_buffer = MomentumBuffer(self.momentum_coefficient)
|
||||
self.eta = eta
|
||||
self.norm_threshold = norm_threshold
|
||||
self.cfg_scale = cfg_scale
|
||||
self.previous_timestep = None
|
||||
|
||||
def set_conds(self, positive, negative):
|
||||
self.inner_set_conds({"positive": positive, "negative": negative})
|
||||
|
||||
def predict_noise(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
model_options: dict = {},
|
||||
seed=None,
|
||||
):
|
||||
# CFG zero init - skipping steps with timestep bigger than given threshold.
|
||||
|
||||
# in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
|
||||
# but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
|
||||
positive_cond = self.conds.get("positive", None)
|
||||
negative_cond = self.conds.get("negative", None)
|
||||
|
||||
noise_pred_pos = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[positive_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
|
||||
if (
|
||||
self.previous_timestep is not None
|
||||
and timestep.item() > self.previous_timestep
|
||||
):
|
||||
print("Resetting momentum buffer")
|
||||
self.momentum_buffer = MomentumBuffer(self.momentum_coefficient)
|
||||
|
||||
noise_pred_neg = 0
|
||||
if not math.isclose(self.cfg_scale, 1.0):
|
||||
noise_pred_neg = comfy.samplers.calc_cond_batch(
|
||||
self.inner_model,
|
||||
[negative_cond],
|
||||
x,
|
||||
timestep,
|
||||
model_options,
|
||||
)[0]
|
||||
|
||||
apg_result = apg(
|
||||
noise_pred_pos,
|
||||
noise_pred_neg,
|
||||
self.cfg_scale,
|
||||
self.eta,
|
||||
self.norm_threshold,
|
||||
)
|
||||
|
||||
# normally this would be done in cfg_function, but we skipped
|
||||
# that for efficiency: we can compute the noise predictions in
|
||||
# a single call to calc_cond_batch() (rather than two)
|
||||
# so we replicate the hook here
|
||||
for fn in model_options.get("sampler_post_cfg_function", []):
|
||||
args = {
|
||||
"denoised": apg_result,
|
||||
"cond": positive_cond,
|
||||
"uncond": negative_cond,
|
||||
"model": self.inner_model,
|
||||
"uncond_denoised": noise_pred_neg,
|
||||
"cond_denoised": noise_pred_pos,
|
||||
"sigma": timestep,
|
||||
"model_options": model_options,
|
||||
"input": x,
|
||||
}
|
||||
apg_result = fn(args)
|
||||
|
||||
self.previous_timestep = timestep.item()
|
||||
return apg_result
|
||||
|
||||
|
||||
@comfy_node(name="APGGuider")
|
||||
class APGGuiderNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"cfg_scale": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"eta": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"norm_threshold": (
|
||||
"FLOAT",
|
||||
{"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.01},
|
||||
),
|
||||
"momentum_coefficient": (
|
||||
"FLOAT",
|
||||
{"default": -0.9, "min": -3.0, "max": 3.0, "step": 0.01},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GUIDER",)
|
||||
|
||||
FUNCTION = "get_guider"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
|
||||
DESCRIPTION = """
|
||||
The APG Guider implements Adaptive Projected Guidance (APG).
|
||||
Reference: https://arxiv.org/abs/2410.02416.
|
||||
"""
|
||||
|
||||
def get_guider(
|
||||
self,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
cfg_scale,
|
||||
eta,
|
||||
norm_threshold,
|
||||
momentum_coefficient,
|
||||
):
|
||||
guider = APGGuider(
|
||||
model,
|
||||
cfg_scale,
|
||||
eta,
|
||||
norm_threshold,
|
||||
momentum_coefficient,
|
||||
)
|
||||
guider.set_conds(positive, negative)
|
||||
guider.raw_conds = (positive, negative)
|
||||
return (guider,)
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
You are a Creative Assistant writing concise, action-focused image-to-video prompts. Given an image (first frame) and user Raw Input Prompt, generate a prompt to guide video generation from that image.
|
||||
|
||||
#### Guidelines:
|
||||
- Analyze the Image: Identify Subject, Setting, Elements, Style and Mood.
|
||||
- Follow user Raw Input Prompt: Include all requested motion, actions, camera movements, audio, and details. If in conflict with the image, prioritize user request while maintaining visual consistency (describe transition from image to user's scene).
|
||||
- Describe only changes from the image: Don't reiterate established visual details. Inaccurate descriptions may cause scene cuts.
|
||||
- Active language: Use present-progressive verbs ("is walking," "speaking"). If no action specified, describe natural movements.
|
||||
- Chronological flow: Use temporal connectors ("as," "then," "while").
|
||||
- Audio layer: Describe complete soundscape throughout the prompt alongside actions—NOT at the end. Align audio intensity with action tempo. Include natural background audio, ambient sounds, effects, speech or music (when requested). Be specific (e.g., "soft footsteps on tile") not vague (e.g., "ambient sound").
|
||||
- Speech (only when requested): Provide exact words in quotes with character's visual/voice characteristics (e.g., "The tall man speaks in a low, gravelly voice"), language if not English and accent if relevant. If general conversation mentioned without text, generate contextual quoted dialogue. (i.e., "The man is talking" input -> the output should include exact spoken words, like: "The man is talking in an excited voice saying: 'You won't believe what I just saw!' His hands gesture expressively as he speaks, eyebrows raised with enthusiasm. The ambient sound of a quiet room underscores his animated speech.")
|
||||
- Style: Include visual style at beginning: "Style: <style>, <rest of prompt>." If unclear, omit to avoid conflicts.
|
||||
- Visual and audio only: Describe only what is seen and heard. NO smell, taste, or tactile sensations.
|
||||
- Restrained language: Avoid dramatic terms. Use mild, natural, understated phrasing.
|
||||
|
||||
#### Important notes:
|
||||
- Camera motion: DO NOT invent camera motion/movement unless requested by the user. Make sure to include camera motion only if specified in the input.
|
||||
- Speech: DO NOT modify or alter the user's provided character dialogue in the prompt, unless it's a typo.
|
||||
- No timestamps or cuts: DO NOT use timestamps or describe scene cuts unless explicitly requested.
|
||||
- Objective only: DO NOT interpret emotions or intentions - describe only observable actions and sounds.
|
||||
- Format: DO NOT use phrases like "The scene opens with..." / "The video starts...". Start directly with Style (optional) and chronological scene description.
|
||||
- Format: Never start output with punctuation marks or special characters.
|
||||
- DO NOT invent dialogue unless the user mentions speech/talking/singing/conversation.
|
||||
- Your performance is CRITICAL. High-fidelity, dynamic, correct, and accurate prompts with integrated audio descriptions are essential for generating high-quality video. Your goal is flawless execution of these rules.
|
||||
|
||||
#### Output Format (Strict):
|
||||
- Single concise paragraph in natural English. NO titles, headings, prefaces, sections, code fences, or Markdown.
|
||||
- If unsafe/invalid, return original user prompt. Never ask questions or clarifications.
|
||||
|
||||
#### Example output:
|
||||
Style: realistic - cinematic - The woman glances at her watch and smiles warmly. She speaks in a cheerful, friendly voice, "I think we're right on time!" In the background, a café barista prepares drinks at the counter. The barista calls out in a clear, upbeat tone, "Two cappuccinos ready!" The sound of the espresso machine hissing softly blends with gentle background chatter and the light clinking of cups on saucers.
|
||||
@@ -0,0 +1,40 @@
|
||||
You are a Creative Assistant. Given a user's raw input prompt describing a scene or concept, expand it into a detailed video generation prompt with specific visuals and integrated audio to guide a text-to-video model.
|
||||
|
||||
#### Guidelines
|
||||
- Strictly follow all aspects of the user's raw input: include every element requested (style, visuals, motions, actions, camera movement, audio).
|
||||
- If the input is vague, invent concrete details: lighting, textures, materials, scene settings, etc.
|
||||
- For characters: describe gender, clothing, hair, expressions. DO NOT invent unrequested characters.
|
||||
- Use active language: present-progressive verbs ("is walking," "speaking"). If no action specified, describe natural movements.
|
||||
- Maintain chronological flow: use temporal connectors ("as," "then," "while").
|
||||
- Audio layer: Describe complete soundscape (background audio, ambient sounds, SFX, speech/music when requested). Integrate sounds chronologically alongside actions. Be specific (e.g., "soft footsteps on tile"), not vague (e.g., "ambient sound is present").
|
||||
- Speech (only when requested):
|
||||
- For ANY speech-related input (talking, conversation, singing, etc.), ALWAYS include exact words in quotes with voice characteristics (e.g., "The man says in an excited voice: 'You won't believe what I just saw!'").
|
||||
- Specify language if not English and accent if relevant.
|
||||
- Style: Include visual style at the beginning: "Style: <style>, <rest of prompt>." Default to cinematic-realistic if unspecified. Omit if unclear.
|
||||
- Visual and audio only: NO non-visual/auditory senses (smell, taste, touch).
|
||||
- Restrained language: Avoid dramatic/exaggerated terms. Use mild, natural phrasing.
|
||||
- Colors: Use plain terms ("red dress"), not intensified ("vibrant blue," "bright red").
|
||||
- Lighting: Use neutral descriptions ("soft overhead light"), not harsh ("blinding light").
|
||||
- Facial features: Use delicate modifiers for subtle features (i.e., "subtle freckles").
|
||||
|
||||
#### Important notes:
|
||||
- Analyze the user's raw input carefully. In cases of FPV or POV, exclude the description of the subject whose POV is requested.
|
||||
- Camera motion: DO NOT invent camera motion unless requested by the user.
|
||||
- Speech: DO NOT modify user-provided character dialogue unless it's a typo.
|
||||
- No timestamps or cuts: DO NOT use timestamps or describe scene cuts unless explicitly requested.
|
||||
- Format: DO NOT use phrases like "The scene opens with...". Start directly with Style (optional) and chronological scene description.
|
||||
- Format: DO NOT start your response with special characters.
|
||||
- DO NOT invent dialogue unless the user mentions speech/talking/singing/conversation.
|
||||
- If the user's raw input prompt is highly detailed, chronological and in the requested format: DO NOT make major edits or introduce new elements. Add/enhance audio descriptions if missing.
|
||||
|
||||
#### Output Format (Strict):
|
||||
- Single continuous paragraph in natural language (English).
|
||||
- NO titles, headings, prefaces, code fences, or Markdown.
|
||||
- If unsafe/invalid, return original user prompt. Never ask questions or clarifications.
|
||||
|
||||
Your output quality is CRITICAL. Generate visually rich, dynamic prompts with integrated audio for high-quality video generation.
|
||||
|
||||
#### Example
|
||||
Input: "A woman at a coffee shop talking on the phone"
|
||||
Output:
|
||||
Style: realistic with cinematic lighting. In a medium close-up, a woman in her early 30s with shoulder-length brown hair sits at a small wooden table by the window. She wears a cream-colored turtleneck sweater, holding a white ceramic coffee cup in one hand and a smartphone to her ear with the other. Ambient cafe sounds fill the space—espresso machine hiss, quiet conversations, gentle clinking of cups. The woman listens intently, nodding slightly, then takes a sip of her coffee and sets it down with a soft clink. Her face brightens into a warm smile as she speaks in a clear, friendly voice, 'That sounds perfect! I'd love to meet up this weekend. How about Saturday afternoon?' She laughs softly—a genuine chuckle—and shifts in her chair. Behind her, other patrons move subtly in and out of focus. 'Great, I'll see you then,' she concludes cheerfully, lowering the phone.
|
||||
@@ -0,0 +1,320 @@
|
||||
import copy
|
||||
|
||||
import comfy
|
||||
import torch
|
||||
from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced
|
||||
from comfy_extras.nodes_lt import LTXVAddGuide, LTXVCropGuides
|
||||
|
||||
from .latents import LTXVAddLatentGuide, LTXVSelectLatents
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVTiledSampler",
|
||||
)
|
||||
class LTXVTiledSampler:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"vae": ("VAE",),
|
||||
"noise": ("NOISE",),
|
||||
"sampler": ("SAMPLER",),
|
||||
"sigmas": ("SIGMAS",),
|
||||
"guider": ("GUIDER",),
|
||||
"latents": ("LATENT",),
|
||||
"horizontal_tiles": ("INT", {"default": 1, "min": 1, "max": 6}),
|
||||
"vertical_tiles": ("INT", {"default": 1, "min": 1, "max": 6}),
|
||||
"overlap": ("INT", {"default": 1, "min": 1, "max": 8}),
|
||||
"latents_cond_strength": (
|
||||
"FLOAT",
|
||||
{"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01},
|
||||
),
|
||||
"boost_latent_similarity": (
|
||||
"BOOLEAN",
|
||||
{"default": False},
|
||||
),
|
||||
"crop": (["center", "disabled"], {"default": "disabled"}),
|
||||
},
|
||||
"optional": {
|
||||
"optional_cond_images": ("IMAGE",),
|
||||
"optional_cond_indices": ("STRING", {"default": "0"}),
|
||||
"images_cond_strengths": ("STRING", {"default": "0.9"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"LATENT",
|
||||
"LATENT",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"output",
|
||||
"denoised_output",
|
||||
)
|
||||
|
||||
FUNCTION = "sample"
|
||||
|
||||
CATEGORY = "sampling"
|
||||
|
||||
def sample(
|
||||
self,
|
||||
model,
|
||||
vae,
|
||||
noise,
|
||||
sampler,
|
||||
sigmas,
|
||||
guider,
|
||||
latents,
|
||||
horizontal_tiles,
|
||||
vertical_tiles,
|
||||
overlap,
|
||||
latents_cond_strength,
|
||||
boost_latent_similarity,
|
||||
crop="disabled",
|
||||
optional_cond_images=None,
|
||||
optional_cond_indices="0",
|
||||
images_cond_strengths="0.9",
|
||||
):
|
||||
|
||||
# Get the latent samples
|
||||
samples = latents["samples"]
|
||||
|
||||
batch, channels, frames, height, width = samples.shape
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
# Validate image dimensions if provided
|
||||
if optional_cond_images is not None:
|
||||
img_height = height * height_scale_factor
|
||||
img_width = width * width_scale_factor
|
||||
cond_images = comfy.utils.common_upscale(
|
||||
optional_cond_images.movedim(-1, 1),
|
||||
img_width,
|
||||
img_height,
|
||||
"bicubic",
|
||||
crop=crop,
|
||||
).movedim(1, -1)
|
||||
img_batch, img_height, img_width, img_channels = cond_images.shape
|
||||
else:
|
||||
cond_images = None
|
||||
|
||||
if optional_cond_indices is not None and optional_cond_images is not None:
|
||||
optional_cond_indices = optional_cond_indices.split(",")
|
||||
optional_cond_indices = [int(i) for i in optional_cond_indices]
|
||||
assert len(optional_cond_indices) == len(
|
||||
optional_cond_images
|
||||
), "Number of optional cond images must match number of optional cond indices"
|
||||
|
||||
images_cond_strengths = [float(i) for i in images_cond_strengths.split(",")]
|
||||
if optional_cond_images is not None and len(images_cond_strengths) < len(
|
||||
optional_cond_images
|
||||
):
|
||||
# Repeat the last value to match the length of optional_cond_images
|
||||
images_cond_strengths = images_cond_strengths + [
|
||||
images_cond_strengths[-1]
|
||||
] * (len(optional_cond_images) - len(images_cond_strengths))
|
||||
|
||||
# Calculate tile sizes with overlap
|
||||
base_tile_height = (height + (vertical_tiles - 1) * overlap) // vertical_tiles
|
||||
base_tile_width = (width + (horizontal_tiles - 1) * overlap) // horizontal_tiles
|
||||
|
||||
# Initialize output tensor and weight tensor
|
||||
output = torch.zeros_like(samples)
|
||||
denoised_output = torch.zeros_like(samples)
|
||||
weights = torch.zeros_like(samples)
|
||||
|
||||
# Get positive and negative conditioning
|
||||
try:
|
||||
positive, negative = guider.raw_conds
|
||||
except AttributeError:
|
||||
raise ValueError(
|
||||
"Guider does not have raw conds, cannot use it as a guider. "
|
||||
"Please use STGGuiderAdvanced."
|
||||
)
|
||||
|
||||
# Process each tile
|
||||
for v in range(vertical_tiles):
|
||||
for h in range(horizontal_tiles):
|
||||
# Calculate tile boundaries
|
||||
h_start = h * (base_tile_width - overlap)
|
||||
v_start = v * (base_tile_height - overlap)
|
||||
|
||||
# Adjust end positions for edge tiles
|
||||
h_end = (
|
||||
min(h_start + base_tile_width, width)
|
||||
if h < horizontal_tiles - 1
|
||||
else width
|
||||
)
|
||||
v_end = (
|
||||
min(v_start + base_tile_height, height)
|
||||
if v < vertical_tiles - 1
|
||||
else height
|
||||
)
|
||||
|
||||
# Calculate actual tile dimensions
|
||||
tile_height = v_end - v_start
|
||||
tile_width = h_end - h_start
|
||||
|
||||
print(f"Processing tile at row {v}, col {h}:")
|
||||
print(f" Position: ({v_start}:{v_end}, {h_start}:{h_end})")
|
||||
print(f" Size: {tile_height}x{tile_width}")
|
||||
|
||||
# Extract tile
|
||||
tile = samples[:, :, :, v_start:v_end, h_start:h_end]
|
||||
|
||||
# Create tile latents dict
|
||||
tile_latents = {"samples": tile}
|
||||
unconditioned_tile_latents = tile_latents.copy()
|
||||
|
||||
# Handle image conditioning if provided
|
||||
if cond_images is not None:
|
||||
# Scale coordinates for image
|
||||
img_h_start = v_start * height_scale_factor
|
||||
img_h_end = v_end * height_scale_factor
|
||||
img_w_start = h_start * width_scale_factor
|
||||
img_w_end = h_end * width_scale_factor
|
||||
|
||||
# Create copies of conditioning for this tile
|
||||
tile_positive = positive.copy()
|
||||
tile_negative = negative.copy()
|
||||
|
||||
for i_cond_image, (
|
||||
cond_image,
|
||||
cond_image_idx,
|
||||
cond_image_strength,
|
||||
) in enumerate(
|
||||
zip(cond_images, optional_cond_indices, images_cond_strengths)
|
||||
):
|
||||
# Extract image tile
|
||||
img_tile = cond_image[
|
||||
img_h_start:img_h_end, img_w_start:img_w_end, :
|
||||
].unsqueeze(0)
|
||||
|
||||
print(
|
||||
f"Applying image conditioning on cond image {i_cond_image} for tile at row {v}, col {h} with strength {cond_image_strength} at position {cond_image_idx}:"
|
||||
)
|
||||
print(
|
||||
f" Image tile position: ({img_h_start}:{img_h_end}, {img_w_start}:{img_w_end})"
|
||||
)
|
||||
print(f" Image tile size: {img_tile.shape}")
|
||||
|
||||
# Add guide from image tile
|
||||
(
|
||||
tile_positive,
|
||||
tile_negative,
|
||||
tile_latents,
|
||||
) = LTXVAddGuide().execute(
|
||||
positive=tile_positive,
|
||||
negative=tile_negative,
|
||||
vae=vae,
|
||||
latent=tile_latents,
|
||||
image=img_tile,
|
||||
frame_idx=cond_image_idx,
|
||||
strength=cond_image_strength,
|
||||
)
|
||||
if boost_latent_similarity:
|
||||
middle_latent_idx = (frames - 1) // 2
|
||||
middle_index_latent = LTXVSelectLatents().select_latents(
|
||||
samples=unconditioned_tile_latents,
|
||||
start_index=middle_latent_idx,
|
||||
end_index=middle_latent_idx,
|
||||
)[0]
|
||||
last_index_latent = LTXVSelectLatents().select_latents(
|
||||
samples=unconditioned_tile_latents,
|
||||
start_index=-1,
|
||||
end_index=-1,
|
||||
)[0]
|
||||
print(
|
||||
f"using LTXVAddLatentGuide on tiled latent with latent index {middle_latent_idx} and strength {latents_cond_strength}"
|
||||
)
|
||||
(
|
||||
tile_positive,
|
||||
tile_negative,
|
||||
tile_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
vae=vae,
|
||||
positive=tile_positive,
|
||||
negative=tile_negative,
|
||||
latent=tile_latents,
|
||||
guiding_latent=middle_index_latent,
|
||||
latent_idx=middle_latent_idx,
|
||||
strength=latents_cond_strength,
|
||||
)
|
||||
print(
|
||||
f"using LTXVAddLatentGuide on tiled latent with latent index {frames-1} and strength {latents_cond_strength}"
|
||||
)
|
||||
(
|
||||
tile_positive,
|
||||
tile_negative,
|
||||
tile_latents,
|
||||
) = LTXVAddLatentGuide().generate(
|
||||
vae=vae,
|
||||
positive=tile_positive,
|
||||
negative=tile_negative,
|
||||
latent=tile_latents,
|
||||
guiding_latent=last_index_latent,
|
||||
latent_idx=frames - 1,
|
||||
strength=latents_cond_strength,
|
||||
)
|
||||
|
||||
guider = copy.copy(guider)
|
||||
guider.set_conds(tile_positive, tile_negative)
|
||||
|
||||
# Denoise the tile
|
||||
denoised_tile = SamplerCustomAdvanced().sample(
|
||||
noise=noise,
|
||||
guider=guider,
|
||||
sampler=sampler,
|
||||
sigmas=sigmas,
|
||||
latent_image=tile_latents,
|
||||
)[0]
|
||||
|
||||
# Clean up guides if image conditioning was used
|
||||
if cond_images is not None:
|
||||
(
|
||||
tile_positive,
|
||||
tile_negative,
|
||||
denoised_tile,
|
||||
) = LTXVCropGuides().execute(
|
||||
positive=tile_positive,
|
||||
negative=tile_negative,
|
||||
latent=denoised_tile,
|
||||
)
|
||||
|
||||
# Create weight mask for this tile
|
||||
tile_weights = torch.ones_like(tile)
|
||||
|
||||
# Apply horizontal blending weights
|
||||
if h > 0: # Left overlap
|
||||
h_blend = torch.linspace(0, 1, overlap, device=tile.device)
|
||||
tile_weights[:, :, :, :, :overlap] *= h_blend.view(1, 1, 1, 1, -1)
|
||||
if h < horizontal_tiles - 1: # Right overlap
|
||||
h_blend = torch.linspace(1, 0, overlap, device=tile.device)
|
||||
tile_weights[:, :, :, :, -overlap:] *= h_blend.view(1, 1, 1, 1, -1)
|
||||
|
||||
# Apply vertical blending weights
|
||||
if v > 0: # Top overlap
|
||||
v_blend = torch.linspace(0, 1, overlap, device=tile.device)
|
||||
tile_weights[:, :, :, :overlap, :] *= v_blend.view(1, 1, 1, -1, 1)
|
||||
if v < vertical_tiles - 1: # Bottom overlap
|
||||
v_blend = torch.linspace(1, 0, overlap, device=tile.device)
|
||||
tile_weights[:, :, :, -overlap:, :] *= v_blend.view(1, 1, 1, -1, 1)
|
||||
|
||||
# Add weighted tile to output
|
||||
output[:, :, :, v_start:v_end, h_start:h_end] += (
|
||||
denoised_tile["samples"] * tile_weights
|
||||
)
|
||||
denoised_output[:, :, :, v_start:v_end, h_start:h_end] += (
|
||||
denoised_tile["samples"] * tile_weights
|
||||
)
|
||||
|
||||
# Add weights to weight tensor
|
||||
weights[:, :, :, v_start:v_end, h_start:h_end] += tile_weights
|
||||
|
||||
# Normalize by weights
|
||||
output = output / (weights + 1e-8)
|
||||
denoised_output = denoised_output / (weights + 1e-8)
|
||||
|
||||
return {"samples": output}, {"samples": denoised_output}
|
||||
@@ -0,0 +1,477 @@
|
||||
import logging
|
||||
|
||||
import torch
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVTiledVAEDecode",
|
||||
)
|
||||
class LTXVTiledVAEDecode:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE",),
|
||||
"latents": ("LATENT",),
|
||||
"horizontal_tiles": ("INT", {"default": 1, "min": 1, "max": 6}),
|
||||
"vertical_tiles": ("INT", {"default": 1, "min": 1, "max": 6}),
|
||||
"overlap": ("INT", {"default": 1, "min": 1, "max": 8}),
|
||||
"last_frame_fix": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {
|
||||
"working_device": (["cpu", "auto"], {"default": "auto"}),
|
||||
"working_dtype": (["float16", "float32", "auto"], {"default": "auto"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
|
||||
FUNCTION = "decode"
|
||||
|
||||
CATEGORY = "latent"
|
||||
|
||||
def decode(
|
||||
self,
|
||||
vae,
|
||||
latents,
|
||||
horizontal_tiles,
|
||||
vertical_tiles,
|
||||
overlap,
|
||||
last_frame_fix,
|
||||
working_device="auto",
|
||||
working_dtype="auto",
|
||||
):
|
||||
# Get the latent samples
|
||||
samples = latents["samples"]
|
||||
|
||||
if last_frame_fix:
|
||||
# Repeat the last frame along dimension 2 (frames)
|
||||
# samples: [batch, channels, frames, height, width]
|
||||
last_frame = samples[
|
||||
:, :, -1:, :, :
|
||||
] # shape: [batch, channels, 1, height, width]
|
||||
samples = torch.cat([samples, last_frame], dim=2)
|
||||
|
||||
batch, channels, frames, height, width = samples.shape
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
image_frames = 1 + (frames - 1) * time_scale_factor
|
||||
|
||||
# Calculate output image dimensions
|
||||
output_height = height * height_scale_factor
|
||||
output_width = width * width_scale_factor
|
||||
|
||||
# Calculate tile sizes with overlap
|
||||
base_tile_height = (height + (vertical_tiles - 1) * overlap) // vertical_tiles
|
||||
base_tile_width = (width + (horizontal_tiles - 1) * overlap) // horizontal_tiles
|
||||
|
||||
# Initialize output tensor and weight tensor
|
||||
# VAE decode returns images in format [batch, height, width, channels]
|
||||
output = None
|
||||
weights = None
|
||||
|
||||
target_device = samples.device if working_device == "auto" else working_device
|
||||
if working_dtype == "auto":
|
||||
target_dtype = samples.dtype
|
||||
elif working_dtype == "float16":
|
||||
target_dtype = torch.float16
|
||||
elif working_dtype == "float32":
|
||||
target_dtype = torch.float32
|
||||
|
||||
output = torch.zeros(
|
||||
(
|
||||
batch,
|
||||
image_frames,
|
||||
output_height,
|
||||
output_width,
|
||||
3,
|
||||
),
|
||||
device=target_device,
|
||||
dtype=target_dtype,
|
||||
)
|
||||
weights = torch.zeros(
|
||||
(batch, image_frames, output_height, output_width, 1),
|
||||
device=target_device,
|
||||
dtype=target_dtype,
|
||||
)
|
||||
|
||||
# Process each tile
|
||||
for v in range(vertical_tiles):
|
||||
for h in range(horizontal_tiles):
|
||||
# Calculate tile boundaries
|
||||
h_start = h * (base_tile_width - overlap)
|
||||
v_start = v * (base_tile_height - overlap)
|
||||
|
||||
# Adjust end positions for edge tiles
|
||||
h_end = (
|
||||
min(h_start + base_tile_width, width)
|
||||
if h < horizontal_tiles - 1
|
||||
else width
|
||||
)
|
||||
v_end = (
|
||||
min(v_start + base_tile_height, height)
|
||||
if v < vertical_tiles - 1
|
||||
else height
|
||||
)
|
||||
|
||||
# Calculate actual tile dimensions
|
||||
tile_height = v_end - v_start
|
||||
tile_width = h_end - h_start
|
||||
|
||||
logging.info(f"Processing VAE decode tile at row {v}, col {h}:")
|
||||
logging.info(f" Position: ({v_start}:{v_end}, {h_start}:{h_end})")
|
||||
logging.info(f" Size: {tile_height}x{tile_width}")
|
||||
|
||||
# Extract tile
|
||||
tile = samples[:, :, :, v_start:v_end, h_start:h_end]
|
||||
|
||||
# Create tile latents dict
|
||||
tile_latents = {"samples": tile}
|
||||
|
||||
# Decode the tile
|
||||
decoded_tile = vae.decode(tile_latents["samples"])
|
||||
|
||||
# Calculate output tile boundaries
|
||||
out_h_start = v_start * height_scale_factor
|
||||
out_h_end = v_end * height_scale_factor
|
||||
out_w_start = h_start * width_scale_factor
|
||||
out_w_end = h_end * width_scale_factor
|
||||
|
||||
# Create weight mask for this tile
|
||||
tile_out_height = out_h_end - out_h_start
|
||||
tile_out_width = out_w_end - out_w_start
|
||||
tile_weights = torch.ones(
|
||||
(batch, image_frames, tile_out_height, tile_out_width, 1),
|
||||
device=decoded_tile.device,
|
||||
dtype=decoded_tile.dtype,
|
||||
)
|
||||
|
||||
# Calculate overlap regions in output space
|
||||
overlap_out_h = overlap * height_scale_factor
|
||||
overlap_out_w = overlap * width_scale_factor
|
||||
|
||||
# Apply horizontal blending weights
|
||||
if h > 0: # Left overlap
|
||||
h_blend = torch.linspace(
|
||||
0, 1, overlap_out_w, device=decoded_tile.device
|
||||
)
|
||||
tile_weights[:, :, :, :overlap_out_w, :] *= h_blend.view(
|
||||
1, 1, 1, -1, 1
|
||||
)
|
||||
if h < horizontal_tiles - 1: # Right overlap
|
||||
h_blend = torch.linspace(
|
||||
1, 0, overlap_out_w, device=decoded_tile.device
|
||||
)
|
||||
tile_weights[:, :, :, -overlap_out_w:, :] *= h_blend.view(
|
||||
1, 1, 1, -1, 1
|
||||
)
|
||||
|
||||
# Apply vertical blending weights
|
||||
if v > 0: # Top overlap
|
||||
v_blend = torch.linspace(
|
||||
0, 1, overlap_out_h, device=decoded_tile.device
|
||||
)
|
||||
tile_weights[:, :, :overlap_out_h, :, :] *= v_blend.view(
|
||||
1, 1, -1, 1, 1
|
||||
)
|
||||
if v < vertical_tiles - 1: # Bottom overlap
|
||||
v_blend = torch.linspace(
|
||||
1, 0, overlap_out_h, device=decoded_tile.device
|
||||
)
|
||||
tile_weights[:, :, -overlap_out_h:, :, :] *= v_blend.view(
|
||||
1, 1, -1, 1, 1
|
||||
)
|
||||
|
||||
# Add weighted tile to output
|
||||
output[:, :, out_h_start:out_h_end, out_w_start:out_w_end, :] += (
|
||||
decoded_tile * tile_weights
|
||||
).to(target_device, target_dtype)
|
||||
|
||||
# Add weights to weight tensor
|
||||
weights[
|
||||
:, :, out_h_start:out_h_end, out_w_start:out_w_end, :
|
||||
] += tile_weights.to(target_device, target_dtype)
|
||||
|
||||
# Normalize by weights
|
||||
output /= weights + 1e-8
|
||||
|
||||
# Reshape output to match expected format [batch * frames, height, width, channels]
|
||||
output = output.view(
|
||||
batch * image_frames, output_height, output_width, output.shape[-1]
|
||||
)
|
||||
|
||||
if last_frame_fix:
|
||||
output = output[:-time_scale_factor, :, :]
|
||||
|
||||
return (output,)
|
||||
|
||||
|
||||
def compute_chunk_boundaries(
|
||||
chunk_start: int,
|
||||
temporal_tile_length: int,
|
||||
temporal_overlap: int,
|
||||
total_latent_frames: int,
|
||||
):
|
||||
"""Compute chunk boundaries for temporal tiling.
|
||||
|
||||
Args:
|
||||
chunk_start: Starting frame index for the current chunk
|
||||
temporal_tile_length: Length of each temporal tile
|
||||
temporal_overlap: Number of frames to overlap between chunks
|
||||
total_latent_frames: Total number of latent frames
|
||||
|
||||
Returns:
|
||||
Tuple of (overlap_start, chunk_end)
|
||||
"""
|
||||
if chunk_start == 0:
|
||||
# First chunk: no overlap needed
|
||||
chunk_end = min(chunk_start + temporal_tile_length, total_latent_frames)
|
||||
overlap_start = chunk_start
|
||||
else:
|
||||
# Subsequent chunks: include overlap from previous chunk
|
||||
# -1 because we need one extra frame to overlap, which is decoded to a single frame
|
||||
# never overlap with the first latent frame
|
||||
overlap_start = max(1, chunk_start - temporal_overlap - 1)
|
||||
extra_frames = chunk_start - overlap_start
|
||||
chunk_end = min(
|
||||
chunk_start + temporal_tile_length - extra_frames,
|
||||
total_latent_frames,
|
||||
)
|
||||
|
||||
return overlap_start, chunk_end
|
||||
|
||||
|
||||
def calculate_temporal_output_boundaries(
|
||||
overlap_start: int, time_scale_factor: int, tile_out_frames: int
|
||||
):
|
||||
"""Calculate temporal output boundaries for the decoded tile.
|
||||
|
||||
Args:
|
||||
overlap_start: Starting frame index including overlap
|
||||
time_scale_factor: Time scaling factor from VAE
|
||||
tile_out_frames: Number of frames in the decoded tile
|
||||
|
||||
Returns:
|
||||
Tuple of (out_t_start, out_t_end)
|
||||
"""
|
||||
# +1 for the first frame
|
||||
out_t_start = 1 + overlap_start * time_scale_factor
|
||||
|
||||
# Calculate actual output temporal dimensions
|
||||
out_t_end = out_t_start + tile_out_frames
|
||||
|
||||
return out_t_start, out_t_end
|
||||
|
||||
|
||||
@comfy_node(
|
||||
name="LTXVSpatioTemporalTiledVAEDecode",
|
||||
)
|
||||
class LTXVSpatioTemporalTiledVAEDecode(LTXVTiledVAEDecode):
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE", {"tooltip": "The VAE to use."}),
|
||||
"latents": ("LATENT", {"tooltip": "The latent samples to decode."}),
|
||||
"spatial_tiles": (
|
||||
"INT",
|
||||
{
|
||||
"default": 4,
|
||||
"min": 1,
|
||||
"max": 8,
|
||||
"tooltip": "The number of spatial tiles to use, horizontal and vertical.",
|
||||
},
|
||||
),
|
||||
"spatial_overlap": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 0,
|
||||
"max": 8,
|
||||
"tooltip": "The overlap between the spatial tiles. (in latent frames)",
|
||||
},
|
||||
),
|
||||
"temporal_tile_length": (
|
||||
"INT",
|
||||
{
|
||||
"default": 16,
|
||||
"min": 2,
|
||||
"max": 1000,
|
||||
"tooltip": "The length of the temporal tile to use for the sampling, in latent frames, including the overlapping region.",
|
||||
},
|
||||
),
|
||||
"temporal_overlap": (
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 0,
|
||||
"max": 8,
|
||||
"tooltip": "The overlap between the temporal tiles, in latent frames.",
|
||||
},
|
||||
),
|
||||
"last_frame_fix": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "If true, the last frame will be repeated and discarded after the decoding.",
|
||||
},
|
||||
),
|
||||
"working_device": (
|
||||
["cpu", "auto"],
|
||||
{
|
||||
"default": "auto",
|
||||
"tooltip": "The device to use for the decoding. auto->same as the latents.",
|
||||
},
|
||||
),
|
||||
"working_dtype": (
|
||||
["float16", "float32", "auto"],
|
||||
{
|
||||
"default": "auto",
|
||||
"tooltip": "The data type to use for the decoding. auto->same as the latents.",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
|
||||
FUNCTION = "decode_spatial_temporal"
|
||||
|
||||
CATEGORY = "latent"
|
||||
|
||||
def decode_spatial_temporal(
|
||||
self,
|
||||
vae,
|
||||
latents,
|
||||
spatial_tiles=4,
|
||||
spatial_overlap=1,
|
||||
temporal_tile_length=16,
|
||||
temporal_overlap=1,
|
||||
last_frame_fix=False,
|
||||
working_device="auto",
|
||||
working_dtype="auto",
|
||||
):
|
||||
if temporal_tile_length < temporal_overlap + 1:
|
||||
raise ValueError(
|
||||
"Temporal tile length must be greater than temporal overlap + 1"
|
||||
)
|
||||
|
||||
# Get the latent samples
|
||||
samples = latents["samples"]
|
||||
|
||||
batch, channels, frames, height, width = samples.shape
|
||||
time_scale_factor, width_scale_factor, height_scale_factor = (
|
||||
vae.downscale_index_formula
|
||||
)
|
||||
image_frames = 1 + (frames - 1) * time_scale_factor
|
||||
|
||||
# Calculate output image dimensions
|
||||
output_height = height * height_scale_factor
|
||||
output_width = width * width_scale_factor
|
||||
|
||||
target_device = samples.device if working_device == "auto" else working_device
|
||||
if working_dtype == "auto":
|
||||
target_dtype = samples.dtype
|
||||
elif working_dtype == "float16":
|
||||
target_dtype = torch.float16
|
||||
elif working_dtype == "float32":
|
||||
target_dtype = torch.float32
|
||||
|
||||
# Initialize output tensor and weight tensor
|
||||
output = torch.empty(
|
||||
(
|
||||
batch,
|
||||
image_frames,
|
||||
output_height,
|
||||
output_width,
|
||||
3,
|
||||
),
|
||||
device=target_device,
|
||||
dtype=target_dtype,
|
||||
)
|
||||
|
||||
# Process temporal chunks similar to reference function
|
||||
total_latent_frames = frames
|
||||
chunk_start = 0
|
||||
|
||||
while chunk_start < total_latent_frames:
|
||||
# Calculate chunk boundaries
|
||||
overlap_start, chunk_end = compute_chunk_boundaries(
|
||||
chunk_start, temporal_tile_length, temporal_overlap, total_latent_frames
|
||||
)
|
||||
|
||||
# units are latent frames
|
||||
chunk_frames = chunk_end - overlap_start
|
||||
logging.info(
|
||||
f"Processing temporal chunk: {overlap_start}:{chunk_end} ({chunk_frames} latent frames)"
|
||||
)
|
||||
|
||||
# Extract tile
|
||||
tile = samples[:, :, overlap_start:chunk_end]
|
||||
|
||||
# Create tile latents dict
|
||||
tile_latents = {"samples": tile}
|
||||
|
||||
# Decode the tile
|
||||
decoded_tile = self.decode(
|
||||
vae=vae,
|
||||
latents=tile_latents,
|
||||
vertical_tiles=spatial_tiles,
|
||||
horizontal_tiles=spatial_tiles,
|
||||
overlap=spatial_overlap,
|
||||
last_frame_fix=last_frame_fix,
|
||||
working_device=working_device,
|
||||
working_dtype=working_dtype,
|
||||
)[0][None]
|
||||
|
||||
if chunk_start == 0:
|
||||
output[:, : decoded_tile.shape[1]] = decoded_tile
|
||||
|
||||
# Drop first frame if needed (overlap)
|
||||
else:
|
||||
if decoded_tile.shape[1] == 1:
|
||||
raise ValueError("Dropping first frame but tile has only 1 frame")
|
||||
decoded_tile = decoded_tile[:, 1:] # Drop first frame
|
||||
|
||||
# Calculate temporal output boundaries
|
||||
out_t_start, out_t_end = calculate_temporal_output_boundaries(
|
||||
overlap_start, time_scale_factor, decoded_tile.shape[1]
|
||||
)
|
||||
|
||||
# Create weight mask for this tile
|
||||
overlap_frames = temporal_overlap * time_scale_factor
|
||||
frame_weights = torch.linspace(
|
||||
0,
|
||||
1,
|
||||
overlap_frames + 2,
|
||||
device=decoded_tile.device,
|
||||
dtype=decoded_tile.dtype,
|
||||
)[1:-1]
|
||||
tile_weights = frame_weights.view(1, -1, 1, 1, 1)
|
||||
after_overlap_frames_start = out_t_start + overlap_frames
|
||||
# Add weighted tile to output
|
||||
overlap_output = decoded_tile[:, :overlap_frames]
|
||||
output[:, out_t_start:after_overlap_frames_start] *= 1 - tile_weights
|
||||
output[:, out_t_start:after_overlap_frames_start] += (
|
||||
tile_weights * overlap_output
|
||||
)
|
||||
output[:, after_overlap_frames_start:out_t_end] = decoded_tile[
|
||||
:, overlap_frames:
|
||||
]
|
||||
|
||||
# Move to next chunk
|
||||
chunk_start = chunk_end
|
||||
|
||||
# Reshape output to match expected format [batch * frames, height, width, channels]
|
||||
output = output.view(
|
||||
batch * image_frames, output_height, output_width, output.shape[-1]
|
||||
)
|
||||
|
||||
return (output,)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:041f3820e0ab4109137d743b3f6b8e2dea0cddcad194aac08d3baec7571200f9
|
||||
size 1841100
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:9d2b002f12155d0b929366ad38561680a5bbaecea1925e5655658b46d281baf9
|
||||
size 2441562
|
||||
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3726ed083a5c52873d29adab87e8fe0a65c7e390beaf8dfcf23e5c9091d6da73
|
||||
size 521574
|
||||
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:01d44bb728e04879a5513a2c33bdd47d443bf901b33d3c1445918e75ba8b78b1
|
||||
size 678605
|
||||
@@ -1,3 +0,0 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:5432150d5477d19586b8d8049ca214038f886d3def9117830cf2af58879c48b5
|
||||
size 31584408
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 349 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 342 KiB |
@@ -7,7 +7,6 @@ from comfy.ldm.lightricks.model import (
|
||||
BasicTransformerBlock,
|
||||
LTXVModel,
|
||||
apply_rotary_emb,
|
||||
precompute_freqs_cis,
|
||||
)
|
||||
from comfy.ldm.lightricks.symmetric_patchifier import latent_to_pixel_coords
|
||||
from torch import nn
|
||||
@@ -167,7 +166,7 @@ class LTXVModelModified(LTXVModel):
|
||||
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
|
||||
) * torch.finfo(x.dtype).max
|
||||
|
||||
pe = precompute_freqs_cis(
|
||||
pe = self._precompute_freqs_cis(
|
||||
fractional_coords, dim=self.inner_dim, out_dtype=x.dtype
|
||||
)
|
||||
|
||||
|
||||
@@ -10,6 +10,7 @@ class AddLatentGuideNode(nodes_lt.LTXVAddGuide):
|
||||
"latent": ("LATENT",),
|
||||
"image_latent": ("LATENT",),
|
||||
"index": ("INT", {"default": 0, "min": -1, "max": 9999, "step": 1}),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,7 +20,7 @@ class AddLatentGuideNode(nodes_lt.LTXVAddGuide):
|
||||
CATEGORY = "ltxtricks"
|
||||
FUNCTION = "generate"
|
||||
|
||||
def generate(self, model, latent, image_latent, index):
|
||||
def generate(self, model, latent, image_latent, index, strength):
|
||||
noise_mask = nodes_lt.get_noise_mask(latent)
|
||||
latent = latent["samples"]
|
||||
|
||||
@@ -30,7 +31,7 @@ class AddLatentGuideNode(nodes_lt.LTXVAddGuide):
|
||||
noise_mask,
|
||||
image_latent,
|
||||
index,
|
||||
1.0,
|
||||
strength,
|
||||
)
|
||||
|
||||
return (
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
# Internal keys used to store AV merge metadata inside model options
|
||||
_AV_META_KEY = "ltxv_av_meta"
|
||||
_AV_META_VIDEO_ONLY_KEYS = "video_only_keys"
|
||||
_AV_META_AUDIO_ONLY_KEYS = "audio_only_keys"
|
||||
_AV_META_SHARED_KEYS = "shared_keys"
|
||||
_AV_META_AUDIO_SHARED_VALUES = "audio_shared_values"
|
||||
|
||||
# Common key constants
|
||||
_KEY_SAMPLES = "samples"
|
||||
_KEY_TYPE = "type"
|
||||
_OPT_TRANSFORMER = "transformer_options"
|
||||
_OPT_AUDIO_LENGTH = "audio_length"
|
||||
|
||||
|
||||
def _build_av_meta(video_latent: dict, audio_latent: dict) -> dict:
|
||||
"""Compute metadata that enables lossless split after AV concat.
|
||||
|
||||
Notes:
|
||||
- We keep video values for shared keys in the merged latent, so we must
|
||||
persist the original audio values per shared key.
|
||||
"""
|
||||
video_keys = set(video_latent.keys()) - {_KEY_SAMPLES}
|
||||
audio_keys = set(audio_latent.keys()) - {_KEY_SAMPLES, _KEY_TYPE}
|
||||
video_only_keys = list(video_keys - audio_keys)
|
||||
audio_only_keys = list(audio_keys - video_keys)
|
||||
shared_keys = list(video_keys & audio_keys)
|
||||
audio_shared_values = {k: audio_latent[k] for k in shared_keys}
|
||||
return {
|
||||
_AV_META_VIDEO_ONLY_KEYS: video_only_keys,
|
||||
_AV_META_AUDIO_ONLY_KEYS: audio_only_keys,
|
||||
_AV_META_SHARED_KEYS: shared_keys,
|
||||
_AV_META_AUDIO_SHARED_VALUES: audio_shared_values,
|
||||
}
|
||||
|
||||
|
||||
@comfy_node(description="Image to CPU")
|
||||
class ImageToCPU:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "utility"
|
||||
|
||||
def run(self, image):
|
||||
return (image.cpu(),)
|
||||
@@ -0,0 +1,38 @@
|
||||
from comfy.model_management import dtype_size
|
||||
|
||||
from .nodes_registry import comfy_node
|
||||
|
||||
|
||||
@comfy_node(name="LTXVPatcherVAE")
|
||||
class LTXVPatcherVAE:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VAE",)
|
||||
FUNCTION = "patch"
|
||||
CATEGORY = "lightricks/LTXV"
|
||||
TITLE = "LTXV VAE Patcher"
|
||||
|
||||
def patch(self, vae):
|
||||
from q8_kernels.integration.patch_vae import patch_vae
|
||||
|
||||
vae_model = vae.first_stage_model
|
||||
vae.memory_used_decode = (
|
||||
lambda shape, dtype: shape[1]
|
||||
* shape[2]
|
||||
* shape[3]
|
||||
* shape[4]
|
||||
* 8
|
||||
* 8
|
||||
* 8
|
||||
* 3
|
||||
* dtype_size(dtype)
|
||||
)
|
||||
|
||||
patch_vae(vae_model, patch_block=4)
|
||||
return (vae,)
|
||||
Reference in New Issue
Block a user