nodes/mobcat40-ComfyUI-PromptChain
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ComfyUI-PromptChain

Browse Source 
A prompt-authoring and image-iteration suite for ComfyUI.
This commit is contained in:
chrmc40
2026-06-17 19:35:25 -07:00
commit 8e3dceae99
2442 changed files with 634371 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+32
View File
@@ -0,0 +1,32 @@
__pycache__/
*.pyc
node_modules/
wildcards/**
!wildcards/.gitkeep
data/*
!data/tags/
!data/prompts/
!data/models/
!data/templates/
!data/tag-builder/
data/models/_hash_index.json
cache/
*.map
*.db-journal
*.db-wal
*.db-shm
# dev-only verification scratch (top-level one-off ESM/py harnesses, their temp
# evals, and test-run logs) — never ship in the pack. Anchored to repo root and
# to the "_*check"/"_*-eval" shapes so it can't catch __init__.py.
/_*.mjs
/_*check.py
/_*-eval.mjs
/regress.txt
# personal deploy helper (hardcoded local paths) — kept local, never shipped
/deploy.sh
# internal planning docs live at the working-dir root (C:\comfyui\dev-promptchain),
# never inside the shipped node
/dev-promptchain/
# dev-only probes/harnesses/tests moved out of scripts/ (which keeps only the
# runtime-imported modules) — kept local, never shipped
/_harness/
LICENSE
+661
View File
@@ -0,0 +1,661 @@
GNU AFFERO GENERAL PUBLIC LICENSE
Version 3, 19 November 2007
Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
Everyone is permitted to copy and distribute verbatim copies
of this license document, but changing it is not allowed.
Preamble
The GNU Affero General Public License is a free, copyleft license for
software and other kinds of works, specifically designed to ensure
cooperation with the community in the case of network server software.
The licenses for most software and other practical works are designed
to take away your freedom to share and change the works. By contrast,
our General Public Licenses are intended to guarantee your freedom to
share and change all versions of a program--to make sure it remains free
software for all its users.
When we speak of free software, we are referring to freedom, not
price. Our General Public Licenses are designed to make sure that you
have the freedom to distribute copies of free software (and charge for
them if you wish), that you receive source code or can get it if you
want it, that you can change the software or use pieces of it in new
free programs, and that you know you can do these things.
Developers that use our General Public Licenses protect your rights
with two steps: (1) assert copyright on the software, and (2) offer
you this License which gives you legal permission to copy, distribute
and/or modify the software.
A secondary benefit of defending all users' freedom is that
improvements made in alternate versions of the program, if they
receive widespread use, become available for other developers to
incorporate. Many developers of free software are heartened and
encouraged by the resulting cooperation. However, in the case of
software used on network servers, this result may fail to come about.
The GNU General Public License permits making a modified version and
letting the public access it on a server without ever releasing its
source code to the public.
The GNU Affero General Public License is designed specifically to
ensure that, in such cases, the modified source code becomes available
to the community. It requires the operator of a network server to
provide the source code of the modified version running there to the
users of that server. Therefore, public use of a modified version, on
a publicly accessible server, gives the public access to the source
code of the modified version.
An older license, called the Affero General Public License and
published by Affero, was designed to accomplish similar goals. This is
a different license, not a version of the Affero GPL, but Affero has
released a new version of the Affero GPL which permits relicensing under
this license.
The precise terms and conditions for copying, distribution and
modification follow.
TERMS AND CONDITIONS
0. Definitions.
"This License" refers to version 3 of the GNU Affero General Public License.
"Copyright" also means copyright-like laws that apply to other kinds of
works, such as semiconductor masks.
"The Program" refers to any copyrightable work licensed under this
License. Each licensee is addressed as "you". "Licensees" and
"recipients" may be individuals or organizations.
To "modify" a work means to copy from or adapt all or part of the work
in a fashion requiring copyright permission, other than the making of an
exact copy. The resulting work is called a "modified version" of the
earlier work or a work "based on" the earlier work.
A "covered work" means either the unmodified Program or a work based
on the Program.
To "propagate" a work means to do anything with it that, without
permission, would make you directly or secondarily liable for
infringement under applicable copyright law, except executing it on a
computer or modifying a private copy. Propagation includes copying,
distribution (with or without modification), making available to the
public, and in some countries other activities as well.
To "convey" a work means any kind of propagation that enables other
parties to make or receive copies. Mere interaction with a user through
a computer network, with no transfer of a copy, is not conveying.
An interactive user interface displays "Appropriate Legal Notices"
to the extent that it includes a convenient and prominently visible
feature that (1) displays an appropriate copyright notice, and (2)
tells the user that there is no warranty for the work (except to the
extent that warranties are provided), that licensees may convey the
work under this License, and how to view a copy of this License. If
the interface presents a list of user commands or options, such as a
menu, a prominent item in the list meets this criterion.
1. Source Code.
The "source code" for a work means the preferred form of the work
for making modifications to it. "Object code" means any non-source
form of a work.
A "Standard Interface" means an interface that either is an official
standard defined by a recognized standards body, or, in the case of
interfaces specified for a particular programming language, one that
is widely used among developers working in that language.
The "System Libraries" of an executable work include anything, other
than the work as a whole, that (a) is included in the normal form of
packaging a Major Component, but which is not part of that Major
Component, and (b) serves only to enable use of the work with that
Major Component, or to implement a Standard Interface for which an
implementation is available to the public in source code form. A
"Major Component", in this context, means a major essential component
(kernel, window system, and so on) of the specific operating system
(if any) on which the executable work runs, or a compiler used to
produce the work, or an object code interpreter used to run it.
The "Corresponding Source" for a work in object code form means all
the source code needed to generate, install, and (for an executable
work) run the object code and to modify the work, including scripts to
control those activities. However, it does not include the work's
System Libraries, or general-purpose tools or generally available free
programs which are used unmodified in performing those activities but
which are not part of the work. For example, Corresponding Source
includes interface definition files associated with source files for
the work, and the source code for shared libraries and dynamically
linked subprograms that the work is specifically designed to require,
such as by intimate data communication or control flow between those
subprograms and other parts of the work.
The Corresponding Source need not include anything that users
can regenerate automatically from other parts of the Corresponding
Source.
The Corresponding Source for a work in source code form is that
same work.
2. Basic Permissions.
All rights granted under this License are granted for the term of
copyright on the Program, and are irrevocable provided the stated
conditions are met. This License explicitly affirms your unlimited
permission to run the unmodified Program. The output from running a
covered work is covered by this License only if the output, given its
content, constitutes a covered work. This License acknowledges your
rights of fair use or other equivalent, as provided by copyright law.
You may make, run and propagate covered works that you do not
convey, without conditions so long as your license otherwise remains
in force. You may convey covered works to others for the sole purpose
of having them make modifications exclusively for you, or provide you
with facilities for running those works, provided that you comply with
the terms of this License in conveying all material for which you do
not control copyright. Those thus making or running the covered works
for you must do so exclusively on your behalf, under your direction
and control, on terms that prohibit them from making any copies of
your copyrighted material outside their relationship with you.
Conveying under any other circumstances is permitted solely under
the conditions stated below. Sublicensing is not allowed; section 10
makes it unnecessary.
3. Protecting Users' Legal Rights From Anti-Circumvention Law.
No covered work shall be deemed part of an effective technological
measure under any applicable law fulfilling obligations under article
11 of the WIPO copyright treaty adopted on 20 December 1996, or
similar laws prohibiting or restricting circumvention of such
measures.
When you convey a covered work, you waive any legal power to forbid
circumvention of technological measures to the extent such circumvention
is effected by exercising rights under this License with respect to
the covered work, and you disclaim any intention to limit operation or
modification of the work as a means of enforcing, against the work's
users, your or third parties' legal rights to forbid circumvention of
technological measures.
4. Conveying Verbatim Copies.
You may convey verbatim copies of the Program's source code as you
receive it, in any medium, provided that you conspicuously and
appropriately publish on each copy an appropriate copyright notice;
keep intact all notices stating that this License and any
non-permissive terms added in accord with section 7 apply to the code;
keep intact all notices of the absence of any warranty; and give all
recipients a copy of this License along with the Program.
You may charge any price or no price for each copy that you convey,
and you may offer support or warranty protection for a fee.
5. Conveying Modified Source Versions.
You may convey a work based on the Program, or the modifications to
produce it from the Program, in the form of source code under the
terms of section 4, provided that you also meet all of these conditions:
a) The work must carry prominent notices stating that you modified
it, and giving a relevant date.
b) The work must carry prominent notices stating that it is
released under this License and any conditions added under section
7. This requirement modifies the requirement in section 4 to
"keep intact all notices".
c) You must license the entire work, as a whole, under this
License to anyone who comes into possession of a copy. This
License will therefore apply, along with any applicable section 7
additional terms, to the whole of the work, and all its parts,
regardless of how they are packaged. This License gives no
permission to license the work in any other way, but it does not
invalidate such permission if you have separately received it.
d) If the work has interactive user interfaces, each must display
Appropriate Legal Notices; however, if the Program has interactive
interfaces that do not display Appropriate Legal Notices, your
work need not make them do so.
A compilation of a covered work with other separate and independent
works, which are not by their nature extensions of the covered work,
and which are not combined with it such as to form a larger program,
in or on a volume of a storage or distribution medium, is called an
"aggregate" if the compilation and its resulting copyright are not
used to limit the access or legal rights of the compilation's users
beyond what the individual works permit. Inclusion of a covered work
in an aggregate does not cause this License to apply to the other
parts of the aggregate.
6. Conveying Non-Source Forms.
You may convey a covered work in object code form under the terms
of sections 4 and 5, provided that you also convey the
machine-readable Corresponding Source under the terms of this License,
in one of these ways:
a) Convey the object code in, or embodied in, a physical product
(including a physical distribution medium), accompanied by the
Corresponding Source fixed on a durable physical medium
customarily used for software interchange.
b) Convey the object code in, or embodied in, a physical product
(including a physical distribution medium), accompanied by a
written offer, valid for at least three years and valid for as
long as you offer spare parts or customer support for that product
model, to give anyone who possesses the object code either (1) a
copy of the Corresponding Source for all the software in the
product that is covered by this License, on a durable physical
medium customarily used for software interchange, for a price no
more than your reasonable cost of physically performing this
conveying of source, or (2) access to copy the
Corresponding Source from a network server at no charge.
c) Convey individual copies of the object code with a copy of the
written offer to provide the Corresponding Source. This
alternative is allowed only occasionally and noncommercially, and
only if you received the object code with such an offer, in accord
with subsection 6b.
d) Convey the object code by offering access from a designated
place (gratis or for a charge), and offer equivalent access to the
Corresponding Source in the same way through the same place at no
further charge. You need not require recipients to copy the
Corresponding Source along with the object code. If the place to
copy the object code is a network server, the Corresponding Source
may be on a different server (operated by you or a third party)
that supports equivalent copying facilities, provided you maintain
clear directions next to the object code saying where to find the
Corresponding Source. Regardless of what server hosts the
Corresponding Source, you remain obligated to ensure that it is
available for as long as needed to satisfy these requirements.
e) Convey the object code using peer-to-peer transmission, provided
you inform other peers where the object code and Corresponding
Source of the work are being offered to the general public at no
charge under subsection 6d.
A separable portion of the object code, whose source code is excluded
from the Corresponding Source as a System Library, need not be
included in conveying the object code work.
A "User Product" is either (1) a "consumer product", which means any
tangible personal property which is normally used for personal, family,
or household purposes, or (2) anything designed or sold for incorporation
into a dwelling. In determining whether a product is a consumer product,
doubtful cases shall be resolved in favor of coverage. For a particular
product received by a particular user, "normally used" refers to a
typical or common use of that class of product, regardless of the status
of the particular user or of the way in which the particular user
actually uses, or expects or is expected to use, the product. A product
is a consumer product regardless of whether the product has substantial
commercial, industrial or non-consumer uses, unless such uses represent
the only significant mode of use of the product.
"Installation Information" for a User Product means any methods,
procedures, authorization keys, or other information required to install
and execute modified versions of a covered work in that User Product from
a modified version of its Corresponding Source. The information must
suffice to ensure that the continued functioning of the modified object
code is in no case prevented or interfered with solely because
modification has been made.
If you convey an object code work under this section in, or with, or
specifically for use in, a User Product, and the conveying occurs as
part of a transaction in which the right of possession and use of the
User Product is transferred to the recipient in perpetuity or for a
fixed term (regardless of how the transaction is characterized), the
Corresponding Source conveyed under this section must be accompanied
by the Installation Information. But this requirement does not apply
if neither you nor any third party retains the ability to install
modified object code on the User Product (for example, the work has
been installed in ROM).
The requirement to provide Installation Information does not include a
requirement to continue to provide support service, warranty, or updates
for a work that has been modified or installed by the recipient, or for
the User Product in which it has been modified or installed. Access to a
network may be denied when the modification itself materially and
adversely affects the operation of the network or violates the rules and
protocols for communication across the network.
Corresponding Source conveyed, and Installation Information provided,
in accord with this section must be in a format that is publicly
documented (and with an implementation available to the public in
source code form), and must require no special password or key for
unpacking, reading or copying.
7. Additional Terms.
"Additional permissions" are terms that supplement the terms of this
License by making exceptions from one or more of its conditions.
Additional permissions that are applicable to the entire Program shall
be treated as though they were included in this License, to the extent
that they are valid under applicable law. If additional permissions
apply only to part of the Program, that part may be used separately
under those permissions, but the entire Program remains governed by
this License without regard to the additional permissions.
When you convey a copy of a covered work, you may at your option
remove any additional permissions from that copy, or from any part of
it. (Additional permissions may be written to require their own
removal in certain cases when you modify the work.) You may place
additional permissions on material, added by you to a covered work,
for which you have or can give appropriate copyright permission.
Notwithstanding any other provision of this License, for material you
add to a covered work, you may (if authorized by the copyright holders of
that material) supplement the terms of this License with terms:
a) Disclaiming warranty or limiting liability differently from the
terms of sections 15 and 16 of this License; or
b) Requiring preservation of specified reasonable legal notices or
author attributions in that material or in the Appropriate Legal
Notices displayed by works containing it; or
c) Prohibiting misrepresentation of the origin of that material, or
requiring that modified versions of such material be marked in
reasonable ways as different from the original version; or
d) Limiting the use for publicity purposes of names of licensors or
authors of the material; or
e) Declining to grant rights under trademark law for use of some
trade names, trademarks, or service marks; or
f) Requiring indemnification of licensors and authors of that
material by anyone who conveys the material (or modified versions of
it) with contractual assumptions of liability to the recipient, for
any liability that these contractual assumptions directly impose on
those licensors and authors.
All other non-permissive additional terms are considered "further
restrictions" within the meaning of section 10. If the Program as you
received it, or any part of it, contains a notice stating that it is
governed by this License along with a term that is a further
restriction, you may remove that term. If a license document contains
a further restriction but permits relicensing or conveying under this
License, you may add to a covered work material governed by the terms
of that license document, provided that the further restriction does
not survive such relicensing or conveying.
If you add terms to a covered work in accord with this section, you
must place, in the relevant source files, a statement of the
additional terms that apply to those files, or a notice indicating
where to find the applicable terms.
Additional terms, permissive or non-permissive, may be stated in the
form of a separately written license, or stated as exceptions;
the above requirements apply either way.
8. Termination.
You may not propagate or modify a covered work except as expressly
provided under this License. Any attempt otherwise to propagate or
modify it is void, and will automatically terminate your rights under
this License (including any patent licenses granted under the third
paragraph of section 11).
However, if you cease all violation of this License, then your
license from a particular copyright holder is reinstated (a)
provisionally, unless and until the copyright holder explicitly and
finally terminates your license, and (b) permanently, if the copyright
holder fails to notify you of the violation by some reasonable means
prior to 60 days after the cessation.
Moreover, your license from a particular copyright holder is
reinstated permanently if the copyright holder notifies you of the
violation by some reasonable means, this is the first time you have
received notice of violation of this License (for any work) from that
copyright holder, and you cure the violation prior to 30 days after
your receipt of the notice.
Termination of your rights under this section does not terminate the
licenses of parties who have received copies or rights from you under
this License. If your rights have been terminated and not permanently
reinstated, you do not qualify to receive new licenses for the same
material under section 10.
9. Acceptance Not Required for Having Copies.
You are not required to accept this License in order to receive or
run a copy of the Program. Ancillary propagation of a covered work
occurring solely as a consequence of using peer-to-peer transmission
to receive a copy likewise does not require acceptance. However,
nothing other than this License grants you permission to propagate or
modify any covered work. These actions infringe copyright if you do
not accept this License. Therefore, by modifying or propagating a
covered work, you indicate your acceptance of this License to do so.
10. Automatic Licensing of Downstream Recipients.
Each time you convey a covered work, the recipient automatically
receives a license from the original licensors, to run, modify and
propagate that work, subject to this License. You are not responsible
for enforcing compliance by third parties with this License.
An "entity transaction" is a transaction transferring control of an
organization, or substantially all assets of one, or subdividing an
organization, or merging organizations. If propagation of a covered
work results from an entity transaction, each party to that
transaction who receives a copy of the work also receives whatever
licenses to the work the party's predecessor in interest had or could
give under the previous paragraph, plus a right to possession of the
Corresponding Source of the work from the predecessor in interest, if
the predecessor has it or can get it with reasonable efforts.
You may not impose any further restrictions on the exercise of the
rights granted or affirmed under this License. For example, you may
not impose a license fee, royalty, or other charge for exercise of
rights granted under this License, and you may not initiate litigation
(including a cross-claim or counterclaim in a lawsuit) alleging that
any patent claim is infringed by making, using, selling, offering for
sale, or importing the Program or any portion of it.
11. Patents.
A "contributor" is a copyright holder who authorizes use under this
License of the Program or a work on which the Program is based. The
work thus licensed is called the contributor's "contributor version".
A contributor's "essential patent claims" are all patent claims
owned or controlled by the contributor, whether already acquired or
hereafter acquired, that would be infringed by some manner, permitted
by this License, of making, using, or selling its contributor version,
but do not include claims that would be infringed only as a
consequence of further modification of the contributor version. For
purposes of this definition, "control" includes the right to grant
patent sublicenses in a manner consistent with the requirements of
this License.
Each contributor grants you a non-exclusive, worldwide, royalty-free
patent license under the contributor's essential patent claims, to
make, use, sell, offer for sale, import and otherwise run, modify and
propagate the contents of its contributor version.
In the following three paragraphs, a "patent license" is any express
agreement or commitment, however denominated, not to enforce a patent
(such as an express permission to practice a patent or covenant not to
sue for patent infringement). To "grant" such a patent license to a
party means to make such an agreement or commitment not to enforce a
patent against the party.
If you convey a covered work, knowingly relying on a patent license,
and the Corresponding Source of the work is not available for anyone
to copy, free of charge and under the terms of this License, through a
publicly available network server or other readily accessible means,
then you must either (1) cause the Corresponding Source to be so
available, or (2) arrange to deprive yourself of the benefit of the
patent license for this particular work, or (3) arrange, in a manner
consistent with the requirements of this License, to extend the patent
license to downstream recipients. "Knowingly relying" means you have
actual knowledge that, but for the patent license, your conveying the
covered work in a country, or your recipient's use of the covered work
in a country, would infringe one or more identifiable patents in that
country that you have reason to believe are valid.
If, pursuant to or in connection with a single transaction or
arrangement, you convey, or propagate by procuring conveyance of, a
covered work, and grant a patent license to some of the parties
receiving the covered work authorizing them to use, propagate, modify
or convey a specific copy of the covered work, then the patent license
you grant is automatically extended to all recipients of the covered
work and works based on it.
A patent license is "discriminatory" if it does not include within
the scope of its coverage, prohibits the exercise of, or is
conditioned on the non-exercise of one or more of the rights that are
specifically granted under this License. You may not convey a covered
work if you are a party to an arrangement with a third party that is
in the business of distributing software, under which you make payment
to the third party based on the extent of your activity of conveying
the work, and under which the third party grants, to any of the
parties who would receive the covered work from you, a discriminatory
patent license (a) in connection with copies of the covered work
conveyed by you (or copies made from those copies), or (b) primarily
for and in connection with specific products or compilations that
contain the covered work, unless you entered into that arrangement,
or that patent license was granted, prior to 28 March 2007.
Nothing in this License shall be construed as excluding or limiting
any implied license or other defenses to infringement that may
otherwise be available to you under applicable patent law.
12. No Surrender of Others' Freedom.
If conditions are imposed on you (whether by court order, agreement or
otherwise) that contradict the conditions of this License, they do not
excuse you from the conditions of this License. If you cannot convey a
covered work so as to satisfy simultaneously your obligations under this
License and any other pertinent obligations, then as a consequence you may
not convey it at all. For example, if you agree to terms that obligate you
to collect a royalty for further conveying from those to whom you convey
the Program, the only way you could satisfy both those terms and this
License would be to refrain entirely from conveying the Program.
13. Remote Network Interaction; Use with the GNU General Public License.
Notwithstanding any other provision of this License, if you modify the
Program, your modified version must prominently offer all users
interacting with it remotely through a computer network (if your version
supports such interaction) an opportunity to receive the Corresponding
Source of your version by providing access to the Corresponding Source
from a network server at no charge, through some standard or customary
means of facilitating copying of software. This Corresponding Source
shall include the Corresponding Source for any work covered by version 3
of the GNU General Public License that is incorporated pursuant to the
following paragraph.
Notwithstanding any other provision of this License, you have
permission to link or combine any covered work with a work licensed
under version 3 of the GNU General Public License into a single
combined work, and to convey the resulting work. The terms of this
License will continue to apply to the part which is the covered work,
but the work with which it is combined will remain governed by version
3 of the GNU General Public License.
14. Revised Versions of this License.
The Free Software Foundation may publish revised and/or new versions of
the GNU Affero General Public License from time to time. Such new versions
will be similar in spirit to the present version, but may differ in detail to
address new problems or concerns.
Each version is given a distinguishing version number. If the
Program specifies that a certain numbered version of the GNU Affero General
Public License "or any later version" applies to it, you have the
option of following the terms and conditions either of that numbered
version or of any later version published by the Free Software
Foundation. If the Program does not specify a version number of the
GNU Affero General Public License, you may choose any version ever published
by the Free Software Foundation.
If the Program specifies that a proxy can decide which future
versions of the GNU Affero General Public License can be used, that proxy's
public statement of acceptance of a version permanently authorizes you
to choose that version for the Program.
Later license versions may give you additional or different
permissions. However, no additional obligations are imposed on any
author or copyright holder as a result of your choosing to follow a
later version.
15. Disclaimer of Warranty.
THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
16. Limitation of Liability.
IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
SUCH DAMAGES.
17. Interpretation of Sections 15 and 16.
If the disclaimer of warranty and limitation of liability provided
above cannot be given local legal effect according to their terms,
reviewing courts shall apply local law that most closely approximates
an absolute waiver of all civil liability in connection with the
Program, unless a warranty or assumption of liability accompanies a
copy of the Program in return for a fee.
END OF TERMS AND CONDITIONS
How to Apply These Terms to Your New Programs
If you develop a new program, and you want it to be of the greatest
possible use to the public, the best way to achieve this is to make it
free software which everyone can redistribute and change under these terms.
To do so, attach the following notices to the program. It is safest
to attach them to the start of each source file to most effectively
state the exclusion of warranty; and each file should have at least
the "copyright" line and a pointer to where the full notice is found.
<one line to give the program's name and a brief idea of what it does.>
Copyright (C) <year> <name of author>
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.
You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
Also add information on how to contact you by electronic and paper mail.
If your software can interact with users remotely through a computer
network, you should also make sure that it provides a way for users to
get its source. For example, if your program is a web application, its
interface could display a "Source" link that leads users to an archive
of the code. There are many ways you could offer source, and different
solutions will be better for different programs; see section 13 for the
specific requirements.
You should also get your employer (if you work as a programmer) or school,
if any, to sign a "copyright disclaimer" for the program, if necessary.
For more information on this, and how to apply and follow the GNU AGPL, see
<https://www.gnu.org/licenses/>.
README.md
+117
View File
@@ -0,0 +1,117 @@
<h1 align="left"><img src="docs/spacer.webp" width="10" alt=""><img src="docs/wordmark.webp" alt="ComfyUI-PromptChain" width="460"></h1>
<p align="center">
<img src="docs/hero.webp" alt="PromptChain — 3D posing, structured prompt authoring, and layered editing in ComfyUI" width="100%">
</p>
> ⚠️ **Beta.** This is an early public release — things may change and break. Feedback welcome.
A prompt-authoring and image-iteration suite for [ComfyUI](https://github.com/comfyanonymous/ComfyUI). PromptChain turns the raw prompt box into a model-aware, structured workflow — and carries that structure through posing, generating, viewing, editing, and upscaling.
## Prompt authoring
- **Prompt Chain node** — a real code editor on the node: syntax highlighting, comments, multi-line editing, and a collapsible output panel showing the compiled positive/negative prompt, a live server console, and this node's gallery.
- **Chain modes** — link multiple Prompt Chain nodes and choose how their text merges: **Combine** all, **Randomize** one at random, **Switch** to a chosen branch, or **Iterate** through options across queued runs — shown live in the toolbar.
- **Wildcards** — inline `::Label::a|b` option sets that randomize, cycle, or switch, plus `__file__` wildcards whose per-token badge sets the mode inline.
- **Autocomplete** — booru tag completion as you type, `@` to pull in curated characters/styles/poses (auto-cascading their outfit and pose), and `$` to insert a regional prompt block for any figure in your 3D pose.
- **Fast editing** — weight nudging with Ctrl+↑/↓ on a tag, `(text:weight)` group, or `<lora:…:weight>`; comment, font-size, comma-split, and panel-toggle shortcuts; a right-click menu for Tag Builder, saved prompts, and image-prompt recovery.
- **Drag an image to merge its prompts** — drop any ComfyUI-generated image onto the editor and its positive/negative prompts merge in, deduplicated. A LoadImage node can also pull an image's original prompts into the editor via keywords.
- **Saved prompt library** — store reusable snippets scoped to a checkpoint and drop them in from the right-click menu.
- **Fullscreen editor** — open the whole prompt chain in a distraction-free overlay that edits every linked node side by side.
## Tag Builder
- **Visual prompt composer** — assemble a prompt from curated, model-aware tag sets instead of typing. Open it from the editor, pick a subject, dress it, pose it, and scene it; the tags write straight into your prompt.
- **Multi-subject cards** — each subject is its own card with appearance, clothing, pose, expression, action, and prop slots, so multi-character prompts stay organized.
- **Character library** — browse and search a large library of named characters (anime, original, video game, V-Tuber, archetype, fantasy, creatures); picking one drops in its canonical identity tags.
- **Curated picks with thumbnails** — pick from image-backed sets for appearance, clothing, pose, expression, action, scene, and 16 prop categories.
- **Customizers** — refine clothing (color / material / condition / pattern), fantasy traits (wings / horns / tails / ears), and props & furniture (materials, patterns, attach to a subject as an interaction).
- **Round-trips** — open the builder on an existing prompt and it parses your styles, characters, outfits, poses, and props back into chips, preserving your original wording so visual and hand edits stay in sync.
- **Model-aware styles** — a Styles tab of render presets matched to your loaded checkpoint's architecture and family appears when a model is detected.
## Model awareness
- **Checkpoint recognition** — models are fingerprinted and looked up on CivitAI by hash, so PromptChain shows the real name and applies the right presets, prompt defaults, sampler recommendations, and ControlNet routing for the detected family. Unknown models are recognized in the background at startup with live progress.
- **In-place model swap** — pick a compatible checkpoint from the model panel and PromptChain swaps it into your existing nodes and reapplies its settings; custom wiring (like a regional couple) stays put.
- **Model panel** — a clickable model chip on each node opens a searchable picker and a settings modal (Info / Settings / Prompts / Templates / Versions) for tuning recommended ranges, scoped prompt snippets, workflow templates, and spotting newer releases.
- **Version inheritance** — a new version of a model you already set up inherits your trigger words, prompt style, sampler ranges, and detailer settings automatically.
- **Built-in downloads** — search CivitAI or browse a curated model catalog and download checkpoints into the right folder with a live progress bar and size estimates.
## 3D Poser
- **Pose Studio node** — pose full 3D mannequins in the browser and emit depth + OpenPose control maps that stay in sync with your prompt sections; the pose is saved into the workflow for cache reuse.
- **Per-figure masks** — alongside the control map, the node outputs a per-character silhouette mask (and per-named-prop region) to drive regional conditioning.
- **Figures & bodies** — stage multiple people, start from body templates, and reshape anyone live with a **MakeHuman-derived slider engine** (height, weight, muscle, gender, age, proportions, and more) — applied instantly with no rebake.
- **Creatures** — pose a horse or a cat as first-class rigged figures with their own IK.
- **Props** — drop in matte primitives and a large categorized library (furniture, weapons, vehicles, tools, environment), spawn pre-arranged prop sets, wear calibrated shoes, or **import your own `.glb`/`.obj`** — and even fit a skeleton inside an import to make it posable.
- **Build geometry** — **draw a 3D shape** from a freehand outline (flat / extruded / revolved / inflated), or **trace an image** into inflated solids that keep its layout.
- **Mesh sculpt** — reshape a body or prop freehand with a grab / inflate / smooth brush, including sculpting while the figure is posed.
- **Posing tools** — move/rotate/scale gizmo, per-finger hand posing, foot pinning, prop linking, object grouping, mirror/reset, full undo/redo, and per-prop `$` prompt regions.
- **Pose presets** — apply bundled poses or save your own, filtered to the active figure's rig.
- **Workspace** — dock the live Poser into a panel, go fullscreen, toggle a ground floor, and export/import whole scenes as JSON.
## Multi-character & regional control
- **Regional (Attention Couple)** — per-figure silhouette masks from the 3D Poser pin each character's prompt block to that figure, so characters and outfits stop bleeding into each other.
- **Regional Conditioning** — carries each character's prompt as native masked conditioning that survives tiled samplers, so an upscale tile over a figure samples with that figure's own prompt.
- **Regional Detailer** — mask-correct face detailing that pairs each face with its own character prompt (detects faces or uses the Poser's exact head positions; works non-regionally too, and installs its own detector).
## Generate & view
- **Generation gallery** — every image a node makes collects into a built-in gallery, newest first, with justified / grid / list views, selection, keyboard navigation, and bulk delete.
- **Lineage grouping** — upscales, inpaints, and edits collapse into their source image's card with the latest on top and a count badge; re-generating from the same source spawns a separate card. Toggleable in settings.
- **Image viewer** — a full-screen viewer with a timeline strip, a lineage family tree for stepping between an image and its variants, a side-by-side compare slider, and a panel that recovers the prompt, model, seed, and settings from any PromptChain image.
- **Safe deletes & recovery** — deleting purges every copy on disk and tombstones the image so self-healing history can't revive it; if a source file goes missing, the viewer offers a hash-verified "Locate file" re-link.
## Edit, inpaint & upscale
Everything below lives behind **Edit** — the single doorway to post-processing any generated image.
- **Layered editor** — a Photoshop-style editor with non-destructive layers (opacity, blend modes, reorder, merge, flatten), brush/eraser, paint bucket, spot heal, healing brush, clone stamp, smooth, move-with-snap, rectangle/lasso selections with feather, crop, **Liquify**, and **Camera-Raw-style** tonal/color adjustments — with full undo history. Save back onto the same image (keeping its lineage) or export a new one.
- **AI selection** — one-click masking: **Select Subject** (BiRefNet) segments the main subject, **Object Select** (SAM2) turns any click into that object's mask, and pose-derived masks isolate an individual figure in a multi-character scene.
- **Region → Inpaint** — send a layer or selection to a full inpaint pass and get it back as a layer. Modes for Basic, Depth (structure-locked), Regional (each figure's own prompt), or both; a choice of engines (the image's own model, a standalone SDXL / FLUX.1, or a Qwen Edit instruction model); and **character lock** via PuLID or Style Reference referenced from the image itself.
- **Region → Upscale** — targeted upscale that can grow the canvas so new detail lands pixel-for-pixel, or keep the canvas and fit the re-detail back in place. Basic / Depth / Regional modes, scale and denoise controls, recommended sampler markers, character lock, and a choice of engines — a deterministic **UltraSharp** climb by default, an opt-in **SeedVR2** restore for degraded sources, or AI re-detail.
- **Background upscales** — queue an upscale from any viewed image without touching your open workflow; it runs in the background with live tile preview, progress, ETA, and a real Cancel, then records into lineage.
- **Re-pose** — pose a detached 3D mannequin from Edit and re-render the whole image into the new pose while keeping identity and outfit, using a Qwen Edit AnyPose recipe or a FLUX.2 Klein depth recipe.
## Asset browser
- **Sidebar file manager** — a desktop-style browser docked in ComfyUI's sidebar for your **Workflows**, **Input**, and **Output** folders, with breadcrumbs, history navigation, sort/group, and type-to-select.
- **Recent feed** — flatten every subfolder into one newest-first stream with infinite scroll and full-subtree search, so your latest renders are always on top.
- **Thumbnails & previews** — inline image/video thumbnails (grid or list, Ctrl+scroll to zoom), saved workflows showing their last render, favorites, and a properties view with embedded generation metadata.
- **File ops** — cut / copy / paste / move (drag-and-drop or shortcuts with conflict resolution), rename, delete, multi-select, and drag-and-drop upload of files or whole folders. New outputs appear live as renders finish.
- **Find Duplicates** — scan a folder for visually duplicate images via perceptual hashing and clean them up.
- **New workflow in a few clicks** — pick a checkpoint from a searchable, architecture-grouped list and a template wired for it; the file is auto-named and you land on the canvas ready to render.
## Onboarding & installation
- **First-run walkthrough** — a guided Welcome / Setup / Extras splash flips on recommended ComfyUI options, validates a CivitAI key, and offers to install optional features.
- **Self-installing features** — Style Reference, Upscaler, Face Detailer, ControlNet, PuLID, SeedVR2, and AI Selection install their own nodes and model files on first use, with live progress and a restart — no Manager round-trip. A health dashboard in **Settings → PromptChain → Install** shows status, sizes, and updates.
- **Bundled ControlNet preprocessors** — license-clean OpenPose, Depth, Canny, Tile, Luminance, Scribble, Soft Edge, and Line Art preprocessors ship with the pack, so full ControlNet workflows work without `comfyui_controlnet_aux`.
- **In-app updater** — an About panel checks the repo, applies updates, and restarts ComfyUI in place.
To install, clone into your ComfyUI `custom_nodes` directory and restart ComfyUI:
```bash
cd ComfyUI/custom_nodes
git clone https://github.com/mobcat40/ComfyUI-PromptChain.git
```
Add the **Prompt Chain** node from the node menu — onboarding takes it from there.
## Nodes
| Node | Purpose |
| --- | --- |
| Prompt Chain | Structured prompt editor and compiler (chain modes, wildcards, regions) |
| PromptChain 3D Poser | 3D posing → depth / OpenPose maps + per-figure masks |
| Prompt Chain Regional (Attention Couple) | Per-character regional attention from per-figure masks |
| Prompt Chain Regional Conditioning | Tile-safe per-region conditioning that survives upscaling |
| Prompt Chain Regional Detailer | Mask-correct per-region face detailing |
| Prompt Chain Masked Detail | Crop-and-upscale re-render of a masked region |
| Prompt Chain Background Mask (Depth) | Depth-based background mask for bokeh-preserving upscales |
| Prompt Chain OpenPose / Depth Anything / Canny / Tile / Luminance / Scribble / Soft Edge / Line Art | Bundled, license-clean ControlNet image preprocessors |
## License
[AGPL-3.0](LICENSE) — © 2026 mobcat40. If you run a modified PromptChain as part of a hosted service, you must offer your users the modified source.
__init__.py
+226
View File
@@ -0,0 +1,226 @@
from __future__ import annotations
import asyncio as _asyncio
import logging as _logging
import sys as _sys
import threading as _threading
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
from .nodes.promptchain import PromptChainNode
from .nodes.pose_studio import PoseStudioNode
from .nodes.attention_couple import AttentionCoupleNode
from .nodes.regional_detailer import RegionalDetailerNode, MaskedDetailNode
from .nodes.regional_conditioning import RegionalConditioningNode
from .nodes.defocus_mask import DefocusMaskNode
# Timestamps on every PromptChain log line. ComfyUI's default formatter
# emits bare messages with no clock — fine for crash-style logs but
# useless for measuring patch latency or correlating events. We attach
# our own StreamHandler to the `promptchain` namespace with HH:MM:SS.mmm
# timestamps and disable propagation so messages don't double-log.
_pc_logger = _logging.getLogger("promptchain")
if not any(getattr(h, "_pc_timestamped", False) for h in _pc_logger.handlers):
_pc_handler = _logging.StreamHandler(_sys.stdout)
_pc_handler.setFormatter(_logging.Formatter(
"%(asctime)s.%(msecs)03d %(levelname).1s [%(name)s] %(message)s",
datefmt="%H:%M:%S",
))
_pc_handler._pc_timestamped = True # marker so re-import doesn't dup
_pc_logger.addHandler(_pc_handler)
_pc_logger.setLevel(_logging.INFO)
_pc_logger.propagate = False
from .core.fingerprint import scan_models, list_models
from .core.tags import get_store as get_tag_store
from .core import model_settings
from .core.shared import send_ws
from .core.model_api import detect_model, recognition_state
# Route modules — register their endpoints on import
from .core import tag_builder # noqa: F401
from .core import browse_api # noqa: F401
from .core import template_api # noqa: F401
from .core import tag_api # noqa: F401
from .core import wildcard_api # noqa: F401
from .core import history_api # noqa: F401
from .core import model_api # noqa: F401
from .core import civitai_api # noqa: F401
from .core import system_api # noqa: F401
from .core import node_install_api # noqa: F401
from .core import pose_api # noqa: F401
from .core import inpaint_files # noqa: F401 (age-sweeps input/promptchain_inpaint at startup)
from .core import upscale_api # noqa: F401
from .core import edit_api # noqa: F401
from .core import subject_api # noqa: F401
from .core import install_sections # noqa: F401
from .core import ai_api # noqa: F401
from .core import ai_agent # noqa: F401
# ── native preprocessor nodes ────────────────────────────────────
# pcr_cnaux (DepthAnything/TEED-based) and pcr_rtmpose (rtmlib) are OUR OWN
# preprocessor nodes — the license-clean replacements for comfyui_controlnet_aux.
# They're written against the V1 `NODE_CLASS_MAPPINGS` convention, but PromptChain
# registers via the V3 `comfy_entrypoint` below and ComfyUI's loader is
# `if NODE_CLASS_MAPPINGS … elif comfy_entrypoint` (mutually exclusive, see
# comfyui/nodes.py), so a module-level NODE_CLASS_MAPPINGS here would shadow our
# own nodes. Instead we inject these maps into ComfyUI's global registry at import
# time (during exec_module, before the loader reaches our entrypoint). Their ids
# are unique (PromptChain_*), so there is no collision to guard against.
# Third-party packs are NOT registered here — they live in bundled_packs/ and are
# copied into the user's custom_nodes/ at install time. Wrapped so an import error
# in one can never take down PromptChain's own nodes.
def _register_vendored_nodes():
import nodes as _comfy
native = []
try:
from .vendor.pcr_cnaux import NODE_CLASS_MAPPINGS as m, NODE_DISPLAY_NAME_MAPPINGS as d
native.append(("pcr_cnaux", m, d))
except Exception:
_pc_logger.warning("native pcr_cnaux failed to load", exc_info=True)
try:
from .vendor.pcr_rtmpose import NODE_CLASS_MAPPINGS as m, NODE_DISPLAY_NAME_MAPPINGS as d
native.append(("pcr_rtmpose", m, d))
except Exception:
_pc_logger.warning("native pcr_rtmpose failed to load", exc_info=True)
for name, class_map, display_map in native:
_comfy.NODE_CLASS_MAPPINGS.update(class_map)
_comfy.NODE_DISPLAY_NAME_MAPPINGS.update(display_map or {})
_pc_logger.info("registered native preprocessors '%s' (%d nodes)", name, len(class_map))
_register_vendored_nodes()
WEB_DIRECTORY = "./js"
# ── boot-time preloading ─────────────────────────────────────────
# Load heavy data in background threads so the first API request
# doesn't block the aiohttp event loop (which freezes all nodes).
def _preload():
scan_models()
get_tag_store().load_all()
# Heal any CivitAI-model-id groups whose versions ended up with
# different model_name strings (creator renamed the CivitAI page
# between recognitions). Runs every boot; idempotent when data
# is already consistent.
try:
model_settings.normalize_civitai_model_names()
except Exception as e:
print(f"[PromptChain] normalize_civitai_model_names failed: {e}")
# Warm up bucket-search embeddings so the first Prompt Generator
# request doesn't pay the model-download + index-build cost (~5-15s).
# Best-effort; if transformers isn't installed the search endpoint
# degrades to empty results and the rest of the pipeline still works.
try:
from .core import bucket_search
bucket_search.warmup()
except Exception as e:
print(f"[PromptChain] bucket_search warmup failed: {e}")
# Same model as bucket_search; this just builds the small modifier
# index (~7 rows) using the cached weights. First /ai/patch request
# then doesn't pay the embed-the-modifiers cost.
try:
from .core import modifier_search
modifier_search.warmup()
except Exception as e:
print(f"[PromptChain] modifier_search warmup failed: {e}")
# Danbooru tag-wiki retrieval. Disk-cached index — typically < 1s
# to load. Cold first-build is ~5-10 min on CPU and only happens
# after a wiki ingest run; we do that out-of-band, not at boot.
try:
from .core import tag_search
tag_search.warmup()
except Exception as e:
print(f"[PromptChain] tag_search warmup failed: {e}")
try:
from .core import style_search
style_search.warmup()
except Exception as e:
print(f"[PromptChain] style_search warmup failed: {e}")
_recognize_unknown_models()
def _recognize_unknown_models():
"""Process unrecognized models: full SHA256 -> CivitAI -> save settings."""
models = list_models()
unrecognized = [
m for m in models
if not model_settings.load(m["hash"])
and not model_settings.is_catalog_filename(m["filename"])
]
if not unrecognized:
print(f"[PromptChain] All {len(models)} model(s) already recognized")
return
print(f"[PromptChain] {len(unrecognized)} unrecognized model(s), starting CivitAI detection...")
recognition_state["running"] = True
recognition_state["total"] = len(unrecognized)
recognition_state["done"] = 0
send_ws("promptchain_recognition_start", {
"total": len(unrecognized),
})
loop = _asyncio.new_event_loop()
try:
for model in unrecognized:
recognition_state["current"] = model["filename"]
print(f"[PromptChain] Detecting: {model['filename']}...")
try:
# Timeout per model: if CivitAI or DNS hangs, don't stall
# the whole boot preload. 30s is generous for a single
# hash lookup + metadata fetch.
result = loop.run_until_complete(
_asyncio.wait_for(
detect_model(model["hash"], model["filepath"], model.get("architecture", "")),
timeout=30,
)
)
except _asyncio.TimeoutError:
print(f"[PromptChain] Timeout detecting {model['filename']}, skipping")
result = {"status": "timeout"}
status = result.get("status", "not_found")
display = result.get("settings", {}).get("display_name", "")
if status == "found":
print(f"[PromptChain] Found: {display} (arch={result['settings'].get('architecture')}, family={result['settings'].get('family', '')})")
else:
print(f"[PromptChain] {status}")
recognition_state["done"] += 1
send_ws("promptchain_model_recognized", {
"hash": model["hash"],
"filename": model["filename"],
"status": status,
"settings": result.get("settings"),
"done": recognition_state["done"],
"total": recognition_state["total"],
})
finally:
loop.close()
recognition_state["running"] = False
recognition_state["current"] = None
print(f"[PromptChain] Recognition complete: {recognition_state['done']}/{recognition_state['total']}")
send_ws("promptchain_recognition_done", {
"total": recognition_state["total"],
"done": recognition_state["done"],
})
_threading.Thread(target=_preload, daemon=True).start()
# ── extension ────────────────────────────────────────────────────
class PromptChainExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [PromptChainNode, PoseStudioNode, AttentionCoupleNode, RegionalDetailerNode, MaskedDetailNode, RegionalConditioningNode, DefocusMaskNode]
async def comfy_entrypoint() -> PromptChainExtension:
return PromptChainExtension()
build-src/build.mjs
+34
View File
@@ -0,0 +1,34 @@
import { build } from "esbuild";
import { resolve, dirname } from "path";
import { fileURLToPath } from "url";
const __dirname = dirname(fileURLToPath(import.meta.url));
const bundles = [
{
entry: "codemirror-entry.mjs",
out: "codemirror.bundle.js",
globalName: "PromptChainCM",
},
{
entry: "three-entry.mjs",
out: "three.bundle.js",
globalName: "PromptChainThree",
},
];
await Promise.all(
bundles.map(({ entry, out, globalName }) =>
build({
entryPoints: [resolve(__dirname, entry)],
bundle: true,
format: "iife",
globalName,
outfile: resolve(__dirname, "..", "js", out),
minify: true,
target: ["es2020"],
}).then(() => {
console.log(`${out} built`);
})
)
);
build-src/codemirror-entry.mjs
+954
View File
@@ -0,0 +1,954 @@
// CodeMirror 6 bundle for PromptChain
// Exports: core CM6 + custom prompt language, theme, linter, keybindings, wildcard system
// -- Core CM6 --
import { EditorView, keymap, Decoration, ViewPlugin, WidgetType } from "@codemirror/view";
import { EditorState, StateField, StateEffect, Compartment } from "@codemirror/state";
import { history, historyKeymap, defaultKeymap } from "@codemirror/commands";
import { bracketMatching } from "@codemirror/language";
import { StreamLanguage, syntaxHighlighting, HighlightStyle } from "@codemirror/language";
import { linter, diagnosticCount, forceLinting } from "@codemirror/lint";
import { autocompletion, startCompletion, acceptCompletion, completionStatus, pickedCompletion } from "@codemirror/autocomplete";
import { tooltips } from "@codemirror/view";
import { tags } from "@lezer/highlight";
// -- Prompt Language (StreamLanguage) --
const promptLanguage = StreamLanguage.define({
startState() {
return { inBlockComment: false, inNegativeRegion: false, inScriptBlock: false };
},
copyState(state) {
return {
inBlockComment: state.inBlockComment,
inNegativeRegion: state.inNegativeRegion,
inScriptBlock: state.inScriptBlock,
};
},
token(stream, state) {
// script blocks
if (state.inScriptBlock && stream.match(/<\/SCRIPT>/i)) {
state.inScriptBlock = false;
return "moduleKeyword";
}
if (state.inScriptBlock) {
if (stream.match(/\{PROMPT\}/i)) return "operator";
stream.next();
return "typeName";
}
if (stream.match(/<SCRIPT>/i)) {
state.inScriptBlock = true;
return "moduleKeyword";
}
// negative prompt region — reset at ::Label:: boundaries so
// per-option negatives don't bleed into the next option
if (state.inNegativeRegion && stream.sol() && stream.match(/::/, false)) {
state.inNegativeRegion = false;
}
if (stream.match("Negative Prompt:")) {
state.inNegativeRegion = true;
return "deleted";
}
if (state.inNegativeRegion) {
stream.next();
return "deleted";
}
// block comments
if (state.inBlockComment) {
while (!stream.eol()) {
if (stream.match("*/")) {
state.inBlockComment = false;
return "comment";
}
stream.next();
}
return "comment";
}
if (stream.match("/*")) {
state.inBlockComment = true;
while (!stream.eol()) {
if (stream.match("*/")) {
state.inBlockComment = false;
return "comment";
}
stream.next();
}
return "comment";
}
// line comments
if (stream.match("//")) {
stream.skipToEnd();
return "comment";
}
// # comments only at line start (with optional leading whitespace),
// matching the Python compiler's strip_comments behavior.
if (stream.match("#", false)) {
const before = stream.string.slice(0, stream.pos);
if (/^\s*$/.test(before)) {
stream.match("#");
stream.skipToEnd();
return "comment";
}
}
// ::labels::
if (stream.match(/::[^:\n]+::/)) return "regexp";
// lora
if (stream.match(/<lora:[^>]+>/i)) return "meta";
// embeddings
if (stream.match(/<embedding:[^>]+>/i) || stream.match(/embedding:\w+/i)) return "className";
// wildcards (supports hierarchical paths like __name/key__)
if (stream.match(/__[a-zA-Z0-9_/.-]+__/)) return "variableName";
// BREAK — case-sensitive, matches the convention used by A1111,
// dfl-clip-with-break, and asagi4 prompt-control encoders.
if (stream.match(/\bBREAK\b/)) return "keyword";
// parentheses (emphasis)
if (stream.peek() === "(") { stream.next(); return "atom"; }
if (stream.peek() === ")") { stream.next(); return "atom"; }
// weight colon + number
if (stream.match(/:(?=[\d.]+\))/)) return "punctuation";
if (stream.match(/[\d.]+(?=\))/)) return "number";
// brackets (de-emphasis)
if (stream.peek() === "[") { stream.next(); return "tagName"; }
if (stream.peek() === "]") { stream.next(); return "tagName"; }
// standalone numbers
if (stream.match(/\b\d+\.?\d*\b/)) return "number";
// operators: comma, pipe, braces
if (stream.match(/[,|{}]/)) return "operator";
stream.next();
return null;
},
});
// -- Prompt Theme --
const promptTheme = EditorView.theme({
"&": {
backgroundColor: "rgba(0, 0, 0, 0.4)",
color: "#ffffffe3",
fontSize: "12px",
borderRadius: "4px",
cursor: "text",
},
".cm-content": {
caretColor: "#fff",
fontFamily: "inherit",
cursor: "text",
},
".cm-cursor": { borderLeftColor: "#fff" },
"&.cm-focused .cm-selectionBackground, .cm-selectionBackground": {
backgroundColor: "rgba(74, 158, 255, 0.3)",
},
".cm-gutters": { display: "none" },
".cm-scroller": {
overflow: "auto",
cursor: "text",
userSelect: "text",
WebkitUserSelect: "text",
},
".cm-line": { cursor: "text" },
}, { dark: true });
// -- Highlight Style --
const highlightStyle = HighlightStyle.define([
{ tag: tags.comment, color: "#6a6a6a", fontStyle: "italic" },
{ tag: tags.keyword, color: "#F92672", fontWeight: "bold" }, // BREAK
{ tag: tags.string, color: "#e4d3ff" }, // weighted
{ tag: tags.operator, color: "#F92672", fontWeight: "bold" }, // comma, pipe
{ tag: tags.meta, color: "#66D9EF" }, // LoRA
{ tag: tags.className, color: "#AE81FF" }, // embeddings
{ tag: tags.variableName, color: "#E6DB74" }, // wildcards
{ tag: tags.atom, color: "#A6E22E" }, // emphasis ()
{ tag: tags.tagName, color: "#FD971F" }, // de-emphasis []
{ tag: tags.number, color: "#FD971F" }, // numbers
{ tag: tags.punctuation, color: "#888888" }, // colon in weights
{ tag: tags.deleted, color: "#e74c3c" }, // negative region
{ tag: tags.moduleKeyword, color: "#BB86FC", fontWeight: "bold" }, // <SCRIPT>
{ tag: tags.typeName, color: "#D4BBFF" }, // script content
{ tag: tags.regexp, color: "#8274d9", fontWeight: "bold" }, // ::Label::
]);
// -- Prompt Linter --
const promptLinter = linter((view) => {
const diagnostics = [];
const doc = view.state.doc.toString();
// invalid weight format (multiple decimals)
for (const match of doc.matchAll(/:\d*\.\.+\d*|\:\d+\.\d+\.\d*/g)) {
diagnostics.push({
from: match.index, to: match.index + match[0].length,
severity: "error", message: "Invalid weight format - multiple decimal points",
});
}
// missing weight value
for (const match of doc.matchAll(/\([^():]+:\s*\)/g)) {
diagnostics.push({
from: match.index, to: match.index + match[0].length,
severity: "error", message: "Missing weight value after colon",
});
}
// unmatched parentheses
let parenDepth = 0;
let lastOpen = -1;
for (let i = 0; i < doc.length; i++) {
if (doc[i] === "(") {
if (parenDepth === 0) lastOpen = i;
parenDepth++;
} else if (doc[i] === ")") {
parenDepth--;
if (parenDepth < 0) {
diagnostics.push({
from: i, to: i + 1,
severity: "error", message: "Unmatched closing parenthesis",
});
parenDepth = 0;
}
}
}
if (parenDepth > 0 && lastOpen >= 0) {
diagnostics.push({
from: lastOpen, to: lastOpen + 1,
severity: "error", message: "Unclosed parenthesis",
});
}
// extreme weights
for (const match of doc.matchAll(/\([^():]+:(-?\d+\.?\d*)\)/g)) {
const weight = parseFloat(match[1]);
if (weight > 5 || weight < -5) {
const colonIdx = match[0].lastIndexOf(":");
const numStart = match.index + colonIdx + 1;
const numEnd = match.index + match[0].length - 1;
diagnostics.push({
from: numStart, to: numEnd,
severity: "warning", message: `Unusual weight value (${weight}) - typically between -2 and 2`,
});
}
}
// missing parentheses around weight
for (const match of doc.matchAll(/(?<!\()(\b\w+):(\d+\.?\d*)(?!\))/g)) {
const before = doc.slice(0, match.index);
const openCount = (before.match(/\(/g) || []).length;
const closeCount = (before.match(/\)/g) || []).length;
if (openCount === closeCount) {
diagnostics.push({
from: match.index, to: match.index + match[0].length,
severity: "warning", message: "Missing parentheses - use (word:1.5) instead of word:1.5",
});
}
}
// regional block ($name { ... }) syntax. Grammar (region-highlight.js):
// $ + a letter-led word name + optional space + a plain `{`, closed by the
// next `}`. Catch the breakages that make a block silently fail to bind.
// space after the $ — "$ mannequin {" is not recognized as a region.
for (const match of doc.matchAll(/\$\s+[A-Za-z][\w-]*\s*\{/g)) {
diagnostics.push({
from: match.index, to: match.index + match[0].length,
severity: "error", message: "No space after $ in a regional block — write $name {",
});
}
// decorative / full-width opening brace instead of a plain `{`.
for (const match of doc.matchAll(/\$\s*[A-Za-z][\w-]*\s*([｛﹛❴⦃])/g)) {
diagnostics.push({
from: match.index, to: match.index + match[0].length,
severity: "error", message: "Use a plain { to open a regional block (not a decorative brace)",
});
}
// unclosed block — a valid opener with no `}` after it (a full-width `}`
// closer trips this too, which is the intended nudge).
for (const match of doc.matchAll(/\$[A-Za-z]\w*\s*\{/g)) {
const openBrace = match.index + match[0].length - 1;
if (doc.indexOf("}", openBrace + 1) === -1) {
diagnostics.push({
from: match.index, to: openBrace + 1,
severity: "error", message: "Unclosed regional block — add a closing }",
});
}
}
// wildcard validation: check __name__ references against API
for (const match of doc.matchAll(/__([a-zA-Z0-9_/.-]+)__/g)) {
const name = match[1];
const cached = getWildcardSync(name);
if (cached === null) {
fetchWildcard(name, view);
} else if (!cached.exists) {
const msg = cached.error === "parse_error"
? `Wildcard "${name}" found but failed to parse (check ${cached.format || "file"} syntax)`
: `Wildcard "${name}" not found in ${cached.folder}`;
diagnostics.push({
from: match.index, to: match.index + match[0].length,
severity: "error", message: msg,
});
}
}
return diagnostics;
}, { delay: 750 });
// -- Lint Styles --
const lintStyles = EditorView.baseTheme({
".cm-diagnostic": {
padding: "3px 6px 3px 8px", marginLeft: "-1px",
display: "block", whiteSpace: "pre-wrap",
},
".cm-diagnostic-error": { borderLeft: "5px solid #e74c3c" },
".cm-diagnostic-warning": { borderLeft: "5px solid #f0ad4e" },
".cm-lintRange-error": {
backgroundImage: "none",
textDecoration: "underline wavy #e74c3c", textDecorationSkipInk: "none",
},
".cm-lintRange-warning": {
backgroundImage: "none",
textDecoration: "underline wavy #f0ad4e", textDecorationSkipInk: "none",
},
".cm-tooltip-lint": {
backgroundColor: "#1e1e1e", border: "1px solid #444",
borderRadius: "4px", padding: "0", zIndex: "100000",
},
});
// -- Wildcard Cache & Linting --
const wildcardCache = new Map();
let wildcardCacheVersionValue = 0;
const wildcardRefreshEffect = StateEffect.define();
const wildcardCacheVersion = StateField.define({
create() { return wildcardCacheVersionValue; },
update(value, tr) {
for (const e of tr.effects) {
if (e.is(wildcardRefreshEffect)) return wildcardCacheVersionValue;
}
return value;
},
});
async function fetchWildcard(name, view) {
if (wildcardCache.has(name)) return wildcardCache.get(name);
try {
const response = await fetch(`/promptchain/wildcard?name=${encodeURIComponent(name)}`);
const data = await response.json();
wildcardCache.set(name, data);
wildcardCacheVersionValue++;
if (view) view.dispatch({ effects: wildcardRefreshEffect.of(null) });
return data;
} catch (e) {
const fallback = { exists: false, name, folder: "wildcards" };
wildcardCache.set(name, fallback);
return fallback;
}
}
function getWildcardSync(name) {
return wildcardCache.get(name) || null;
}
// -- Wildcard List Cache --
let wildcardListCache = null;
let wildcardListFetching = false;
let wildcardListFetchedAt = 0;
const LIST_REFRESH_INTERVAL = 5000; // throttle refreshes to 5s
fetchWildcardList();
async function fetchWildcardList() {
if (wildcardListFetching) return wildcardListCache;
// throttle: skip if fetched recently
if (wildcardListCache && Date.now() - wildcardListFetchedAt < LIST_REFRESH_INTERVAL) {
return wildcardListCache;
}
wildcardListFetching = true;
try {
const response = await fetch("/promptchain/wildcard/list");
const data = await response.json();
const newList = data.wildcards || [];
wildcardListCache = newList;
wildcardListFetchedAt = Date.now();
// merge list data into wildcard cache (don't clear individually-fetched entries)
for (const wc of newList) {
wildcardCache.set(wc.name, {
exists: wc.count > 0,
name: wc.name,
folder: "wildcards",
count: wc.count,
format: wc.format,
sections: wc.sections || [],
error: wc.count > 0 ? null : "parse_error",
});
}
return wildcardListCache;
} catch (e) {
if (!wildcardListCache) wildcardListCache = [];
return wildcardListCache;
} finally {
wildcardListFetching = false;
}
}
// -- Shared helpers --
function isInsideWildcard(textBefore) {
const dunderCount = textBefore.split("__").length - 1;
return dunderCount > 0 && dunderCount % 2 !== 0;
}
function buildSectionOptions(wildcardName, sections, sectionPrefix) {
return sections
.filter(s => s.toLowerCase().startsWith(sectionPrefix))
.map(s => ({
label: s,
apply: wildcardName + "/" + s + "__",
type: "property",
detail: "section",
}));
}
function findMatchingFile(prefix) {
if (!wildcardListCache) return null;
let best = null;
const lowerPrefix = prefix.toLowerCase();
for (const wc of wildcardListCache) {
const namePrefix = wc.name.toLowerCase() + "/";
if (lowerPrefix.startsWith(namePrefix) && (!best || wc.name.length > best.name.length)) {
best = wc;
}
}
return best;
}
// -- Wildcard Autocomplete --
function wildcardCompletionSource(context) {
const pos = context.pos;
const line = context.state.doc.lineAt(pos);
const textBefore = line.text.slice(0, pos - line.from);
if (!isInsideWildcard(textBefore)) return null;
const lastDunder = textBefore.lastIndexOf("__");
const prefix = textBefore.slice(lastDunder + 2);
const from = line.from + lastDunder + 2;
// Section completion (prefix contains /)
if (prefix.includes("/") && wildcardListCache) {
const matched = findMatchingFile(prefix);
if (matched && matched.sections && matched.sections.length > 0) {
const sectionPrefix = prefix.slice(matched.name.length + 1).toLowerCase();
return { from, filter: false, options: buildSectionOptions(matched.name, matched.sections, sectionPrefix) };
}
return null;
}
// File name completion
const list = wildcardListCache;
if (!list) {
fetchWildcardList();
return null;
}
fetchWildcardList(); // background refresh
const lowerPrefix = prefix.toLowerCase();
// show inline sections only when narrowed to one file, to avoid flooding the dropdown
const matchCount = lowerPrefix
? list.filter(w => w.name.toLowerCase().startsWith(lowerPrefix)).length
: list.length;
const options = [];
for (const wc of list) {
const name = wc.name;
if (wc.count === 0 && (!wc.sections || wc.sections.length === 0)) continue;
if (lowerPrefix && !name.toLowerCase().startsWith(lowerPrefix)) continue;
const hasSections = wc.sections && wc.sections.length > 0;
if (hasSections) {
options.push({
label: name + "/",
apply: (view, _completion, from, to) => {
const insert = name + "/";
view.dispatch({
changes: { from, to, insert },
selection: { anchor: from + insert.length },
});
setTimeout(() => startCompletion(view), 50);
},
type: "variable",
detail: `${wc.format} · ${wc.sections.length} sections`,
});
} else {
options.push({
label: name,
apply: name + "__",
type: "variable",
detail: `${wc.format} · ${wc.count} options`,
});
}
// show subsections only when narrowed to a single file
if (lowerPrefix && hasSections && matchCount === 1) {
for (const section of wc.sections) {
const fullName = name + "/" + section;
if (!fullName.toLowerCase().startsWith(lowerPrefix)) continue;
options.push({
label: fullName,
apply: fullName + "__",
type: "property",
detail: "section",
});
}
}
}
options.sort((a, b) => a.label.localeCompare(b.label));
return { from, filter: false, options };
}
const wildcardAutocomplete = autocompletion({
override: [wildcardCompletionSource],
activateOnTyping: true,
maxRenderedOptions: 25,
icons: false,
});
// Compartment allows reconfiguring the autocomplete after editor creation
// to add additional sources (e.g. tag DB) alongside wildcards.
const autocompleteCompartment = new Compartment();
function reconfigureAutocomplete(view, sources) {
view.dispatch({
effects: autocompleteCompartment.reconfigure(
autocompletion({
override: sources,
activateOnTyping: true,
maxRenderedOptions: 25,
icons: false,
})
),
});
}
// Re-trigger autocomplete on backspace when inside a __wildcard__ context
const autocompleteOnDelete = EditorView.updateListener.of((update) => {
if (!update.docChanged) return;
if (update.changes.newLength >= update.changes.length) return;
const pos = update.state.selection.main.head;
const line = update.state.doc.lineAt(pos);
const before = line.text.slice(0, pos - line.from);
if (isInsideWildcard(before)) {
setTimeout(() => startCompletion(update.view), 50);
}
});
// -- Autocomplete Theme --
const autocompleteTheme = EditorView.baseTheme({
".cm-tooltip-autocomplete": {
backgroundColor: "#1e1e1e !important",
border: "1px solid #444 !important",
borderRadius: "4px !important",
boxShadow: "0 4px 12px rgba(0,0,0,0.5) !important",
zIndex: "100000 !important",
maxHeight: "300px !important",
overflow: "hidden !important",
pointerEvents: "auto !important",
},
".cm-tooltip-autocomplete > ul": {
fontFamily: "system-ui, sans-serif !important",
fontSize: "12px !important",
maxHeight: "300px !important",
overflowY: "auto !important",
pointerEvents: "auto !important",
},
".cm-tooltip-autocomplete > ul > li": {
padding: "4px 8px !important",
borderBottom: "1px solid #333 !important",
pointerEvents: "auto !important",
cursor: "pointer !important",
},
".cm-tooltip-autocomplete > ul > li:last-child": {
borderBottom: "none !important",
},
".cm-tooltip-autocomplete > ul > li[aria-selected]": {
backgroundColor: "#2d4f7c !important",
color: "#fff !important",
},
".cm-completionLabel": {
color: "#e0e0e0 !important",
},
".cm-completionDetail": {
color: "#888 !important",
marginLeft: "8px !important",
fontSize: "11px !important",
},
});
// -- Keybindings --
const INDENT = " ";
function indentLines(view) {
const { state } = view;
const sel = state.selection.main;
if (sel.from === sel.to) {
view.dispatch(state.replaceSelection(INDENT));
} else {
const from = state.doc.lineAt(sel.from).from;
const to = state.doc.lineAt(sel.to).to;
let changes = [];
for (let pos = from; pos <= to; ) {
const line = state.doc.lineAt(pos);
changes.push({ from: line.from, insert: INDENT });
pos = line.to + 1;
}
view.dispatch({ changes });
}
return true;
}
function dedentLines(view) {
const { state } = view;
const sel = state.selection.main;
const from = state.doc.lineAt(sel.from).from;
const to = state.doc.lineAt(sel.to).to;
let changes = [];
for (let pos = from; pos <= to; ) {
const line = state.doc.lineAt(pos);
let removeCount = 0;
for (let i = 0; i < INDENT.length && i < line.text.length; i++) {
if (line.text[i] === " ") removeCount++;
else break;
}
if (removeCount > 0) {
changes.push({ from: line.from, to: line.from + removeCount, insert: "" });
}
pos = line.to + 1;
}
if (changes.length > 0) view.dispatch({ changes });
return true;
}
function formatPromptText(text) {
// Label-aware formatting: ::Label:: lines are kept flush-left and the
// content below each label is indented one INDENT level. Sections are
// separated by blank lines. `//` line-comments directly above a label
// (no blank line between) are treated as that label's preamble and
// stay flush-left with it. If the text has no labels this falls
// through to the single-block branch below.
const lines = text.split("\n");
// Match ::Label:: at line start (allowing leading whitespace) with any
// inline body after it on the same line. The inline body is prepended
// to the label's body lines so `::Label::inline` formats the same as
// `::Label::\ninline`.
const labelRE = /^\s*(::[^:]+::)(.*)$/;
const commentRE = /^\s*\/\//;
const labelMatches = [];
for (let i = 0; i < lines.length; i++) {
const m = lines[i].match(labelRE);
if (m) labelMatches.push({ idx: i, label: m[1], inline: m[2] });
}
if (labelMatches.length === 0) return formatPromptBlock(text);
// For each label, walk upward over contiguous comment lines with no
// blank line between — those attach as the label's preamble.
const sectionStarts = labelMatches.map(({ idx }) => {
let start = idx;
while (start > 0 && commentRE.test(lines[start - 1])) start--;
return start;
});
const sections = [];
if (sectionStarts[0] > 0) {
sections.push({ kind: "prelude", lines: lines.slice(0, sectionStarts[0]) });
}
for (let s = 0; s < labelMatches.length; s++) {
const { idx: labelIdx, label, inline } = labelMatches[s];
const preamble = lines.slice(sectionStarts[s], labelIdx);
const bodyEnd = s + 1 < labelMatches.length ? sectionStarts[s + 1] : lines.length;
const body = lines.slice(labelIdx + 1, bodyEnd);
if (inline.trim()) body.unshift(inline);
sections.push({ kind: "label", preamble, label, body });
}
return sections.map(sec => {
if (sec.kind === "prelude") {
return sec.lines.join("\n").replace(/^\s+|\s+$/g, "");
}
const parts = [];
if (sec.preamble.length) parts.push(sec.preamble.map(l => l.trim()).join("\n"));
parts.push(sec.label);
const rawBody = sec.body.join("\n").trim();
if (rawBody) {
const formatted = formatPromptBlock(rawBody);
const indented = formatted
.split("\n")
.map(l => (l === "" ? "" : INDENT + l))
.join("\n");
parts.push(indented);
}
return parts.join("\n");
}).filter(s => s).join("\n\n");
}
function formatPromptBlock(text) {
if (!text.includes("{")) {
return text.split("\n").map((line) => line.trimStart()).join("\n").trim();
}
text = text.replace(/[^\S\n]+/g, " ");
text = text.replace(/\n\s*\n/g, "\n");
text = text.replace(/\s+/g, " ").trim();
// Returns true if string contains nested braces like {a|{b|c}}
const hasNestedBrackets = (str) => {
let depth = 0;
for (const char of str) {
if (char === "{") { depth++; if (depth > 0) return true; }
else if (char === "}") depth--;
}
return false;
};
const splitByPipes = (str) => {
let depth = 0, options = [], start = 0;
for (let i = 0; i < str.length; i++) {
if (str[i] === "{") depth++;
else if (str[i] === "}") depth--;
else if (str[i] === "|" && depth === 0) {
options.push(str.slice(start, i).trim());
start = i + 1;
}
}
options.push(str.slice(start).trim());
return options;
};
const formatWithDepth = (str, depth) => {
let result = "", i = 0;
while (i < str.length) {
if (str[i] === "{") {
let braceDepth = 1, j = i + 1;
while (j < str.length && braceDepth > 0) {
if (str[j] === "{") braceDepth++;
else if (str[j] === "}") braceDepth--;
j++;
}
const content = str.slice(i + 1, j - 1);
const options = splitByPipes(content);
const nested = hasNestedBrackets(content);
if (nested || options.length > 3) {
result += "{\n";
for (let k = 0; k < options.length; k++) {
const prefix = k === 0 ? "" : "| ";
result += INDENT.repeat(depth + 1) + prefix + formatWithDepth(options[k].trim(), depth + 1) + "\n";
}
result += INDENT.repeat(depth) + "}";
} else {
result += "{" + content + "}";
}
i = j;
} else {
result += str[i];
i++;
}
}
return result;
};
return formatWithDepth(text, 0);
}
// Regional blocks ($name { ... }) format to the house style: the `$name {`
// opener and `}` closer each on their own line, the body indented one INDENT
// level. Blocks are pulled out first so the wildcard-brace formatter never
// touches a region's `{`; the text around them formats as a normal prompt.
function formatRegionInner(inner) {
const lines = inner.split("\n").map((l) => l.trim().replace(/[ \t]{2,}/g, " "));
const out = [];
for (const l of lines) {
if (l === "" && (out.length === 0 || out[out.length - 1] === "")) continue;
out.push(l);
}
while (out.length && out[out.length - 1] === "") out.pop();
return out.map((l) => (l === "" ? "" : INDENT + l)).join("\n");
}
function formatWithRegions(text) {
const OPEN = /\$[A-Za-z]\w*\s*\{/g;
let result = "";
let last = 0;
let m;
while ((m = OPEN.exec(text)) !== null) {
const openBrace = m.index + m[0].length - 1;
const close = text.indexOf("}", openBrace + 1);
if (close === -1) continue; // unterminated — leave it for the linter to flag
const name = m[0].match(/^\$[A-Za-z]\w*/)[0];
const before = text.slice(last, m.index);
if (before.trim()) result += formatPromptText(before).trim() + "\n\n";
const inner = text.slice(openBrace + 1, close);
result += `${name} {\n${formatRegionInner(inner)}\n}\n\n`;
last = close + 1;
OPEN.lastIndex = close + 1;
}
const tail = text.slice(last);
if (tail.trim()) result += formatPromptText(tail).trim();
return result.trim();
}
function formatPrompt(view) {
const text = view.state.doc.toString();
const formatted = formatWithRegions(text);
if (formatted !== text) {
view.dispatch({ changes: { from: 0, to: text.length, insert: formatted } });
}
return true;
}
// -- Indent-aware word wrap --
// Wrapped continuations align to the line's leading whitespace.
const indentWrap = ViewPlugin.fromClass(class {
constructor(view) { this.decorations = this.build(view); }
update(update) {
if (update.docChanged || update.viewportChanged) this.decorations = this.build(update.view);
}
build(view) {
const decos = [];
for (const { from, to } of view.visibleRanges) {
let pos = from;
while (pos <= to) {
const line = view.state.doc.lineAt(pos);
const match = line.text.match(/^(\s+)/);
if (match) {
const indent = match[1].length;
decos.push(Decoration.line({
attributes: { style: `padding-left: ${indent}ch; text-indent: -${indent}ch;` }
}).range(line.from));
}
pos = line.to + 1;
}
}
return Decoration.set(decos);
}
}, { decorations: v => v.decorations });
// -- Minimal Setup --
const minimalSetup = [
EditorView.lineWrapping,
indentWrap,
EditorView.contentAttributes.of({ draggable: "false" }),
bracketMatching(),
history(),
wildcardCacheVersion,
promptLinter,
lintStyles,
autocompleteCompartment.of(wildcardAutocomplete),
autocompleteOnDelete,
autocompleteTheme,
tooltips({
position: "fixed",
}),
keymap.of([
// Custom bindings come BEFORE defaultKeymap because CM6 takes the first
// matching keymap entry. defaultKeymap binds Mod-Enter (insertBlankLine)
// and Tab (insertTab), which would otherwise win over our queue and
// accept-completion bindings.
// Tab accepts the active completion if one is open (matches editor
// conventions like VS Code), otherwise indents.
{ key: "Tab", run: (view) => completionStatus(view.state) === "active" ? acceptCompletion(view) : indentLines(view) },
{ key: "Shift-Tab", run: dedentLines },
{ key: "Mod-]", run: indentLines },
{ key: "Mod-[", run: dedentLines },
{ key: "Shift-Alt-f", run: formatPrompt },
// Ctrl/Cmd+S: route to ComfyUI's save command instead of letting the
// browser's "Save Page As" dialog through. ComfyUI's keybindingService
// bails on contenteditable targets, so we own the binding here.
{
key: "Mod-s",
preventDefault: true,
run: () => {
window.app?.extensionManager?.command?.execute?.("Comfy.SaveWorkflow");
return true;
},
},
// Ctrl/Cmd+Shift+Enter: queue at front. Listed before Mod-Enter so the
// more specific binding wins.
{
key: "Mod-Shift-Enter",
preventDefault: true,
run: () => {
window.app?.extensionManager?.command?.execute?.("Comfy.QueuePromptFront");
return true;
},
},
// Ctrl/Cmd+Enter: queue the workflow from inside the editor. Same
// contenteditable-bypass reason as Mod-s above.
{
key: "Mod-Enter",
preventDefault: true,
run: () => {
window.app?.extensionManager?.command?.execute?.("Comfy.QueuePrompt");
return true;
},
},
// Ctrl/Cmd+Alt+Enter: interrupt the current generation.
{
key: "Mod-Alt-Enter",
preventDefault: true,
run: () => {
window.app?.extensionManager?.command?.execute?.("Comfy.Interrupt");
return true;
},
},
...defaultKeymap,
...historyKeymap,
]),
];
// -- Exports --
export {
// core
EditorView,
EditorState,
minimalSetup,
// language/syntax
promptLanguage,
promptTheme,
highlightStyle,
syntaxHighlighting,
// linting
promptLinter,
lintStyles,
diagnosticCount,
forceLinting,
// autocomplete
autocompletion,
wildcardCompletionSource,
reconfigureAutocomplete,
startCompletion,
pickedCompletion,
// wildcard system
wildcardCacheVersion,
wildcardRefreshEffect,
wildcardAutocomplete,
fetchWildcard,
getWildcardSync,
fetchWildcardList,
// state/view primitives (for extensions)
Decoration,
ViewPlugin,
StateField,
StateEffect,
WidgetType,
Compartment,
// utilities
formatPromptText,
};
build-src/package-lock.json
Generated
+627
View File
@@ -0,0 +1,627 @@
{
"name": "promptchain-codemirror",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "promptchain-codemirror",
"dependencies": {
"@codemirror/autocomplete": "^6.20.1",
"@codemirror/commands": "^6.10.3",
"@codemirror/language": "^6.12.2",
"@codemirror/lint": "^6.9.5",
"@codemirror/search": "^6.6.0",
"@codemirror/state": "^6.6.0",
"@codemirror/theme-one-dark": "^6.1.3",
"@codemirror/view": "^6.40.0",
"@lezer/highlight": "^1.2.3",
"esbuild": "^0.27.4",
"three": "^0.169.0"
}
},
"node_modules/@codemirror/autocomplete": {
"version": "6.20.1",
"resolved": "https://registry.npmjs.org/@codemirror/autocomplete/-/autocomplete-6.20.1.tgz",
"integrity": "sha512-1cvg3Vz1dSSToCNlJfRA2WSI4ht3K+WplO0UMOgmUYPivCyy2oueZY6Lx7M9wThm7SDUBViRmuT+OG/i8+ON9A==",
"license": "MIT",
"dependencies": {
"@codemirror/language": "^6.0.0",
"@codemirror/state": "^6.0.0",
"@codemirror/view": "^6.17.0",
"@lezer/common": "^1.0.0"
}
},
"node_modules/@codemirror/commands": {
"version": "6.10.3",
"resolved": "https://registry.npmjs.org/@codemirror/commands/-/commands-6.10.3.tgz",
"integrity": "sha512-JFRiqhKu+bvSkDLI+rUhJwSxQxYb759W5GBezE8Uc8mHLqC9aV/9aTC7yJSqCtB3F00pylrLCwnyS91Ap5ej4Q==",
"license": "MIT",
"dependencies": {
"@codemirror/language": "^6.0.0",
"@codemirror/state": "^6.6.0",
"@codemirror/view": "^6.27.0",
"@lezer/common": "^1.1.0"
}
},
"node_modules/@codemirror/language": {
"version": "6.12.2",
"resolved": "https://registry.npmjs.org/@codemirror/language/-/language-6.12.2.tgz",
"integrity": "sha512-jEPmz2nGGDxhRTg3lTpzmIyGKxz3Gp3SJES4b0nAuE5SWQoKdT5GoQ69cwMmFd+wvFUhYirtDTr0/DRHpQAyWg==",
"license": "MIT",
"dependencies": {
"@codemirror/state": "^6.0.0",
"@codemirror/view": "^6.23.0",
"@lezer/common": "^1.5.0",
"@lezer/highlight": "^1.0.0",
"@lezer/lr": "^1.0.0",
"style-mod": "^4.0.0"
}
},
"node_modules/@codemirror/lint": {
"version": "6.9.5",
"resolved": "https://registry.npmjs.org/@codemirror/lint/-/lint-6.9.5.tgz",
"integrity": "sha512-GElsbU9G7QT9xXhpUg1zWGmftA/7jamh+7+ydKRuT0ORpWS3wOSP0yT1FOlIZa7mIJjpVPipErsyvVqB9cfTFA==",
"license": "MIT",
"dependencies": {
"@codemirror/state": "^6.0.0",
"@codemirror/view": "^6.35.0",
"crelt": "^1.0.5"
}
},
"node_modules/@codemirror/search": {
"version": "6.6.0",
"resolved": "https://registry.npmjs.org/@codemirror/search/-/search-6.6.0.tgz",
"integrity": "sha512-koFuNXcDvyyotWcgOnZGmY7LZqEOXZaaxD/j6n18TCLx2/9HieZJ5H6hs1g8FiRxBD0DNfs0nXn17g872RmYdw==",
"license": "MIT",
"dependencies": {
"@codemirror/state": "^6.0.0",
"@codemirror/view": "^6.37.0",
"crelt": "^1.0.5"
}
},
"node_modules/@codemirror/state": {
"version": "6.6.0",
"resolved": "https://registry.npmjs.org/@codemirror/state/-/state-6.6.0.tgz",
"integrity": "sha512-4nbvra5R5EtiCzr9BTHiTLc+MLXK2QGiAVYMyi8PkQd3SR+6ixar/Q/01Fa21TBIDOZXgeWV4WppsQolSreAPQ==",
"license": "MIT",
"dependencies": {
"@marijn/find-cluster-break": "^1.0.0"
}
},
"node_modules/@codemirror/theme-one-dark": {
"version": "6.1.3",
"resolved": "https://registry.npmjs.org/@codemirror/theme-one-dark/-/theme-one-dark-6.1.3.tgz",
"integrity": "sha512-NzBdIvEJmx6fjeremiGp3t/okrLPYT0d9orIc7AFun8oZcRk58aejkqhv6spnz4MLAevrKNPMQYXEWMg4s+sKA==",
"license": "MIT",
"dependencies": {
"@codemirror/language": "^6.0.0",
"@codemirror/state": "^6.0.0",
"@codemirror/view": "^6.0.0",
"@lezer/highlight": "^1.0.0"
}
},
"node_modules/@codemirror/view": {
"version": "6.40.0",
"resolved": "https://registry.npmjs.org/@codemirror/view/-/view-6.40.0.tgz",
"integrity": "sha512-WA0zdU7xfF10+5I3HhUUq3kqOx3KjqmtQ9lqZjfK7jtYk4G72YW9rezcSywpaUMCWOMlq+6E0pO1IWg1TNIhtg==",
"license": "MIT",
"dependencies": {
"@codemirror/state": "^6.6.0",
"crelt": "^1.0.6",
"style-mod": "^4.1.0",
"w3c-keyname": "^2.2.4"
}
},
"node_modules/@esbuild/aix-ppc64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/aix-ppc64/-/aix-ppc64-0.27.4.tgz",
"integrity": "sha512-cQPwL2mp2nSmHHJlCyoXgHGhbEPMrEEU5xhkcy3Hs/O7nGZqEpZ2sUtLaL9MORLtDfRvVl2/3PAuEkYZH0Ty8Q==",
"cpu": [
"ppc64"
],
"license": "MIT",
"optional": true,
"os": [
"aix"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/android-arm": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/android-arm/-/android-arm-0.27.4.tgz",
"integrity": "sha512-X9bUgvxiC8CHAGKYufLIHGXPJWnr0OCdR0anD2e21vdvgCI8lIfqFbnoeOz7lBjdrAGUhqLZLcQo6MLhTO2DKQ==",
"cpu": [
"arm"
],
"license": "MIT",
"optional": true,
"os": [
"android"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/android-arm64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/android-arm64/-/android-arm64-0.27.4.tgz",
"integrity": "sha512-gdLscB7v75wRfu7QSm/zg6Rx29VLdy9eTr2t44sfTW7CxwAtQghZ4ZnqHk3/ogz7xao0QAgrkradbBzcqFPasw==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"android"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/android-x64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/android-x64/-/android-x64-0.27.4.tgz",
"integrity": "sha512-PzPFnBNVF292sfpfhiyiXCGSn9HZg5BcAz+ivBuSsl6Rk4ga1oEXAamhOXRFyMcjwr2DVtm40G65N3GLeH1Lvw==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"android"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/darwin-arm64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/darwin-arm64/-/darwin-arm64-0.27.4.tgz",
"integrity": "sha512-b7xaGIwdJlht8ZFCvMkpDN6uiSmnxxK56N2GDTMYPr2/gzvfdQN8rTfBsvVKmIVY/X7EM+/hJKEIbbHs9oA4tQ==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/darwin-x64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/darwin-x64/-/darwin-x64-0.27.4.tgz",
"integrity": "sha512-sR+OiKLwd15nmCdqpXMnuJ9W2kpy0KigzqScqHI3Hqwr7IXxBp3Yva+yJwoqh7rE8V77tdoheRYataNKL4QrPw==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"darwin"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/freebsd-arm64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/freebsd-arm64/-/freebsd-arm64-0.27.4.tgz",
"integrity": "sha512-jnfpKe+p79tCnm4GVav68A7tUFeKQwQyLgESwEAUzyxk/TJr4QdGog9sqWNcUbr/bZt/O/HXouspuQDd9JxFSw==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"freebsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/freebsd-x64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/freebsd-x64/-/freebsd-x64-0.27.4.tgz",
"integrity": "sha512-2kb4ceA/CpfUrIcTUl1wrP/9ad9Atrp5J94Lq69w7UwOMolPIGrfLSvAKJp0RTvkPPyn6CIWrNy13kyLikZRZQ==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"freebsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-arm": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/linux-arm/-/linux-arm-0.27.4.tgz",
"integrity": "sha512-aBYgcIxX/wd5n2ys0yESGeYMGF+pv6g0DhZr3G1ZG4jMfruU9Tl1i2Z+Wnj9/KjGz1lTLCcorqE2viePZqj4Eg==",
"cpu": [
"arm"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-arm64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/linux-arm64/-/linux-arm64-0.27.4.tgz",
"integrity": "sha512-7nQOttdzVGth1iz57kxg9uCz57dxQLHWxopL6mYuYthohPKEK0vU0C3O21CcBK6KDlkYVcnDXY099HcCDXd9dA==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-ia32": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/linux-ia32/-/linux-ia32-0.27.4.tgz",
"integrity": "sha512-oPtixtAIzgvzYcKBQM/qZ3R+9TEUd1aNJQu0HhGyqtx6oS7qTpvjheIWBbes4+qu1bNlo2V4cbkISr8q6gRBFA==",
"cpu": [
"ia32"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-loong64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/linux-loong64/-/linux-loong64-0.27.4.tgz",
"integrity": "sha512-8mL/vh8qeCoRcFH2nM8wm5uJP+ZcVYGGayMavi8GmRJjuI3g1v6Z7Ni0JJKAJW+m0EtUuARb6Lmp4hMjzCBWzA==",
"cpu": [
"loong64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-mips64el": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/linux-mips64el/-/linux-mips64el-0.27.4.tgz",
"integrity": "sha512-1RdrWFFiiLIW7LQq9Q2NES+HiD4NyT8Itj9AUeCl0IVCA459WnPhREKgwrpaIfTOe+/2rdntisegiPWn/r/aAw==",
"cpu": [
"mips64el"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-ppc64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/linux-ppc64/-/linux-ppc64-0.27.4.tgz",
"integrity": "sha512-tLCwNG47l3sd9lpfyx9LAGEGItCUeRCWeAx6x2Jmbav65nAwoPXfewtAdtbtit/pJFLUWOhpv0FpS6GQAmPrHA==",
"cpu": [
"ppc64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-riscv64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/linux-riscv64/-/linux-riscv64-0.27.4.tgz",
"integrity": "sha512-BnASypppbUWyqjd1KIpU4AUBiIhVr6YlHx/cnPgqEkNoVOhHg+YiSVxM1RLfiy4t9cAulbRGTNCKOcqHrEQLIw==",
"cpu": [
"riscv64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-s390x": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/linux-s390x/-/linux-s390x-0.27.4.tgz",
"integrity": "sha512-+eUqgb/Z7vxVLezG8bVB9SfBie89gMueS+I0xYh2tJdw3vqA/0ImZJ2ROeWwVJN59ihBeZ7Tu92dF/5dy5FttA==",
"cpu": [
"s390x"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/linux-x64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/linux-x64/-/linux-x64-0.27.4.tgz",
"integrity": "sha512-S5qOXrKV8BQEzJPVxAwnryi2+Iq5pB40gTEIT69BQONqR7JH1EPIcQ/Uiv9mCnn05jff9umq/5nqzxlqTOg9NA==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"linux"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/netbsd-arm64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/netbsd-arm64/-/netbsd-arm64-0.27.4.tgz",
"integrity": "sha512-xHT8X4sb0GS8qTqiwzHqpY00C95DPAq7nAwX35Ie/s+LO9830hrMd3oX0ZMKLvy7vsonee73x0lmcdOVXFzd6Q==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"netbsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/netbsd-x64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/netbsd-x64/-/netbsd-x64-0.27.4.tgz",
"integrity": "sha512-RugOvOdXfdyi5Tyv40kgQnI0byv66BFgAqjdgtAKqHoZTbTF2QqfQrFwa7cHEORJf6X2ht+l9ABLMP0dnKYsgg==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"netbsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/openbsd-arm64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/openbsd-arm64/-/openbsd-arm64-0.27.4.tgz",
"integrity": "sha512-2MyL3IAaTX+1/qP0O1SwskwcwCoOI4kV2IBX1xYnDDqthmq5ArrW94qSIKCAuRraMgPOmG0RDTA74mzYNQA9ow==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"openbsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/openbsd-x64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/openbsd-x64/-/openbsd-x64-0.27.4.tgz",
"integrity": "sha512-u8fg/jQ5aQDfsnIV6+KwLOf1CmJnfu1ShpwqdwC0uA7ZPwFws55Ngc12vBdeUdnuWoQYx/SOQLGDcdlfXhYmXQ==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"openbsd"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/openharmony-arm64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/openharmony-arm64/-/openharmony-arm64-0.27.4.tgz",
"integrity": "sha512-JkTZrl6VbyO8lDQO3yv26nNr2RM2yZzNrNHEsj9bm6dOwwu9OYN28CjzZkH57bh4w0I2F7IodpQvUAEd1mbWXg==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"openharmony"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/sunos-x64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/sunos-x64/-/sunos-x64-0.27.4.tgz",
"integrity": "sha512-/gOzgaewZJfeJTlsWhvUEmUG4tWEY2Spp5M20INYRg2ZKl9QPO3QEEgPeRtLjEWSW8FilRNacPOg8R1uaYkA6g==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"sunos"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/win32-arm64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/win32-arm64/-/win32-arm64-0.27.4.tgz",
"integrity": "sha512-Z9SExBg2y32smoDQdf1HRwHRt6vAHLXcxD2uGgO/v2jK7Y718Ix4ndsbNMU/+1Qiem9OiOdaqitioZwxivhXYg==",
"cpu": [
"arm64"
],
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/win32-ia32": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/win32-ia32/-/win32-ia32-0.27.4.tgz",
"integrity": "sha512-DAyGLS0Jz5G5iixEbMHi5KdiApqHBWMGzTtMiJ72ZOLhbu/bzxgAe8Ue8CTS3n3HbIUHQz/L51yMdGMeoxXNJw==",
"cpu": [
"ia32"
],
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@esbuild/win32-x64": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/@esbuild/win32-x64/-/win32-x64-0.27.4.tgz",
"integrity": "sha512-+knoa0BDoeXgkNvvV1vvbZX4+hizelrkwmGJBdT17t8FNPwG2lKemmuMZlmaNQ3ws3DKKCxpb4zRZEIp3UxFCg==",
"cpu": [
"x64"
],
"license": "MIT",
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">=18"
}
},
"node_modules/@lezer/common": {
"version": "1.5.1",
"resolved": "https://registry.npmjs.org/@lezer/common/-/common-1.5.1.tgz",
"integrity": "sha512-6YRVG9vBkaY7p1IVxL4s44n5nUnaNnGM2/AckNgYOnxTG2kWh1vR8BMxPseWPjRNpb5VtXnMpeYAEAADoRV1Iw==",
"license": "MIT"
},
"node_modules/@lezer/highlight": {
"version": "1.2.3",
"resolved": "https://registry.npmjs.org/@lezer/highlight/-/highlight-1.2.3.tgz",
"integrity": "sha512-qXdH7UqTvGfdVBINrgKhDsVTJTxactNNxLk7+UMwZhU13lMHaOBlJe9Vqp907ya56Y3+ed2tlqzys7jDkTmW0g==",
"license": "MIT",
"dependencies": {
"@lezer/common": "^1.3.0"
}
},
"node_modules/@lezer/lr": {
"version": "1.4.8",
"resolved": "https://registry.npmjs.org/@lezer/lr/-/lr-1.4.8.tgz",
"integrity": "sha512-bPWa0Pgx69ylNlMlPvBPryqeLYQjyJjqPx+Aupm5zydLIF3NE+6MMLT8Yi23Bd9cif9VS00aUebn+6fDIGBcDA==",
"license": "MIT",
"dependencies": {
"@lezer/common": "^1.0.0"
}
},
"node_modules/@marijn/find-cluster-break": {
"version": "1.0.2",
"resolved": "https://registry.npmjs.org/@marijn/find-cluster-break/-/find-cluster-break-1.0.2.tgz",
"integrity": "sha512-l0h88YhZFyKdXIFNfSWpyjStDjGHwZ/U7iobcK1cQQD8sejsONdQtTVU+1wVN1PBw40PiiHB1vA5S7VTfQiP9g==",
"license": "MIT"
},
"node_modules/crelt": {
"version": "1.0.6",
"resolved": "https://registry.npmjs.org/crelt/-/crelt-1.0.6.tgz",
"integrity": "sha512-VQ2MBenTq1fWZUH9DJNGti7kKv6EeAuYr3cLwxUWhIu1baTaXh4Ib5W2CqHVqib4/MqbYGJqiL3Zb8GJZr3l4g==",
"license": "MIT"
},
"node_modules/esbuild": {
"version": "0.27.4",
"resolved": "https://registry.npmjs.org/esbuild/-/esbuild-0.27.4.tgz",
"integrity": "sha512-Rq4vbHnYkK5fws5NF7MYTU68FPRE1ajX7heQ/8QXXWqNgqqJ/GkmmyxIzUnf2Sr/bakf8l54716CcMGHYhMrrQ==",
"hasInstallScript": true,
"license": "MIT",
"bin": {
"esbuild": "bin/esbuild"
},
"engines": {
"node": ">=18"
},
"optionalDependencies": {
"@esbuild/aix-ppc64": "0.27.4",
"@esbuild/android-arm": "0.27.4",
"@esbuild/android-arm64": "0.27.4",
"@esbuild/android-x64": "0.27.4",
"@esbuild/darwin-arm64": "0.27.4",
"@esbuild/darwin-x64": "0.27.4",
"@esbuild/freebsd-arm64": "0.27.4",
"@esbuild/freebsd-x64": "0.27.4",
"@esbuild/linux-arm": "0.27.4",
"@esbuild/linux-arm64": "0.27.4",
"@esbuild/linux-ia32": "0.27.4",
"@esbuild/linux-loong64": "0.27.4",
"@esbuild/linux-mips64el": "0.27.4",
"@esbuild/linux-ppc64": "0.27.4",
"@esbuild/linux-riscv64": "0.27.4",
"@esbuild/linux-s390x": "0.27.4",
"@esbuild/linux-x64": "0.27.4",
"@esbuild/netbsd-arm64": "0.27.4",
"@esbuild/netbsd-x64": "0.27.4",
"@esbuild/openbsd-arm64": "0.27.4",
"@esbuild/openbsd-x64": "0.27.4",
"@esbuild/openharmony-arm64": "0.27.4",
"@esbuild/sunos-x64": "0.27.4",
"@esbuild/win32-arm64": "0.27.4",
"@esbuild/win32-ia32": "0.27.4",
"@esbuild/win32-x64": "0.27.4"
}
},
"node_modules/style-mod": {
"version": "4.1.3",
"resolved": "https://registry.npmjs.org/style-mod/-/style-mod-4.1.3.tgz",
"integrity": "sha512-i/n8VsZydrugj3Iuzll8+x/00GH2vnYsk1eomD8QiRrSAeW6ItbCQDtfXCeJHd0iwiNagqjQkvpvREEPtW3IoQ==",
"license": "MIT"
},
"node_modules/three": {
"version": "0.169.0",
"resolved": "https://registry.npmjs.org/three/-/three-0.169.0.tgz",
"integrity": "sha512-Ed906MA3dR4TS5riErd4QBsRGPcx+HBDX2O5yYE5GqJeFQTPU+M56Va/f/Oph9X7uZo3W3o4l2ZhBZ6f6qUv0w==",
"license": "MIT"
},
"node_modules/w3c-keyname": {
"version": "2.2.8",
"resolved": "https://registry.npmjs.org/w3c-keyname/-/w3c-keyname-2.2.8.tgz",
"integrity": "sha512-dpojBhNsCNN7T82Tm7k26A6G9ML3NkhDsnw9n/eoxSRlVBB4CEtIQ/KTCLI2Fwf3ataSXRhYFkQi3SlnFwPvPQ==",
"license": "MIT"
}
}
}
build-src/package.json
+21
View File
@@ -0,0 +1,21 @@
{
"name": "promptchain-codemirror",
"private": true,
"type": "module",
"scripts": {
"build": "node build.mjs"
},
"dependencies": {
"@codemirror/autocomplete": "^6.20.1",
"@codemirror/commands": "^6.10.3",
"@codemirror/language": "^6.12.2",
"@codemirror/lint": "^6.9.5",
"@codemirror/search": "^6.6.0",
"@codemirror/state": "^6.6.0",
"@codemirror/theme-one-dark": "^6.1.3",
"@codemirror/view": "^6.40.0",
"@lezer/highlight": "^1.2.3",
"esbuild": "^0.27.4",
"three": "^0.169.0"
}
}
build-src/three-entry.mjs
+12
View File
@@ -0,0 +1,12 @@
// Three.js bundle for PromptChain Pose Studio.
// Built to js/three.bundle.js as an IIFE exposing window.PromptChainThree,
// lazy-loaded only when a Pose Studio node mounts (Three is ~600KB — it must
// stay out of the main extension bundle). Mirrors codemirror.bundle.js.
import * as THREE from "three";
import { OrbitControls } from "three/addons/controls/OrbitControls.js";
import { TransformControls } from "three/addons/controls/TransformControls.js";
import { GLTFLoader } from "three/addons/loaders/GLTFLoader.js";
import { OBJLoader } from "three/addons/loaders/OBJLoader.js";
export { THREE, OrbitControls, TransformControls, GLTFLoader, OBJLoader };
bundled_packs/ComfyUI-PuLID-Flux2/LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2026 iFayens
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
bundled_packs/ComfyUI-PuLID-Flux2/__init__.py
+4
View File
@@ -0,0 +1,4 @@
from .pulid_flux2 import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
from .pulid_flux2 import PuLIDFlux2
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "PuLIDFlux2"]
bundled_packs/ComfyUI-PuLID-Flux2/pulid_flux2.py
+581
View File
@@ -0,0 +1,581 @@
"""
ComfyUI-PuLID-Flux2 - VERSION UNIFIÉE KLEIN + DEV32B
Compatible automatiquement Klein 4B / 9B et Flux.2 Dev 32B
Version optimisée avec adaptation dynamique à chaque run
"""
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Tuple
import warnings
import comfy.model_management
import folder_paths
# Configuration des dossiers
PULID_DIR = os.path.join(folder_paths.models_dir, "pulid")
INSIGHTFACE_DIR = os.path.join(folder_paths.models_dir, "insightface")
os.makedirs(PULID_DIR, exist_ok=True)
os.makedirs(INSIGHTFACE_DIR, exist_ok=True)
# ============================================================================
# CACHE GLOBAL
# ============================================================================
_MODEL_CACHE = {
"eva_clip": None,
"insightface": None,
"eva_device": None,
"insightface_device": None,
}
def get_cached_model(model_type: str, device: torch.device, loader_func):
cache_key = f"{model_type}"
device_key = f"{model_type}_device"
if _MODEL_CACHE[cache_key] is not None and _MODEL_CACHE[device_key] == device:
return _MODEL_CACHE[cache_key]
model = loader_func(device)
_MODEL_CACHE[cache_key] = model
_MODEL_CACHE[device_key] = device
return model
# ============================================================================
# Classes d'attention
# ============================================================================
class PerceiverAttentionCA(nn.Module):
def __init__(self, dim: int = 4096, dim_head: int = 64, heads: int = 16):
super().__init__()
self.heads = heads
self.dim_head = dim_head
inner_dim = dim_head * heads
self.norm1 = nn.LayerNorm(dim)
self.norm2 = nn.LayerNorm(dim)
self.to_q = nn.Linear(dim, inner_dim, bias=False)
self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
self.to_out = nn.Linear(inner_dim, dim, bias=False)
def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
B, N, _ = x.shape
target_dtype = self.norm1.weight.dtype
x = x.to(target_dtype)
context = context.to(target_dtype)
x_n = self.norm1(x)
ctx = self.norm2(context)
q = self.to_q(x_n)
kv = self.to_kv(ctx)
k, v = kv.chunk(2, dim=-1)
def reshape(t):
return t.view(B, -1, self.heads, self.dim_head).transpose(1, 2)
q, k, v = reshape(q), reshape(k), reshape(v)
attn_out = F.scaled_dot_product_attention(q, k, v)
attn_out = attn_out.transpose(1, 2).contiguous().view(B, N, -1)
return self.to_out(attn_out)
class IDFormer(nn.Module):
def __init__(self, dim: int = 4096, num_tokens: int = 4):
super().__init__()
self.num_tokens = num_tokens
self.proj = nn.Sequential(
nn.Linear(512 + 768, dim),
nn.GELU(),
nn.Linear(dim, dim * num_tokens),
)
self.latents = nn.Parameter(torch.randn(1, num_tokens, dim) * 0.02)
self.layers = nn.ModuleList([PerceiverAttentionCA(dim=dim) for _ in range(4)])
self.norm = nn.LayerNorm(dim)
def forward(self, id_embed: torch.Tensor, clip_embed: torch.Tensor) -> torch.Tensor:
B = id_embed.shape[0]
clip_embed = clip_embed - clip_embed.mean(dim=-1, keepdim=True)
clip_embed = clip_embed * 1.0
combined = torch.cat([id_embed, clip_embed], dim=-1)
tokens = self.proj(combined).view(B, self.num_tokens, -1)
latents = self.latents.expand(B, -1, -1)
for layer in self.layers:
latents = latents + layer(latents, tokens)
return self.norm(latents)
class PuLIDFlux2(nn.Module):
"""Modèle PuLID – peut être créé avec n’importe quelle dim"""
def __init__(self, dim: int = 4096):
super().__init__()
self.dim = dim
self.id_former = IDFormer(dim=dim)
self.double_ca = nn.ModuleList([PerceiverAttentionCA(dim=dim) for _ in range(12)]) # plus flexible
self.single_ca = nn.ModuleList([PerceiverAttentionCA(dim=dim) for _ in range(60)]) # supporte Dev32B (48 single)
@classmethod
def from_pretrained(cls, path: str):
state = torch.load(path, map_location="cpu", weights_only=True)
dim = state["id_former.latents"].shape[-1] if "id_former.latents" in state else 4096
model = cls(dim=dim)
model.load_state_dict(state, strict=False)
return model
# ============================================================================
# Utilitaires
# ============================================================================
def get_flux_inner(model):
if hasattr(model, "model"):
model = model.model
if hasattr(model, "diffusion_model"):
model = model.diffusion_model
return model
def detect_flux_variant(model) -> Tuple[str, int, int, int]:
"""Retourne variante, dim cachée, nb double blocks, nb single blocks"""
dm = get_flux_inner(model)
double_blocks = getattr(dm, "transformer_blocks", None) or getattr(dm, "double_blocks", [])
single_blocks = getattr(dm, "single_transformer_blocks", None) or getattr(dm, "single_blocks", [])
n_double, n_single = len(double_blocks), len(single_blocks)
if n_double <= 6 and n_single <= 22:
return "klein_4b", 3072, n_double, n_single
elif n_double <= 10 and n_single <= 28:
return "klein_9b", 4096, n_double, n_single
elif n_single >= 40: # Flux.2 Dev 32B
return "flux2_dev", 6144, n_double, n_single
else:
print(f"[PuLID-Flux2] Variante inconnue ({n_double}d/{n_single}s) → fallback klein_9b")
return "klein_9b", 4096, n_double, n_single
def load_eva_clip(device):
"""Charge le modèle EVA-CLIP pour l'extraction des features"""
try:
import open_clip
model, _, _ = open_clip.create_model_and_transforms(
"EVA02-L-14-336",
pretrained="merged2b_s6b_b61k",
)
visual = model.visual
visual.eval().to(device)
return visual
except Exception as e:
warnings.warn(f"[PuLID] Erreur lors du chargement d'EVA-CLIP: {e}")
return None
def get_ca_index(block_idx: int, total_blocks: int, num_ca: int) -> int:
"""Distribution intelligente des couches CA sur tous les blocks (même Dev32B)"""
if total_blocks <= num_ca:
return block_idx
return int(block_idx * num_ca / total_blocks)
def get_scale_factors(block_idx: int, total_blocks: int, block_type: str = "double", variant: str = "klein") -> float:
progress = block_idx / max(total_blocks, 1)
if block_type == "double":
if progress < 0.4:
return 8.0
elif progress < 0.7:
return 5.0
else:
return 3.0
else: # single
if progress < 0.3:
return 6.5
elif progress < 0.6:
return 4.5
elif progress < 0.85:
return 3.0
else:
return 1.8
def patch_flux(model, pulid_module, id_tokens, strength, debug=False):
dm = get_flux_inner(model)
double_blocks = getattr(dm, "transformer_blocks", None) or getattr(dm, "double_blocks", [])
single_blocks = getattr(dm, "single_transformer_blocks", None) or getattr(dm, "single_blocks", [])
original_double = {}
original_single = {}
variant, _, n_double, n_single = detect_flux_variant(dm)
# Patch double blocks
for idx, block in enumerate(double_blocks):
original_double[idx] = block.forward
def make_double_patch(block_idx):
def patched(img, txt, vec, **kwargs):
out_img, out_txt = original_double[block_idx](img, txt, vec, **kwargs)
ca_idx = get_ca_index(block_idx, n_double, len(pulid_module.double_ca))
ca = pulid_module.double_ca[ca_idx]
factor = get_scale_factors(block_idx, n_double, "double", variant)
correction = F.normalize(ca(out_img, id_tokens), p=2, dim=-1)
out_img = out_img + strength * factor * correction
if debug:
print(f"[Double {block_idx}/{n_double-1}] ca={ca_idx} factor={factor:.1f}")
return out_img, out_txt
return patched
block.forward = make_double_patch(idx)
# Patch single blocks
for idx, block in enumerate(single_blocks):
original_single[idx] = block.forward
def make_single_patch(block_idx):
def patched(x, vec, pe, *args, **kwargs):
try:
out = original_single[block_idx](x, vec, pe, args[0] if args else None, **kwargs)
except:
out = original_single[block_idx](x, vec, pe, **kwargs)
ca_idx = get_ca_index(block_idx, n_single, len(pulid_module.single_ca))
ca = pulid_module.single_ca[ca_idx]
factor = get_scale_factors(block_idx, n_single, "single", variant)
correction = F.normalize(ca(out, id_tokens), p=2, dim=-1)
out = out + strength * factor * correction
if debug:
print(f"[Single {block_idx}/{n_single-1}] ca={ca_idx} factor={factor:.1f}")
return out
return patched
block.forward = make_single_patch(idx)
def unpatch():
for idx, block in enumerate(double_blocks):
if idx in original_double:
block.forward = original_double[idx]
for idx, block in enumerate(single_blocks):
if idx in original_single:
block.forward = original_single[idx]
return unpatch
# ============================================================================
# Nodes ComfyUI (inchangés sauf ApplyPuLIDFlux2)
# ============================================================================
class PuLIDInsightFaceLoader:
"""Charge InsightFace pour la détection et l'analyse des visages"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"provider": (["CPU", "CUDA", "ROCM"],),
}
}
RETURN_TYPES = ("INSIGHTFACE",)
FUNCTION = "load"
CATEGORY = "PuLID-Flux2"
def load(self, provider):
def _load_insightface(device):
try:
from insightface.app import FaceAnalysis
providers = {
"CPU": "CPUExecutionProvider",
"CUDA": "CUDAExecutionProvider",
"ROCM": "ROCMExecutionProvider",
}
app = FaceAnalysis(
name="antelopev2",
root=INSIGHTFACE_DIR,
providers=[providers.get(provider, "CPUExecutionProvider")]
)
app.prepare(ctx_id=0, det_size=(640, 640))
return app
except Exception as e:
raise RuntimeError(f"[PuLID] Erreur lors du chargement d'InsightFace: {e}")
device = comfy.model_management.get_torch_device()
model = get_cached_model("insightface", device, _load_insightface)
print(f"✅ InsightFace chargé (provider={provider}, cache={'HIT' if _MODEL_CACHE['insightface'] else 'MISS'})")
return (model,)
class PuLIDEVACLIPLoader:
"""Charge EVA-CLIP pour l'extraction des features visuelles"""
@classmethod
def INPUT_TYPES(cls):
return {"required": {}}
RETURN_TYPES = ("EVA_CLIP",)
FUNCTION = "load"
CATEGORY = "PuLID-Flux2"
def load(self):
device = comfy.model_management.get_torch_device()
model = get_cached_model("eva_clip", device, load_eva_clip)
if model is None:
raise RuntimeError("[PuLID] EVA-CLIP non disponible. Vérifiez l'installation d'open_clip.")
print(f"✅ EVA-CLIP chargé (cache={'HIT' if _MODEL_CACHE['eva_clip'] else 'MISS'})")
return (model,)
class PuLIDModelLoader:
"""Charge le modèle PuLID depuis un fichier .safetensors ou .pt"""
@classmethod
def INPUT_TYPES(cls):
if not os.path.exists(PULID_DIR):
os.makedirs(PULID_DIR, exist_ok=True)
files = [f for f in os.listdir(PULID_DIR) if f.endswith((".safetensors", ".pt", ".bin"))]
if not files:
files = ["__create_new__"]
return {
"required": {
"pulid_file": (files,),
}
}
RETURN_TYPES = ("PULID_MODEL",)
FUNCTION = "load"
CATEGORY = "PuLID-Flux2"
def load(self, pulid_file):
if pulid_file == "__create_new__":
print("⚠️ Création d'un nouveau modèle PuLID (non entraîné)")
return (PuLIDFlux2(dim=4096),)
path = os.path.join(PULID_DIR, pulid_file)
if not os.path.exists(path):
raise FileNotFoundError(f"[PuLID] Fichier introuvable: {path}")
try:
if path.endswith(".safetensors"):
from safetensors.torch import load_file
state = load_file(path, device="cpu")
dim = state["id_former.latents"].shape[-1]
model = PuLIDFlux2(dim=dim)
model.load_state_dict(state, strict=False)
else:
model = PuLIDFlux2.from_pretrained(path)
model.eval()
print(f"✅ Modèle PuLID chargé: {pulid_file} (dim={model.dim})")
return (model,)
except Exception as e:
raise RuntimeError(f"[PuLID] Erreur lors du chargement du modèle: {e}")
class ApplyPuLIDFlux2:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("MODEL",),
"pulid_model": ("PULID_MODEL",),
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
"eva_clip": ("EVA_CLIP",),
"face_analysis": ("INSIGHTFACE",),
"image": ("IMAGE",),
},
"optional": {
"face_index": ("INT", {"default": 0, "min": 0, "max": 9}),
"debug_mode": ("BOOLEAN", {"default": False}),
}
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply"
CATEGORY = "PuLID-Flux2"
def apply(self, model, pulid_model, strength, eva_clip, face_analysis, image,
face_index=0, debug_mode=False):
device = comfy.model_management.get_torch_device()
dtype = torch.bfloat16
# Extraction visage + embeddings (identique à ton code original)
img_np = (image[0].numpy() * 255).astype(np.uint8)
faces = face_analysis.get(img_np)
if not faces:
print("⚠️ [PuLID] AUCUN VISAGE → retour sans modification")
return (model,)
faces = sorted(faces, key=lambda f: (f.bbox[2]-f.bbox[0])*(f.bbox[3]-f.bbox[1]), reverse=True)
if face_index >= len(faces):
face_index = 0
face = faces[face_index]
id_embed = torch.from_numpy(face.embedding).unsqueeze(0).to(device, dtype=dtype)
id_embed = F.normalize(id_embed, dim=-1)
# Crop + EVA-CLIP (identique)
x1, y1, x2, y2 = face.bbox.astype(int)
margin = int(max(x2-x1, y2-y1) * 0.2)
x1, y1 = max(0, x1-margin), max(0, y1-margin)
x2, y2 = min(img_np.shape[1], x2+margin), min(img_np.shape[0], y2+margin)
face_crop = image[:1, y1:y2, x1:x2, :]
if face_crop.shape[1] == 0 or face_crop.shape[2] == 0:
face_crop = image[:1]
face_crop = F.interpolate(face_crop.permute(0,3,1,2), size=(336,336), mode="bilinear").to(device)
mean = torch.tensor([0.48145466, 0.4578275, 0.40821073], device=device).view(1,3,1,1)
std = torch.tensor([0.26862954, 0.26130258, 0.27577711], device=device).view(1,3,1,1)
face_crop = (face_crop - mean) / std
with torch.no_grad():
clip_out = eva_clip(face_crop.float())
if isinstance(clip_out, (list, tuple)):
clip_out = clip_out[0]
if clip_out.dim() == 3:
clip_out = clip_out[:, 0, :]
clip_embed = clip_out.to(device, dtype=dtype)
# ====================== PARTIE ADAPTATIVE ======================
pulid_model = pulid_model.to(device, dtype=dtype)
with torch.no_grad():
id_tokens = pulid_model.id_former(id_embed, clip_embed)
id_tokens = F.normalize(id_tokens, p=2, dim=-1)
work_model = model.clone()
dm = get_flux_inner(work_model)
# Undo any previous monkey-patch before applying a new one.
# patch_flux overwrites block.forward on layer instances that are
# shared across runs (model.clone() does not deep-copy them).
# Without this cleanup the PuLID corrections stack exponentially
# and after 3-4 generations the residual stream degenerates
# (banding, ghost face, eventually background only).
# Fixes #17.
if hasattr(dm, "_pulid_unpatcher"):
try:
dm._pulid_unpatcher()
except Exception as _e:
print(f"[PuLID] unpatch warning: {_e}")
delattr(dm, "_pulid_unpatcher")
variant, flux_dim, n_double, n_single = detect_flux_variant(dm)
# Projection si dimension différente (Klein → Dev)
if id_tokens.shape[-1] != flux_dim:
print(f"[PuLID] 🔄 Projection {id_tokens.shape[-1]} → {flux_dim} ({variant})")
proj = nn.Linear(id_tokens.shape[-1], flux_dim, bias=False).to(device, dtype=dtype)
nn.init.normal_(proj.weight, std=0.02)
id_tokens = proj(id_tokens)
id_tokens = F.normalize(id_tokens, p=2, dim=-1)
# Création d'un injector avec les bonnes dimensions pour Dev
injector = PuLIDFlux2(dim=flux_dim).to(device, dtype=dtype)
else:
injector = pulid_model
if debug_mode:
print(f"[PuLID Debug] {variant} | {n_double}d / {n_single}s | dim={flux_dim}")
# Patch via wrapper: applied just before each sampler call, removed after.
# Unlike direct patching on the shared PyTorch module, this approach
# ensures bypass works correctly: a bypassed node produces no work_model,
# so no wrapper is registered and no patch is applied.
original_wrapper = work_model.model_options.get("model_function_wrapper", None)
captured = {"injector": injector, "id_tokens": id_tokens, "strength": strength, "debug": debug_mode}
def pulid_wrapper(model_function, kwargs):
unpatch = patch_flux(
get_flux_inner(work_model.model),
captured["injector"],
captured["id_tokens"],
captured["strength"],
debug=captured["debug"]
)
try:
if original_wrapper is not None:
return original_wrapper(model_function, kwargs)
return model_function(kwargs["input"], kwargs["timestep"], **kwargs["c"])
finally:
unpatch()
work_model.set_model_unet_function_wrapper(pulid_wrapper)
emoji = "🟢" if strength >= 1.0 else "🟡" if strength > 0 else "⚪"
print(f"{emoji} PuLID Flux2 | {variant.upper()} | strength={strength:.2f} | face={face_index}")
return (work_model,)
class PuLIDFacePreview:
"""Affiche les visages détectés avec leur index et informations de confiance"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"face_analysis": ("INSIGHTFACE",),
"image": ("IMAGE",),
},
"optional": {
"show_confidence": ("BOOLEAN", {"default": True}),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "preview"
CATEGORY = "PuLID-Flux2"
OUTPUT_NODE = True
def preview(self, face_analysis, image, show_confidence=True):
try:
import cv2
img_np = (image[0].numpy() * 255).astype(np.uint8).copy()
img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
faces = face_analysis.get(img_np)
if not faces:
# Ajouter un message "No faces detected"
h, w = img_bgr.shape[:2]
cv2.putText(img_bgr, "No faces detected", (w//2-100, h//2),
cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0,0,255), 2)
else:
for i, face in enumerate(faces):
x1, y1, x2, y2 = face.bbox.astype(int)
# Rectangle vert autour du visage
cv2.rectangle(img_bgr, (x1, y1), (x2, y2), (0,255,0), 2)
# Label avec index
label = f"Face {i}"
if show_confidence and hasattr(face, 'det_score'):
label += f" ({face.det_score:.2f})"
cv2.putText(img_bgr, label, (x1, y1-8),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
# Afficher la taille du visage
size = f"{x2-x1}x{y2-y1}"
cv2.putText(img_bgr, size, (x1, y2+20),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0,255,0), 1)
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
out = torch.from_numpy(img_rgb.astype(np.float32) / 255.0).unsqueeze(0)
print(f"[PuLID Preview] Visages détectés: {len(faces)}")
return (out,)
except Exception as e:
print(f"[PuLID Preview] Erreur: {e}")
return (image,)
# ============================================================================
# Enregistrement des nodes
# ============================================================================
NODE_CLASS_MAPPINGS = {
"PuLIDInsightFaceLoader": PuLIDInsightFaceLoader,
"PuLIDEVACLIPLoader": PuLIDEVACLIPLoader,
"PuLIDModelLoader": PuLIDModelLoader,
"ApplyPuLIDFlux2": ApplyPuLIDFlux2,
"PuLIDFacePreview": PuLIDFacePreview,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"PuLIDInsightFaceLoader": "Load InsightFace (PuLID)",
"PuLIDEVACLIPLoader": "Load EVA-CLIP (PuLID)",
"PuLIDModelLoader": "Load PuLID ✦ Flux.2",
"ApplyPuLIDFlux2": "Apply PuLID ✦ Flux.2",
"PuLIDFacePreview": "PuLID — Face Preview",
}
bundled_packs/ComfyUI-PuLID-Flux2/requirements.txt
+19
View File
@@ -0,0 +1,19 @@
# ⚠️ Do NOT run pip install -r requirements.txt if you already have a working ComfyUI
# Install only:
# pip install insightface onnxruntime-gpu open-clip-torch safetensors ml_dtypes==0.3.2
insightface>=0.7.3
onnxruntime-gpu>=1.16.0
# If onnxruntime-gpu fails (CPU only): pip install onnxruntime
open-clip-torch>=3.2.0
# EVA02-CLIP-L-14-336 downloads automatically on first run (~800MB)
# DO NOT install eva_clip from GitHub — that package is broken and not required
safetensors>=0.4.0
opencv-python
numpy<2.0.0
# numpy<2.0 required for InsightFace + open_clip compatibility
ml_dtypes==0.3.2
# Pin ml_dtypes to avoid: AttributeError: module 'ml_dtypes' has no attribute 'float4_e2m1fn'
bundled_packs/ComfyUI-PuLID-Flux2/workflows/PuLID_Flux2_(Klein).json
+1
View File
File diff suppressed because one or more lines are too long
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/LICENSE
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"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.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"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.
"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).
"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.
"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."
"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.
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.
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.
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) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(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) 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
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.
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.
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.
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,
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.
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.
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.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
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.
Copyright 2025 seed
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
http://www.apache.org/licenses/LICENSE-2.0
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.
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/__init__.py
+9
View File
@@ -0,0 +1,9 @@
"""
ComfyUI-SeedVR2_VideoUpscaler
Official SeedVR2 integration for ComfyUI
"""
from .src.optimization.compatibility import ensure_triton_compat # noqa: F401
from .src.interfaces import comfy_entrypoint, SeedVR2Extension
__all__ = ["comfy_entrypoint", "SeedVR2Extension"]
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/configs_3b/main.yaml
+91
View File
@@ -0,0 +1,91 @@
__object__:
path: projects.video_diffusion_sr.train
name: VideoDiffusionTrainer
dit:
model:
__object__:
path: "dit_3b.nadit"
name: "NaDiT"
args: "as_params"
vid_in_channels: 33
vid_out_channels: 16
vid_dim: 2560
vid_out_norm: fusedrms
txt_in_dim: 5120
txt_in_norm: fusedln
txt_dim: ${.vid_dim}
emb_dim: ${eval:'6 * ${.vid_dim}'}
heads: 20
head_dim: 128 # llm-like
expand_ratio: 4
norm: fusedrms
norm_eps: 1.0e-05
ada: single
qk_bias: False
qk_norm: fusedrms
patch_size: [1, 2, 2]
num_layers: 32 # llm-like
mm_layers: 10
mlp_type: swiglu
msa_type: None
block_type: ${eval:'${.num_layers} * ["mmdit_sr"]'} # space-full
window: ${eval:'${.num_layers} * [(4,3,3)]'} # space-full
window_method: ${eval:'${.num_layers} // 2 * ["720pwin_by_size_bysize","720pswin_by_size_bysize"]'} # space-full
rope_type: mmrope3d
rope_dim: 128
compile: False
gradient_checkpoint: True
fsdp:
sharding_strategy: _HYBRID_SHARD_ZERO2
ema:
decay: 0.9998
vae:
model:
__object__:
path: "video_vae_v3.modules.attn_video_vae"
name: "VideoAutoencoderKLWrapper"
args: "as_params"
freeze_encoder: False
gradient_checkpoint: True # Disabled to prevent VRAM leaks in inference
slicing:
split_size: 4
memory_device: same
memory_limit:
conv_max_mem: 0.5
norm_max_mem: 0.5
checkpoint: ema_vae_fp16.safetensors
scaling_factor: 0.9152
compile: False
grouping: False
dtype: float16
diffusion:
schedule:
type: lerp
T: 1000.0
sampler:
type: euler
prediction_type: v_lerp
timesteps:
training:
type: logitnormal
loc: 0.0
scale: 1.0
sampling:
type: uniform_trailing
steps: 50
transform: True
loss:
type: v_lerp
cfg:
scale: 7.5
rescale: 0
condition:
i2v: 0.0
v2v: 0.0
sr: 1.0
noise_scale: 0.25
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/configs_7b/main.yaml
+88
View File
@@ -0,0 +1,88 @@
__object__:
path: projects.video_diffusion_sr.train
name: VideoDiffusionTrainer
dit:
model:
__object__:
path: "dit_7b.nadit"
name: "NaDiT"
args: "as_params"
vid_in_channels: 33
vid_out_channels: 16
vid_dim: 3072
txt_in_dim: 5120
txt_dim: ${.vid_dim}
emb_dim: ${eval:'6 * ${.vid_dim}'}
heads: 24
head_dim: 128 # llm-like
expand_ratio: 4
norm: fusedrms
norm_eps: 1e-5
ada: single
qk_bias: False
qk_rope: True
qk_norm: fusedrms
patch_size: [1, 2, 2]
num_layers: 36 # llm-like
shared_mlp: False
shared_qkv: False
mlp_type: normal
block_type: ${eval:'${.num_layers} * ["mmdit_sr"]'} # space-full
window: ${eval:'${.num_layers} * [(4,3,3)]'} # space-full
window_method: ${eval:'${.num_layers} // 2 * ["720pwin_by_size_bysize","720pswin_by_size_bysize"]'} # space-full
compile: False
gradient_checkpoint: True
fsdp:
sharding_strategy: _HYBRID_SHARD_ZERO2
ema:
decay: 0.9998
vae:
model:
__object__:
path: "video_vae_v3.modules.attn_video_vae"
name: "VideoAutoencoderKLWrapper"
args: "as_params"
freeze_encoder: False
# gradient_checkpoint: True
slicing:
split_size: 4
memory_device: same
memory_limit:
conv_max_mem: 0.5
norm_max_mem: 0.5
checkpoint: ema_vae_fp16.safetensors
scaling_factor: 0.9152
compile: False
grouping: False
dtype: float16
diffusion:
schedule:
type: lerp
T: 1000.0
sampler:
type: euler
prediction_type: v_lerp
timesteps:
training:
type: logitnormal
loc: 0.0
scale: 1.0
sampling:
type: uniform_trailing
steps: 50
transform: True
loss:
type: v_lerp
cfg:
scale: 7.5
rescale: 0
condition:
i2v: 0.0
v2v: 0.0
sr: 1.0
noise_scale: 0.25
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_4K_image_upscale.json
+1
View File
File diff suppressed because one or more lines are too long
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_HD_video_upscale.json
+1
View File
File diff suppressed because one or more lines are too long
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_simple_image_upscale.json
+1
View File
File diff suppressed because one or more lines are too long
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/inference_cli.py
+1712
View File
File diff suppressed because it is too large Load Diff
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/neg_emb.pt
BIN
View File
Binary file not shown.
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/pos_emb.pt
BIN
View File
Binary file not shown.
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/pyproject.toml
+40
View File
@@ -0,0 +1,40 @@
[project]
name = "seedvr2_videoupscaler"
description = "SeedVR2 official ComfyUI integration: ByteDance-Seed's one-step diffusion-based video/image upscaling with memory-efficient inference"
version = "2.5.24"
authors = [
{name = "numz"},
{name = "adrientoupet"}
]
license = {file = "LICENSE"}
classifiers = [
"Operating System :: OS Independent"
]
dependencies = [
"torch",
"torchvision",
"safetensors",
"numpy",
"tqdm",
"psutil",
"einops",
"omegaconf>=2.3.0",
"diffusers>=0.33.1",
"peft>=0.17.0",
"rotary_embedding_torch>=0.5.3",
"opencv-python",
"gguf",
"matplotlib"
]
[project.urls]
Repository = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler"
Documentation = "https://www.youtube.com/@AInVFX"
"Bug Tracker" = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/issues"
Forum = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/discussions"
[tool.comfy]
PublisherId = "ainvfx"
DisplayName = "ComfyUI-SeedVR2_VideoUpscaler"
Icon = "https://raw.githubusercontent.com/numz/ComfyUI-SeedVR2_VideoUpscaler/refs/heads/main/docs/seedvr_logo.png"
includes = []
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/requirements.txt
+14
View File
@@ -0,0 +1,14 @@
torch
torchvision
safetensors
numpy
tqdm
psutil
einops
omegaconf>=2.3.0
diffusers>=0.33.1
peft>=0.17.0
rotary_embedding_torch>=0.5.3
opencv-python
gguf
matplotlib
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/__init__.py
View File
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/__init__.py
View File
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/cache.py
+47
View File
@@ -0,0 +1,47 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Callable
class Cache:
"""Caching reusable args for faster inference"""
def __init__(self, disable=False, prefix="", cache=None):
self.cache = cache if cache is not None else {}
self.disable = disable
self.prefix = prefix
def __call__(self, key: str, fn: Callable):
if self.disable:
return fn()
key = self.prefix + key
try:
result = self.cache[key]
except KeyError:
result = fn()
self.cache[key] = result
return result
def namespace(self, namespace: str):
return Cache(
disable=self.disable,
prefix=self.prefix + namespace + ".",
cache=self.cache,
)
def get(self, key: str):
key = self.prefix + key
return self.cache[key]
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/config.py
+134
View File
@@ -0,0 +1,134 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Configuration utility functions
"""
import importlib
from typing import Any, Callable, List, Union
from omegaconf import DictConfig, ListConfig, OmegaConf
from ..utils.model_registry import MODEL_CLASSES
try:
OmegaConf.register_new_resolver("eval", eval)
except Exception as e:
if "already registered" not in str(e):
raise
def load_config(path: str, argv: List[str] = None) -> Union[DictConfig, ListConfig]:
"""
Load a configuration. Will resolve inheritance.
"""
#print(path)
config = OmegaConf.load(path)
if argv is not None:
config_argv = OmegaConf.from_dotlist(argv)
config = OmegaConf.merge(config, config_argv)
config = resolve_recursive(config, resolve_inheritance)
return config
def resolve_recursive(
config: Any,
resolver: Callable[[Union[DictConfig, ListConfig]], Union[DictConfig, ListConfig]],
) -> Any:
config = resolver(config)
if isinstance(config, DictConfig):
for k in config.keys():
v = config.get(k)
if isinstance(v, (DictConfig, ListConfig)):
config[k] = resolve_recursive(v, resolver)
if isinstance(config, ListConfig):
for i in range(len(config)):
v = config.get(i)
if isinstance(v, (DictConfig, ListConfig)):
config[i] = resolve_recursive(v, resolver)
return config
def resolve_inheritance(config: Union[DictConfig, ListConfig]) -> Any:
"""
Recursively resolve inheritance if the config contains:
__inherit__: path/to/parent.yaml or a ListConfig of such paths.
"""
if isinstance(config, DictConfig):
inherit = config.pop("__inherit__", None)
if inherit:
inherit_list = inherit if isinstance(inherit, ListConfig) else [inherit]
parent_config = None
for parent_path in inherit_list:
assert isinstance(parent_path, str)
parent_config = (
load_config(parent_path)
if parent_config is None
else OmegaConf.merge(parent_config, load_config(parent_path))
)
if len(config.keys()) > 0:
config = OmegaConf.merge(parent_config, config)
else:
config = parent_config
return config
def import_item(path: str, name: str) -> Any:
"""
Import a python item, checking model registry first.
Args:
path: Module path
name: Class/function name to import
Returns:
Imported object
"""
# Simple lookup with path as key
if path in MODEL_CLASSES:
return MODEL_CLASSES[path]
# Fallback to dynamic import for everything else
try:
return getattr(importlib.import_module(path), name)
except (ImportError, AttributeError) as e:
raise ImportError(f"Could not import '{name}' from '{path}': {e}")
def create_object(config: DictConfig) -> Any:
"""
Create an object from config.
The config is expected to contains the following:
__object__:
path: path.to.module
name: MyClass
args: as_config | as_params (default to as_config)
"""
item = import_item(
path=config.__object__.path,
name=config.__object__.name,
)
args = config.__object__.get("args", "as_config")
if args == "as_config":
return item(config)
if args == "as_params":
config = OmegaConf.to_object(config)
config.pop("__object__")
return item(**config)
raise NotImplementedError(f"Unknown args type: {args}")
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/decorators.py
+130
View File
@@ -0,0 +1,130 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Decorators.
"""
import functools
import threading
import time
from typing import Callable
import torch
from .distributed import barrier_if_distributed, get_global_rank, get_local_rank
from .logger import get_logger
logger = get_logger(__name__)
def log_on_entry(func: Callable) -> Callable:
"""
Functions with this decorator will log the function name at entry.
When using multiple decorators, this must be applied innermost to properly capture the name.
"""
def log_on_entry_wrapper(*args, **kwargs):
logger.info(f"Entering {func.__name__}")
return func(*args, **kwargs)
return log_on_entry_wrapper
def barrier_on_entry(func: Callable) -> Callable:
"""
Functions with this decorator will start executing when all ranks are ready to enter.
"""
def barrier_on_entry_wrapper(*args, **kwargs):
barrier_if_distributed()
return func(*args, **kwargs)
return barrier_on_entry_wrapper
def _conditional_execute_wrapper_factory(execute: bool, func: Callable) -> Callable:
"""
Helper function for local_rank_zero_only and global_rank_zero_only.
"""
def conditional_execute_wrapper(*args, **kwargs):
# Only execute if needed.
result = func(*args, **kwargs) if execute else None
# All GPUs must wait.
barrier_if_distributed()
# Return results.
return result
return conditional_execute_wrapper
def _asserted_wrapper_factory(condition: bool, func: Callable, err_msg: str = "") -> Callable:
"""
Helper function for some functions with special constraints,
especially functions called by other global_rank_zero_only / local_rank_zero_only ones,
in case they are wrongly invoked in other scenarios.
"""
def asserted_execute_wrapper(*args, **kwargs):
assert condition, err_msg
result = func(*args, **kwargs)
return result
return asserted_execute_wrapper
def local_rank_zero_only(func: Callable) -> Callable:
"""
Functions with this decorator will only execute on local rank zero.
"""
return _conditional_execute_wrapper_factory(get_local_rank() == 0, func)
def global_rank_zero_only(func: Callable) -> Callable:
"""
Functions with this decorator will only execute on global rank zero.
"""
return _conditional_execute_wrapper_factory(get_global_rank() == 0, func)
def assert_only_global_rank_zero(func: Callable) -> Callable:
"""
Functions with this decorator are only accessible to processes with global rank zero.
"""
return _asserted_wrapper_factory(
get_global_rank() == 0, func, err_msg="Not accessible to processes with global_rank != 0"
)
def assert_only_local_rank_zero(func: Callable) -> Callable:
"""
Functions with this decorator are only accessible to processes with local rank zero.
"""
return _asserted_wrapper_factory(
get_local_rank() == 0, func, err_msg="Not accessible to processes with local_rank != 0"
)
def new_thread(func: Callable) -> Callable:
"""
Functions with this decorator will run in a new thread.
The function will return the thread, which can be joined to wait for completion.
"""
def new_thread_wrapper(*args, **kwargs):
thread = threading.Thread(target=func, args=args, kwargs=kwargs)
thread.start()
return thread
return new_thread_wrapper
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/__init__.py
+56
View File
@@ -0,0 +1,56 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Diffusion package.
"""
from .config import (
create_sampler_from_config,
create_sampling_timesteps_from_config,
create_schedule_from_config,
)
from .samplers.base import Sampler
from .samplers.euler import EulerSampler
from .schedules.base import Schedule
from .schedules.lerp import LinearInterpolationSchedule
from .timesteps.base import SamplingTimesteps, Timesteps
from .timesteps.sampling.trailing import UniformTrailingSamplingTimesteps
from .types import PredictionType, SamplingDirection
from .utils import classifier_free_guidance, classifier_free_guidance_dispatcher, expand_dims
__all__ = [
# Configs
"create_sampler_from_config",
"create_sampling_timesteps_from_config",
"create_schedule_from_config",
# Schedules
"Schedule",
"DiscreteVariancePreservingSchedule",
"LinearInterpolationSchedule",
# Samplers
"Sampler",
"EulerSampler",
# Timesteps
"Timesteps",
"SamplingTimesteps",
# Types
"PredictionType",
"SamplingDirection",
"UniformTrailingSamplingTimesteps",
# Utils
"classifier_free_guidance",
"classifier_free_guidance_dispatcher",
"expand_dims",
]
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/config.py
+75
View File
@@ -0,0 +1,75 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Utility functions for creating schedules and samplers from config.
"""
import torch
from omegaconf import DictConfig
from .samplers.base import Sampler
from .samplers.euler import EulerSampler
from .schedules.base import Schedule
from .schedules.lerp import LinearInterpolationSchedule
from .timesteps.base import SamplingTimesteps
from .timesteps.sampling.trailing import UniformTrailingSamplingTimesteps
def create_schedule_from_config(
config: DictConfig,
device: torch.device,
dtype: torch.dtype = torch.float32,
) -> Schedule:
"""
Create a schedule from configuration.
"""
if config.type == "lerp":
return LinearInterpolationSchedule(T=config.get("T", 1.0))
raise NotImplementedError
def create_sampler_from_config(
config: DictConfig,
schedule: Schedule,
timesteps: SamplingTimesteps,
) -> Sampler:
"""
Create a sampler from configuration.
"""
if config.type == "euler":
return EulerSampler(
schedule=schedule,
timesteps=timesteps,
prediction_type=config.prediction_type,
)
raise NotImplementedError
def create_sampling_timesteps_from_config(
config: DictConfig,
schedule: Schedule,
device: torch.device,
dtype: torch.dtype = torch.float32,
) -> SamplingTimesteps:
if config.type == "uniform_trailing":
return UniformTrailingSamplingTimesteps(
T=schedule.T,
steps=config.steps,
shift=config.get("shift", 1.0),
device=device,
dtype=dtype,
)
raise NotImplementedError
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/samplers/base.py
+108
View File
@@ -0,0 +1,108 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Sampler base class.
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Callable
import torch
from tqdm import tqdm
from ..schedules.base import Schedule
from ..timesteps.base import SamplingTimesteps
from ..types import PredictionType, SamplingDirection
from ..utils import assert_schedule_timesteps_compatible
@dataclass
class SamplerModelArgs:
x_t: torch.Tensor
t: torch.Tensor
i: int
class Sampler(ABC):
"""
Samplers are ODE/SDE solvers.
"""
def __init__(
self,
schedule: Schedule,
timesteps: SamplingTimesteps,
prediction_type: PredictionType,
return_endpoint: bool = True,
):
assert_schedule_timesteps_compatible(
schedule=schedule,
timesteps=timesteps,
)
self.schedule = schedule
self.timesteps = timesteps
self.prediction_type = prediction_type
self.return_endpoint = return_endpoint
@abstractmethod
def sample(
self,
x: torch.Tensor,
f: Callable[[SamplerModelArgs], torch.Tensor],
) -> torch.Tensor:
"""
Generate a new sample given the the intial sample x and score function f.
"""
def get_next_timestep(
self,
t: torch.Tensor,
) -> torch.Tensor:
"""
Get the next sample timestep.
Support multiple different timesteps t in a batch.
If no more steps, return out of bound value -1 or T+1.
"""
T = self.timesteps.T
steps = len(self.timesteps)
curr_idx = self.timesteps.index(t)
next_idx = curr_idx + 1
bound = -1 if self.timesteps.direction == SamplingDirection.backward else T + 1
s = self.timesteps[next_idx.clamp_max(steps - 1)]
s = s.where(next_idx < steps, bound)
return s
def get_endpoint(
self,
pred: torch.Tensor,
x_t: torch.Tensor,
t: torch.Tensor,
) -> torch.Tensor:
"""
Get to the endpoint of the probability flow.
"""
x_0, x_T = self.schedule.convert_from_pred(pred, self.prediction_type, x_t, t)
return x_0 if self.timesteps.direction == SamplingDirection.backward else x_T
def get_progress_bar(self):
"""
Get progress bar for sampling.
"""
return tqdm(
iterable=range(len(self.timesteps) - (0 if self.return_endpoint else 1)),
dynamic_ncols=True,
desc=self.__class__.__name__,
)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/samplers/euler.py
+99
View File
@@ -0,0 +1,99 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Euler ODE solver.
"""
from typing import Callable
import torch
from einops import rearrange
from torch.nn import functional as F
from ..types import PredictionType
from ..utils import expand_dims
from .base import Sampler, SamplerModelArgs
class EulerSampler(Sampler):
"""
The Euler method is the simplest ODE solver.
<https://en.wikipedia.org/wiki/Euler_method>
"""
def sample(
self,
x: torch.Tensor,
f: Callable[[SamplerModelArgs], torch.Tensor],
) -> torch.Tensor:
timesteps = self.timesteps.timesteps
progress = self.get_progress_bar()
i = 0
# Keep native dtype throughout sampling
# The DiT model already handles dtype internally via compatibility wrapper
for t, s in zip(timesteps[:-1], timesteps[1:]):
pred = f(SamplerModelArgs(x, t, i))
# Next step
x = self.step_to(pred, x, t, s)
# Clean up temporary tensors
del pred
i += 1
progress.update()
if self.return_endpoint:
t = timesteps[-1]
pred = f(SamplerModelArgs(x, t, i))
x = self.get_endpoint(pred, x, t)
del pred
progress.update()
return x
def step(
self,
pred: torch.Tensor,
x_t: torch.Tensor,
t: torch.Tensor,
) -> torch.Tensor:
"""
Step to the next timestep.
"""
return self.step_to(pred, x_t, t, self.get_next_timestep(t))
def step_to(
self,
pred: torch.Tensor,
x_t: torch.Tensor,
t: torch.Tensor,
s: torch.Tensor,
) -> torch.Tensor:
"""
Steps from x_t at timestep t to x_s at timestep s. Returns x_s.
"""
t = expand_dims(t, x_t.ndim)
s = expand_dims(s, x_t.ndim)
T = self.schedule.T
# Step from x_t to x_s.
pred_x_0, pred_x_T = self.schedule.convert_from_pred(pred, self.prediction_type, x_t, t)
pred_x_s = self.schedule.forward(pred_x_0, pred_x_T, s.clamp(0, T))
# Clamp x_s to x_0 and x_T if s is out of bound.
pred_x_s = pred_x_s.where(s >= 0, pred_x_0)
pred_x_s = pred_x_s.where(s <= T, pred_x_T)
return pred_x_s
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/schedules/base.py
+131
View File
@@ -0,0 +1,131 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Schedule base class.
"""
from abc import ABC, abstractmethod, abstractproperty
from typing import Tuple, Union
import torch
from ..types import PredictionType
from ..utils import expand_dims
class Schedule(ABC):
"""
Diffusion schedules are uniquely defined by T, A, B:
x_t = A(t) * x_0 + B(t) * x_T, where t in [0, T]
Schedules can be continuous or discrete.
"""
@abstractproperty
def T(self) -> Union[int, float]:
"""
Maximum timestep inclusive.
Schedule is continuous if float, discrete if int.
"""
@abstractmethod
def A(self, t: torch.Tensor) -> torch.Tensor:
"""
Interpolation coefficient A.
Returns tensor with the same shape as t.
"""
@abstractmethod
def B(self, t: torch.Tensor) -> torch.Tensor:
"""
Interpolation coefficient B.
Returns tensor with the same shape as t.
"""
# ----------------------------------------------------
def snr(self, t: torch.Tensor) -> torch.Tensor:
"""
Signal to noise ratio.
Returns tensor with the same shape as t.
"""
return (self.A(t) ** 2) / (self.B(t) ** 2)
def isnr(self, snr: torch.Tensor) -> torch.Tensor:
"""
Inverse signal to noise ratio.
Returns tensor with the same shape as snr.
Subclass may implement.
"""
raise NotImplementedError
# ----------------------------------------------------
def is_continuous(self) -> bool:
"""
Whether the schedule is continuous.
"""
return isinstance(self.T, float)
def forward(self, x_0: torch.Tensor, x_T: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""
Diffusion forward function.
"""
t = expand_dims(t, x_0.ndim)
return self.A(t) * x_0 + self.B(t) * x_T
def convert_from_pred(
self, pred: torch.Tensor, pred_type: PredictionType, x_t: torch.Tensor, t: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Convert from prediction. Return predicted x_0 and x_T.
"""
t = expand_dims(t, x_t.ndim)
A_t = self.A(t)
B_t = self.B(t)
if pred_type == PredictionType.x_T:
pred_x_T = pred
pred_x_0 = (x_t - B_t * pred_x_T) / A_t
elif pred_type == PredictionType.x_0:
pred_x_0 = pred
pred_x_T = (x_t - A_t * pred_x_0) / B_t
elif pred_type == PredictionType.v_cos:
pred_x_0 = A_t * x_t - B_t * pred
pred_x_T = A_t * pred + B_t * x_t
elif pred_type == PredictionType.v_lerp:
pred_x_0 = (x_t - B_t * pred) / (A_t + B_t)
pred_x_T = (x_t + A_t * pred) / (A_t + B_t)
else:
raise NotImplementedError
return pred_x_0, pred_x_T
def convert_to_pred(
self, x_0: torch.Tensor, x_T: torch.Tensor, t: torch.Tensor, pred_type: PredictionType
) -> torch.FloatTensor:
"""
Convert to prediction target given x_0 and x_T.
"""
if pred_type == PredictionType.x_T:
return x_T
if pred_type == PredictionType.x_0:
return x_0
if pred_type == PredictionType.v_cos:
t = expand_dims(t, x_0.ndim)
return self.A(t) * x_T - self.B(t) * x_0
if pred_type == PredictionType.v_lerp:
return x_T - x_0
raise NotImplementedError
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/schedules/lerp.py
+55
View File
@@ -0,0 +1,55 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Linear interpolation schedule (lerp).
"""
from typing import Union
import torch
from .base import Schedule
class LinearInterpolationSchedule(Schedule):
"""
Linear interpolation schedule (lerp) is proposed by flow matching and rectified flow.
It leads to straighter probability flow theoretically. It is also used by Stable Diffusion 3.
<https://arxiv.org/abs/2209.03003>
<https://arxiv.org/abs/2210.02747>
x_t = (1 - t) * x_0 + t * x_T
Can be either continuous or discrete.
"""
def __init__(self, T: Union[int, float] = 1.0):
self._T = T
@property
def T(self) -> Union[int, float]:
return self._T
def A(self, t: torch.Tensor) -> torch.Tensor:
return 1 - (t / self.T)
def B(self, t: torch.Tensor) -> torch.Tensor:
return t / self.T
# ----------------------------------------------------
def isnr(self, snr: torch.Tensor) -> torch.Tensor:
t = self.T / (1 + snr**0.5)
t = t if self.is_continuous() else t.round().int()
return t
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/timesteps/base.py
+72
View File
@@ -0,0 +1,72 @@
from abc import ABC, abstractmethod
from typing import Sequence, Union
import torch
from ..types import SamplingDirection
class Timesteps(ABC):
"""
Timesteps base class.
"""
def __init__(self, T: Union[int, float]):
assert T > 0
self._T = T
@property
def T(self) -> Union[int, float]:
"""
Maximum timestep inclusive.
int if discrete, float if continuous.
"""
return self._T
def is_continuous(self) -> bool:
"""
Whether the schedule is continuous.
"""
return isinstance(self.T, float)
class SamplingTimesteps(Timesteps):
"""
Sampling timesteps.
It defines the discretization of sampling steps.
"""
def __init__(
self,
T: Union[int, float],
timesteps: torch.Tensor,
direction: SamplingDirection,
):
assert timesteps.ndim == 1
super().__init__(T)
self.timesteps = timesteps
self.direction = direction
def __len__(self) -> int:
"""
Number of sampling steps.
"""
return len(self.timesteps)
def __getitem__(self, idx: Union[int, torch.IntTensor]) -> torch.Tensor:
"""
The timestep at the sampling step.
Returns a scalar tensor if idx is int,
or tensor of the same size if idx is a tensor.
"""
return self.timesteps[idx]
def index(self, t: torch.Tensor) -> torch.Tensor:
"""
Find index by t.
Return index of the same shape as t.
Index is -1 if t not found in timesteps.
"""
i, j = t.reshape(-1, 1).eq(self.timesteps).nonzero(as_tuple=True)
idx = torch.full_like(t, fill_value=-1, dtype=torch.int)
idx.view(-1)[i] = j.int()
return idx
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/timesteps/sampling/trailing.py
+50
View File
@@ -0,0 +1,50 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
import torch
from ...types import SamplingDirection
from ..base import SamplingTimesteps
class UniformTrailingSamplingTimesteps(SamplingTimesteps):
"""
Uniform trailing sampling timesteps.
Defined in (https://arxiv.org/abs/2305.08891)
Shift is proposed in SD3 for RF schedule.
Defined in (https://arxiv.org/pdf/2403.03206) eq.23
"""
def __init__(
self,
T: int,
steps: int,
shift: float = 1.0,
device: torch.device = "cpu",
dtype: torch.dtype = torch.float32,
):
# Create trailing timesteps with specified dtype
timesteps = torch.arange(1.0, 0.0, -1.0 / steps, device='cpu').to(device=device, dtype=dtype)
# Shift timesteps.
timesteps = shift * timesteps / (1 + (shift - 1) * timesteps)
# Scale to T range.
if isinstance(T, float):
timesteps = timesteps * T
else:
timesteps = timesteps.mul(T + 1).sub(1).round().int()
super().__init__(T=T, timesteps=timesteps, direction=SamplingDirection.backward)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/types.py
+59
View File
@@ -0,0 +1,59 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Type definitions.
"""
from enum import Enum
class PredictionType(str, Enum):
"""
x_0:
Predict data sample.
x_T:
Predict noise sample.
Proposed by DDPM (https://arxiv.org/abs/2006.11239)
Proved problematic by zsnr paper (https://arxiv.org/abs/2305.08891)
v_cos:
Predict velocity dx/dt based on the cosine schedule (A_t * x_T - B_t * x_0).
Proposed by progressive distillation (https://arxiv.org/abs/2202.00512)
v_lerp:
Predict velocity dx/dt based on the lerp schedule (x_T - x_0).
Proposed by rectified flow (https://arxiv.org/abs/2209.03003)
"""
x_0 = "x_0"
x_T = "x_T"
v_cos = "v_cos"
v_lerp = "v_lerp"
class SamplingDirection(str, Enum):
"""
backward: Sample from x_T to x_0 for data generation.
forward: Sample from x_0 to x_T for noise inversion.
"""
backward = "backward"
forward = "forward"
@staticmethod
def reverse(direction):
if direction == SamplingDirection.backward:
return SamplingDirection.forward
if direction == SamplingDirection.forward:
return SamplingDirection.backward
raise NotImplementedError
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/utils.py
+84
View File
@@ -0,0 +1,84 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Utility functions.
"""
from typing import Callable
import torch
def expand_dims(tensor: torch.Tensor, ndim: int):
"""
Expand tensor to target ndim. New dims are added to the right.
For example, if the tensor shape was (8,), target ndim is 4, return (8, 1, 1, 1).
"""
shape = tensor.shape + (1,) * (ndim - tensor.ndim)
return tensor.reshape(shape)
def assert_schedule_timesteps_compatible(schedule, timesteps):
"""
Check if schedule and timesteps are compatible.
"""
if schedule.T != timesteps.T:
raise ValueError("Schedule and timesteps must have the same T.")
if schedule.is_continuous() != timesteps.is_continuous():
raise ValueError("Schedule and timesteps must have the same continuity.")
def classifier_free_guidance(
pos: torch.Tensor,
neg: torch.Tensor,
scale: float,
rescale: float = 0.0,
):
"""
Apply classifier-free guidance.
"""
# Classifier-free guidance (https://arxiv.org/abs/2207.12598)
cfg = neg + scale * (pos - neg)
# Classifier-free guidance rescale (https://arxiv.org/pdf/2305.08891.pdf)
if rescale != 0.0:
pos_std = pos.std(dim=list(range(1, pos.ndim)), keepdim=True)
cfg_std = cfg.std(dim=list(range(1, cfg.ndim)), keepdim=True)
factor = pos_std / cfg_std
factor = rescale * factor + (1 - rescale)
cfg *= factor
return cfg
def classifier_free_guidance_dispatcher(
pos: Callable,
neg: Callable,
scale: float,
rescale: float = 0.0,
):
"""
Optionally execute models depending on classifer-free guidance scale.
"""
# If scale is 1, no need to execute neg model.
if scale == 1.0:
return pos()
# Otherwise, execute both pos nad neg models and apply cfg.
return classifier_free_guidance(
pos=pos(),
neg=neg(),
scale=scale,
rescale=rescale,
)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/__init__.py
+37
View File
@@ -0,0 +1,37 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Distributed package.
"""
from .basic import (
barrier_if_distributed,
convert_to_ddp,
get_device,
get_global_rank,
get_local_rank,
get_world_size,
init_torch,
)
__all__ = [
"barrier_if_distributed",
"convert_to_ddp",
"get_device",
"get_global_rank",
"get_local_rank",
"get_world_size",
"init_torch",
]
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/advanced.py
+223
View File
@@ -0,0 +1,223 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Advanced distributed functions for sequence parallel.
"""
from typing import Optional, List, TYPE_CHECKING
import torch
import torch.distributed as dist
# Conditional imports for distributed training features (not needed for inference)
try:
from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
from torch.distributed.fsdp import ShardingStrategy
_FSDP_AVAILABLE = True
except (ImportError, AttributeError):
# AMD ROCm and other builds may not have full FSDP support
DeviceMesh = None
init_device_mesh = None
ShardingStrategy = None
_FSDP_AVAILABLE = False
from .basic import get_global_rank, get_world_size
_DATA_PARALLEL_GROUP = None
_SEQUENCE_PARALLEL_GROUP = None
_SEQUENCE_PARALLEL_CPU_GROUP = None
_MODEL_SHARD_CPU_INTER_GROUP = None
_MODEL_SHARD_CPU_INTRA_GROUP = None
_MODEL_SHARD_INTER_GROUP = None
_MODEL_SHARD_INTRA_GROUP = None
_SEQUENCE_PARALLEL_GLOBAL_RANKS = None
def get_data_parallel_group() -> Optional[dist.ProcessGroup]:
"""
Get data parallel process group.
"""
return _DATA_PARALLEL_GROUP
def get_sequence_parallel_group() -> Optional[dist.ProcessGroup]:
"""
Get sequence parallel process group.
"""
return _SEQUENCE_PARALLEL_GROUP
def get_sequence_parallel_cpu_group() -> Optional[dist.ProcessGroup]:
"""
Get sequence parallel CPU process group.
"""
return _SEQUENCE_PARALLEL_CPU_GROUP
def get_data_parallel_rank() -> int:
"""
Get data parallel rank.
"""
group = get_data_parallel_group()
return dist.get_rank(group) if group else get_global_rank()
def get_data_parallel_world_size() -> int:
"""
Get data parallel world size.
"""
group = get_data_parallel_group()
return dist.get_world_size(group) if group else get_world_size()
def get_sequence_parallel_rank() -> int:
"""
Get sequence parallel rank.
"""
group = get_sequence_parallel_group()
return dist.get_rank(group) if group else 0
def get_sequence_parallel_world_size() -> int:
"""
Get sequence parallel world size.
"""
group = get_sequence_parallel_group()
return dist.get_world_size(group) if group else 1
def get_model_shard_cpu_intra_group() -> Optional[dist.ProcessGroup]:
"""
Get the CPU intra process group of model sharding.
"""
return _MODEL_SHARD_CPU_INTRA_GROUP
def get_model_shard_cpu_inter_group() -> Optional[dist.ProcessGroup]:
"""
Get the CPU inter process group of model sharding.
"""
return _MODEL_SHARD_CPU_INTER_GROUP
def get_model_shard_intra_group() -> Optional[dist.ProcessGroup]:
"""
Get the GPU intra process group of model sharding.
"""
return _MODEL_SHARD_INTRA_GROUP
def get_model_shard_inter_group() -> Optional[dist.ProcessGroup]:
"""
Get the GPU inter process group of model sharding.
"""
return _MODEL_SHARD_INTER_GROUP
def init_sequence_parallel(sequence_parallel_size: int):
"""
Initialize sequence parallel.
"""
global _DATA_PARALLEL_GROUP
global _SEQUENCE_PARALLEL_GROUP
global _SEQUENCE_PARALLEL_CPU_GROUP
global _SEQUENCE_PARALLEL_GLOBAL_RANKS
assert dist.is_initialized()
world_size = dist.get_world_size()
rank = dist.get_rank()
data_parallel_size = world_size // sequence_parallel_size
for i in range(data_parallel_size):
start_rank = i * sequence_parallel_size
end_rank = (i + 1) * sequence_parallel_size
ranks = range(start_rank, end_rank)
group = dist.new_group(ranks)
cpu_group = dist.new_group(ranks, backend="gloo")
if rank in ranks:
_SEQUENCE_PARALLEL_GROUP = group
_SEQUENCE_PARALLEL_CPU_GROUP = cpu_group
_SEQUENCE_PARALLEL_GLOBAL_RANKS = list(ranks)
def init_model_shard_group(
*,
sharding_strategy: ShardingStrategy,
device_mesh: Optional[DeviceMesh] = None,
):
"""
Initialize process group of model sharding.
"""
if not _FSDP_AVAILABLE:
raise RuntimeError(
"FSDP features are not available in this PyTorch build. "
"Model sharding requires torch.distributed.fsdp support."
)
global _MODEL_SHARD_INTER_GROUP
global _MODEL_SHARD_INTRA_GROUP
global _MODEL_SHARD_CPU_INTER_GROUP
global _MODEL_SHARD_CPU_INTRA_GROUP
assert dist.is_initialized()
world_size = dist.get_world_size()
if device_mesh is not None:
num_shards_per_group = device_mesh.shape[1]
elif sharding_strategy == ShardingStrategy.NO_SHARD:
num_shards_per_group = 1
elif sharding_strategy in [
ShardingStrategy.HYBRID_SHARD,
ShardingStrategy._HYBRID_SHARD_ZERO2,
]:
num_shards_per_group = torch.cuda.device_count()
else:
num_shards_per_group = world_size
num_groups = world_size // num_shards_per_group
device_mesh = (num_groups, num_shards_per_group)
gpu_mesh_2d = init_device_mesh("cuda", device_mesh, mesh_dim_names=("inter", "intra"))
cpu_mesh_2d = init_device_mesh("cpu", device_mesh, mesh_dim_names=("inter", "intra"))
_MODEL_SHARD_INTER_GROUP = gpu_mesh_2d.get_group("inter")
_MODEL_SHARD_INTRA_GROUP = gpu_mesh_2d.get_group("intra")
_MODEL_SHARD_CPU_INTER_GROUP = cpu_mesh_2d.get_group("inter")
_MODEL_SHARD_CPU_INTRA_GROUP = cpu_mesh_2d.get_group("intra")
def get_sequence_parallel_global_ranks() -> List[int]:
"""
Get all global ranks of the sequence parallel process group
that the caller rank belongs to.
"""
if _SEQUENCE_PARALLEL_GLOBAL_RANKS is None:
return [dist.get_rank()]
return _SEQUENCE_PARALLEL_GLOBAL_RANKS
def get_next_sequence_parallel_rank() -> int:
"""
Get the next global rank of the sequence parallel process group
that the caller rank belongs to.
"""
sp_global_ranks = get_sequence_parallel_global_ranks()
sp_rank = get_sequence_parallel_rank()
sp_size = get_sequence_parallel_world_size()
return sp_global_ranks[(sp_rank + 1) % sp_size]
def get_prev_sequence_parallel_rank() -> int:
"""
Get the previous global rank of the sequence parallel process group
that the caller rank belongs to.
"""
sp_global_ranks = get_sequence_parallel_global_ranks()
sp_rank = get_sequence_parallel_rank()
sp_size = get_sequence_parallel_world_size()
return sp_global_ranks[(sp_rank + sp_size - 1) % sp_size]
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/basic.py
+87
View File
@@ -0,0 +1,87 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Distributed basic functions.
"""
import os
from datetime import timedelta
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from ...optimization.memory_manager import is_mps_available
def get_global_rank() -> int:
"""
Get the global rank, the global index of the GPU.
"""
return int(os.environ.get("RANK", "0"))
def get_local_rank() -> int:
"""
Get the local rank, the local index of the GPU.
"""
return int(os.environ.get("LOCAL_RANK", "0"))
def get_world_size() -> int:
"""
Get the world size, the total amount of GPUs.
"""
return int(os.environ.get("WORLD_SIZE", "1"))
def get_device() -> torch.device:
"""
Get current rank device.
"""
if is_mps_available():
return torch.device("mps")
return torch.device("cuda", get_local_rank())
def barrier_if_distributed(*args, **kwargs):
"""
Synchronizes all processes if under distributed context.
"""
if dist.is_initialized():
return dist.barrier(*args, **kwargs)
def init_torch(cudnn_benchmark=True, timeout=timedelta(seconds=600)):
"""
Common PyTorch initialization configuration.
"""
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = cudnn_benchmark
torch.cuda.set_device(get_local_rank())
dist.init_process_group(
backend="nccl",
rank=get_global_rank(),
world_size=get_world_size(),
timeout=timeout,
)
def convert_to_ddp(module: torch.nn.Module, **kwargs) -> DistributedDataParallel:
return DistributedDataParallel(
module=module,
device_ids=[get_local_rank()],
output_device=get_local_rank(),
**kwargs,
)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/meta_init_utils.py
+40
View File
@@ -0,0 +1,40 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
import torch
from rotary_embedding_torch import RotaryEmbedding
from torch import nn
__all__ = ["meta_non_persistent_buffer_init_fn"]
def meta_non_persistent_buffer_init_fn(module: nn.Module) -> nn.Module:
"""
Used for materializing `non-persistent tensor buffers` while model resuming.
Since non-persistent tensor buffers are not saved in state_dict,
when initializing model with meta device, user should materialize those buffers manually.
Currently, only `rope.dummy` is this special case.
"""
with torch.no_grad():
for submodule in module.modules():
if not isinstance(submodule, RotaryEmbedding):
continue
for buffer_name, buffer in submodule.named_buffers(recurse=False):
if buffer.is_meta and "dummy" in buffer_name:
materialized_buffer = torch.zeros_like(buffer, device="cpu")
setattr(submodule, buffer_name, materialized_buffer)
assert not any(b.is_meta for n, b in module.named_buffers())
return module
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/ops.py
+495
View File
@@ -0,0 +1,495 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Distributed ops for supporting sequence parallel.
"""
from collections import defaultdict
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import torch
import torch.distributed as dist
from torch import Tensor
from ..cache import Cache
from .advanced import (
get_sequence_parallel_group,
get_sequence_parallel_rank,
get_sequence_parallel_world_size,
)
from .basic import get_device
_SEQ_DATA_BUF = defaultdict(lambda: [None, None, None])
_SEQ_DATA_META_SHAPES = defaultdict()
_SEQ_DATA_META_DTYPES = defaultdict()
_SEQ_DATA_ASYNC_COMMS = defaultdict(list)
_SYNC_BUFFER = defaultdict(dict)
def single_all_to_all(
local_input: Tensor,
scatter_dim: int,
gather_dim: int,
group: dist.ProcessGroup,
async_op: bool = False,
):
"""
A function to do all-to-all on a tensor
"""
seq_world_size = dist.get_world_size(group)
prev_scatter_dim = scatter_dim
if scatter_dim != 0:
local_input = local_input.transpose(0, scatter_dim)
if gather_dim == 0:
gather_dim = scatter_dim
scatter_dim = 0
inp_shape = list(local_input.shape)
inp_shape[scatter_dim] = inp_shape[scatter_dim] // seq_world_size
input_t = local_input.reshape(
[seq_world_size, inp_shape[scatter_dim]] + inp_shape[scatter_dim + 1 :]
).contiguous()
output = torch.empty_like(input_t)
comm = dist.all_to_all_single(output, input_t, group=group, async_op=async_op)
if async_op:
# let user's code transpose & reshape
return output, comm, prev_scatter_dim
# first dim is seq_world_size, so we can split it directly
output = torch.cat(output.split(1), dim=gather_dim + 1).squeeze(0)
if prev_scatter_dim:
output = output.transpose(0, prev_scatter_dim).contiguous()
return output
def _all_to_all(
local_input: Tensor,
scatter_dim: int,
gather_dim: int,
group: dist.ProcessGroup,
):
seq_world_size = dist.get_world_size(group)
input_list = [
t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)
]
output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)]
dist.all_to_all(output_list, input_list, group=group)
return torch.cat(output_list, dim=gather_dim).contiguous()
class SeqAllToAll(torch.autograd.Function):
@staticmethod
def forward(
ctx: Any,
group: dist.ProcessGroup,
local_input: Tensor,
scatter_dim: int,
gather_dim: int,
async_op: bool,
) -> Tensor:
ctx.group = group
ctx.scatter_dim = scatter_dim
ctx.gather_dim = gather_dim
ctx.async_op = async_op
if async_op:
output, comm, prev_scatter_dim = single_all_to_all(
local_input, scatter_dim, gather_dim, group, async_op=async_op
)
ctx.prev_scatter_dim = prev_scatter_dim
return output, comm
return _all_to_all(local_input, scatter_dim, gather_dim, group)
@staticmethod
def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]:
if ctx.async_op:
input_t = torch.cat(grad_output[0].split(1), dim=ctx.gather_dim + 1).squeeze(0)
if ctx.prev_scatter_dim:
input_t = input_t.transpose(0, ctx.prev_scatter_dim)
else:
input_t = grad_output[0]
return (
None,
_all_to_all(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group),
None,
None,
None,
)
class Slice(torch.autograd.Function):
@staticmethod
def forward(ctx: Any, group: dist.ProcessGroup, local_input: Tensor, dim: int) -> Tensor:
ctx.group = group
ctx.rank = dist.get_rank(group)
seq_world_size = dist.get_world_size(group)
ctx.seq_world_size = seq_world_size
ctx.dim = dim
dim_size = local_input.shape[dim]
return local_input.split(dim_size // seq_world_size, dim=dim)[ctx.rank].contiguous()
@staticmethod
def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor, None]:
dim_size = list(grad_output.size())
split_size = dim_size[0]
dim_size[0] = dim_size[0] * ctx.seq_world_size
output = torch.empty(dim_size, dtype=grad_output.dtype, device=torch.cuda.current_device())
dist._all_gather_base(output, grad_output, group=ctx.group)
return (None, torch.cat(output.split(split_size), dim=ctx.dim), None)
class Gather(torch.autograd.Function):
@staticmethod
def forward(
ctx: Any,
group: dist.ProcessGroup,
local_input: Tensor,
dim: int,
grad_scale: Optional[bool] = False,
) -> Tensor:
ctx.group = group
ctx.rank = dist.get_rank(group)
ctx.dim = dim
ctx.grad_scale = grad_scale
seq_world_size = dist.get_world_size(group)
ctx.seq_world_size = seq_world_size
dim_size = list(local_input.size())
split_size = dim_size[0]
ctx.part_size = dim_size[dim]
dim_size[0] = dim_size[0] * seq_world_size
output = torch.empty(dim_size, dtype=local_input.dtype, device=torch.cuda.current_device())
dist._all_gather_base(output, local_input.contiguous(), group=ctx.group)
return torch.cat(output.split(split_size), dim=dim)
@staticmethod
def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor]:
if ctx.grad_scale:
grad_output = grad_output * ctx.seq_world_size
return (
None,
grad_output.split(ctx.part_size, dim=ctx.dim)[ctx.rank].contiguous(),
None,
None,
)
def gather_seq_scatter_heads_qkv(
qkv_tensor: Tensor,
*,
seq_dim: int,
qkv_shape: Optional[Tensor] = None,
cache: Cache = Cache(disable=True),
restore_shape: bool = True,
):
"""
A func to sync splited qkv tensor
qkv_tensor: the tensor we want to do alltoall with. The last dim must
be the projection_idx, which we will split into 3 part. After
spliting, the gather idx will be projecttion_idx + 1
seq_dim: gather_dim for all2all comm
restore_shape: if True, output will has the same shape length as input
"""
group = get_sequence_parallel_group()
if not group:
return qkv_tensor
world = get_sequence_parallel_world_size()
orig_shape = qkv_tensor.shape
scatter_dim = qkv_tensor.dim()
bef_all2all_shape = list(orig_shape)
qkv_proj_dim = bef_all2all_shape[-1]
bef_all2all_shape = bef_all2all_shape[:-1] + [3, qkv_proj_dim // 3]
qkv_tensor = qkv_tensor.view(bef_all2all_shape)
qkv_tensor = SeqAllToAll.apply(group, qkv_tensor, scatter_dim, seq_dim, False)
if restore_shape:
out_shape = list(orig_shape)
out_shape[seq_dim] *= world
out_shape[-1] = qkv_proj_dim // world
qkv_tensor = qkv_tensor.view(out_shape)
# remove padding
if qkv_shape is not None:
unpad_dim_size = cache(
"unpad_dim_size", lambda: torch.sum(torch.prod(qkv_shape, dim=-1)).item()
)
if unpad_dim_size % world != 0:
padding_size = qkv_tensor.size(seq_dim) - unpad_dim_size
qkv_tensor = _unpad_tensor(qkv_tensor, seq_dim, padding_size)
return qkv_tensor
def slice_inputs(x: Tensor, dim: int, padding: bool = True):
"""
A func to slice the input sequence in sequence parallel
"""
group = get_sequence_parallel_group()
if group is None:
return x
sp_rank = get_sequence_parallel_rank()
sp_world = get_sequence_parallel_world_size()
dim_size = x.shape[dim]
unit = (dim_size + sp_world - 1) // sp_world
if padding and dim_size % sp_world:
padding_size = sp_world - (dim_size % sp_world)
x = _pad_tensor(x, dim, padding_size)
slc = [slice(None)] * len(x.shape)
slc[dim] = slice(unit * sp_rank, unit * (sp_rank + 1))
return x[slc]
def remove_seqeunce_parallel_padding(x: Tensor, dim: int, unpad_dim_size: int):
"""
A func to remove the padding part of the tensor based on its original shape
"""
group = get_sequence_parallel_group()
if group is None:
return x
sp_world = get_sequence_parallel_world_size()
if unpad_dim_size % sp_world == 0:
return x
padding_size = sp_world - (unpad_dim_size % sp_world)
assert (padding_size + unpad_dim_size) % sp_world == 0
return _unpad_tensor(x, dim=dim, padding_size=padding_size)
def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int) -> Tensor:
"""
A func to sync attention result with alltoall in sequence parallel
"""
group = get_sequence_parallel_group()
if not group:
return x
dim_size = x.size(seq_dim)
sp_world = get_sequence_parallel_world_size()
if dim_size % sp_world != 0:
padding_size = sp_world - (dim_size % sp_world)
x = _pad_tensor(x, seq_dim, padding_size)
return SeqAllToAll.apply(group, x, seq_dim, head_dim, False)
def gather_seq_scatter_heads(x: Tensor, seq_dim: int, head_dim: int) -> Tensor:
"""
A func to sync embedding input with alltoall in sequence parallel
"""
group = get_sequence_parallel_group()
if not group:
return x
return SeqAllToAll.apply(group, x, head_dim, seq_dim, False)
def scatter_heads(x: Tensor, dim: int) -> Tensor:
"""
A func to split heads before attention in sequence parallel
"""
group = get_sequence_parallel_group()
if not group:
return x
return Slice.apply(group, x, dim)
def gather_heads(x: Tensor, dim: int, grad_scale: Optional[bool] = False) -> Tensor:
"""
A func to gather heads for the attention result in sequence parallel
"""
group = get_sequence_parallel_group()
if not group:
return x
return Gather.apply(group, x, dim, grad_scale)
def gather_outputs(
x: Tensor,
*,
gather_dim: int,
padding_dim: Optional[int] = None,
unpad_shape: Optional[Tensor] = None,
cache: Cache = Cache(disable=True),
scale_grad=True,
):
"""
A func to gather the outputs for the model result in sequence parallel
"""
group = get_sequence_parallel_group()
if not group:
return x
x = Gather.apply(group, x, gather_dim, scale_grad)
if padding_dim is not None:
unpad_dim_size = cache(
"unpad_dim_size", lambda: torch.sum(torch.prod(unpad_shape, dim=1)).item()
)
x = remove_seqeunce_parallel_padding(x, padding_dim, unpad_dim_size)
return x
def _pad_tensor(x: Tensor, dim: int, padding_size: int):
shape = list(x.shape)
shape[dim] = padding_size
pad = torch.zeros(shape, dtype=x.dtype, device=x.device)
return torch.cat([x, pad], dim=dim)
def _unpad_tensor(x: Tensor, dim: int, padding_size):
slc = [slice(None)] * len(x.shape)
slc[dim] = slice(0, -padding_size)
return x[slc]
def _broadcast_data(data, shape, dtype, src, group, async_op):
comms = []
if isinstance(data, (list, tuple)):
for i, sub_shape in enumerate(shape):
comms += _broadcast_data(data[i], sub_shape, dtype[i], src, group, async_op)
elif isinstance(data, dict):
for key, sub_data in data.items():
comms += _broadcast_data(sub_data, shape[key], dtype[key], src, group, async_op)
elif isinstance(data, Tensor):
comms.append(dist.broadcast(data, src=src, group=group, async_op=async_op))
return comms
def _traverse(data: Any, op: Callable) -> Union[None, List, Dict, Any]:
if isinstance(data, (list, tuple)):
return [_traverse(sub_data, op) for sub_data in data]
elif isinstance(data, dict):
return {key: _traverse(sub_data, op) for key, sub_data in data.items()}
elif isinstance(data, Tensor):
return op(data)
else:
return None
def _get_shapes(data):
return _traverse(data, op=lambda x: x.shape)
def _get_dtypes(data):
return _traverse(data, op=lambda x: x.dtype)
def _construct_broadcast_buffer(shapes, dtypes, device):
if isinstance(shapes, torch.Size):
return torch.empty(shapes, dtype=dtypes, device=device)
if isinstance(shapes, (list, tuple)):
buffer = []
for i, sub_shape in enumerate(shapes):
buffer.append(_construct_broadcast_buffer(sub_shape, dtypes[i], device))
elif isinstance(shapes, dict):
buffer = {}
for key, sub_shape in shapes.items():
buffer[key] = _construct_broadcast_buffer(sub_shape, dtypes[key], device)
else:
return None
return buffer
class SPDistForward:
"""A forward tool to sync different result across sp group
Args:
module: a function or module to process users input
sp_step: current training step to judge which rank to broadcast its result to all
name: a distinct str to save meta and async comm
comm_shape: if different ranks have different shape, mark this arg to True
device: the device for current rank, can be empty
"""
def __init__(
self,
name: str,
comm_shape: bool,
device: torch.device = None,
):
self.name = name
self.comm_shape = comm_shape
if device:
self.device = device if isinstance(device, torch.device) else torch.device(device)
else:
# Fallback to standard device detection
self.device = get_device()
def __call__(self, inputs) -> Any:
group = get_sequence_parallel_group()
if not group:
yield inputs
else:
device = self.device
sp_world = get_sequence_parallel_world_size()
sp_rank = get_sequence_parallel_rank()
for local_step in range(sp_world):
src_rank = dist.get_global_rank(group, local_step)
is_src = sp_rank == local_step
local_shapes = []
local_dtypes = []
if local_step == 0:
local_result = inputs
_SEQ_DATA_BUF[self.name][-1] = local_result
local_shapes = _get_shapes(local_result)
local_dtypes = _get_dtypes(local_result)
if self.comm_shape:
group_shapes_lists = [None] * sp_world
dist.all_gather_object(group_shapes_lists, local_shapes, group=group)
_SEQ_DATA_META_SHAPES[self.name] = group_shapes_lists
else:
_SEQ_DATA_META_SHAPES[self.name] = [local_shapes] * sp_world
_SEQ_DATA_META_DTYPES[self.name] = local_dtypes
shapes = _SEQ_DATA_META_SHAPES[self.name][local_step]
dtypes = _SEQ_DATA_META_DTYPES[self.name]
buf_id = local_step % 2
if local_step == 0:
sync_data = (
local_result
if is_src
else _construct_broadcast_buffer(shapes, dtypes, device)
)
_broadcast_data(sync_data, shapes, dtypes, src_rank, group, False)
_SEQ_DATA_BUF[self.name][buf_id] = sync_data
# wait for async comm ops
if _SEQ_DATA_ASYNC_COMMS[self.name]:
for comm in _SEQ_DATA_ASYNC_COMMS[self.name]:
comm.wait()
# before return the sync result, do async broadcast for next batch
if local_step < sp_world - 1:
next_buf_id = 1 - buf_id
shapes = _SEQ_DATA_META_SHAPES[self.name][local_step + 1]
src_rank = dist.get_global_rank(group, local_step + 1)
is_src = sp_rank == local_step + 1
next_sync_data = (
_SEQ_DATA_BUF[self.name][-1]
if is_src
else _construct_broadcast_buffer(shapes, dtypes, device)
)
_SEQ_DATA_ASYNC_COMMS[self.name] = _broadcast_data(
next_sync_data, shapes, dtypes, src_rank, group, True
)
_SEQ_DATA_BUF[self.name][next_buf_id] = next_sync_data
yield _SEQ_DATA_BUF[self.name][buf_id]
sync_inputs = SPDistForward(name="bef_fwd", comm_shape=True)
def sync_data(data, sp_idx, name="tmp"):
group = get_sequence_parallel_group()
if group is None:
return data
# if sp_idx in _SYNC_BUFFER[name]:
# return _SYNC_BUFFER[name][sp_idx]
sp_rank = get_sequence_parallel_rank()
src_rank = dist.get_global_rank(group, sp_idx)
objects = [data] if sp_rank == sp_idx else [None]
dist.broadcast_object_list(objects, src=src_rank, group=group)
# _SYNC_BUFFER[name] = {sp_idx: objects[0]}
return objects[0]
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/half_precision_fixes.py
+160
View File
@@ -0,0 +1,160 @@
"""
Half Precision Compatibility Fixes for SeedVR2
Provides safe wrappers around PyTorch operations that may fail with reduced
precision dtypes (float16, bfloat16, float8). These utilities automatically
detect precision issues and apply temporary conversions when needed, then
restore original dtypes.
"""
import torch
import torch.nn.functional as F
from typing import Optional, Tuple, Union
def safe_pad_operation(
x: torch.Tensor,
padding: Union[Tuple[int, ...], int],
mode: str = 'constant',
value: float = 0.0
) -> torch.Tensor:
"""
Safe padding operation with automatic float16 compatibility handling.
Certain padding modes ('replicate', 'reflect', 'circular') are not implemented
for float16 on some backends. This function automatically detects failures and
applies a temporary float32 conversion, then restores the original dtype.
Args:
x: Input tensor of any dtype
padding: Padding specification (left, right, top, bottom, front, back)
mode: Padding mode - 'constant', 'replicate', 'reflect', or 'circular'
value: Fill value for 'constant' mode (default: 0.0)
Returns:
Padded tensor in original dtype
"""
# Modes that may require float16 compatibility fixes
problematic_modes = ['replicate', 'reflect', 'circular']
if mode in problematic_modes:
try:
return F.pad(x, padding, mode=mode, value=value)
except RuntimeError as e:
if "not implemented for 'Half'" in str(e):
original_dtype = x.dtype
result = F.pad(x.float(), padding, mode=mode, value=value)
return result.to(original_dtype)
else:
raise
else:
# 'constant' and other compatible modes work natively
return F.pad(x, padding, mode=mode, value=value)
def safe_interpolate_operation(
x: torch.Tensor,
size: Optional[Union[int, Tuple[int, ...]]] = None,
scale_factor: Optional[Union[float, Tuple[float, ...]]] = None,
mode: str = 'nearest',
align_corners: Optional[bool] = None,
recompute_scale_factor: Optional[bool] = None
) -> torch.Tensor:
"""
Safe interpolation operation with automatic float16 compatibility handling.
Interpolation modes like 'bilinear', 'bicubic', and 'trilinear' may not be
implemented for float16 on some backends. This function automatically detects
failures and applies a temporary float32 conversion, then restores the original dtype.
Args:
x: Input tensor of any dtype
size: Target output size (height, width) or (depth, height, width)
scale_factor: Multiplier for spatial size (alternative to size)
mode: Interpolation mode - 'nearest', 'linear', 'bilinear', 'bicubic', 'trilinear'
align_corners: If True, align corner pixels of input and output
recompute_scale_factor: Recompute scale_factor for use in interpolation
Returns:
Interpolated tensor in original dtype
"""
# Modes that may require float16 compatibility fixes
problematic_modes = ['bilinear', 'bicubic', 'trilinear']
if mode in problematic_modes:
try:
return F.interpolate(
x,
size=size,
scale_factor=scale_factor,
mode=mode,
align_corners=align_corners,
recompute_scale_factor=recompute_scale_factor
)
except RuntimeError as e:
# Check for float16 incompatibility errors
if ("not implemented for 'Half'" in str(e) or
"compute_indices_weights" in str(e)):
original_dtype = x.dtype
result = F.interpolate(
x.float(),
size=size,
scale_factor=scale_factor,
mode=mode,
align_corners=align_corners,
recompute_scale_factor=recompute_scale_factor
)
return result.to(original_dtype)
else:
raise
else:
# 'nearest' and other compatible modes work natively
return F.interpolate(
x,
size=size,
scale_factor=scale_factor,
mode=mode,
align_corners=align_corners,
recompute_scale_factor=recompute_scale_factor
)
def ensure_float32_precision(
tensor: torch.Tensor,
force_float32: bool = True
) -> Tuple[torch.Tensor, torch.dtype]:
"""
Ensure tensor is in float32 for precision-sensitive operations.
Many numerical operations require full precision to avoid accumulated errors:
- Color space conversions (RGB↔LAB, RGB↔HSV) with matrix multiplications
- Statistical operations (mean, variance, standard deviation)
- Edge detection using derivative filters (Sobel, Canny)
- Histogram matching and CDF computations
- Guided filtering with covariance calculations
This function upgrades reduced precision dtypes (float16, bfloat16, float8)
to float32, while preserving the original dtype for restoration after computation.
Args:
tensor: Input tensor of any dtype
force_float32: If True, convert reduced precision dtypes to float32.
If False, return tensor unchanged (useful for disabling conversion)
Returns:
Tuple of (converted_tensor, original_dtype) for easy restoration:
- converted_tensor: Tensor in float32 (or original dtype if already full precision)
- original_dtype: Original tensor dtype for restoration via .to(original_dtype)
"""
original_dtype = tensor.dtype
# Skip conversion if disabled
if not force_float32:
return tensor, original_dtype
# Convert reduced precision dtypes to float32
if original_dtype not in (torch.float32, torch.float64):
return tensor.float(), original_dtype
# Already full precision - return as-is (no copy)
return tensor, original_dtype
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/logger.py
+44
View File
@@ -0,0 +1,44 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Logging utility functions.
"""
import logging
import sys
from typing import Optional
from .distributed import get_global_rank, get_local_rank, get_world_size
_default_handler = logging.StreamHandler(sys.stdout)
_default_handler.setFormatter(
logging.Formatter(
"%(asctime)s "
+ (f"[Rank:{get_global_rank()}]" if get_world_size() > 1 else "")
+ (f"[LocalRank:{get_local_rank()}]" if get_world_size() > 1 else "")
+ "[%(threadName).12s][%(name)s][%(levelname).5s] "
+ "%(message)s"
)
)
def get_logger(name: Optional[str] = None) -> logging.Logger:
"""
Get a logger.
"""
logger = logging.getLogger(name)
logger.addHandler(_default_handler)
logger.setLevel(logging.INFO)
return logger
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/partition.py
+59
View File
@@ -0,0 +1,59 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Partition utility functions.
"""
from typing import Any, List
def partition_by_size(data: List[Any], size: int) -> List[List[Any]]:
"""
Partition a list by size.
When indivisible, the last group contains fewer items than the target size.
Examples:
- data: [1,2,3,4,5]
- size: 2
- return: [[1,2], [3,4], [5]]
"""
assert size > 0
return [data[i : (i + size)] for i in range(0, len(data), size)]
def partition_by_groups(data: List[Any], groups: int) -> List[List[Any]]:
"""
Partition a list by groups.
When indivisible, some groups may have more items than others.
Examples:
- data: [1,2,3,4,5]
- groups: 2
- return: [[1,3,5], [2,4]]
"""
assert groups > 0
return [data[i::groups] for i in range(groups)]
def shift_list(data: List[Any], n: int) -> List[Any]:
"""
Rotate a list by n elements.
Examples:
- data: [1,2,3,4,5]
- n: 3
- return: [4,5,1,2,3]
"""
return data[(n % len(data)) :] + data[: (n % len(data))]
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/common/seed.py
+30
View File
@@ -0,0 +1,30 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
import random
from typing import Optional
import numpy as np
import torch
from .distributed import get_global_rank
def set_seed(seed: Optional[int], same_across_ranks: bool = False):
"""Function that sets the seed for pseudo-random number generators."""
if seed is not None:
seed += get_global_rank() if not same_across_ranks else 0
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/core/__init__.py
View File
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/core/alpha_upscaling.py
+438
View File
@@ -0,0 +1,438 @@
"""
Alpha Channel Upscaling Module
Provides edge-guided Alpha upscaling that matches RGB structural changes
without diffusion artifacts. Uses joint bilateral filtering and edge detection
to maintain alignment between upscaled RGB and Alpha channels.
"""
import torch
import torch.nn.functional as F
import cv2
import numpy as np
from typing import Optional, Any, List
from ..common.half_precision_fixes import ensure_float32_precision
from ..optimization.memory_manager import manage_tensor
def process_alpha_for_batch(
rgb_samples: List[torch.Tensor],
alpha_original: torch.Tensor,
rgb_original: torch.Tensor,
device: torch.device,
compute_dtype: torch.dtype,
debug: Optional['Debug'] = None
) -> List[torch.Tensor]:
"""
Process Alpha channel for an entire batch with temporal consistency.
Called during postprocess phase when VRAM is available.
Args:
rgb_samples: List of decoded RGB samples from VAE
alpha_original: Original Alpha channel (C, T, H, W)
rgb_original: Original RGB for guidance (C, T, H, W)
device: Target device for processing (typically CUDA)
compute_dtype: Pipeline compute dtype (e.g., bfloat16) for final output.
debug: Debug instance for logging
Returns:
List of RGBA samples with Alpha merged, in compute_dtype
"""
# Move alpha and RGB guidance tensors to processing device (GPU)
alpha_original = manage_tensor(
tensor=alpha_original,
target_device=device,
tensor_name="alpha_original",
debug=debug,
reason="Alpha processing",
indent_level=1
)
rgb_original = manage_tensor(
tensor=rgb_original,
target_device=device,
tensor_name="rgb_original",
debug=debug,
reason="Alpha processing",
indent_level=1
)
# Process each RGB sample individually to handle variable frame counts
processed_samples = []
for rgb_sample in rgb_samples:
# Move RGB sample to processing device before use
rgb_sample = manage_tensor(
tensor=rgb_sample,
target_device=device,
tensor_name="rgb_sample",
debug=debug,
reason="Alpha processing",
indent_level=1
)
# Handle dimension variations between single frames and video sequences
if rgb_sample.ndim == 3:
# Single frame: (C, H, W) -> need to add T dimension
rgb_sample_4d = rgb_sample.unsqueeze(0) # (C, H, W) -> (1, C, H, W)
alpha_4d = alpha_original.unsqueeze(1) if alpha_original.ndim == 3 else alpha_original
rgb_guide_4d = rgb_original.unsqueeze(1) if rgb_original.ndim == 3 else rgb_original
was_single_frame = True
else:
# Video: (T, C, H, W) -> already correct
rgb_sample_4d = rgb_sample
alpha_4d = alpha_original
rgb_guide_4d = rgb_original
was_single_frame = False
# Apply edge-guided upscaling to match RGB resolution while preserving alpha structure
alpha_upscaled = edge_guided_alpha_upscale(
input_alpha=alpha_4d[:, :1, :, :] if alpha_4d.shape[1] > 1 else alpha_4d, # Ensure (T, 1, H, W)
input_rgb=rgb_guide_4d[:, :3, :, :], # Ensure (T, 3, H, W)
upscaled_rgb=rgb_sample_4d[:, :3, :, :], # Ensure (T, 3, H, W)
method='guided',
debug=debug
)
# Convert Alpha from float32 to compute dtype
if alpha_upscaled.dtype != compute_dtype:
alpha_upscaled = manage_tensor(
tensor=alpha_upscaled,
target_device=alpha_upscaled.device,
tensor_name="alpha_upscaled",
dtype=compute_dtype,
debug=debug,
reason="dtype alignment for RGBA concatenation",
indent_level=1
)
# Concatenate RGB and upscaled alpha to create RGBA output (T, 4, H, W)
rgba_sample = torch.cat([rgb_sample_4d[:, :3, :, :], alpha_upscaled], dim=1)
# Restore original format
if was_single_frame:
rgba_sample = rgba_sample.squeeze(0) # (1, 4, H, W) -> (4, H, W)
processed_samples.append(rgba_sample)
# Release memory immediately after processing each sample
del alpha_upscaled
# Clean up
del alpha_original, rgb_original
return processed_samples
def detect_edges_batch(
images: torch.Tensor,
method: str = 'sobel',
debug: Optional['Debug'] = None
) -> torch.Tensor:
"""
Detect edges in a batch of images using Sobel or Canny.
Args:
images: Tensor of shape (T, C, H, W) in range [-1, 1] or [0, 1]
method: 'sobel' or 'canny'
debug: Optional debug instance for logging
Returns:
Edge map tensor of shape (T, 1, H, W) in range [0, 1]
"""
# Convert to float32 for OpenCV processing (will be converted to numpy anyway)
images, images_dtype = ensure_float32_precision(images, force_float32=True)
images_np = images.cpu().numpy()
T, C, H, W = images_np.shape
# Denormalize if needed
if images_np.min() < 0:
images_np = (images_np + 1) / 2
# Convert to 0-255 uint8
images_np = (images_np * 255).clip(0, 255).astype(np.uint8)
edges = []
for t in range(T):
frame = images_np[t].transpose(1, 2, 0)
if C == 3:
gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
else:
gray = frame[..., 0]
if method == 'sobel':
sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
edge = np.sqrt(sobelx**2 + sobely**2)
edge = (edge / edge.max() * 255).astype(np.uint8)
else:
edge = cv2.Canny(gray, 50, 150)
edges.append(edge)
edges = np.stack(edges)
edges = torch.from_numpy(edges).float() / 255.0
edges = edges.unsqueeze(1)
# Move back to images device after CPU numpy processing
edges = manage_tensor(
tensor=edges,
target_device=images.device,
tensor_name="edge_map",
dtype=images_dtype,
debug=debug,
reason="edge detection result",
indent_level=1
)
return edges
def guided_filter_pytorch(guide: torch.Tensor, src: torch.Tensor,
radius: int = 8, eps: float = 0.01) -> torch.Tensor:
"""
Apply guided filter for edge-preserving smoothing.
Args:
guide: Guidance image (T, C, H, W)
src: Input to filter (T, 1, H, W)
radius: Filter radius
eps: Regularization parameter
Returns:
Filtered output (T, 1, H, W)
"""
# Convert to float32 for numerical stability in statistical operations
guide, guide_dtype = ensure_float32_precision(guide, force_float32=True)
src, _ = ensure_float32_precision(src, force_float32=True)
T, C, H, W = guide.shape
# Convert to grayscale if RGB
if C == 3:
guide_gray = guide.mean(dim=1, keepdim=True)
else:
guide_gray = guide
# Apply guided filter
output = _apply_guided_filter(guide_gray, src, radius, eps)
# Restore original dtype after float32 processing
if output.dtype != guide_dtype:
output = manage_tensor(
tensor=output,
target_device=output.device,
tensor_name="filtered_output",
dtype=guide_dtype,
debug=debug,
reason="dtype restoration"
)
return output
def _apply_guided_filter(guide_gray: torch.Tensor, src: torch.Tensor,
radius: int, eps: float) -> torch.Tensor:
"""
Apply guided filter algorithm.
Args:
guide_gray: Grayscale guidance image (T, 1, H, W)
src: Source image to filter (T, 1, H, W)
radius: Filter radius in pixels
eps: Regularization parameter for edge preservation
Returns:
Filtered output tensor (T, 1, H, W)
Note:
Uses box filtering for efficient implementation
"""
def box_filter(x, r):
return F.avg_pool2d(x, kernel_size=2*r+1, stride=1, padding=r)
# Compute means
mean_guide = box_filter(guide_gray, radius)
mean_src = box_filter(src, radius)
# Compute correlations
corr_guide = box_filter(guide_gray * guide_gray, radius)
corr_guide_src = box_filter(guide_gray * src, radius)
# Compute variance and covariance
var_guide = corr_guide - mean_guide * mean_guide
cov_guide_src = corr_guide_src - mean_guide * mean_src
del corr_guide, corr_guide_src
# Linear coefficients
a = cov_guide_src / (var_guide + eps)
b = mean_src - a * mean_guide
del cov_guide_src, var_guide, mean_guide, mean_src
# Average coefficients
mean_a = box_filter(a, radius)
mean_b = box_filter(b, radius)
del a, b
# Apply filter
output = mean_a * guide_gray + mean_b
del mean_a, mean_b
return output
def edge_guided_alpha_upscale(
input_alpha: torch.Tensor,
input_rgb: torch.Tensor,
upscaled_rgb: torch.Tensor,
method: str = 'guided',
debug: Optional['Debug'] = None
) -> torch.Tensor:
"""
Upscale Alpha channel using RGB edge structure as guidance.
Creates tight, smooth transitions aligned with RGB edges.
Args:
input_alpha: Original Alpha (T, 1, H_in, W_in) in [0, 1]
input_rgb: Original RGB (T, 3, H_in, W_in) in [-1, 1] or [0, 1]
upscaled_rgb: Upscaled RGB (T, 3, H_out, W_out) in [-1, 1] or [0, 1]
method: 'guided' or 'bilateral'
debug: Debug instance
Returns:
Upscaled Alpha (T, 1, H_out, W_out) in [0, 1]
"""
T, _, H_out, W_out = upscaled_rgb.shape
device = input_alpha.device
dtype = input_alpha.dtype
# Convert to float32 for numerical stability in edge detection and filtering
input_alpha, alpha_dtype = ensure_float32_precision(input_alpha, force_float32=True)
upscaled_rgb, _ = ensure_float32_precision(upscaled_rgb, force_float32=True)
input_rgb, _ = ensure_float32_precision(input_rgb, force_float32=True)
# Analyze alpha distribution to detect binary masks vs gradient alphas
alpha_flat = input_alpha.flatten()
near_zero = (alpha_flat < 0.1).sum().float()
near_one = (alpha_flat > 0.9).sum().float()
binary_ratio = (near_zero + near_one) / alpha_flat.numel()
is_binary_mask = binary_ratio > 0.95
if debug:
debug.log(f"Alpha type: {'binary mask' if is_binary_mask else 'gradient alpha'}", category="alpha", indent_level=1)
debug.log(f"Binary ratio: {binary_ratio:.2%}", category="alpha", indent_level=1)
debug.log("Creating tight edge-aligned transitions", category="alpha", indent_level=1)
# Normalize RGB from [-1, 1] to [0, 1] for edge detection
rgb_normalized = upscaled_rgb.clone()
if rgb_normalized.min() < 0:
rgb_normalized = (rgb_normalized + 1) / 2
# Detect edges in upscaled RGB
rgb_edges = detect_edges_batch(images=rgb_normalized, method='sobel', debug=debug)
# Step 1: Initial bicubic upscale provides smooth base before edge refinement
# MPS on PyTorch < 2.8 doesn't support bicubic+antialias - use CPU fallback
try:
alpha_upscaled = F.interpolate(
input_alpha,
size=(H_out, W_out),
mode='bicubic',
align_corners=False,
antialias=True
).clamp(0, 1)
except NotImplementedError:
alpha_upscaled = F.interpolate(
input_alpha.cpu(),
size=(H_out, W_out),
mode='bicubic',
align_corners=False,
antialias=True
).to(device).clamp(0, 1)
if is_binary_mask:
if debug:
debug.log("Applying tight edge-aware refinement", category="alpha", indent_level=1)
# Step 2: Single-pass guided filter with small radius for tight edges
alpha_refined = guided_filter_pytorch(
guide=rgb_normalized,
src=alpha_upscaled,
radius=2, # Reduced from 3 for tighter edges
eps=0.002
)
# Step 3: Create tight transition zone using 3x3 max pooling on edge map
edge_map = rgb_edges
transition_zone = F.max_pool2d(edge_map, kernel_size=3, stride=1, padding=1)
# Step 4: Identify solid regions (far from edges) vs transition regions (near edges)
solid_threshold = 0.05 # Reduced from 0.08
is_solid = transition_zone < solid_threshold
# Convert solid regions to binary (0 or 1) based on 0.5 threshold
alpha_binary = (alpha_refined > 0.5).float()
# Step 5: Apply strong contrast enhancement in narrow transition zones to complete gradients
in_edge_region = (transition_zone >= solid_threshold) & (transition_zone < 0.25) # Narrow band
# Apply strong sigmoid (strength=12.0) to push alpha values toward binary endpoints
contrast_enhanced = torch.sigmoid((alpha_refined - 0.5) * 12.0) # Increased from 7.0
# Blend original and enhanced alpha based on edge strength for smooth transitions
edge_strength = torch.clamp(edge_map / 0.25, 0, 1)
alpha_in_edges = alpha_refined * (1 - edge_strength) + contrast_enhanced * edge_strength
# Step 6: Combine solid binary regions with enhanced transition regions
alpha_combined = torch.where(is_solid, alpha_binary, alpha_in_edges)
# Step 7: Snap remaining mid-gray values to binary outside tightest edge zone
very_solid = transition_zone < 0.03 # Even tighter threshold
final_binary = (alpha_combined > 0.5).float()
alpha_final = torch.where(very_solid, final_binary, alpha_combined)
# Step 8: Final cleanup - eliminate mid-gray artifacts, preserving only smooth edges in tight regions
tightest_edges = edge_map > 0.15
mid_gray = (alpha_final > 0.3) & (alpha_final < 0.7)
should_be_binary = mid_gray & ~tightest_edges
alpha_final = torch.where(
should_be_binary,
(alpha_final > 0.5).float(),
alpha_final
)
else:
# For gradient alphas: single guided filter pass preserves smooth transitions
if debug:
debug.log("Applying guided filter for gradient alpha", category="alpha", indent_level=1)
alpha_final = alpha_upscaled.clone()
# Single pass with optimized parameters for gradients
alpha_final = guided_filter_pytorch(
guide=rgb_normalized,
src=alpha_final,
radius=3,
eps=0.002
)
# Clamp output to valid alpha range [0, 1]
alpha_final = alpha_final.clamp(0, 1)
if debug:
final_values = alpha_final.flatten()
exact_zeros = (final_values < 0.02).sum().item()
exact_ones = (final_values > 0.98).sum().item()
smooth_edges = ((final_values >= 0.02) & (final_values <= 0.98)).sum().item()
mid_grays = ((final_values >= 0.35) & (final_values <= 0.65)).sum().item()
total = final_values.numel()
debug.log(f"Pure 0s: {exact_zeros/total*100:.1f}% | Pure 1s: {exact_ones/total*100:.1f}% | Smooth edges: {smooth_edges/total*100:.1f}% | Mid-grays: {mid_grays/total*100:.1f}%", category="alpha", indent_level=1)
return alpha_final
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/core/generation_phases.py
+1478
View File
File diff suppressed because it is too large Load Diff
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/core/generation_utils.py
+827
View File
@@ -0,0 +1,827 @@
"""
Generation Utilities for SeedVR2
This module provides setup, configuration, and utility functions for the generation pipeline.
These are helper functions used to prepare, configure, and support the video upscaling process,
but do not contain the core 4-phase orchestration logic itself.
Setup Functions:
- prepare_video_transforms: Create video transformation pipeline (resize, pad, normalize)
- setup_video_transform: Initialize transforms and compute target dimensions
- setup_generation_context: Initialize context with device configuration
- prepare_runner: Configure VideoDiffusionInfer with all settings
Video Processing Utilities:
- pad_video_temporal: Unified temporal padding with reversed frames (prepend/append, any format)
- blend_overlapping_frames: Smooth blending for temporal overlap between batches
Configuration Helpers:
- load_text_embeddings: Load positive/negative text embeddings for DiT
- calculate_optimal_batch_params: Compute batch processing parameters
- check_interrupt: Check for user interruption
Debugging Utilities:
- _draw_tile_boundaries: Draw tile boundaries for debugging VAE tiling
- ensure_precision_initialized: Log model dtype information
These utilities support the 4-phase pipeline implemented in generation_phases.py.
"""
import os
import torch
from typing import Dict, List, Optional, Tuple, Any, Callable, Union
from torchvision.transforms import Compose, Lambda, Normalize
from .model_configuration import configure_runner
from .infer import VideoDiffusionInfer
from ..data.image.transforms.divisible_crop import DivisiblePad
from ..data.image.transforms.na_resize import NaResize
from ..optimization.compatibility import COMPUTE_DTYPE, BFLOAT16_SUPPORTED
from ..optimization.memory_manager import manage_tensor
from ..utils.constants import get_script_directory
# Get script directory for embeddings
script_directory = get_script_directory()
def prepare_video_transforms(resolution: int, max_resolution: int = 0, debug: Optional['Debug'] = None) -> Compose:
"""
Prepare optimized video transformation pipeline
Args:
resolution (int): Target resolution for shortest edge
max_resolution (int): Maximum resolution for any edge (0 = no limit)
debug (Debug, optional): Debug instance for logging
Returns:
Compose: Configured transformation pipeline
Features:
- Resolution-aware upscaling (no downsampling)
- Optional max resolution constraint on longest edge
- Padding to divisible by 16 (no data loss)
- Proper normalization for model compatibility
- Memory-efficient tensor operations
"""
if debug:
msg = f"Initializing video transformation pipeline for {resolution}px (shortest edge)"
if max_resolution > 0:
msg += f", max {max_resolution}px (any edge)"
debug.log(msg, category="setup", indent_level=1)
return Compose([
NaResize(
resolution=resolution,
mode="side",
# Upsample image, model only trained for high res
downsample_only=False,
max_resolution=max_resolution,
),
Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
DivisiblePad((16, 16)),
Normalize(0.5, 0.5),
Lambda(lambda x: x.permute(1, 0, 2, 3)), # t c h w -> c t h w (faster than Rearrange)
])
def setup_video_transform(ctx: Dict[str, Any], resolution: int, max_resolution: int = 0,
debug: Optional['Debug'] = None,
sample_frame: Optional[torch.Tensor] = None) -> Tuple[int, int, int, int]:
"""
Setup video transformation pipeline and compute target dimensions.
Args:
ctx: Generation context dictionary
resolution: Target resolution for shortest edge
max_resolution: Maximum resolution for any edge (0 = no limit)
debug: Debug instance for logging
sample_frame: Optional sample frame tensor (C, H, W) to compute dimensions
Returns:
(true_height, true_width, padded_height, padded_width) if dimensions computed,
(0, 0, 0, 0) otherwise
"""
# Check if transform exists AND is not None
existing_transform = ctx.get('video_transform')
if existing_transform is not None:
# Transform exists - check if we need to compute dimensions
if 'true_target_dims' in ctx and sample_frame is not None:
# Return cached dimensions + recompute padded from sample
true_h, true_w = ctx['true_target_dims']
transformed = existing_transform(sample_frame)
padded_h, padded_w = transformed.shape[-2:]
if debug:
debug.log("Reusing pre-initialized video transformation pipeline", category="reuse")
return true_h, true_w, padded_h, padded_w
elif debug:
debug.log("Reusing pre-initialized video transformation pipeline", category="reuse")
return 0, 0, 0, 0
# Create transformation pipeline (first time or after cleanup)
ctx['video_transform'] = prepare_video_transforms(resolution, max_resolution, debug)
# Compute dimensions if sample frame provided
if sample_frame is not None:
# Get true target size (after resize, before padding)
temp_transform = Compose([
NaResize(resolution=resolution, mode="side", downsample_only=False, max_resolution=max_resolution),
Lambda(lambda x: torch.clamp(x, 0.0, 1.0))
])
resized_sample = temp_transform(sample_frame)
true_h, true_w = resized_sample.shape[-2:]
# Round to even numbers for video codec compatibility (libx264 requirement)
true_h = (true_h // 2) * 2
true_w = (true_w // 2) * 2
# Cache for later use in trimming
ctx['true_target_dims'] = (true_h, true_w)
# Get padded dimensions
transformed_sample = ctx['video_transform'](sample_frame)
padded_h, padded_w = transformed_sample.shape[-2:]
if debug:
if true_h == padded_h and true_w == padded_w:
debug.log(f"Target dimensions: {true_w}x{true_h} (no padding needed)",
category="setup", indent_level=1)
else:
debug.log(f"Target dimensions: {true_w}x{true_h} (padded to {padded_w}x{padded_h} for processing)",
category="setup", indent_level=1)
del temp_transform, resized_sample, transformed_sample
return true_h, true_w, padded_h, padded_w
return 0, 0, 0, 0
def compute_generation_info(
ctx: Dict[str, Any],
images: torch.Tensor,
resolution: int = 1080,
max_resolution: int = 0,
batch_size: int = 5,
uniform_batch_size: bool = False,
seed: int = 42,
prepend_frames: int = 0,
temporal_overlap: int = 0,
debug: Optional['Debug'] = None
) -> Tuple[torch.Tensor, Dict[str, Any]]:
"""
Compute all generation parameters and dimensions for logging.
Args:
ctx: Generation context dictionary
images: Input frames tensor [T, H, W, C]
resolution: Target resolution for shortest edge
max_resolution: Maximum resolution for any edge (0 = no limit)
batch_size: Frames per batch
uniform_batch_size: Whether to pad final batch to match batch_size
seed: Random seed
prepend_frames: Number of frames to prepend
temporal_overlap: Overlapping frames between batches
debug: Debug instance for logging
Returns:
Tuple of (processed_images, info_dict)
- processed_images: Input images with prepending applied if needed
- info_dict: Information dictionary for logging
"""
# Track input frames before any modifications
input_frames = len(images)
input_h, input_w = images.shape[1], images.shape[2]
channels_info = "RGBA" if images.shape[-1] == 4 else "RGB"
# Apply prepending if requested
if prepend_frames > 0:
images = pad_video_temporal(images, count=prepend_frames, temporal_dim=0, prepend=True, debug=debug)
# Track total frames after prepending
total_frames = len(images)
ctx['total_frames'] = total_frames
# Setup transform and compute dimensions on final frame count
sample_frame = images[0].permute(2, 0, 1).unsqueeze(0)
true_h, true_w, padded_h, padded_w = setup_video_transform(
ctx, resolution, max_resolution, debug, sample_frame
)
del sample_frame
info = {
'input_frames': input_frames,
'input_h': input_h,
'input_w': input_w,
'total_frames': total_frames,
'true_h': true_h,
'true_w': true_w,
'padded_h': padded_h,
'padded_w': padded_w,
'channels_info': channels_info,
'batch_size': batch_size,
'uniform_batch_size': uniform_batch_size,
'seed': seed,
'prepend_frames': prepend_frames,
'temporal_overlap': temporal_overlap,
'resolution': resolution,
'max_resolution': max_resolution
}
return images, info
def log_generation_start(info: Dict[str, Any], debug: Optional['Debug'] = None) -> None:
"""
Log generation start information in a consistent format.
Args:
info: Information dictionary from compute_generation_info()
debug: Debug instance for logging
"""
if debug is None:
return
debug.log("", category="none", force=True)
debug.log("Starting upscaling generation...", category="generation", force=True)
# Build concise parameter info
batch_text = f"Batch size: {info['batch_size']}"
if info.get('uniform_batch_size', False):
batch_text += " (uniform)"
params_info = batch_text
if info['prepend_frames'] > 0:
params_info += f", Prepend frames: {info['prepend_frames']}"
if info['temporal_overlap'] > 0:
params_info += f", Temporal overlap: {info['temporal_overlap']}"
params_info += f", Seed: {info['seed']}, Channels: {info['channels_info']}"
# Build resolution constraint info
res_constraint = f"shortest edge: {info['resolution']}px"
if info['max_resolution'] > 0:
res_constraint += f", max edge: {info['max_resolution']}px"
# Log dimension flow with full context
if info['true_h'] > 0:
frame_text = "frame" if info['input_frames'] <= 1 else "frames"
if info['true_h'] == info['padded_h'] and info['true_w'] == info['padded_w']:
debug.log(
f"Input: {info['input_frames']} {frame_text}, "
f"{info['input_w']}x{info['input_h']}px → Output: {info['true_w']}x{info['true_h']}px "
f"({res_constraint})",
category="generation", force=True, indent_level=1
)
else:
debug.log(
f"Input: {info['input_frames']} {frame_text}, "
f"{info['input_w']}x{info['input_h']}px → Padded: {info['padded_w']}x{info['padded_h']}px → "
f"Output: {info['true_w']}x{info['true_h']}px ({res_constraint})",
category="generation", force=True, indent_level=1
)
debug.log(f"{params_info}", category="generation", force=True, indent_level=1)
def blend_overlapping_frames(prev_tail: torch.Tensor, cur_head: torch.Tensor, overlap: int) -> torch.Tensor:
"""
Blend two overlapping frame sequences in-place.
Args:
prev_tail: Last `overlap` frames from previous batch [overlap, H, W, C]
cur_head: First `overlap` frames from current batch [overlap, H, W, C]
overlap: Number of overlapping frames
Returns:
torch.Tensor: Blended frames [overlap, H, W, C]
"""
device = prev_tail.device
dtype = prev_tail.dtype
# Smooth crossfade with Hann window for overlap >= 3, linear for smaller overlaps
if overlap >= 3:
t = torch.linspace(0.0, 1.0, steps=overlap, device=device, dtype=dtype)
blend_start = 1.0 / 3.0
blend_end = 2.0 / 3.0
u = ((t - blend_start) / (blend_end - blend_start)).clamp(0.0, 1.0)
w_prev_1d = 0.5 + 0.5 * torch.cos(torch.pi * u) # Hann window
else:
w_prev_1d = torch.linspace(1.0, 0.0, steps=overlap, device=device, dtype=dtype)
w_prev = w_prev_1d.view(overlap, 1, 1, 1)
w_cur = 1.0 - w_prev
return prev_tail * w_prev + cur_head * w_cur
def setup_generation_context(
dit_device: Optional[Union[str, torch.device]] = None,
vae_device: Optional[Union[str, torch.device]] = None,
dit_offload_device: Optional[Union[str, torch.device]] = None,
vae_offload_device: Optional[Union[str, torch.device]] = None,
tensor_offload_device: Optional[Union[str, torch.device]] = None,
debug: Optional['Debug'] = None
) -> Dict[str, Any]:
"""
Initialize generation context with device configuration.
Processes device objects, configures environment variables, and creates the
generation context dictionary with all necessary state.
Args:
dit_device: Device for DiT model (str or torch.device, defaults to 'cpu')
vae_device: Device for VAE model (str or torch.device, defaults to 'cpu')
dit_offload_device: Device to offload DiT to when not in use (optional)
vae_offload_device: Device to offload VAE to when not in use (optional)
tensor_offload_device: Device to offload intermediate tensors to (optional)
debug: Debug instance for logging
Returns:
Dict[str, Any]: Generation context dictionary with torch.device objects
"""
# Normalize devices to torch.device objects (follows PyTorch convention)
def _normalize_device(device_spec: Optional[Union[str, torch.device]]) -> torch.device:
"""Convert device specification to torch.device object."""
if device_spec is None:
return torch.device("cpu")
if isinstance(device_spec, torch.device):
return device_spec
return torch.device(device_spec)
dit_device = _normalize_device(dit_device)
vae_device = _normalize_device(vae_device)
dit_offload_device = _normalize_device(dit_offload_device) if dit_offload_device is not None else None
vae_offload_device = _normalize_device(vae_offload_device) if vae_offload_device is not None else None
# MPS unified memory: CPU offload causes sync overhead with no memory benefit
is_mps = dit_device.type == 'mps' or vae_device.type == 'mps'
if is_mps and tensor_offload_device is not None and str(tensor_offload_device) == 'cpu':
tensor_offload_device = None
else:
tensor_offload_device = _normalize_device(tensor_offload_device) if tensor_offload_device is not None else None
# Set LOCAL_RANK to 0 for single-GPU inference mode
# CLI multi-GPU uses CUDA_VISIBLE_DEVICES to restrict visibility per worker
os.environ.setdefault("LOCAL_RANK", "0")
# Detect ComfyUI integration for interrupt support
try:
import comfy.model_management
interrupt_fn = comfy.model_management.throw_exception_if_processing_interrupted
comfyui_available = True
except:
interrupt_fn = None
comfyui_available = False
# Create generation context
ctx = {
'dit_device': dit_device,
'vae_device': vae_device,
'dit_offload_device': dit_offload_device,
'vae_offload_device': vae_offload_device,
'tensor_offload_device': tensor_offload_device,
'compute_dtype': COMPUTE_DTYPE,
'interrupt_fn': interrupt_fn,
'video_transform': None,
'text_embeds': None,
'all_latents': [],
'all_upscaled_latents': [],
'batch_samples': [],
'final_video': None,
'comfyui_available': comfyui_available,
}
if debug:
# Build device configuration summary
offload_info = []
if dit_offload_device:
offload_info.append(f"DiT offload={str(dit_offload_device)}")
if vae_offload_device:
offload_info.append(f"VAE offload={str(vae_offload_device)}")
if tensor_offload_device:
offload_info.append(f"Tensor offload={str(tensor_offload_device)}")
offload_str = ", ".join(offload_info) if offload_info else "none"
debug.log(
f"Generation context initialized: "
f"DiT={str(dit_device)}, VAE={str(vae_device)}, "
f"Offload=[{offload_str}], "
f"LOCAL_RANK={os.environ['LOCAL_RANK']}",
category="setup"
)
if ctx['compute_dtype'] == torch.float32:
reason = "quality"
elif not BFLOAT16_SUPPORTED:
reason = "compatibility (GPU lacks bfloat16 CUBLAS - 7B models unsupported, 3B may have artifacts)"
else:
reason = "performance"
debug.log(f"Unified compute dtype: {ctx['compute_dtype']} across entire pipeline for maximum {reason}", category="precision")
return ctx
def prepare_runner(
dit_model: str,
vae_model: str,
model_dir: str,
debug: 'Debug',
ctx: Dict[str, Any],
dit_cache: bool = False,
vae_cache: bool = False,
dit_id: Optional[int] = None,
vae_id: Optional[int] = None,
block_swap_config: Optional[Dict[str, Any]] = None,
encode_tiled: bool = False,
encode_tile_size: Optional[Tuple[int, int]] = None,
encode_tile_overlap: Optional[Tuple[int, int]] = None,
decode_tiled: bool = False,
decode_tile_size: Optional[Tuple[int, int]] = None,
decode_tile_overlap: Optional[Tuple[int, int]] = None,
tile_debug: str = "false",
attention_mode: str = 'sdpa',
torch_compile_args_dit: Optional[Dict[str, Any]] = None,
torch_compile_args_vae: Optional[Dict[str, Any]] = None
) -> Tuple['VideoDiffusionInfer', Dict[str, Any]]:
"""
Prepare runner with model state management and global cache integration.
Handles model changes and caching logic with independent DiT/VAE caching support.
Args:
dit_model: DiT model filename (e.g., "seedvr2_ema_3b_fp16.safetensors")
vae_model: VAE model filename (e.g., "ema_vae_fp16.safetensors")
model_dir: Base directory containing model files
debug: Debug instance for logging (required)
ctx: Generation context from setup_generation_context
dit_cache: Whether to cache DiT model between runs
vae_cache: Whether to cache VAE model between runs
dit_id: Node instance ID for DiT model caching
vae_id: Node instance ID for VAE model caching
block_swap_config: Optional BlockSwap configuration for DiT memory optimization
encode_tiled: Enable tiled encoding to reduce VRAM during VAE encoding
encode_tile_size: Tile size for encoding (height, width)
encode_tile_overlap: Tile overlap for encoding (height, width)
decode_tiled: Enable tiled decoding to reduce VRAM during VAE decoding
decode_tile_size: Tile size for decoding (height, width)
decode_tile_overlap: Tile overlap for decoding (height, width)
tile_debug: Tile visualization mode (false/encode/decode)
attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
torch_compile_args_dit: Optional torch.compile configuration for DiT model
torch_compile_args_vae: Optional torch.compile configuration for VAE model
Returns:
Tuple['VideoDiffusionInfer', Dict[str, Any]]: Tuple containing:
- VideoDiffusionInfer: Configured runner instance with models loaded and settings applied
- Dict[str, Any]: Cache context dictionary containing cache state and metadata with keys:
- 'global_cache': GlobalModelCache instance
- 'dit_cache', 'vae_cache': Caching enabled flags
- 'dit_id', 'vae_id': Node IDs for cache lookup
- 'cached_dit', 'cached_vae': Cached model instances (if found)
- 'reusing_runner': Flag indicating if runner template was reused
Features:
- Independent DiT and VAE caching for flexible memory management
- Dynamic model reloading when models change
- Optional torch.compile optimization for inference speedup
- Separate encode/decode tiling configuration for optimal performance
- Memory optimization and BlockSwap integration
"""
dit_changed = False
vae_changed = False
# Configure runner
debug.log("Configuring inference runner...", category="runner")
runner, cache_context = configure_runner(
dit_model=dit_model,
vae_model=vae_model,
base_cache_dir=model_dir,
debug=debug,
ctx=ctx,
dit_cache=dit_cache,
vae_cache=vae_cache,
dit_id=dit_id,
vae_id=vae_id,
block_swap_config=block_swap_config,
encode_tiled=encode_tiled,
encode_tile_size=encode_tile_size,
encode_tile_overlap=encode_tile_overlap,
decode_tiled=decode_tiled,
decode_tile_size=decode_tile_size,
decode_tile_overlap=decode_tile_overlap,
tile_debug=tile_debug,
attention_mode=attention_mode,
torch_compile_args_dit=torch_compile_args_dit,
torch_compile_args_vae=torch_compile_args_vae
)
return runner, cache_context
def load_text_embeddings(script_directory: str, device: torch.device,
dtype: torch.dtype, debug: Optional['Debug'] = None) -> Dict[str, List[torch.Tensor]]:
"""
Load and prepare text embeddings for generation
Args:
script_directory (str): Script directory path
device (torch.device): Target device
dtype (torch.dtype): Target dtype
debug: Optional debug instance for logging
Returns:
dict: Text embeddings dictionary
Features:
- Adaptive dtype handling
- Device-optimized loading
- Memory-efficient embedding preparation
- Consistent movement logging
"""
text_pos_embeds = torch.load(os.path.join(script_directory, 'pos_emb.pt'), weights_only=True)
text_neg_embeds = torch.load(os.path.join(script_directory, 'neg_emb.pt'), weights_only=True)
text_pos_embeds = manage_tensor(
tensor=text_pos_embeds,
target_device=device,
tensor_name="text_pos_embeds",
dtype=dtype,
debug=debug,
reason="DiT inference"
)
text_neg_embeds = manage_tensor(
tensor=text_neg_embeds,
target_device=device,
tensor_name="text_neg_embeds",
dtype=dtype,
debug=debug,
reason="DiT inference"
)
return {"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]}
def calculate_optimal_batch_params(total_frames: int, batch_size: int,
temporal_overlap: int) -> Dict[str, Any]:
"""
Calculate batch processing parameters.
Args:
total_frames (int): Total number of frames to process
batch_size (int): Desired batch size
temporal_overlap (int): Number of overlapping frames between batches
Returns:
dict: {
'step': Effective step size between batches,
'temporal_overlap': Adjusted temporal overlap,
'best_batch': Optimal batch size matching video length (4n+1 format)
}
The 4n+1 constraint (1, 5, 9, 13, 17, 21...) is required by the model.
"""
# Calculate step size
step = batch_size - temporal_overlap
if step <= 0:
step = batch_size
temporal_overlap = 0
# Find all valid 4n+1 batch sizes up to total_frames
valid_sizes = [i for i in range(1, total_frames + 1) if i % 4 == 1]
# Best batch: largest valid size ≤ total_frames (maximizes temporal stability)
best_batch = max(valid_sizes) if valid_sizes else 1
return {
'step': step,
'temporal_overlap': temporal_overlap,
'best_batch': best_batch
}
def pad_video_temporal(videos: torch.Tensor, count: int = 0, temporal_dim: int = 1,
prepend: bool = False, debug: Optional['Debug'] = None) -> torch.Tensor:
"""
Extend video with reversed frames for temporal continuity.
Single source of truth for all temporal padding (prepend/append, 4n+1 constraint).
Args:
videos: Video tensor (TCHW or CTHW format)
count: Frames to add (0 = auto-calculate for 4n+1 when prepend=False)
temporal_dim: Temporal axis (0=TCHW, 1=CTHW)
prepend: Add to start (True) or end (False)
debug: Optional debug logger
Returns:
torch.Tensor: Extended video
"""
t = videos.size(temporal_dim)
# Auto-calculate for 4n+1 constraint
if count == 0 and not prepend:
if t % 4 == 1:
return videos
count = ((t - 1) // 4 + 1) * 4 + 1 - t
if count <= 0:
return videos
if debug and prepend:
debug.log(f"Prepending {count} reversed frames to reduce start artifacts",
category="video", indent_level=1)
# Helper to select frames along temporal dimension
def select(start, end):
return videos[start:end] if temporal_dim == 0 else videos[:, start:end]
# Padding exceeds video length: repeat + reverse
if count >= t:
repeat_count = count - t + 1
last = select(-1, None)
# Repeat last frame
if temporal_dim == 0:
repeated = last.repeat(repeat_count, 1, 1, 1)
reversed_frames = select(1, None).flip(temporal_dim) if t > 1 else last[:0]
else:
repeated = last.expand(-1, repeat_count, -1, -1).contiguous()
reversed_frames = select(1, None).flip(temporal_dim) if t > 1 else last[:, :0]
return torch.cat([repeated, reversed_frames, videos] if prepend else
[videos, reversed_frames, repeated], dim=temporal_dim)
# Normal case: extract and reverse frames
if prepend:
reversed_frames = select(1, count+1).flip(temporal_dim)
else:
reversed_frames = select(-count-1, -1).flip(temporal_dim)
return torch.cat([reversed_frames, videos] if prepend else
[videos, reversed_frames], dim=temporal_dim)
def check_interrupt(ctx: Dict[str, Any]) -> None:
"""Single interrupt check to avoid redundant imports"""
if ctx.get('interrupt_fn') is not None:
ctx['interrupt_fn']()
def _draw_tile_boundaries(image: torch.Tensor, debug: 'Debug', tile_boundaries: list, phase: str) -> torch.Tensor:
"""
Draw tile boundary overlays on all frames for debugging (non-destructive).
Args:
image: Image tensor [T, H, W, C] or [H, W, C] in range [0, 1]
debug: Debug instance for logging
tile_boundaries: List of tile boundary info dictionaries
phase: Phase name ('encode' or 'decode') for logging
Returns:
Image with boundary overlays drawn inside tiles on all frames
"""
if not tile_boundaries:
return image
# Try to import required libraries
try:
import cv2
import numpy as np
import random
import colorsys
except ImportError as e:
debug.log(f"Tile debug ignored: missing imports ({e})", level="WARNING", category="video")
return image
# Handle both [T, H, W, C] and [H, W, C]
squeeze_t = False
if image.ndim == 3:
image = image.unsqueeze(0)
squeeze_t = True
T, H, W, C = image.shape
original_dtype = image.dtype
log_frames = f"all {T} frames" if T > 1 else "1 frame"
debug.log(f"Drawing {phase} tile boundaries ({len(tile_boundaries)} tiles) on {log_frames}", category="video", indent_level=1, force=True)
# Scale line thickness and font size based on video width
# Reference points: 512px (min) to 1920px (max)
min_width, max_width = 512, 1920
min_line_thickness, max_line_thickness = 2, 6
min_font_scale, max_font_scale = 0.8, 2.5
min_text_thickness, max_text_thickness = 2, 4
# Calculate scale factor (clamped between 0 and 1)
scale_factor = max(0.0, min(1.0, (W - min_width) / (max_width - min_width)))
# Apply scaling
line_thickness = int(min_line_thickness + scale_factor * (max_line_thickness - min_line_thickness))
font_scale = min_font_scale + scale_factor * (max_font_scale - min_font_scale)
text_thickness = int(min_text_thickness + scale_factor * (max_text_thickness - min_text_thickness))
# Generate high-contrast colors using HSV color space
num_tiles = len(tile_boundaries)
colors = []
for i in range(num_tiles):
hue = (i * 360 / num_tiles) % 360
saturation = 0.9 + (i % 2) * 0.1
brightness = 0.8 + ((i // 2) % 2) * 0.2
r, g, b = colorsys.hsv_to_rgb(hue / 360, saturation, brightness)
colors.append((int(b * 255), int(g * 255), int(r * 255))) # BGR for OpenCV
random.seed(42)
random.shuffle(colors)
# Process all frames
annotated_frames = []
for frame_idx in range(T):
# Convert frame to numpy (handle RGB and RGBA)
img = np.ascontiguousarray((image[frame_idx].float().cpu().numpy() * 255).astype(np.uint8)) # [H, W, C]
# Draw boundary lines inside each tile
for idx, tile_info in enumerate(tile_boundaries):
tile_id = tile_info['id']
x, y = tile_info['x'], tile_info['y']
w, h = tile_info['w'], tile_info['h']
color = colors[idx]
inset = line_thickness // 2
# Draw four edges
cv2.line(img, (x, y + inset), (x + w, y + inset), color, line_thickness)
cv2.line(img, (x, y + h - inset), (x + w, y + h - inset), color, line_thickness)
cv2.line(img, (x + inset, y), (x + inset, y + h), color, line_thickness)
cv2.line(img, (x + w - inset, y), (x + w - inset, y + h), color, line_thickness)
# Draw tile number
text = str(tile_id)
font = cv2.FONT_HERSHEY_SIMPLEX
(text_w, text_h), baseline = cv2.getTextSize(text, font, font_scale, text_thickness)
margin = int(15 * scale_factor) if scale_factor > 0.5 else 8 # Scale margin too
text_x = x + margin
text_y = y + text_h + margin
cv2.putText(img, text, (text_x, text_y), font, font_scale, color, text_thickness, cv2.LINE_AA)
# Convert back to tensor
frame_tensor = torch.from_numpy(img.astype(np.float32) / 255.0).to(device=image.device, dtype=original_dtype)
annotated_frames.append(frame_tensor)
# Stack all frames
image = torch.stack(annotated_frames, dim=0)
if squeeze_t:
image = image.squeeze(0)
return image
def ensure_precision_initialized(
ctx: Dict[str, Any],
runner: 'VideoDiffusionInfer',
debug: Optional['Debug'] = None
) -> None:
"""
Log model dtypes for debugging. Compute dtype is hardcoded in context.
Since compute_dtype is hardcoded to bfloat16 in setup_generation_context(),
this function only logs model dtypes for informational purposes.
Args:
ctx: Generation context dictionary (compute_dtype already set)
runner: VideoDiffusionInfer instance with loaded models
debug: Optional Debug instance for logging
"""
if not debug:
return
try:
# Get model dtypes for informational logging
dit_dtype = None
vae_dtype = None
if runner.dit is not None:
try:
param_device = next(runner.dit.parameters()).device
if param_device.type != 'meta':
dit_dtype = next(runner.dit.parameters()).dtype
except StopIteration:
pass
if runner.vae is not None:
try:
param_device = next(runner.vae.parameters()).device
if param_device.type != 'meta':
vae_dtype = next(runner.vae.parameters()).dtype
except StopIteration:
pass
# Build precision info string
parts = []
if dit_dtype is not None:
parts.append(f"DiT={dit_dtype}")
if vae_dtype is not None:
parts.append(f"VAE={vae_dtype}")
parts.append(f"compute={ctx['compute_dtype']}")
if parts:
debug.log(f"Model precision: {', '.join(parts)}", category="precision")
except Exception as e:
debug.log(f"Could not log model dtypes: {e}", level="WARNING", category="precision", force=True)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/core/infer.py
+395
View File
@@ -0,0 +1,395 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import List, Optional, Tuple, Union
import torch
from einops import rearrange
from omegaconf import DictConfig, ListConfig
from torch import Tensor
from ..common.diffusion import (
classifier_free_guidance_dispatcher,
create_sampler_from_config,
create_sampling_timesteps_from_config,
create_schedule_from_config,
)
from ..common.distributed import (
get_device,
)
from ..optimization.performance import (
optimized_channels_to_last,
optimized_channels_to_second
)
from ..models.dit_3b import na
class VideoDiffusionInfer():
def __init__(self, config: DictConfig, debug: 'Debug',
encode_tiled: bool = False, encode_tile_size: Tuple[int, int] = (512, 512),
encode_tile_overlap: Tuple[int, int] = (64, 64),
decode_tiled: bool = False, decode_tile_size: Tuple[int, int] = (512, 512),
decode_tile_overlap: Tuple[int, int] = (64, 64),
tile_debug: str = "false"):
self.config = config
self.debug = debug
# Store separate encode and decode tiling parameters
self.encode_tiled = encode_tiled
self.encode_tile_size = encode_tile_size
self.encode_tile_overlap = encode_tile_overlap
self.decode_tiled = decode_tiled
self.decode_tile_size = decode_tile_size
self.decode_tile_overlap = decode_tile_overlap
self.tile_debug = tile_debug
def get_condition(self, latent: Tensor, latent_blur: Tensor, task: str) -> Tensor:
t, h, w, c = latent.shape
cond = torch.zeros([t, h, w, c + 1], device=latent.device, dtype=latent.dtype)
if task == "t2v" or t == 1:
# t2i or t2v generation.
if task == "sr":
cond[:, ..., :-1] = latent_blur[:]
cond[:, ..., -1:] = 1.0
return cond
if task == "i2v":
# i2v generation.
cond[:1, ..., :-1] = latent[:1]
cond[:1, ..., -1:] = 1.0
return cond
if task == "v2v":
# v2v frame extension.
cond[:2, ..., :-1] = latent[:2]
cond[:2, ..., -1:] = 1.0
return cond
if task == "sr":
# sr generation.
cond[:, ..., :-1] = latent_blur[:]
cond[:, ..., -1:] = 1.0
return cond
raise NotImplementedError
def configure_diffusion(self, device: Optional[torch.device] = None, dtype=torch.float32):
"""
Configure diffusion schedule and sampler.
Args:
device: Device for schedule tensors. If None, uses get_device()
dtype: Data type for computations
"""
# Use provided device or fallback to standard detection
if device is None:
device = get_device()
elif not isinstance(device, torch.device):
device = torch.device(device)
self.schedule = create_schedule_from_config(
config=self.config.diffusion.schedule,
device=device,
dtype=dtype,
)
self.sampling_timesteps = create_sampling_timesteps_from_config(
config=self.config.diffusion.timesteps.sampling,
schedule=self.schedule,
device=device,
dtype=dtype,
)
self.sampler = create_sampler_from_config(
config=self.config.diffusion.sampler,
schedule=self.schedule,
timesteps=self.sampling_timesteps,
)
# Propagate debug to sampler
if hasattr(self, 'debug'):
self.sampler.debug = self.debug
# -------------------------------- Helper ------------------------------- #
@torch.no_grad()
def vae_encode(self, samples: List[Tensor]) -> List[Tensor]:
"""VAE encode with configured dtype - converts samples to latents with optional tiling"""
use_sample = self.config.vae.get("use_sample", True)
latents = []
if len(samples) > 0:
# Use VAE model's current device
# This ensures consistency with where the VAE model is loaded
try:
device = next(self.vae.parameters()).device
except StopIteration:
# Fallback if VAE has no parameters (shouldn't happen)
device = get_device()
dtype = getattr(torch, self.config.vae.dtype)
scale = self.config.vae.scaling_factor
shift = self.config.vae.get("shifting_factor", 0.0)
if isinstance(scale, ListConfig):
scale = torch.tensor(scale, device=device, dtype=dtype)
if isinstance(shift, ListConfig):
shift = torch.tensor(shift, device=device, dtype=dtype)
# Group samples of the same shape to batches if enabled.
if self.config.vae.grouping:
batches, indices = na.pack(samples)
else:
batches = [sample.unsqueeze(0) for sample in samples]
# VAE process by each group.
for sample in batches:
if hasattr(self.vae, "preprocess"):
sample = self.vae.preprocess(sample)
# Detect VAE model dtype
try:
vae_dtype = next(self.vae.parameters()).dtype
except StopIteration:
vae_dtype = dtype # Fallback
# Use autocast if VAE dtype differs from input dtype
# Skip autocast on MPS (only supports bf16, unified memory = no benefit)
# Instead, explicitly convert input to model dtype
if vae_dtype != sample.dtype:
if device.type == 'mps':
# MPS: explicit dtype conversion instead of autocast
sample = sample.to(vae_dtype)
if use_sample:
latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
tile_overlap=self.encode_tile_overlap).latent
else:
latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
tile_overlap=self.encode_tile_overlap).posterior.mode().squeeze(2)
else:
with torch.autocast(device.type, sample.dtype, enabled=True):
if use_sample:
latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
tile_overlap=self.encode_tile_overlap).latent
else:
latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
tile_overlap=self.encode_tile_overlap).posterior.mode().squeeze(2)
else:
if use_sample:
latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
tile_overlap=self.encode_tile_overlap).latent
else:
# Deterministic vae encode, only used for i2v inference (optionally)
latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
tile_overlap=self.encode_tile_overlap).posterior.mode().squeeze(2)
latent = latent.unsqueeze(2) if latent.ndim == 4 else latent
latent = optimized_channels_to_last(latent)
latent = (latent - shift) * scale
latents.append(latent)
# Ungroup back to individual latent with the original order.
if self.config.vae.grouping:
latents = na.unpack(latents, indices)
else:
latents = [latent.squeeze(0) for latent in latents]
self.debug.log(f"Latents shape: {latents[0].shape}", category="info", indent_level=1)
return latents
@torch.no_grad()
def vae_decode(self, latents: List[Tensor]) -> List[Tensor]:
"""VAE decode with configured dtype - converts latents to samples with optional tiling"""
samples = []
if len(latents) > 0:
# Use VAE model's current device
# This ensures consistency with where the VAE model is loaded
try:
device = next(self.vae.parameters()).device
except StopIteration:
# Fallback if VAE has no parameters (shouldn't happen)
device = get_device()
dtype = getattr(torch, self.config.vae.dtype)
scale = self.config.vae.scaling_factor
shift = self.config.vae.get("shifting_factor", 0.0)
if isinstance(scale, ListConfig):
scale = torch.tensor(scale, device=device, dtype=dtype)
if isinstance(shift, ListConfig):
shift = torch.tensor(shift, device=device, dtype=dtype)
# Group samples of the same shape to batches if enabled.
if self.config.vae.grouping:
latents, indices = na.pack(latents)
else:
latents = [latent.unsqueeze(0) for latent in latents]
self.debug.log(f"Latents shape: {latents[0].shape}", category="info", indent_level=1)
for i, latent in enumerate(latents):
latent = latent / scale + shift
latent = optimized_channels_to_second(latent)
latent = latent.squeeze(2)
# Detect VAE model dtype
try:
vae_dtype = next(self.vae.parameters()).dtype
except StopIteration:
vae_dtype = dtype # Fallback
# Use autocast if VAE dtype differs from latent dtype
# Skip autocast on MPS (only supports bf16, unified memory = no benefit)
if vae_dtype != latent.dtype:
if device.type == 'mps':
# MPS: explicit dtype conversion instead of autocast
latent = latent.to(vae_dtype)
sample = self.vae.decode(
latent,
tiled=self.decode_tiled, tile_size=self.decode_tile_size,
tile_overlap=self.decode_tile_overlap
).sample
else:
with torch.autocast(device.type, latent.dtype, enabled=True):
sample = self.vae.decode(
latent,
tiled=self.decode_tiled, tile_size=self.decode_tile_size,
tile_overlap=self.decode_tile_overlap
).sample
else:
sample = self.vae.decode(
latent,
tiled=self.decode_tiled, tile_size=self.decode_tile_size,
tile_overlap=self.decode_tile_overlap
).sample
if hasattr(self.vae, "postprocess"):
sample = self.vae.postprocess(sample)
samples.append(sample)
if self.config.vae.grouping:
samples = na.unpack(samples, indices)
else:
samples = [sample.squeeze(0) for sample in samples]
return samples
def timestep_transform(self, timesteps: Tensor, latents_shapes: Tensor):
# Skip if not needed.
if not self.config.diffusion.timesteps.get("transform", False):
return timesteps
# Compute resolution.
vt = self.config.vae.model.get("temporal_downsample_factor", 4)
vs = self.config.vae.model.get("spatial_downsample_factor", 8)
frames = (latents_shapes[:, 0] - 1) * vt + 1
heights = latents_shapes[:, 1] * vs
widths = latents_shapes[:, 2] * vs
# Compute shift factor.
def get_lin_function(x1, y1, x2, y2):
m = (y2 - y1) / (x2 - x1)
b = y1 - m * x1
return lambda x: m * x + b
img_shift_fn = get_lin_function(x1=256 * 256, y1=1.0, x2=1024 * 1024, y2=3.2)
vid_shift_fn = get_lin_function(x1=256 * 256 * 37, y1=1.0, x2=1280 * 720 * 145, y2=5.0)
shift = torch.where(
frames > 1,
vid_shift_fn(heights * widths * frames),
img_shift_fn(heights * widths),
)
# Shift timesteps.
timesteps = timesteps / self.schedule.T
timesteps = shift * timesteps / (1 + (shift - 1) * timesteps)
timesteps = timesteps * self.schedule.T
return timesteps
@torch.no_grad()
def inference(
self,
noises: List[Tensor],
conditions: List[Tensor],
texts_pos: Union[List[str], List[Tensor], List[Tuple[Tensor]]],
texts_neg: Union[List[str], List[Tensor], List[Tuple[Tensor]]],
cfg_scale: Optional[float] = None,
) -> List[Tensor]:
assert len(noises) == len(conditions) == len(texts_pos) == len(texts_neg)
batch_size = len(noises)
# Return if empty.
if batch_size == 0:
return []
# Set cfg scale
if cfg_scale is None:
cfg_scale = self.config.diffusion.cfg.scale
# Text embeddings.
assert type(texts_pos[0]) is type(texts_neg[0])
if isinstance(texts_pos[0], str):
text_pos_embeds, text_pos_shapes = self.text_encode(texts_pos)
text_neg_embeds, text_neg_shapes = self.text_encode(texts_neg)
elif isinstance(texts_pos[0], tuple):
text_pos_embeds, text_pos_shapes = [], []
text_neg_embeds, text_neg_shapes = [], []
for pos in zip(*texts_pos):
emb, shape = na.flatten(pos)
text_pos_embeds.append(emb)
text_pos_shapes.append(shape)
for neg in zip(*texts_neg):
emb, shape = na.flatten(neg)
text_neg_embeds.append(emb)
text_neg_shapes.append(shape)
else:
text_pos_embeds, text_pos_shapes = na.flatten(texts_pos)
text_neg_embeds, text_neg_shapes = na.flatten(texts_neg)
# Flatten.
latents, latents_shapes = na.flatten(noises)
latents_cond, _ = na.flatten(conditions)
latents = self.sampler.sample(
x=latents,
f=lambda args: classifier_free_guidance_dispatcher(
pos=lambda: self.dit(
vid=torch.cat([args.x_t, latents_cond], dim=-1),
txt=text_pos_embeds,
vid_shape=latents_shapes,
txt_shape=text_pos_shapes,
timestep=args.t.repeat(batch_size),
).vid_sample,
neg=lambda: self.dit(
vid=torch.cat([args.x_t, latents_cond], dim=-1),
txt=text_neg_embeds,
vid_shape=latents_shapes,
txt_shape=text_neg_shapes,
timestep=args.t.repeat(batch_size),
).vid_sample,
scale=(
cfg_scale
if (args.i + 1) / len(self.sampler.timesteps)
<= self.config.diffusion.cfg.get("partial", 1)
else 1.0
),
rescale=self.config.diffusion.cfg.rescale,
),
)
latents = na.unflatten(latents, latents_shapes)
# Clean up temporary tensors
del latents_cond
del latents_shapes
del text_pos_embeds
del text_neg_embeds
del text_pos_shapes
del text_neg_shapes
return latents
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_cache.py
+239
View File
@@ -0,0 +1,239 @@
"""
Global Model Cache for SeedVR2
Enables independent DiT and VAE model sharing across multiple upscaler node instances
"""
from typing import Dict, Any, Optional, Tuple
from ..optimization.memory_manager import release_model_memory
class GlobalModelCache:
"""
Global cache for sharing DiT and VAE models independently across upscaler instances.
Each upscaler gets its own runner but can share the underlying models.
Also caches runner templates for instant runner creation.
"""
def __init__(self):
# Storage for cached DiT models: node_id -> (model, config)
self._dit_models: Dict[str, Tuple[Any, Dict]] = {}
# Storage for cached VAE models: node_id -> (model, config)
self._vae_models: Dict[str, Tuple[Any, Dict]] = {}
# Storage for runner templates: "dit_id+vae_id" -> runner
self._runner_templates: Dict[str, Any] = {}
def get_dit(self, dit_config: Dict[str, Any], debug: Optional['Debug'] = None) -> Optional[Any]:
"""
Get cached DiT model if available.
Args:
dit_config: Configuration dictionary with 'cache_model' and 'node_id' keys
debug: Optional debug instance for logging
Returns:
Cached DiT model instance if found, None otherwise
"""
if not dit_config.get('cache_model', False):
return None
node_id = dit_config.get('node_id')
if node_id in self._dit_models:
model, stored_config = self._dit_models[node_id]
return model
return None
def get_vae(self, vae_config: Dict[str, Any], debug: Optional['Debug'] = None) -> Optional[Any]:
"""
Get cached VAE model if available.
Args:
vae_config: Configuration dictionary with 'cache_model' and 'node_id' keys
debug: Optional debug instance for logging
Returns:
Cached VAE model instance if found, None otherwise
"""
if not vae_config.get('cache_model', False):
return None
node_id = vae_config.get('node_id')
if node_id in self._vae_models:
model, stored_config = self._vae_models[node_id]
return model
return None
def get_runner(self, dit_id: Optional[int], vae_id: Optional[int],
debug: Optional['Debug'] = None) -> Optional[Any]:
"""
Get cached runner template if available.
Args:
dit_id: DiT node ID for lookup
vae_id: VAE node ID for lookup
debug: Optional debug instance for logging
Returns:
Cached runner template if found, None if either ID is None or not cached
"""
if dit_id is None or vae_id is None:
return None
runner_key = f"{dit_id}+{vae_id}"
if runner_key in self._runner_templates:
return self._runner_templates[runner_key]
return None
def set_dit(self, dit_config: Dict[str, Any], model: Any, model_name: str, debug: Optional['Debug'] = None) -> Optional[str]:
"""
Store DiT model in cache.
Args:
dit_config: Configuration dictionary with 'cache_model' and 'node_id' keys
model: DiT model instance to cache
model_name: Name identifier for the model
debug: Optional debug instance for logging
Returns:
Node ID string if cached successfully, None if caching disabled
"""
if not dit_config.get('cache_model', False):
return None
node_id = dit_config.get('node_id')
self._dit_models[node_id] = (model, dit_config)
if debug:
debug.log(f"DiT model cached in memory (node {node_id}): {model_name}",
category="cache", force=True)
return node_id
def set_vae(self, vae_config: Dict[str, Any], model: Any, model_name: str, debug: Optional['Debug'] = None) -> Optional[str]:
"""
Store VAE model in cache.
Args:
vae_config: Configuration dictionary with 'cache_model' and 'node_id' keys
model: VAE model instance to cache
model_name: Name identifier for the model
debug: Optional debug instance for logging
Returns:
Node ID string if cached successfully, None if caching disabled
"""
if not vae_config.get('cache_model', False):
return None
node_id = vae_config.get('node_id')
self._vae_models[node_id] = (model, vae_config)
if debug:
debug.log(f"VAE model cached in memory (node {node_id}): {model_name}",
category="cache", force=True)
return node_id
def set_runner(self, dit_id: Optional[int], vae_id: Optional[int],
runner: Any, debug: Optional['Debug'] = None) -> Optional[str]:
"""
Store runner template in cache.
Args:
dit_id: DiT node ID for the runner template
vae_id: VAE node ID for the runner template
runner: Runner instance to cache as template
debug: Optional debug instance for logging
Returns:
Runner key string (format: "dit_id+vae_id") if cached successfully,
None if either ID is None or runner already cached
"""
if dit_id is None or vae_id is None:
return None
runner_key = f"{dit_id}+{vae_id}"
if runner_key not in self._runner_templates:
self._runner_templates[runner_key] = runner
if debug:
debug.log(f"Runner template cached in memory: nodes {runner_key}", category="cache", force=True)
return runner_key
return None
def remove_dit(self, dit_config: Dict[str, Any], debug: Optional['Debug'] = None) -> bool:
"""
Remove DiT model from cache if it exists.
Args:
dit_config: Configuration dictionary with 'node_id' key
debug: Optional debug instance for logging
Returns:
True if model was removed, False if not found in cache
Note:
Also removes any runner templates that used this DiT model
"""
node_id = dit_config.get('node_id')
if node_id in self._dit_models:
if debug:
debug.log(f"Removing cached DiT: {node_id}", category="cache", force=True)
model, stored_config = self._dit_models[node_id]
# Release model memory
if model is not None:
release_model_memory(model=model, debug=debug)
del self._dit_models[node_id]
# Remove any runner templates that used this DiT
templates_to_remove = [k for k in self._runner_templates.keys() if k.startswith(str(node_id) + "+")]
for template_key in templates_to_remove:
del self._runner_templates[template_key]
return True
return False
def remove_vae(self, vae_config: Dict[str, Any], debug: Optional['Debug'] = None) -> bool:
"""
Remove VAE model from cache if it exists.
Args:
vae_config: Configuration dictionary with 'node_id' key
debug: Optional debug instance for logging
Returns:
True if model was removed, False if not found in cache
Note:
Also removes any runner templates that used this VAE model
"""
node_id = vae_config.get('node_id')
if node_id in self._vae_models:
if debug:
debug.log(f"Removing cached VAE: {node_id}", category="cache", force=True)
model, stored_config = self._vae_models[node_id]
# Release model memory directly
if model is not None:
release_model_memory(model=model, debug=debug)
del self._vae_models[node_id]
# Remove any runner templates that used this VAE
templates_to_remove = [k for k in self._runner_templates.keys() if k.endswith("+" + str(node_id))]
for template_key in templates_to_remove:
del self._runner_templates[template_key]
return True
return False
# Global singleton instance
_global_cache = GlobalModelCache()
def get_global_cache() -> GlobalModelCache:
"""Get the global model cache instance."""
return _global_cache
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_configuration.py
+1480
View File
File diff suppressed because it is too large Load Diff
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_loader.py
+965
View File
@@ -0,0 +1,965 @@
"""
Model Weight Loading for SeedVR2
This module handles all weight loading operations for DiT and VAE models:
- Loading state dictionaries from multiple formats (SafeTensors, PyTorch, GGUF)
- Materializing models from meta device to target device
- Applying weights with dtype conversion
- GGUF quantized model support with dequantization
- Meta buffer initialization for non-persistent buffers
Key Features:
- Multi-format support: .safetensors, .pth, .gguf files
- Memory-efficient loading with meta device initialization
- Native FP8 weight handling with optimal performance
- GGUF quantization support (Q4_K_M, Q8_0, etc.)
- Automatic dtype conversion for compatibility
- Meta buffer initialization post-materialization
Main Functions:
- load_quantized_state_dict: Load state dict from checkpoint file
- materialize_model: Move model from meta device and load weights
- prepare_model_structure: Create model structure on meta device
GGUF Support:
- apply_gguf_parameters: Apply GGUF weights to model (handles meta and materialized)
- _load_gguf_state: Load GGUF quantized weights from file
- _load_gguf_weights: Apply GGUF weights to model with validation
- _validate_gguf_architecture: Validate GGUF model architecture
- _create_dequantize_method: Create dequantization callable
- _create_gguf_parameter: Create parameter preserving quantization info
- _set_parameter_on_meta_model: Set parameter on meta device model
- _set_parameter_on_materialized_model: Set parameter on materialized model
- _navigate_to_parameter: Navigate to module containing parameter
- _get_tensor_shape: Get logical shape of tensor (handling GGUF)
- _is_quantized_tensor: Check if tensor is GGUF quantized
- _report_parameter_mismatches: Report parameter mismatches
Meta Buffer Initialization:
- initialize_meta_buffers: Initialize meta buffers with timing wrapper
- initialize_meta_buffers_impl: Initialize non-persistent buffers on target device
Standard Loading:
- _load_model_weights: Orchestrate weight loading process
- _load_standard_weights: Apply SafeTensors/PyTorch weights
- _convert_state_dtype: Convert weight dtypes
- _log_weight_stats: Log weight statistics
This module is used by model_configuration for weight loading during materialization.
"""
import os
import torch
from omegaconf import OmegaConf
from typing import Dict, Any, Optional, Tuple, Union, Callable
# Import SafeTensors with fallback
try:
from safetensors.torch import load_file as load_safetensors_file
SAFETENSORS_AVAILABLE = True
except ImportError:
SAFETENSORS_AVAILABLE = False
from .infer import VideoDiffusionInfer
from ..common.config import create_object
from ..optimization.compatibility import (
GGUF_AVAILABLE,
GGMLQuantizationType,
validate_gguf_availability
)
# GGUF-specific imports (only when available)
if GGUF_AVAILABLE:
import gguf
import traceback
from ..optimization.gguf_dequant import dequantize_tensor
from ..optimization.gguf_ops import replace_linear_with_quantized
from ..utils.constants import get_script_directory, suppress_tensor_warnings
# Get script directory for config paths
script_directory = get_script_directory()
def load_quantized_state_dict(checkpoint_path: str, device: torch.device = torch.device("cpu"),
debug: Optional['Debug'] = None) -> Dict[str, torch.Tensor]:
"""
Load model state dict from checkpoint with support for multiple formats.
Handles .safetensors, .gguf, and .pth files. GGUF models support quantization
for memory-efficient loading. Validates required libraries are installed.
Args:
checkpoint_path: Path to checkpoint file
device: Target device for tensor placement (torch.device object, defaults to CPU)
debug: Optional Debug instance for logging
Returns:
dict: State dictionary loaded with appropriate format handler
Notes:
- SafeTensors files use optimized loading with direct device placement
- PyTorch files use memory-mapped loading to reduce RAM usage
"""
device_str = str(device)
if checkpoint_path.endswith('.safetensors'):
if not SAFETENSORS_AVAILABLE:
error_msg = (
f"Cannot load {os.path.basename(checkpoint_path)}\n"
f"SafeTensors library is required but not installed.\n"
f"Please install it with: pip install safetensors"
)
if debug:
debug.log(error_msg, level="ERROR", category="dit", force=True)
debug.log("This is a one-time installation that will enable loading of .safetensors files",
level="INFO", category="info", force=True)
raise ImportError(error_msg)
# Try direct device loading first (optimal path)
try:
state = load_safetensors_file(checkpoint_path, device=device_str)
except RuntimeError as e:
# MPS allocator fallback: some PyTorch/macOS versions have issues with
# direct MPS loading (allocation failures, watermark errors, etc.)
error_msg = str(e).lower()
is_mps_alloc_error = device.type == "mps" and any(
keyword in error_msg for keyword in ["watermark", "allocat", "memory"]
)
if is_mps_alloc_error:
# Transparent fallback - only log if debug enabled
if debug:
debug.log("Using CPU intermediate loading for MPS compatibility",
category="info", indent_level=1)
state = load_safetensors_file(checkpoint_path, device="cpu")
# Tensors will be moved to MPS during model.load_state_dict()
else:
# Re-raise if it's a different error (file corruption, etc.)
raise
elif checkpoint_path.endswith('.gguf'):
validate_gguf_availability(f"load {os.path.basename(checkpoint_path)}", debug)
state = _load_gguf_state(
checkpoint_path=checkpoint_path,
device=device,
debug=debug,
handle_prefix="model.diffusion_model."
)
elif checkpoint_path.endswith('.pth'):
state = torch.load(checkpoint_path, map_location=device_str, mmap=True, weights_only=True)
else:
raise ValueError(f"Unsupported checkpoint format. Expected .safetensors or .pth, got: {checkpoint_path}")
return state
def _load_gguf_state(checkpoint_path: str, device: torch.device, debug: Optional['Debug'] = None,
handle_prefix: str = "model.diffusion_model.") -> Dict[str, torch.Tensor]:
"""
Load GGUF state dict
Args:
checkpoint_path: Path to GGUF file
device: Target device (torch.device object)
debug: Debug instance
handle_prefix: Prefix to strip from tensor names
Returns:
State dictionary with loaded tensors
"""
reader = gguf.GGUFReader(checkpoint_path)
# Filter and strip prefix
has_prefix = False
if handle_prefix is not None:
prefix_len = len(handle_prefix)
tensor_names = set(tensor.name for tensor in reader.tensors)
has_prefix = any(s.startswith(handle_prefix) for s in tensor_names)
tensors = []
for tensor in reader.tensors:
sd_key = tensor_name = tensor.name
if has_prefix:
if not tensor_name.startswith(handle_prefix):
continue
sd_key = tensor_name[prefix_len:]
tensors.append((sd_key, tensor))
state_dict = {}
total_tensors = len(reader.tensors)
device_str = str(device)
debug.log(f"Loading {total_tensors} tensors to {str(device_str)}...", category="dit")
# Suppress expected warnings: GGUF tensors are read-only numpy arrays that trigger warnings when converted
suppress_tensor_warnings()
for i, (sd_key, tensor) in enumerate(tensors):
tensor_name = tensor.name
# Create tensor directly on target device to avoid CPU->GPU copy overhead
# For meta-initialized models, this directly materializes to the target device
torch_tensor = torch.from_numpy(tensor.data).to(device, non_blocking=False)
# Get original shape from metadata or infer from tensor shape
shape = _get_tensor_logical_shape(reader, tensor_name)
if shape is None:
shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
# Handle tensors based on quantization type
if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
# For unquantized tensors, just reshape
torch_tensor = torch_tensor.view(*shape)
else:
# For quantized tensors, keep them quantized but track original shape
torch_tensor = GGUFTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape, debug=debug)
state_dict[sd_key] = torch_tensor
# Progress reporting
if (i + 1) % 100 == 0:
debug.log(f"Loaded {i+1}/{total_tensors} tensors...", category="dit", indent_level=1)
debug.log(f"Successfully loaded {len(state_dict)} tensors to {device_str}", category="success")
return state_dict
def _get_tensor_logical_shape(reader: 'gguf.GGUFReader', tensor_name: str) -> Optional[torch.Size]:
"""
Extract the logical (unquantized) shape from GGUF metadata
"""
field_key = f"comfy.gguf.orig_shape.{tensor_name}"
field = reader.get_field(field_key)
if field is None:
return None
# Has original shape metadata, so we try to decode it.
if len(field.types) != 2 or field.types[0] != gguf.GGUFValueType.ARRAY or field.types[1] != gguf.GGUFValueType.INT32:
raise TypeError(f"Bad original shape metadata for {field_key}: Expected ARRAY of INT32, got {field.types}")
return torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data))
class GGUFTensor(torch.Tensor):
"""
Tensor wrapper for GGUF quantized tensors that preserves quantization info
"""
def __init__(self, *args, tensor_type, tensor_shape, **kwargs):
super().__init__()
self.tensor_type = tensor_type
self.tensor_shape = tensor_shape
def __new__(cls, *args, tensor_type, tensor_shape, debug, **kwargs):
# Create tensor with requires_grad=False to avoid gradient issues
tensor = super().__new__(cls, *args, **kwargs)
tensor.requires_grad_(False)
tensor.tensor_type = tensor_type
tensor.tensor_shape = tensor_shape
tensor.debug = debug
return tensor
def to(self, *args, **kwargs):
new = super().to(*args, **kwargs)
new.tensor_type = getattr(self, "tensor_type", None)
new.tensor_shape = getattr(self, "tensor_shape", self.tensor_shape if hasattr(self, "tensor_shape") else new.shape)
new.debug = getattr(self, "debug", None)
new.requires_grad_(False) # Ensure no gradients
return new
@property
def shape(self):
# Always return the logical tensor shape, not the quantized data shape
if hasattr(self, "tensor_shape"):
return self.tensor_shape
else:
# Fallback to actual data shape if tensor_shape is not available
return self.size()
def size(self, *args):
# Override size() to also return logical shape
if hasattr(self, "tensor_shape") and len(args) == 0:
return self.tensor_shape
elif hasattr(self, "tensor_shape") and len(args) == 1:
return self.tensor_shape[args[0]]
else:
return super().size(*args)
def dequantize(self, device=None, dtype=torch.float16, dequant_dtype=None):
"""Dequantize this tensor to its original shape"""
if device is None:
device = self.device
# Suppress expected warning when converting from GGUFTensor subclass to regular tensor
suppress_tensor_warnings()
# Check if already unquantized
if self.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
# Return regular tensor, not GGUFTensor
result = self.to(device, dtype)
if isinstance(result, GGUFTensor):
# Convert to regular tensor to avoid __torch_function__ calls
result = torch.tensor(result, dtype=dtype, device=device, requires_grad=False)
return result
# Try fast dequantization with crash protection
try:
result = dequantize_tensor(self, dtype, dequant_dtype)
final_result = result.to(device)
# Ensure we return a regular tensor, not GGUFTensor
if isinstance(final_result, GGUFTensor):
final_result = torch.tensor(final_result.data, dtype=dtype, device=device, requires_grad=False)
return final_result
except Exception as e:
self.debug.log(f"Fast dequantization failed: {e}", level="WARNING", category="dit", force=True)
self.debug.log(f"Falling back to numpy dequantization", level="WARNING", category="dit", force=True)
# Fallback to numpy (slower but reliable)
try:
numpy_data = self.cpu().numpy()
dequantized = gguf.quants.dequantize(numpy_data, self.tensor_type)
result = torch.from_numpy(dequantized).to(device, dtype)
result.requires_grad_(False)
final_result = result.reshape(self.tensor_shape)
# from_numpy already returns a regular tensor, no conversion needed
return final_result
except Exception as e:
self.debug.log(f"Numpy fallback also failed: {e}", level="WARNING", category="dit", force=True)
self.debug.log(f"Tensor type: {self.tensor_type}", level="WARNING", category="dit", force=True, indent_level=1)
self.debug.log(f"Shape: {self.shape}", level="WARNING", category="dit", force=True, indent_level=1)
self.debug.log(f"Target shape: {self.tensor_shape}", level="WARNING", category="dit", force=True, indent_level=1)
traceback.print_exc()
# Return regular tensor as last resort
result = self.to(device, dtype)
if isinstance(result, GGUFTensor):
result = torch.tensor(result.data, dtype=dtype, device=device, requires_grad=False)
return result
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
"""Override torch function calls to automatically dequantize"""
if kwargs is None:
kwargs = {}
# Find the GGUFTensor instance(s) in args
gguf_tensors = [arg for arg in args if isinstance(arg, cls)]
if not gguf_tensors:
return super().__torch_function__(func, types, args, kwargs)
# Use the first GGUFTensor instance for attribute access
self = gguf_tensors[0]
# Check if the tensor is fully constructed and still quantized
tensor_type = getattr(self, 'tensor_type', None)
if tensor_type is None:
# Tensor is either being constructed or already dequantized
return super().__torch_function__(func, types, args, kwargs)
# Check if tensor is already unquantized (F32/F16)
if tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}:
return super().__torch_function__(func, types, args, kwargs)
# Check if debug exists before using it
debug = getattr(self, 'debug', None)
# Handle linear operations specially
if func == torch.nn.functional.linear:
if len(args) >= 2 and isinstance(args[1], cls): # weight is the second argument
try:
weight_tensor = args[1]
dequantized_weight = weight_tensor.dequantize(device=args[0].device, dtype=args[0].dtype)
new_args = (args[0], dequantized_weight) + args[2:]
return func(*new_args, **kwargs)
except Exception as e:
if debug:
debug.log(f"Error in linear dequantization: {e}", level="WARNING", category="dit", force=True)
debug.log(f"Function: {func}", level="WARNING", category="dit", force=True, indent_level=1)
debug.log(f"Args: {[arg.shape if hasattr(arg, 'shape') else type(arg) for arg in args]}", level="WARNING", category="dit", force=True, indent_level=1)
raise
# Handle matrix multiplication operations that need dequantization
if func in {torch.matmul, torch.mm, torch.bmm, torch.addmm, torch.addmv,
torch.addr, torch.baddbmm, torch.chain_matmul}:
try:
new_args = []
for arg in args:
if isinstance(arg, cls):
new_args.append(arg.dequantize())
else:
new_args.append(arg)
return func(*tuple(new_args), **kwargs)
except Exception as e:
if debug:
debug.log(f"Error in {func.__name__} dequantization: {e}", level="WARNING", category="dit", force=True)
raise
# Handle conv2d/conv3d operations (critical for GGUF VAE models)
# Conv3d layers (InflatedCausalConv3d) are not replaced by layer replacement
if func in {torch.nn.functional.conv2d, torch.nn.functional.conv3d}:
if len(args) >= 2 and isinstance(args[1], cls): # weight is second arg
try:
weight_tensor = args[1]
dequantized_weight = weight_tensor.dequantize(device=args[0].device, dtype=args[0].dtype)
new_args = (args[0], dequantized_weight) + args[2:]
return func(*new_args, **kwargs)
except Exception as e:
if debug:
debug.log(f"Error in conv dequantization: {e}", level="WARNING", category="dit", force=True)
raise
# For ALL other operations, delegate to parent WITHOUT dequantization
# This includes .cpu(), .to(), .device, .dtype, .shape, etc.
return super().__torch_function__(func, types, args, kwargs)
def prepare_model_structure(
runner: VideoDiffusionInfer,
model_type: str,
checkpoint_path: str,
config: OmegaConf,
debug: 'Debug',
block_swap_config: Optional[Dict[str, Any]] = None
) -> VideoDiffusionInfer:
"""
Prepare model structure on meta device without loading weights.
This uses zero memory as meta device doesn't allocate real memory.
Args:
runner: VideoDiffusionInfer instance
model_type: "dit" or "vae"
checkpoint_path: Path to checkpoint (stored for later loading)
config: Model configuration
debug: Debug instance for logging (required)
block_swap_config: BlockSwap config (stored for DiT, optional)
Returns:
runner: Updated runner with model structure on meta device
"""
if debug is None:
raise ValueError(f"Debug instance required for prepare_model_structure")
is_dit = (model_type == "dit")
model_type_upper = "DiT" if is_dit else "VAE"
model_config = config.dit.model if is_dit else config.vae.model
# Always create on meta device for zero memory usage
debug.log(f"Creating {model_type_upper} model structure on meta device",
category=model_type, force=True)
debug.start_timer(f"{model_type}_structure")
with torch.device("meta"):
model = create_object(model_config)
debug.end_timer(f"{model_type}_structure", f"{model_type_upper} structure created")
# Store model and config for later materialization
if is_dit:
runner.dit = model
runner._dit_checkpoint = checkpoint_path
runner._dit_block_swap_config = block_swap_config
else:
runner.vae = model
runner._vae_checkpoint = checkpoint_path
return runner
def materialize_model(runner: VideoDiffusionInfer, model_type: str, device: torch.device,
config: OmegaConf, debug: 'Debug') -> None:
"""
Materialize model weights from checkpoint to memory.
Call this right before the model is needed.
Args:
runner: Runner with model structure on meta device
model_type: "dit" or "vae"
device: Target device for inference (torch.device object)
config: Full configuration
debug: Debug instance
"""
if debug is None:
raise ValueError(f"Debug instance required for materialize_model")
is_dit = (model_type == "dit")
model_type_upper = "DiT" if is_dit else "VAE"
# Get model and checkpoint path
if is_dit:
model = runner.dit
checkpoint_path = runner._dit_checkpoint
block_swap_config = runner._dit_block_swap_config
override_dtype = getattr(runner, '_dit_dtype_override', None)
else:
model = runner.vae
checkpoint_path = runner._vae_checkpoint
block_swap_config = None
override_dtype = getattr(runner, '_vae_dtype_override', None)
# Check if already materialized
if model is None:
debug.log(f"No {model_type_upper} model structure found", level="WARNING", category=model_type, force=True)
return
param_device = next(model.parameters()).device
if param_device.type != 'meta':
debug.log(f"{model_type_upper} already materialized on {model.device}", category=model_type)
return
# Determine target device for materialization
offload_device_str = None
if hasattr(runner, f'_{model_type}_offload_device'):
offload_device_str = getattr(runner, f'_{model_type}_offload_device')
# If offload_device is set and not "none", materialize to offload device
if offload_device_str and offload_device_str != "none":
target_device = torch.device(offload_device_str)
offload_reason = " (offload device)"
else:
# Otherwise materialize to inference device
target_device = device
offload_reason = ""
# Start materialization
debug.start_timer(f"{model_type}_materialize")
# Load weights (this materializes from meta to target device)
model = _load_model_weights(model, checkpoint_path, target_device, True,
model_type_upper, offload_reason, debug, override_dtype)
# Apply model-specific configurations (includes BlockSwap and torch.compile)
# Import here to avoid circular dependency
from .model_configuration import apply_model_specific_config
model = apply_model_specific_config(model, runner, config, is_dit, debug)
debug.end_timer(f"{model_type}_materialize", f"{model_type_upper} materialized")
# Clean up checkpoint paths (no longer needed after weights are loaded)
# Note: Config attributes (_dit_block_swap_config, _dit_compile_args) are preserved
# for configuration change detection on subsequent runs
if is_dit:
runner._dit_checkpoint = None
runner._dit_dtype_override = None
else:
runner._vae_checkpoint = None
runner._vae_dtype_override = None
def _load_model_weights(model: torch.nn.Module, checkpoint_path: str, target_device: torch.device,
used_meta: bool, model_type: str, cpu_reason: str,
debug: Optional['Debug'] = None, override_dtype: Optional[torch.dtype] = None) -> torch.nn.Module:
"""
Load model weights from checkpoint file with optimized GGUF support.
For meta-initialized models, materializes to target device.
For standard models, loads weights and applies state dict.
Args:
model: Model instance (may be on meta device)
checkpoint_path: Path to checkpoint file
target_device: Target device for weights (torch.device object)
used_meta: Whether model was created on meta device
model_type: Model type string for logging
cpu_reason: Reason string if using CPU
debug: Debug instance
override_dtype: Optional dtype override for weights
Returns:
Model with loaded weights
"""
model_type_lower = model_type.lower()
# Log loading action
action = "Materializing" if used_meta else "Loading"
target_device_str = str(target_device).upper()
debug.log(f"{action} {model_type} weights to {target_device_str}{cpu_reason}: {checkpoint_path}",
category=model_type_lower, force=True)
# Load state dict from file
debug.start_timer(f"{model_type_lower}_weights_load")
state = load_quantized_state_dict(checkpoint_path, target_device, debug)
debug.end_timer(f"{model_type_lower}_weights_load", f"{model_type} weights loaded from file")
# Apply dtype conversion if requested
if override_dtype is not None:
state = _convert_state_dtype(state, override_dtype, model_type, debug)
# Log weight statistics
_log_weight_stats(state, used_meta, model_type, debug)
# Handle GGUF or standard loading
if checkpoint_path.endswith('.gguf'):
model = _load_gguf_weights(model, state, used_meta, model_type_lower, debug)
else:
model = _load_standard_weights(model, state, used_meta, model_type, model_type_lower, debug)
# Clean up state dict
del state
# Initialize meta buffers if needed
if used_meta:
initialize_meta_buffers(model, target_device, debug)
return model
def _convert_state_dtype(state: Dict[str, torch.Tensor], target_dtype: torch.dtype,
model_type: str, debug: Optional['Debug'] = None) -> Dict[str, torch.Tensor]:
"""Convert floating point tensors in state dict to target dtype."""
debug.log(f"Converting {model_type} weights to {target_dtype} during loading", category="precision")
debug.start_timer(f"{model_type.lower()}_dtype_convert")
for key in state:
if torch.is_tensor(state[key]) and state[key].is_floating_point():
state[key] = state[key].to(target_dtype)
debug.end_timer(f"{model_type.lower()}_dtype_convert", f"{model_type} weights converted to {target_dtype}")
return state
def _log_weight_stats(state: Dict[str, torch.Tensor], used_meta: bool, model_type: str, debug: Optional['Debug'] = None) -> None:
"""Log statistics about loaded weights."""
num_params = len(state)
total_size_mb = sum(p.nelement() * p.element_size() for p in state.values()) / (1024 * 1024)
action = "Materializing" if used_meta else "Applying"
debug.log(f"{action} {model_type}: {num_params} parameters, {total_size_mb:.2f}MB total",
category=model_type.lower())
def apply_gguf_parameters(model: torch.nn.Module, state: Dict[str, torch.Tensor],
model_state: Dict[str, torch.Tensor], debug: Optional['Debug'] = None) -> Dict[str, Any]:
"""
Apply GGUF parameters to model, handling both meta and materialized models.
Returns:
Statistics dictionary with loaded count, quantized count, and parameter names
"""
loaded_names = set()
quantized_count = 0
for name, param in state.items():
if name not in model_state:
continue
model_param = model_state[name]
param_shape = _get_tensor_shape(param)
if param_shape != model_param.shape:
debug.log(f"Unexpected shape mismatch for {name}: {param_shape} vs {model_param.shape}",
level="ERROR", category="dit", force=True)
raise ValueError(f"Shape mismatch for parameter {name}")
# Apply parameter based on device type
with torch.no_grad():
if model_param.device.type == 'meta':
_set_parameter_on_meta_model(model, name, param, debug)
else:
_set_parameter_on_materialized_model(model, name, param, debug)
loaded_names.add(name)
if _is_quantized_tensor(param):
quantized_count += 1
return {
'loaded': len(loaded_names),
'quantized': quantized_count,
'loaded_names': loaded_names
}
def _set_parameter_on_meta_model(model: torch.nn.Module, param_name: str,
param_value: torch.Tensor, debug: Optional['Debug'] = None) -> None:
"""Set parameter on meta device model."""
module, attr_name = _navigate_to_parameter(model, param_name)
new_param = _create_gguf_parameter(param_value, debug)
setattr(module, attr_name, new_param)
def _set_parameter_on_materialized_model(model: torch.nn.Module, param_name: str,
param_value: torch.Tensor, debug: Optional['Debug'] = None) -> None:
"""Set parameter on already materialized model."""
module, attr_name = _navigate_to_parameter(model, param_name)
if _is_quantized_tensor(param_value):
# For quantized tensors, replace with wrapped parameter
new_param = _create_gguf_parameter(param_value, debug)
setattr(module, attr_name, new_param)
else:
# For regular tensors, just copy
existing_param = getattr(module, attr_name)
existing_param.copy_(param_value)
def _navigate_to_parameter(model: torch.nn.Module, param_path: str) -> Tuple[torch.nn.Module, str]:
"""
Navigate to the module containing a parameter.
Args:
model: Root model
param_path: Dot-separated path to parameter
Returns:
Tuple of (parent module, parameter name)
"""
path_parts = param_path.split('.')
module = model
# Navigate to parent module
for part in path_parts[:-1]:
module = getattr(module, part)
return module, path_parts[-1]
def _create_gguf_parameter(tensor: torch.Tensor, debug: Optional['Debug'] = None) -> torch.nn.Parameter:
"""
Create a parameter from a GGUF tensor, preserving quantization info.
Args:
tensor: GGUF tensor (may be quantized)
debug: Debug instance for logging
Returns:
Parameter with GGUF attributes and dequantize method if quantized
"""
param = torch.nn.Parameter(tensor, requires_grad=False)
# Preserve GGUF attributes if present
if hasattr(tensor, 'tensor_type'):
param.tensor_type = tensor.tensor_type
param.tensor_shape = tensor.tensor_shape
# Add dequantize method for runtime dequantization
param.gguf_dequantize = _create_dequantize_method(tensor, debug)
return param
def _get_tensor_shape(tensor: torch.Tensor) -> torch.Size:
"""Get the logical shape of a tensor (handling GGUF quantized tensors)."""
if hasattr(tensor, 'tensor_shape'):
return tensor.tensor_shape
return tensor.shape
def _is_quantized_tensor(tensor: torch.Tensor) -> bool:
"""Check if a tensor is GGUF quantized."""
return hasattr(tensor, 'tensor_type') and hasattr(tensor, 'tensor_shape')
def _report_parameter_mismatches(state: Dict[str, torch.Tensor],
model_state: Dict[str, torch.Tensor],
loaded_names: set, debug: Optional['Debug'] = None) -> None:
"""Report any parameter mismatches between GGUF and model."""
# Check for unmatched GGUF parameters
unmatched = [name for name in state if name not in model_state]
if unmatched:
debug.log(f"Warning: {len(unmatched)} parameters from GGUF not found in model",
level="WARNING", category="dit", force=True)
debug.log(f"First few unmatched: {unmatched[:5]}", level="WARNING", category="dit", force=True, indent_level=1)
# Check for missing model parameters
missing = [name for name in model_state if name not in loaded_names]
if missing:
debug.log(f"Warning: {len(missing)} model parameters not loaded from GGUF",
level="WARNING", category="dit", force=True)
debug.log(f"First few missing: {missing[:5]}", level="WARNING", category="dit", force=True, indent_level=1)
def initialize_meta_buffers(model: torch.nn.Module, target_device: torch.device, debug: Optional['Debug'] = None) -> None:
"""Initialize meta buffers with timing."""
debug.start_timer("buffer_init")
initialized = initialize_meta_buffers_impl(model, target_device, debug)
if initialized > 0:
debug.log(f"Initialized {initialized} non-persistent buffers", category="success")
debug.end_timer("buffer_init", "Buffer initialization")
def initialize_meta_buffers_impl(model: torch.nn.Module, target_device: torch.device, debug: Optional['Debug'] = None) -> int:
"""
Initialize any buffers still on meta device after materialization.
Non-persistent buffers aren't included in state_dict and remain on meta
device after load_state_dict. This function moves them to the target device.
Args:
model: Model potentially containing meta device buffers
target_device: Target device for initialization (torch.device object)
debug: Debug instance for logging
Returns:
Number of buffers initialized
"""
initialized_count = 0
# Simply initialize all meta device buffers to zeros on target device
for name, buffer in model.named_buffers():
if buffer is not None and buffer.device.type == 'meta':
# Get the module that owns this buffer
module_path = name.rsplit('.', 1)[0] if '.' in name else ''
buffer_name = name.rsplit('.', 1)[1] if '.' in name else name
# Get the actual module
if module_path:
module = model
for part in module_path.split('.'):
module = getattr(module, part)
else:
module = model
# Create a zero tensor of the same shape on target device
# This is safe for all non-persistent buffers (caches, dummy tensors, etc.)
initialized_buffer = torch.zeros_like(buffer, device=target_device)
module.register_buffer(buffer_name, initialized_buffer, persistent=False)
initialized_count += 1
return initialized_count
def _load_standard_weights(model: torch.nn.Module, state: Dict[str, torch.Tensor],
used_meta: bool, model_type: str, model_type_lower: str,
debug: Optional['Debug'] = None) -> torch.nn.Module:
"""Load standard (non-GGUF) weights into model."""
debug.start_timer(f"{model_type_lower}_state_apply")
model.load_state_dict(state, strict=False, assign=True)
action = "materialized" if used_meta else "applied"
debug.end_timer(f"{model_type_lower}_state_apply", f"{model_type} weights {action}")
if used_meta:
debug.log(f"{model_type} materialized directly from meta with loaded weights", category=model_type_lower)
else:
debug.log(f"{model_type} weights applied", category=model_type_lower)
return model
def _load_gguf_weights(model: torch.nn.Module, state: Dict[str, torch.Tensor],
used_meta: bool, model_type_lower: str, debug: Optional['Debug'] = None) -> torch.nn.Module:
"""
Load GGUF quantized weights into model with architecture validation.
Args:
model: Target model
state: GGUF state dict with quantized tensors
used_meta: Whether model was initialized on meta device
model_type_lower: Lowercase model type for logging
debug: Debug instance
Returns:
Model with GGUF weights loaded
"""
debug.log("Loading GGUF weights", category="dit")
# Get model state dict for validation
model_state = model.state_dict()
# Validate architecture compatibility
_validate_gguf_architecture(state, model_state, debug)
# Load GGUF parameters
stats = apply_gguf_parameters(model, state, model_state, debug)
# Log results
debug.log(f"GGUF loading complete: {stats['loaded']} parameters loaded", category="success")
debug.log(f"Quantized parameters: {stats['quantized']}", category="info")
# Report any mismatches
_report_parameter_mismatches(state, model_state, stats['loaded_names'], debug)
# Replace Linear/Conv2d layers with quantized versions for optimal precision handling
if stats['quantized'] > 0:
debug.log("Replacing layers with GGUF-optimized versions for precision handling", category="dit")
replacements, quant_types = replace_linear_with_quantized(model, debug=debug)
if replacements > 0:
debug.log(f"Replaced {replacements} layers with GGUF-optimized versions", category="success")
# Show actual quantization types found and precision strategy
if quant_types:
qtypes_str = ', '.join([f"{qtype}:{count}" for qtype, count in quant_types.items()])
debug.log(
f"GGUF precision path: {qtypes_str} → FP16 (preserve) → BF16/FP32 (compute)",
category="precision"
)
else:
debug.log(
"GGUF precision: Dequantizing to FP16 first, then converting to compute dtype",
category="precision"
)
else:
debug.log("Warning: No layers were replaced despite having quantized parameters",
level="WARNING", category="dit", force=True)
return model
def _validate_gguf_architecture(state: Dict[str, torch.Tensor],
model_state: Dict[str, torch.Tensor], debug: Optional['Debug'] = None) -> None:
"""
Validate GGUF model architecture matches target model.
Raises:
ValueError: If architecture mismatch is detected
"""
key_params = [
"blocks.0.attn.proj_qkv.vid.weight",
"blocks.0.attn.proj_qkv.txt.weight",
"blocks.0.mlp.vid.proj_in.weight"
]
for key in key_params:
if key in state and key in model_state:
model_shape = model_state[key].shape
gguf_shape = _get_tensor_shape(state[key])
if model_shape != gguf_shape:
# Check if it's just a quantization difference
if hasattr(state[key], 'tensor_shape') and state[key].tensor_shape == model_shape:
continue
raise ValueError(
f"GGUF model architecture mismatch: This GGUF model is incompatible with the current architecture.\n\n"
f"Detected mismatch:\n"
f" Parameter: {key}\n"
f" Expected shape: {model_shape}\n"
f" GGUF shape: {gguf_shape}\n\n"
f"Possible solutions:\n"
f"1. Use a GGUF model that matches the current architecture\n"
f"2. Try using a regular FP16 model instead\n"
f"3. Verify you're using the correct model variant (3B vs 7B)"
)
debug.log(f"Architecture check complete, no shape mismatch", category="success")
def _create_dequantize_method(tensor: torch.Tensor, debug: Optional['Debug'] = None) -> callable:
"""
Create a dequantization method for a GGUF tensor.
Args:
tensor: GGUF quantized tensor with tensor_type and tensor_shape attributes
debug: Debug instance
Returns:
Callable dequantization method
"""
def dequantize(device: Optional[torch.device] = None,
dtype: torch.dtype = torch.float16) -> torch.Tensor:
"""Dequantize GGUF tensor on demand."""
if hasattr(tensor, 'dequantize'):
return tensor.dequantize(device, dtype)
try:
# Fallback to manual dequantization using gguf library
numpy_data = tensor.cpu().numpy()
dequantized = gguf.quants.dequantize(numpy_data, tensor.tensor_type)
result = torch.from_numpy(dequantized).to(device or tensor.device, dtype)
result.requires_grad_(False)
return result.reshape(tensor.tensor_shape)
except Exception as e:
if debug:
debug.log(f"Warning: Could not dequantize tensor: {e}", level="WARNING", category="dit", force=True)
return tensor.to(device or tensor.device, dtype)
return dequantize
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/area_resize.py
+141
View File
@@ -0,0 +1,141 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
import math
import random
from typing import Union
import torch
from PIL import Image
from torchvision.transforms import functional as TVF
from torchvision.transforms.functional import InterpolationMode
from ....optimization.memory_manager import is_mps_available
class AreaResize:
def __init__(
self,
max_area: float,
downsample_only: bool = False,
interpolation: InterpolationMode = InterpolationMode.BICUBIC,
):
self.max_area = max_area
self.downsample_only = downsample_only
self.interpolation = interpolation
if is_mps_available():
self.interpolation = InterpolationMode.BILINEAR
def __call__(self, image: Union[torch.Tensor, Image.Image]):
if isinstance(image, torch.Tensor):
height, width = image.shape[-2:]
elif isinstance(image, Image.Image):
width, height = image.size
else:
raise NotImplementedError
scale = math.sqrt(self.max_area / (height * width))
# keep original height and width for small pictures.
scale = 1 if scale >= 1 and self.downsample_only else scale
resized_height, resized_width = round(height * scale), round(width * scale)
antialias = not (isinstance(image, torch.Tensor) and image.device.type == 'mps')
return TVF.resize(
image,
size=(resized_height, resized_width),
interpolation=self.interpolation,
antialias=antialias,
)
class AreaRandomCrop:
def __init__(
self,
max_area: float,
):
self.max_area = max_area
def get_params(self, input_size, output_size):
"""Get parameters for ``crop`` for a random crop.
Args:
img (PIL Image): Image to be cropped.
output_size (tuple): Expected output size of the crop.
Returns:
tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
"""
# w, h = _get_image_size(img)
h, w = input_size
th, tw = output_size
if w <= tw and h <= th:
return 0, 0, h, w
i = random.randint(0, h - th)
j = random.randint(0, w - tw)
return i, j, th, tw
def __call__(self, image: Union[torch.Tensor, Image.Image]):
if isinstance(image, torch.Tensor):
height, width = image.shape[-2:]
elif isinstance(image, Image.Image):
width, height = image.size
else:
raise NotImplementedError
resized_height = math.sqrt(self.max_area / (width / height))
resized_width = (width / height) * resized_height
# print('>>>>>>>>>>>>>>>>>>>>>')
# print((height, width))
# print( (resized_height, resized_width))
resized_height, resized_width = round(resized_height), round(resized_width)
i, j, h, w = self.get_params((height, width), (resized_height, resized_width))
image = TVF.crop(image, i, j, h, w)
return image
class ScaleResize:
def __init__(
self,
scale: float,
):
self.scale = scale
def __call__(self, image: Union[torch.Tensor, Image.Image]):
if isinstance(image, torch.Tensor):
height, width = image.shape[-2:]
interpolation_mode = InterpolationMode.BILINEAR
antialias = True if image.ndim == 4 else "warn"
elif isinstance(image, Image.Image):
width, height = image.size
interpolation_mode = InterpolationMode.LANCZOS
antialias = "warn"
else:
raise NotImplementedError
scale = self.scale
# keep original height and width for small pictures
resized_height, resized_width = round(height * scale), round(width * scale)
image = TVF.resize(
image,
size=(resized_height, resized_width),
interpolation=interpolation_mode,
antialias=antialias,
)
return image
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/divisible_crop.py
+80
View File
@@ -0,0 +1,80 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Union
import torch
from PIL import Image
from torchvision.transforms import functional as TVF
class DivisibleCrop:
def __init__(self, factor):
if not isinstance(factor, tuple):
factor = (factor, factor)
self.height_factor, self.width_factor = factor[0], factor[1]
def __call__(self, image: Union[torch.Tensor, Image.Image]):
if isinstance(image, torch.Tensor):
height, width = image.shape[-2:]
elif isinstance(image, Image.Image):
width, height = image.size
else:
raise NotImplementedError
cropped_height = height - (height % self.height_factor)
cropped_width = width - (width % self.width_factor)
image = TVF.center_crop(img=image, output_size=(cropped_height, cropped_width))
return image
class DivisiblePad:
"""
Pad image to make dimensions divisible by a factor.
Pads with black (0) to avoid data loss.
"""
def __init__(self, factor):
if not isinstance(factor, tuple):
factor = (factor, factor)
self.height_factor, self.width_factor = factor[0], factor[1]
def __call__(self, image: Union[torch.Tensor, Image.Image]):
if isinstance(image, torch.Tensor):
height, width = image.shape[-2:]
elif isinstance(image, Image.Image):
width, height = image.size
else:
raise NotImplementedError
# Calculate padding needed
pad_height = (self.height_factor - (height % self.height_factor)) % self.height_factor
pad_width = (self.width_factor - (width % self.width_factor)) % self.width_factor
if pad_height == 0 and pad_width == 0:
return image
# Pad symmetrically (or bottom/right)
if isinstance(image, torch.Tensor):
# Pad format: (left, right, top, bottom)
padding = (0, pad_width, 0, pad_height)
image = torch.nn.functional.pad(image, padding, mode='constant', value=0.0)
elif isinstance(image, Image.Image):
new_width = width + pad_width
new_height = height + pad_height
result = Image.new(image.mode, (new_width, new_height), (0, 0, 0))
result.paste(image, (0, 0))
image = result
return image
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/na_resize.py
+54
View File
@@ -0,0 +1,54 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
import torch
from typing import Literal
from torchvision.transforms import CenterCrop, Compose, InterpolationMode, Resize
from .area_resize import AreaResize
from .side_resize import SideResize
from ....optimization.memory_manager import is_mps_available
def NaResize(
resolution: int,
mode: Literal["area", "side"],
downsample_only: bool,
max_resolution: int = 0,
interpolation: InterpolationMode = InterpolationMode.BICUBIC,
):
Interpolation = InterpolationMode.BILINEAR if is_mps_available() else interpolation
if mode == "area":
return AreaResize(
max_area=resolution**2,
downsample_only=downsample_only,
interpolation=Interpolation,
)
if mode == "side":
return SideResize(
size=resolution,
max_size=max_resolution,
downsample_only=downsample_only,
interpolation=Interpolation,
)
if mode == "square":
return Compose(
[
Resize(
size=resolution,
interpolation=Interpolation,
),
CenterCrop(resolution),
]
)
raise ValueError(f"Unknown resize mode: {mode}")
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/side_resize.py
+76
View File
@@ -0,0 +1,76 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Union
import torch
from PIL import Image
from torchvision.transforms import InterpolationMode
from torchvision.transforms import functional as TVF
from ....optimization.memory_manager import is_mps_available
class SideResize:
def __init__(
self,
size: int,
max_size: int = 0,
downsample_only: bool = False,
interpolation: InterpolationMode = InterpolationMode.BICUBIC,
):
self.size = size
self.max_size = max_size
self.downsample_only = downsample_only
self.interpolation = interpolation
if is_mps_available():
self.interpolation = InterpolationMode.BILINEAR
def __call__(self, image: Union[torch.Tensor, Image.Image]):
"""
Resize image with shortest edge set to size, optionally limiting longest edge.
Args:
image (PIL Image or Tensor): Image to be scaled.
Returns:
PIL Image or Tensor: Rescaled image with shortest edge = size,
and no edge exceeding max_size (if max_size > 0).
"""
if isinstance(image, torch.Tensor):
height, width = image.shape[-2:]
elif isinstance(image, Image.Image):
width, height = image.size
else:
raise NotImplementedError
if self.downsample_only and min(width, height) < self.size:
size = min(width, height)
else:
size = self.size
# Resize to shortest edge (disable antialias only for MPS tensors - not supported)
antialias = not (isinstance(image, torch.Tensor) and image.device.type == 'mps')
resized = TVF.resize(image, size, self.interpolation, antialias=antialias)
# Apply max_size constraint if specified
if self.max_size > 0:
if isinstance(resized, torch.Tensor):
h, w = resized.shape[-2:]
else:
w, h = resized.size
if max(h, w) > self.max_size:
scale = self.max_size / max(h, w)
new_h, new_w = round(h * scale), round(w * scale)
resized = TVF.resize(resized, (new_h, new_w), self.interpolation, antialias=antialias)
return resized
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/__init__.py
+39
View File
@@ -0,0 +1,39 @@
"""
SeedVR2 ComfyUI Nodes
Central registry for all SeedVR2 nodes
"""
from comfy_api.latest import ComfyExtension, io
from .video_upscaler import SeedVR2VideoUpscaler
from .dit_model_loader import SeedVR2LoadDiTModel
from .vae_model_loader import SeedVR2LoadVAEModel
from .torch_compile_settings import SeedVR2TorchCompileSettings
class SeedVR2Extension(ComfyExtension):
"""SeedVR2 ComfyUI Extension"""
async def get_node_list(self) -> list[type[io.ComfyNode]]:
"""Return list of all SeedVR2 nodes"""
return [
SeedVR2VideoUpscaler,
SeedVR2LoadDiTModel,
SeedVR2LoadVAEModel,
SeedVR2TorchCompileSettings,
]
async def comfy_entrypoint() -> ComfyExtension:
"""ComfyUI V3 entry point"""
return SeedVR2Extension()
__all__ = [
'SeedVR2VideoUpscaler',
'SeedVR2LoadDiTModel',
'SeedVR2LoadVAEModel',
'SeedVR2TorchCompileSettings',
'SeedVR2Extension',
'comfy_entrypoint',
]
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/dit_model_loader.py
+180
View File
@@ -0,0 +1,180 @@
"""
SeedVR2 DiT Model Loader Node
Configure DiT (Diffusion Transformer) model with memory optimization
"""
from comfy_api.latest import io
from comfy_execution.utils import get_executing_context
from typing import Dict, Any, Tuple
from ..utils.model_registry import get_available_dit_models, DEFAULT_DIT
from ..optimization.memory_manager import get_device_list
class SeedVR2LoadDiTModel(io.ComfyNode):
"""
Configure DiT (Diffusion Transformer) model loader with memory optimization
Provides configuration for:
- Model selection and device placement
- BlockSwap memory optimization for limited VRAM
- Model caching between runs
- Optional torch.compile integration
Returns:
SEEDVR2_DIT configuration dictionary for main upscaler node
"""
@classmethod
def define_schema(cls) -> io.Schema:
devices = get_device_list()
dit_models = get_available_dit_models()
return io.Schema(
node_id="SeedVR2LoadDiTModel",
display_name="SeedVR2 (Down)Load DiT Model",
category="SEEDVR2",
description=(
"Load and configure SeedVR2 DiT (Diffusion Transformer) model for video upscaling. "
"Supports BlockSwap memory optimization for low VRAM systems, model caching for batch processing, "
"multi-GPU offloading, and torch.compile acceleration. \n\n"
"Connect to Video Upscaler node."
),
inputs=[
io.Combo.Input("model",
options=dit_models,
default=DEFAULT_DIT,
tooltip=(
"DiT (Diffusion Transformer) model for video upscaling.\n"
"Models automatically download on first use.\n"
"Additional models can be added to the ComfyUI models folder."
)
),
io.Combo.Input("device",
options=devices,
default=devices[0],
tooltip="GPU device for DiT model inference (upscaling phase)"
),
io.Int.Input("blocks_to_swap",
default=0,
min=0,
max=36,
step=1,
optional=True,
tooltip=(
"Number of transformer blocks to swap between devices for VRAM optimization.\n"
"• 0: Disabled (default)\n"
"• 3B model: 0-32 blocks\n"
"• 7B model: 0-36 blocks\n"
"\n"
"Requires offload_device to be set and different from device.\n"
"Not available on macOS (unified memory architecture)."
)
),
io.Boolean.Input("swap_io_components",
default=False,
optional=True,
tooltip=(
"Offload input/output embeddings and normalization layers to reduce VRAM.\n"
"Requires offload_device to be set and different from device.\n"
"Not available on macOS (unified memory architecture)."
)
),
io.Combo.Input("offload_device",
options=get_device_list(include_none=True, include_cpu=True),
default="none",
optional=True,
tooltip=(
"Device to offload DiT model when not actively processing.\n"
"• 'none': Keep model on inference device (default, fastest)\n"
"• 'cpu': Offload to system RAM (reduces VRAM usage)\n"
"• 'cuda:X': Offload to another GPU (good balance if available)\n"
"\n"
"Required for BlockSwap (blocks_to_swap or swap_io_components)."
)
),
io.Boolean.Input("cache_model",
default=False,
optional=True,
tooltip=(
"Keep DiT model loaded on offload_device between workflow runs.\n"
"Useful for batch processing to avoid repeated loading.\n"
"Requires offload_device to be set."
)
),
io.Combo.Input("attention_mode",
options=["sdpa", "flash_attn_2", "flash_attn_3", "sageattn_2", "sageattn_3"],
default="sdpa",
optional=True,
tooltip=(
"Attention computation backend:\n"
"• sdpa: PyTorch scaled_dot_product_attention (default, stable, always available)\n"
"• flash_attn_2: Flash Attention 2 (Ampere+, requires flash-attn package)\n"
"• flash_attn_3: Flash Attention 3 (Hopper+, requires flash-attn with FA3 support)\n"
"• sageattn_2: SageAttention 2 (requires sageattention package)\n"
"• sageattn_3: SageAttention 3 (Blackwell/RTX 50xx only, requires sageattn3 package)\n"
"\n"
"SDPA is recommended - stable and works everywhere.\n"
"Flash Attention and SageAttention provide speedup through optimized CUDA kernels on compatible GPUs."
)
),
io.Custom("TORCH_COMPILE_ARGS").Input("torch_compile_args",
optional=True,
tooltip=(
"Optional torch.compile optimization settings from SeedVR2 Torch Compile Settings node.\n"
"Provides 20-40% speedup with compatible PyTorch 2.0+ and Triton installation."
)
),
],
outputs=[
io.Custom("SEEDVR2_DIT").Output(
tooltip="DiT model configuration containing model path, device settings, BlockSwap parameters, and compilation options. Connect to Video Upscaler node."
)
]
)
@classmethod
def execute(cls, model: str, device: str, offload_device: str = "none",
cache_model: bool = False, blocks_to_swap: int = 0,
swap_io_components: bool = False, attention_mode: str = "sdpa",
torch_compile_args: Dict[str, Any] = None) -> io.NodeOutput:
"""
Create DiT model configuration for SeedVR2 main node
Args:
model: Model filename to load
device: Target device for model execution
offload_device: Device to offload model to when not in use
cache_model: Whether to keep model loaded between runs
blocks_to_swap: Number of transformer blocks to swap (requires offload_device != device)
swap_io_components: Whether to offload I/O components (requires offload_device != device)
attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
torch_compile_args: Optional torch.compile configuration from settings node
Returns:
NodeOutput containing configuration dictionary for SeedVR2 main node
Raises:
ValueError: If cache_model is enabled but offload_device is not set
"""
# Validate cache_model configuration
if cache_model and offload_device == "none":
raise ValueError(
"Model caching (cache_model=True) requires offload_device to be set. "
f"Current: offload_device='{offload_device}'. "
"Please set offload_device to specify where the cached DiT model should be stored "
"(e.g., 'cpu' or another device). Set cache_model=False if you don't want to cache the model."
)
config = {
"model": model,
"device": device,
"offload_device": offload_device,
"cache_model": cache_model,
"blocks_to_swap": blocks_to_swap,
"swap_io_components": swap_io_components,
"attention_mode": attention_mode,
"torch_compile_args": torch_compile_args,
"node_id": get_executing_context().node_id,
}
return io.NodeOutput(config)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/torch_compile_settings.py
+121
View File
@@ -0,0 +1,121 @@
"""
SeedVR2 Torch Compile Settings Node
Configure torch.compile optimization for DiT and VAE models
"""
from comfy_api.latest import io
from typing import Dict, Any, Tuple
class SeedVR2TorchCompileSettings(io.ComfyNode):
"""Configure torch.compile optimization for DiT and VAE models"""
@classmethod
def define_schema(cls) -> io.Schema:
return io.Schema(
node_id="SeedVR2TorchCompileSettings",
display_name="SeedVR2 Torch Compile Settings",
category="SEEDVR2",
description=(
"Configure SeedVR2 torch.compile optimization for 20-40% DiT speedup and 15-25% VAE speedup. "
"Trades longer first-run compilation time for faster inference.\n\n"
"Connect to DiT and/or VAE model loaders. Requires PyTorch 2.0+ and Triton for inductor backend."
),
inputs=[
io.Combo.Input("backend",
options=["inductor", "cudagraphs"],
default="inductor",
tooltip=(
"Compilation backend:\n"
"• inductor: Full optimization with Triton kernel generation and fusion (recommended)\n"
"• cudagraphs: Lightweight wrapper using CUDA graphs, no kernel optimization"
)
),
io.Combo.Input("mode",
options=["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"],
default="default",
tooltip=(
"Optimization level (compilation time vs runtime performance):\n"
"• default: Fast compilation with good speedup (recommended for development)\n"
"• reduce-overhead: Lower overhead, optimized for smaller models\n"
"• max-autotune: Slowest compilation, best runtime performance (recommended for production)\n"
"• max-autotune-no-cudagraphs: Like max-autotune but without CUDA graphs"
)
),
io.Boolean.Input("fullgraph",
default=False,
tooltip=(
"Compile entire model as single graph without breaks.\n"
"• False: Allow graph breaks for better compatibility (default)\n"
"• True: Enforce no breaks for maximum optimization (may fail with dynamic shapes)"
)
),
io.Boolean.Input("dynamic",
default=False,
tooltip=(
"Handle varying input shapes without recompilation.\n"
"• False: Specialize for exact input shapes (default)\n"
"• True: Create dynamic kernels that adapt to shape variations\n"
"\n"
"Enable when processing different resolutions or batch sizes."
)
),
io.Int.Input("dynamo_cache_size_limit",
default=64,
min=0,
max=1024,
step=1,
tooltip=(
"Maximum cached compiled versions per function (default: 64).\n"
"Controls how many shape variations to compile before stopping.\n"
"\n"
"• Increase: When processing many different input shapes (more memory usage)\n"
"• Decrease: When recompilation cost outweighs benefits (faster fallback to eager)"
)
),
io.Int.Input("dynamo_recompile_limit",
default=128,
min=0,
max=1024,
step=1,
tooltip=(
"Maximum recompilation attempts before fallback to eager mode (default: 128).\n"
"Safety limit to prevent infinite compilation loops.\n"
"\n"
"Only increase if you see 'hit config.recompile_limit' warnings and have bounded shape variations."
)
),
],
outputs=[
io.Custom("TORCH_COMPILE_ARGS").Output(
tooltip="torch.compile optimization settings including backend, mode, and Dynamo configuration. Connect to DiT and/or VAE model loader nodes."
)
]
)
@classmethod
def execute(cls, backend: str, mode: str, fullgraph: bool, dynamic: bool,
dynamo_cache_size_limit: int, dynamo_recompile_limit: int) -> io.NodeOutput:
"""
Create torch.compile configuration for model optimization
Args:
backend: Compilation backend ("inductor" or "cudagraphs")
mode: Optimization mode ("default", "reduce-overhead", "max-autotune", etc.)
fullgraph: Whether to compile entire model as single graph
dynamic: Whether to handle varying input shapes without recompilation
dynamo_cache_size_limit: Maximum cached compiled versions per function
dynamo_recompile_limit: Maximum recompilation attempts before fallback
Returns:
NodeOutput containing torch.compile configuration dictionary
"""
compile_args = {
"backend": backend,
"mode": mode,
"fullgraph": fullgraph,
"dynamic": dynamic,
"dynamo_cache_size_limit": dynamo_cache_size_limit,
"dynamo_recompile_limit": dynamo_recompile_limit,
}
return io.NodeOutput(compile_args)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/vae_model_loader.py
+219
View File
@@ -0,0 +1,219 @@
"""
SeedVR2 VAE Model Loader Node
Configure VAE (Variational Autoencoder) model with tiling support
"""
from comfy_api.latest import io
from comfy_execution.utils import get_executing_context
from typing import Dict, Any, Tuple
from ..utils.model_registry import get_available_vae_models, DEFAULT_VAE
from ..optimization.memory_manager import get_device_list
class SeedVR2LoadVAEModel(io.ComfyNode):
"""
Configure VAE (Variational Autoencoder) model loader with tiling support
Provides configuration for:
- Model selection and device placement
- Tiled encoding/decoding for VRAM reduction
- Tile size and overlap control
- Model caching between runs
- Optional torch.compile integration
Returns:
SEEDVR2_VAE configuration dictionary for main upscaler node
"""
@classmethod
def define_schema(cls) -> io.Schema:
devices = get_device_list()
vae_models = get_available_vae_models()
return io.Schema(
node_id="SeedVR2LoadVAEModel",
display_name="SeedVR2 (Down)Load VAE Model",
category="SEEDVR2",
description=(
"Load and configure SeedVR2 VAE (Variational Autoencoder) for encoding/decoding video frames to/from latent space. "
"Supports tiled processing to handle high resolutions on limited VRAM, model caching, "
"multi-GPU offloading, and torch.compile acceleration. \n\n"
"Connect to Video Upscaler node."
),
inputs=[
io.Combo.Input("model",
options=vae_models,
default=DEFAULT_VAE,
tooltip=(
"VAE (Variational Autoencoder) model for encoding/decoding.\n"
"Models automatically download on first use.\n"
"Additional models can be added to the ComfyUI models folder."
)
),
io.Combo.Input("device",
options=devices,
default=devices[0],
tooltip="GPU device for VAE model inference (encoding/decoding phases)"
),
io.Boolean.Input("encode_tiled",
default=False,
optional=True,
tooltip="Enable tiled encoding to reduce VRAM usage during the encoding phase"
),
io.Int.Input("encode_tile_size",
default=1024,
min=64,
step=32,
optional=True,
tooltip=(
"Encoding tile size in pixels (default: 1024).\n"
"Applied to both height and width.\n"
"Lower values reduce VRAM usage but may increase processing time.\n"
"Only used when encode_tiled is enabled."
)
),
io.Int.Input("encode_tile_overlap",
default=128,
min=0,
step=32,
optional=True,
tooltip=(
"Pixel overlap between encoding tiles (default: 128).\n"
"Reduces visible seams between tiles through blending.\n"
"Higher values improve quality but slow processing.\n"
"Only used when encode_tiled is enabled."
)
),
io.Boolean.Input("decode_tiled",
default=False,
optional=True,
tooltip="Enable tiled decoding to reduce VRAM usage during the decoding phase"
),
io.Int.Input("decode_tile_size",
default=1024,
min=64,
step=32,
optional=True,
tooltip=(
"Decoding tile size in pixels (default: 1024).\n"
"Applied to both height and width.\n"
"Lower values reduce VRAM usage but may increase processing time.\n"
"Only used when decode_tiled is enabled."
)
),
io.Int.Input("decode_tile_overlap",
default=128,
min=0,
step=32,
optional=True,
tooltip=(
"Pixel overlap between decoding tiles (default: 128).\n"
"Reduces visible seams between tiles through blending.\n"
"Higher values improve quality but slow processing.\n"
"Only used when decode_tiled is enabled."
)
),
io.Combo.Input("tile_debug",
options=["false", "encode", "decode"],
default="false",
optional=True,
tooltip=(
"Tile debug visualization mode:\n"
"• 'false': No visualization overlay (default)\n"
"• 'encode': Show encoding tile boundaries\n"
"• 'decode': Show decoding tile boundaries\n"
"\n"
"Only works when respective tiling is enabled."
)
),
io.Combo.Input("offload_device",
options=get_device_list(include_none=True, include_cpu=True),
default="none",
optional=True,
tooltip=(
"Device to offload VAE model when not actively processing.\n"
"• 'none': Keep model on inference device (default, fastest)\n"
"• 'cpu': Offload to system RAM (reduces VRAM usage)\n"
"• 'cuda:X': Offload to another GPU (good balance if available)"
)
),
io.Boolean.Input("cache_model",
default=False,
optional=True,
tooltip=(
"Keep VAE model loaded on offload_device between workflow runs.\n"
"Useful for batch processing to avoid repeated loading.\n"
"Requires offload_device to be set."
)
),
io.Custom("TORCH_COMPILE_ARGS").Input("torch_compile_args",
optional=True,
tooltip=(
"Optional torch.compile optimization settings from SeedVR2 Torch Compile Settings node.\n"
"Provides 15-25% speedup with compatible PyTorch 2.0+ and Triton installation."
)
),
],
outputs=[
io.Custom("SEEDVR2_VAE").Output(
tooltip="VAE model configuration containing model path, device settings, tiling parameters, and compilation options. Connect to Video Upscaler node."
)
]
)
@classmethod
def execute(cls, model: str, device: str, offload_device: str = "none",
cache_model: bool = False, encode_tiled: bool = False,
encode_tile_size: int = 512, encode_tile_overlap: int = 64,
decode_tiled: bool = False, decode_tile_size: int = 512,
decode_tile_overlap: int = 64, tile_debug: str = "false",
torch_compile_args: Dict[str, Any] = None
) -> io.NodeOutput:
"""
Create VAE model configuration for SeedVR2 main node
Args:
model: Model filename to load
device: Target device for model execution
offload_device: Device to offload model to when not in use
cache_model: Whether to keep model loaded between runs
encode_tiled: Enable tiled encoding
encode_tile_size: Tile size for encoding
encode_tile_overlap: Tile overlap for encoding
decode_tiled: Enable tiled decoding
decode_tile_size: Tile size for decoding
decode_tile_overlap: Tile overlap for decoding
tile_debug: Tile visualization mode (false/encode/decode)
torch_compile_args: Optional torch.compile configuration from settings node
Returns:
NodeOutput containing configuration dictionary for SeedVR2 main node
Raises:
ValueError: If cache_model is enabled but offload_device is invalid
"""
# Validate cache_model configuration
if cache_model and offload_device == "none":
raise ValueError(
"Model caching (cache_model=True) requires offload_device to be set. "
f"Current: offload_device='{offload_device}'. "
"Please set offload_device to specify where the cached VAE model should be stored "
"(e.g., 'cpu' or another device). Set cache_model=False if you don't want to cache the model."
)
config = {
"model": model,
"device": device,
"offload_device": offload_device,
"cache_model": cache_model,
"encode_tiled": encode_tiled,
"encode_tile_size": encode_tile_size,
"encode_tile_overlap": encode_tile_overlap,
"decode_tiled": decode_tiled,
"decode_tile_size": decode_tile_size,
"decode_tile_overlap": decode_tile_overlap,
"tile_debug": tile_debug,
"torch_compile_args": torch_compile_args,
"node_id": get_executing_context().node_id,
}
return io.NodeOutput(config)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/video_upscaler.py
+580
View File
@@ -0,0 +1,580 @@
"""
SeedVR2 Video Upscaler Node
Main ComfyUI node for high-quality video upscaling using diffusion models
"""
import torch
from comfy_api.latest import io
from typing import Tuple, Dict, Any, Optional
from ..utils.constants import get_base_cache_dir, __version__
from ..utils.downloads import download_weight
from ..utils.debug import Debug
from ..core.generation_phases import (
encode_all_batches,
upscale_all_batches,
decode_all_batches,
postprocess_all_batches
)
from ..core.generation_utils import (
setup_generation_context,
prepare_runner,
compute_generation_info,
log_generation_start,
load_text_embeddings,
script_directory
)
from ..optimization.memory_manager import (
cleanup_text_embeddings,
complete_cleanup,
get_device_list
)
# Import ComfyUI progress reporting
try:
from comfy.utils import ProgressBar
except ImportError:
ProgressBar = None
class SeedVR2VideoUpscaler(io.ComfyNode):
"""
SeedVR2 Video Upscaler ComfyUI Node
High-quality video upscaling using diffusion models with support for:
- Multiple model variants (3B/7B, FP16/FP8)
- Adaptive VRAM management
- Advanced dtype compatibility
- Optimized inference pipeline
- Real-time progress reporting
"""
@classmethod
def define_schema(cls) -> io.Schema:
return io.Schema(
node_id="SeedVR2VideoUpscaler",
display_name=f"SeedVR2 Video Upscaler (v{__version__})",
category="SEEDVR2",
description=(
"SeedVR2 main upscaling node: processes video frames using DiT and VAE models with diffusion-based enhancement. "
"Handles RGB/RGBA formats, maintains temporal consistency across frames, applies color correction, "
"and manages VRAM through intelligent tensor offloading. \n\n"
"Requires DiT and VAE model configurations."
),
inputs=[
io.Image.Input("image",
tooltip=(
"Input video frames as image batch.\n"
"Accepts both RGB (3-channel) and RGBA (4-channel) formats.\n"
"Output format will match input format."
)
),
io.Custom("SEEDVR2_DIT").Input("dit",
tooltip="DiT model configuration from SeedVR2 (Down)Load DiT Model node"
),
io.Custom("SEEDVR2_VAE").Input("vae",
tooltip="VAE model configuration from SeedVR2 (Down)Load VAE Model node"
),
io.Int.Input("seed",
default=42,
min=0,
max=2**32 - 1,
step=1,
tooltip=(
"Random seed for reproducible generation (default: 42).\n"
"Same seed with same inputs produces identical output."
)
),
io.Int.Input("resolution",
default=1080,
min=16,
max=16384,
step=2,
tooltip=(
"Target resolution for the shortest edge in pixels (default: 1080).\n"
"Automatically maintains aspect ratio of input.\n"
"Even values required for optimal processing."
)
),
io.Int.Input("max_resolution",
default=0,
min=0,
max=16384,
step=2,
tooltip=(
"Maximum resolution limit for any dimension (default: 0, no limit).\n"
"If any edge exceeds this after applying resolution,\n"
"both dimensions are scaled down proportionally.\n"
"Useful to prevent excessive VRAM usage on extreme aspect ratios."
)
),
io.Int.Input("batch_size",
default=5,
min=1,
max=16384,
step=4,
tooltip=(
"Number of frames processed together per batch (default: 5).\n"
"Must follow pattern 4n+1: 1, 5, 9, 13, 17, 21, ...\n"
"\n"
"• Higher values: Better temporal consistency and faster processing\n"
"• Lower values: Reduced VRAM usage\n"
"\n"
"Ideally match to shot length for best quality."
)
),
io.Boolean.Input("uniform_batch_size",
default=False,
tooltip=(
"Pad final batch to match batch_size (default: False).\n"
"Prevents temporal artifacts caused by small final batch.\n"
"Add extra compute but recommended for optimal quality."
)
),
io.Int.Input("temporal_overlap",
default=0,
min=0,
max=16,
step=1,
optional=True,
tooltip=(
"Overlapping frames between consecutive batches (default: 0, disabled).\n"
"Improves temporal consistency across batch boundaries through blending.\n"
"Values 1-4 work well for most content."
)
),
io.Int.Input("prepend_frames",
default=0,
min=0,
max=32,
step=1,
optional=True,
tooltip=(
"Number of frames to prepend (reversed from start) before processing (default: 0).\n"
"Helps reduce artifacts at video beginning.\n"
"Prepended frames are automatically removed from final output."
)
),
io.Combo.Input("color_correction",
options=["lab", "wavelet", "wavelet_adaptive", "hsv", "adain", "none"],
default="lab",
tooltip=(
"Corrects color shifts in upscaled output to match original input (default: lab).\n"
"The upscaling process may alter colors; this applies color grading to restore them.\n"
"\n"
"• lab: Perceptual color matching with detail preservation (recommended)\n"
"• wavelet: Frequency-based natural colors, preserves fine details\n"
"• wavelet_adaptive: Wavelet base with targeted saturation correction\n"
"• hsv: Hue-conditional saturation matching\n"
"• adain: Statistical style transfer approach\n"
"• none: No color correction applied"
)
),
io.Float.Input("input_noise_scale",
default=0.0,
min=0.0,
max=1.0,
step=0.001,
optional=True,
tooltip=(
"Input noise injection scale (default: 0.0, disabled).\n"
"Adds controlled variation to input images before encoding.\n"
"Range: 0.0 (no noise) to 1.0 (maximum noise).\n"
"Can help with certain types of artifacts."
)
),
io.Float.Input("latent_noise_scale",
default=0.0,
min=0.0,
max=1.0,
step=0.001,
optional=True,
tooltip=(
"Latent space noise injection scale (default: 0.0, disabled).\n"
"Adds controlled variation during the diffusion process.\n"
"Range: 0.0 (no noise) to 1.0 (maximum noise).\n"
"Can soften details if input_noise_scale doesn't help."
)
),
io.Combo.Input("offload_device",
options=get_device_list(include_none=True, include_cpu=True),
default="cpu",
optional=True,
tooltip=(
"Device for storing intermediate tensors between processing phases (default: cpu).\n"
"• 'none': Keep all tensors on inference device (fastest but highest VRAM usage)\n"
"• 'cpu': Offload to system RAM (recommended for long videos, slower transfers)\n"
"• 'cuda:X': Offload to another GPU (good balance if available, faster than CPU)"
)
),
io.Boolean.Input("enable_debug",
default=False,
optional=True,
tooltip=(
"Enable detailed debug logging (default: False).\n"
"Shows memory usage, timing information, and processing details.\n"
"Useful for troubleshooting errors and performance issues."
)
),
],
outputs=[
io.Image.Output(
tooltip="Upscaled video frames with color correction applied. Format (RGB/RGBA) matches input. Range [0, 1] normalized for ComfyUI compatibility."
)
]
)
@classmethod
def execute(cls, image: torch.Tensor, dit: Dict[str, Any], vae: Dict[str, Any],
seed: int, resolution: int = 1080, max_resolution: int = 0, batch_size: int = 5,
uniform_batch_size: bool = False, temporal_overlap: int = 0, prepend_frames: int = 0,
color_correction: str = "wavelet", input_noise_scale: float = 0.0,
latent_noise_scale: float = 0.0, offload_device: str = "none",
enable_debug: bool = False) -> io.NodeOutput:
"""
Execute SeedVR2 video upscaling with progress reporting
Main entry point for ComfyUI node execution.
Automatically detects and preserves input format (RGB or RGBA). Handles model downloads,
configuration unpacking, and delegates to upscaling pipeline.
Args:
image: Input video frames as tensor (N, H, W, C) in [0, 1] range
dit: DiT model configuration from SeedVR2LoadDiTModel node
vae: VAE model configuration from SeedVR2LoadVAEModel node
seed: Random seed for reproducible generation
resolution: Target resolution for shortest edge (maintains aspect ratio)
max_resolution: Maximum resolution for any edge (0 = no limit)
batch_size: Frames per batch (minimum 5 for temporal consistency)
uniform_batch_size: Whether to pad final batch to match batch_size
temporal_overlap: Overlapping frames between batches (0-16)
prepend_frames: Frames to prepend (0-32) to reduce initial artifacts.
color_correction: Color correction method
input_noise_scale: Input noise injection scale [0.0-1.0]
latent_noise_scale: Latent noise injection scale [0.0-1.0]
offload_device: Device to offload intermediate tensors
enable_debug: Enable detailed logging and memory tracking
Returns:
NodeOutput containing upscaled video tensor (N, H', W', C) in [0, 1] range
Raises:
ValueError: If model files cannot be downloaded or configuration is invalid
RuntimeError: If generation fails
"""
# Initialize debug (stateless - stored in local variable)
debug = Debug(enabled=enable_debug)
# Track execution state in local variables (not instance)
runner = None
ctx = None
pbar = None
# Define progress callback as local closure
def progress_callback(current_step: int, total_steps: int,
current_frames: int, phase_name: str) -> None:
"""
Update progress bar based on pipeline phase
Args:
current_step: Current step within phase
total_steps: Total steps in phase
current_frames: Number of frames being processed
phase_name: Name of current phase
"""
if pbar is None:
return
phase_weights = {
"Phase 1: Encoding": 0.2,
"Phase 2: Upscaling": 0.25,
"Phase 3: Decoding": 0.5,
"Phase 4: Post-processing": 0.05
}
phase_offset = {
"Phase 1: Encoding": 0.0,
"Phase 2: Upscaling": 0.2,
"Phase 3: Decoding": 0.45,
"Phase 4: Post-processing": 0.95
}
# Extract phase key from phase_name
phase_key = phase_name.split(" (")[0] if " (" in phase_name else phase_name
# Get weight and offset
weight = phase_weights.get(phase_key, 1.0)
offset = phase_offset.get(phase_key, 0.0)
# Calculate weighted progress
phase_progress = (current_step / total_steps) if total_steps > 0 else 0
overall_progress = offset + (phase_progress * weight)
# Update the progress bar with the overall progress
progress_value = int(overall_progress * 100)
pbar.update_absolute(progress_value, 100)
# Define cleanup as local function
def cleanup(dit_cache: bool = False, vae_cache: bool = False) -> None:
"""Cleanup resources after upscaling"""
nonlocal runner, ctx
# Use complete_cleanup for all cleanup operations
if runner:
complete_cleanup(runner=runner, debug=debug,
dit_cache=dit_cache, vae_cache=vae_cache)
# Delete runner only if neither model is cached
if not (dit_cache or vae_cache):
runner = None
# Clean up context text embeddings if they exist
if ctx:
cleanup_text_embeddings(ctx, debug)
ctx = None
# Extract configuration from dict inputs
dit_model = dit["model"]
vae_model = vae["model"]
dit_device = torch.device(dit["device"])
vae_device = torch.device(vae["device"])
dit_id = dit["node_id"]
vae_id = vae["node_id"]
# OPTIONAL inputs - use .get() with defaults
dit_cache = dit.get("cache_model", False)
attention_mode = dit.get("attention_mode", "sdpa")
vae_cache = vae.get("cache_model", False)
# BlockSwap configuration - construct from individual values
blocks_to_swap = dit.get("blocks_to_swap", 0)
swap_io_components = dit.get("swap_io_components", False)
dit_offload_str = dit.get("offload_device", "none")
block_swap_config = None
if blocks_to_swap > 0 or swap_io_components:
block_swap_config = {
"blocks_to_swap": blocks_to_swap,
"swap_io_components": swap_io_components,
}
if dit_offload_str != "none":
block_swap_config["offload_device"] = torch.device(dit_offload_str)
# Device configuration for offloading - convert "none" to None, else torch.device
vae_offload_str = vae.get("offload_device", "none")
dit_offload_device = torch.device(dit_offload_str) if dit_offload_str != "none" else None
vae_offload_device = torch.device(vae_offload_str) if vae_offload_str != "none" else None
tensor_offload_device = torch.device(offload_device) if offload_device != "none" else None
# VAE tiling configuration
encode_tiled = vae.get("encode_tiled", False)
encode_tile_size = vae.get("encode_tile_size", 512)
encode_tile_overlap = vae.get("encode_tile_overlap", 64)
decode_tiled = vae.get("decode_tiled", False)
decode_tile_size = vae.get("decode_tile_size", 512)
decode_tile_overlap = vae.get("decode_tile_overlap", 64)
tile_debug = vae.get("tile_debug", False)
# TorchCompile args (optional connection, can be None)
dit_torch_compile_args = dit.get("torch_compile_args")
vae_torch_compile_args = vae.get("torch_compile_args")
# Print header
debug.print_header()
debug.start_timer("total_execution", force=True)
debug.log("━━━━━━━━━ Model Preparation ━━━━━━━━━", category="none")
# Initial memory state
debug.log_memory_state("Before model preparation", show_tensors=False, detailed_tensors=False)
debug.start_timer("model_preparation")
# Check if download succeeded
debug.log("Checking and downloading models if needed...", category="download")
if not download_weight(dit_model=dit_model, vae_model=vae_model, debug=debug):
raise RuntimeError(
f"Failed to download required model files. "
f"DiT model: {dit_model}, VAE model: {vae_model}. "
"Please check the console output above for specific file failures and manual download instructions."
)
try:
# Initialize ComfyUI progress bar if available
if ProgressBar is not None:
pbar = ProgressBar(100)
# Setup generation context with device configuration
ctx = setup_generation_context(
dit_device=dit_device,
vae_device=vae_device,
dit_offload_device=dit_offload_device,
vae_offload_device=vae_offload_device,
tensor_offload_device=tensor_offload_device,
debug=debug
)
# Prepare runner with model state management and global cache
runner, cache_context = prepare_runner(
dit_model=dit_model,
vae_model=vae_model,
model_dir=get_base_cache_dir(),
debug=debug,
ctx=ctx,
dit_cache=dit_cache,
vae_cache=vae_cache,
dit_id=dit_id,
vae_id=vae_id,
block_swap_config=block_swap_config,
encode_tiled=encode_tiled,
encode_tile_size=(encode_tile_size, encode_tile_size),
encode_tile_overlap=(encode_tile_overlap, encode_tile_overlap),
decode_tiled=decode_tiled,
decode_tile_size=(decode_tile_size, decode_tile_size),
decode_tile_overlap=(decode_tile_overlap, decode_tile_overlap),
tile_debug=tile_debug,
attention_mode=attention_mode,
torch_compile_args_dit=dit_torch_compile_args,
torch_compile_args_vae=vae_torch_compile_args
)
# Store cache context in ctx for use in generation phases
ctx['cache_context'] = cache_context
# Preload text embeddings before Phase 1 to avoid sync stall in Phase 2
ctx['text_embeds'] = load_text_embeddings(script_directory, ctx['dit_device'], ctx['compute_dtype'], debug)
debug.log("Loaded text embeddings for DiT", category="dit")
debug.log_memory_state("After model preparation", show_tensors=False, detailed_tensors=False)
debug.end_timer("model_preparation", "Model preparation", force=True, show_breakdown=True)
# Compute generation info and log start (handles prepending internally)
image, gen_info = compute_generation_info(
ctx=ctx,
images=image,
resolution=resolution,
max_resolution=max_resolution,
batch_size=batch_size,
uniform_batch_size=uniform_batch_size,
seed=seed,
prepend_frames=prepend_frames,
temporal_overlap=temporal_overlap,
debug=debug
)
# Log generation start in consistent format
log_generation_start(gen_info, debug)
debug.start_timer("generation")
# Phase 1: Encode
ctx = encode_all_batches(
runner,
ctx=ctx,
images=image,
debug=debug,
batch_size=batch_size,
uniform_batch_size=uniform_batch_size,
seed=seed,
progress_callback=progress_callback,
temporal_overlap=temporal_overlap,
resolution=resolution,
max_resolution=max_resolution,
input_noise_scale=input_noise_scale,
color_correction=color_correction
)
# Phase 2: Upscale
ctx = upscale_all_batches(
runner,
ctx=ctx,
debug=debug,
progress_callback=progress_callback,
seed=seed,
latent_noise_scale=latent_noise_scale,
cache_model=dit_cache
)
# Phase 3: Decode
ctx = decode_all_batches(
runner,
ctx=ctx,
debug=debug,
progress_callback=progress_callback,
cache_model=vae_cache
)
# Phase 4: Post-processing
ctx = postprocess_all_batches(
ctx=ctx,
debug=debug,
progress_callback=progress_callback,
color_correction=color_correction,
prepend_frames=prepend_frames,
temporal_overlap=temporal_overlap,
batch_size=batch_size
)
sample = ctx['final_video']
debug.log("", category="none", force=True)
# Ensure CPU tensor in float32 for maximum ComfyUI compatibility
if torch.is_tensor(sample):
if sample.is_cuda or sample.is_mps:
sample = sample.cpu()
if sample.dtype != torch.float32:
src_dtype = sample.dtype
try:
sample = sample.to(torch.float32)
debug.log(f"Converted output from {src_dtype} to float32", category="precision")
except Exception as e:
debug.log(f"Could not convert to float32: {e}. Output is {src_dtype}, compatibility with other nodes not guaranteed",
level="WARNING", category="precision", force=True)
debug.log("Upscaling completed successfully!", category="success", force=True)
debug.end_timer("generation", "Video generation")
# Final cleanup
debug.start_timer("final_cleanup")
cleanup(dit_cache=dit_cache, vae_cache=vae_cache)
debug.end_timer("final_cleanup", "Final cleanup")
debug.log_memory_state("After all phases complete", show_tensors=False, detailed_tensors=False)
# Final peak vram summary
debug.log_peak_memory_summary()
# Final timing summary
debug.log("", category="none")
debug.log("────────────────────────", category="none")
child_times = {
"Model preparation": debug.timer_durations.get("model_preparation", 0),
"Video generation": debug.timer_durations.get("generation", 0),
"Final cleanup": debug.timer_durations.get("final_cleanup", 0)
}
if "phase1_encoding" in debug.timer_durations:
child_times[" Phase 1: VAE encoding"] = debug.timer_durations.get("phase1_encoding", 0)
if "phase2_upscaling" in debug.timer_durations:
child_times[" Phase 2: DiT upscaling"] = debug.timer_durations.get("phase2_upscaling", 0)
if "phase3_decoding" in debug.timer_durations:
child_times[" Phase 3: VAE decoding"] = debug.timer_durations.get("phase3_decoding", 0)
if "phase4_postprocessing" in debug.timer_durations:
child_times[" Phase 4: Post-processing"] = debug.timer_durations.get("phase4_postprocessing", 0)
total_execution_time = debug.end_timer("total_execution", "Total execution", show_breakdown=True, custom_children=child_times)
if total_execution_time > 0:
fps = gen_info['total_frames'] / total_execution_time
debug.log(f"Average FPS: {fps:.2f} frames/sec", category="timing", force=True)
# Print footer
debug.print_footer()
debug.clear_history()
pbar = None
ctx = None
# V3-compatible return with optional UI preview
return io.NodeOutput(sample)
except Exception as e:
cleanup(dit_cache=dit_cache, vae_cache=vae_cache)
raise e
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/attention.py
+148
View File
@@ -0,0 +1,148 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
import torch
import torch.nn.functional as F
# Import flash/sage attn with automatic fallback from compatibility layer
from ...optimization.compatibility import (
call_flash_attn_2_varlen, call_flash_attn_3_varlen,
call_sage_attn_2_varlen, call_sage_attn_3_varlen
)
from torch import nn
def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=None, max_seqlen_k=None, dropout_p=0.0, softmax_scale=None, causal=False, deterministic=False):
"""
A PyTorch-based implementation of variable-length attention to replace flash_attn_varlen_func.
It processes each sequence in the batch individually.
NOTE: max_seqlen_q and max_seqlen_k are accepted for API compatibility but not used.
PyTorch's scaled_dot_product_attention automatically handles variable sequence lengths.
COMPILE OPTIMIZATION: Uses torch.tensor_split to avoid .item() graph breaks
"""
# Split q, k, v using cumulative sequence lengths
# NOTE: torch.tensor_split requires int64 dtype and CPU device (PyTorch requirements)
q_splits = list(torch.tensor_split(q, cu_seqlens_q[1:-1].long().cpu(), dim=0))
k_splits = list(torch.tensor_split(k, cu_seqlens_k[1:-1].long().cpu(), dim=0))
v_splits = list(torch.tensor_split(v, cu_seqlens_k[1:-1].long().cpu(), dim=0))
# Process each sequence
output_splits = []
for q_i, k_i, v_i in zip(q_splits, k_splits, v_splits):
# Reshape for torch's scaled_dot_product_attention which expects (batch, heads, seq, dim).
# Here, we treat each sequence as a batch of 1.
q_i = q_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_q, head_dim)
k_i = k_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim)
v_i = v_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim)
# Use PyTorch's built-in scaled dot-product attention.
output_i = F.scaled_dot_product_attention(
q_i, k_i, v_i,
dropout_p=dropout_p if not deterministic else 0.0,
is_causal=causal
)
# Reshape the output back to the original format (seq_len, heads, head_dim)
output_i = output_i.squeeze(0).permute(1, 0, 2)
output_splits.append(output_i)
# Concatenate all outputs
return torch.cat(output_splits, dim=0)
class TorchAttention(nn.Module):
def tflops(self, args, kwargs, output) -> float:
assert len(args) == 0 or len(args) > 2, "query, key should both provided by args / kwargs"
q = kwargs.get("query") or args[0]
k = kwargs.get("key") or args[1]
b, h, sq, d = q.shape
b, h, sk, d = k.shape
return b * h * (4 * d * (sq / 1e6) * (sk / 1e6))
def forward(self, *args, **kwargs):
return F.scaled_dot_product_attention(*args, **kwargs)
class FlashAttentionVarlen(nn.Module):
"""
Variable-length attention with configurable backend.
Supported backends:
- sdpa: PyTorch SDPA (fully compilable, always available)
- flash_attn_2: Flash Attention 2 (Ampere+)
- flash_attn_3: Flash Attention 3 (Hopper+)
- sageattn_2: SageAttention 2
- sageattn_3: SageAttention 3 (Blackwell/RTX 50xx)
All non-SDPA backends use @torch._dynamo.disable wrapper (C++ extensions).
"""
def __init__(self, attention_mode: str = 'sdpa', compute_dtype: torch.dtype = None):
"""
Initialize with specified attention backend.
Args:
attention_mode: 'sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3'
compute_dtype: Compute dtype for attention (set by pipeline, defaults to None for auto-detection)
"""
super().__init__()
self.attention_mode = attention_mode
self.compute_dtype = compute_dtype
def tflops(self, args, kwargs, output) -> float:
cu_seqlens_q = kwargs["cu_seqlens_q"]
cu_seqlens_k = kwargs["cu_seqlens_k"]
_, h, d = output.shape
seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]) / 1e6
seqlens_k = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]) / 1e6
return h * (4 * d * (seqlens_q * seqlens_k).sum())
def forward(self, q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
kwargs["deterministic"] = torch.are_deterministic_algorithms_enabled()
# Convert to pipeline compute_dtype if configured (handles FP8 → fp16/bf16)
if self.compute_dtype is not None and q.dtype != self.compute_dtype:
q = q.to(self.compute_dtype)
k = k.to(self.compute_dtype)
v = v.to(self.compute_dtype)
if self.attention_mode == 'flash_attn_3':
return call_flash_attn_3_varlen(
q, k, v, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, **kwargs
)
elif self.attention_mode == 'flash_attn_2':
return call_flash_attn_2_varlen(
q, k, v, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, **kwargs
)
elif self.attention_mode == 'sageattn_3':
return call_sage_attn_3_varlen(
q, k, v, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, **kwargs
)
elif self.attention_mode == 'sageattn_2':
return call_sage_attn_2_varlen(
q, k, v, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, **kwargs
)
else:
# PyTorch SDPA
return pytorch_varlen_attention(
q, k, v, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, **kwargs
)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/embedding.py
+62
View File
@@ -0,0 +1,62 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Optional, Union
import torch
from diffusers.models.embeddings import get_timestep_embedding
from torch import nn
def emb_add(emb1: torch.Tensor, emb2: Optional[torch.Tensor]):
return emb1 if emb2 is None else emb1 + emb2
class TimeEmbedding(nn.Module):
def __init__(
self,
sinusoidal_dim: int,
hidden_dim: int,
output_dim: int,
):
super().__init__()
self.sinusoidal_dim = sinusoidal_dim
self.proj_in = nn.Linear(sinusoidal_dim, hidden_dim)
self.proj_hid = nn.Linear(hidden_dim, hidden_dim)
self.proj_out = nn.Linear(hidden_dim, output_dim)
self.act = nn.SiLU()
def forward(
self,
timestep: Union[int, float, torch.IntTensor, torch.FloatTensor],
device: torch.device,
dtype: torch.dtype,
) -> torch.FloatTensor:
if not torch.is_tensor(timestep):
timestep = torch.tensor([timestep], device=device, dtype=dtype)
if timestep.ndim == 0:
timestep = timestep[None]
emb = get_timestep_embedding(
timesteps=timestep,
embedding_dim=self.sinusoidal_dim,
flip_sin_to_cos=False,
downscale_freq_shift=0,
)
emb = emb.to(dtype)
emb = self.proj_in(emb)
emb = self.act(emb)
emb = self.proj_hid(emb)
emb = self.act(emb)
emb = self.proj_out(emb)
return emb
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/mlp.py
+62
View File
@@ -0,0 +1,62 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Optional
import torch
import torch.nn.functional as F
from torch import nn
def get_mlp(mlp_type: Optional[str] = "normal"):
if mlp_type == "normal":
return MLP
elif mlp_type == "swiglu":
return SwiGLUMLP
class MLP(nn.Module):
def __init__(
self,
dim: int,
expand_ratio: int,
):
super().__init__()
self.proj_in = nn.Linear(dim, dim * expand_ratio)
self.act = nn.GELU("tanh")
self.proj_out = nn.Linear(dim * expand_ratio, dim)
def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
x = self.proj_in(x)
x = self.act(x)
x = self.proj_out(x)
return x
class SwiGLUMLP(nn.Module):
def __init__(
self,
dim: int,
expand_ratio: int,
multiple_of: int = 256,
):
super().__init__()
hidden_dim = int(2 * dim * expand_ratio / 3)
hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
self.proj_in_gate = nn.Linear(dim, hidden_dim, bias=False)
self.proj_out = nn.Linear(hidden_dim, dim, bias=False)
self.proj_in = nn.Linear(dim, hidden_dim, bias=False)
def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
x = self.proj_out(F.silu(self.proj_in_gate(x)) * self.proj_in(x))
return x
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/mm.py
+74
View File
@@ -0,0 +1,74 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Tuple
import torch
from torch import nn
@dataclass
class MMArg:
vid: Any
txt: Any
def get_args(key: str, args: List[Any]) -> List[Any]:
return [getattr(v, key) if isinstance(v, MMArg) else v for v in args]
def get_kwargs(key: str, kwargs: Dict[str, Any]) -> Dict[str, Any]:
return {k: getattr(v, key) if isinstance(v, MMArg) else v for k, v in kwargs.items()}
class MMModule(nn.Module):
def __init__(
self,
module: Callable[..., nn.Module],
*args,
shared_weights: bool = False,
vid_only: bool = False,
**kwargs,
):
super().__init__()
self.shared_weights = shared_weights
self.vid_only = vid_only
if self.shared_weights:
assert get_args("vid", args) == get_args("txt", args)
assert get_kwargs("vid", kwargs) == get_kwargs("txt", kwargs)
self.all = module(*get_args("vid", args), **get_kwargs("vid", kwargs))
else:
self.vid = module(*get_args("vid", args), **get_kwargs("vid", kwargs))
self.txt = (
module(*get_args("txt", args), **get_kwargs("txt", kwargs))
if not vid_only
else None
)
def forward(
self,
vid: torch.FloatTensor,
txt: torch.FloatTensor,
*args,
**kwargs,
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
]:
vid_module = self.vid if not self.shared_weights else self.all
vid = vid_module(vid, *get_args("vid", args), **get_kwargs("vid", kwargs))
if not self.vid_only:
txt_module = self.txt if not self.shared_weights else self.all
txt = txt_module(txt, *get_args("txt", args), **get_kwargs("txt", kwargs))
return vid, txt
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/modulation.py
+121
View File
@@ -0,0 +1,121 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Callable, List, Optional
import torch
from einops import rearrange
from torch import nn
from ...common.cache import Cache
from ...common.distributed.ops import slice_inputs
# (dim: int, emb_dim: int)
ada_layer_type = Callable[[int, int], nn.Module]
def get_ada_layer(ada_layer: str) -> ada_layer_type:
if ada_layer == "single":
return AdaSingle
raise NotImplementedError(f"{ada_layer} is not supported")
def expand_dims(x: torch.Tensor, dim: int, ndim: int):
"""
Expand tensor "x" to "ndim" by adding empty dims at "dim".
Example: x is (b d), target ndim is 5, add dim at 1, return (b 1 1 1 d).
"""
shape = x.shape
shape = shape[:dim] + (1,) * (ndim - len(shape)) + shape[dim:]
return x.reshape(shape)
class AdaSingle(nn.Module):
def __init__(
self,
dim: int,
emb_dim: int,
layers: List[str],
modes: List[str] = ["in", "out"],
):
assert emb_dim == 6 * dim, "AdaSingle requires emb_dim == 6 * dim"
super().__init__()
self.dim = dim
self.emb_dim = emb_dim
self.layers = layers
for l in layers:
if "in" in modes:
self.register_parameter(f"{l}_shift", nn.Parameter(torch.randn(dim) / dim**0.5))
self.register_parameter(
f"{l}_scale", nn.Parameter(torch.randn(dim) / dim**0.5 + 1)
)
if "out" in modes:
self.register_parameter(f"{l}_gate", nn.Parameter(torch.randn(dim) / dim**0.5))
def forward(
self,
hid: torch.FloatTensor, # b ... c
emb: torch.FloatTensor, # b d
layer: str,
mode: str,
cache: Cache = Cache(disable=True),
branch_tag: str = "",
hid_len: Optional[torch.LongTensor] = None, # b
) -> torch.FloatTensor:
idx = self.layers.index(layer)
emb = rearrange(emb, "b (d l g) -> b d l g", l=len(self.layers), g=3)[..., idx, :]
emb = expand_dims(emb, 1, hid.ndim + 1)
if hid_len is not None:
emb = cache(
f"emb_repeat_{idx}_{branch_tag}",
lambda: slice_inputs(
torch.repeat_interleave(emb, hid_len, dim=0),
dim=0,
),
)
shiftA, scaleA, gateA = emb.unbind(-1)
shiftB, scaleB, gateB = (
getattr(self, f"{layer}_shift", None),
getattr(self, f"{layer}_scale", None),
getattr(self, f"{layer}_gate", None),
)
# Handle potential FP8 parameters - convert to input computation dtype
if hasattr(torch, 'float8_e4m3fn'):
fp8_types = (torch.float8_e4m3fn, torch.float8_e5m2)
# Use input tensor's dtype as target (respects pipeline precision)
target_dtype = hid.dtype
# Convert FP8 parameters to match input dtype for arithmetic operations
if shiftB is not None and shiftB.dtype in fp8_types:
shiftB = shiftB.to(target_dtype)
if scaleB is not None and scaleB.dtype in fp8_types:
scaleB = scaleB.to(target_dtype)
if gateB is not None and gateB.dtype in fp8_types:
gateB = gateB.to(target_dtype)
if mode == "in":
return hid.mul_(scaleA + scaleB).add_(shiftA + shiftB)
if mode == "out":
if gateB is not None:
return hid.mul_(gateA + gateB)
else:
# If no gate parameter, just use the embedding gate
return hid.mul_(gateA)
raise NotImplementedError
def extra_repr(self) -> str:
return f"dim={self.dim}, emb_dim={self.emb_dim}, layers={self.layers}"
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/na.py
+641
View File
@@ -0,0 +1,641 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Native Resolution Transformer (NaDiT) tensor manipulation utilities.
TORCH.COMPILE OPTIMIZED VERSION
================================
This module has been optimized for torch.compile compatibility by eliminating
data-dependent operations that cause graph breaks:
Key Changes from Original:
- Replaced all .tolist() calls with pure tensor operations
- Minimized .item() calls (only used where required by einops API)
- Replaced list comprehensions with tensor-based splitting
- Added _tensor_split helper for compile-friendly splitting
- Proper device management to ensure tensors stay on correct devices
"""
from itertools import chain
from typing import Callable, Dict, List, Tuple
import einops
import torch
def _tensor_split(tensor: torch.Tensor, lengths: torch.LongTensor, dim: int = 0) -> List[torch.Tensor]:
"""
Optimized compile-friendly split using torch.tensor_split.
Uses PyTorch's native C++ implementation for fast eager mode execution
while remaining fully compatible with torch.compile symbolic tracing.
Args:
tensor: Input tensor to split
lengths: Tensor of split lengths (1D)
dim: Dimension along which to split (default: 0)
Returns:
List of split tensors
"""
if lengths.numel() == 0:
return []
if lengths.numel() == 1:
return [tensor]
# Calculate split indices: torch.tensor_split splits BEFORE each index
# So we need cumsum[:-1] to get the split points
# NOTE: torch.tensor_split requires indices on CPU (PyTorch requirement)
split_indices = lengths.cumsum(0)[:-1].cpu()
# Use torch.tensor_split - native C++ implementation
# Compilable: accepts tensor indices (symbolic shapes work)
# Fast: uses optimized CUDA/CPU kernels instead of Python loops
return list(torch.tensor_split(tensor, split_indices, dim=dim))
def flatten(
hid: List[torch.FloatTensor], # List of (*** c)
) -> Tuple[
torch.FloatTensor, # (L c)
torch.LongTensor, # (b n)
]:
"""
Flatten a list of tensors into a single tensor and track their shapes.
Converts a list of tensors with potentially different spatial shapes but
same feature dimension into a flattened tensor and shape metadata.
Args:
hid: List of tensors, each with shape (d1, d2, ..., dn, c)
Returns:
Tuple of:
- Flattened tensor of shape (L, c) where L = sum of all spatial dimensions
- Shape tensor of shape (b, n) tracking original spatial dimensions
COMPILE OPTIMIZATION: Uses tensor operations on correct device from input
"""
assert len(hid) > 0
device = hid[0].device
# Stack shape metadata - ensure tensors are created on correct device
shapes = []
for x in hid:
shape_tensor = torch.tensor(x.shape[:-1], dtype=torch.long, device=device)
shapes.append(shape_tensor)
shape = torch.stack(shapes)
# Flatten and concatenate
hid = torch.cat([x.flatten(0, -2) for x in hid])
return hid, shape
def unflatten(
hid: torch.FloatTensor, # (L c) or (L ... c)
hid_shape: torch.LongTensor, # (b n)
) -> List[torch.Tensor]: # List of (*** c) or (*** ... c)
"""
Unflatten a tensor back to a list using shape metadata.
Inverse operation of flatten(), reconstructing original tensor shapes.
Args:
hid: Flattened tensor of shape (L, c) or (L, ..., c)
hid_shape: Shape metadata tensor of shape (b, n)
Returns:
List of unflattened tensors with original spatial dimensions
COMPILE OPTIMIZATION:
- Uses optimized _tensor_split with torch.tensor_split (fast)
- .cpu().numpy() conversion needed for torch.compile compatibility
(unflatten() requires concrete Python ints, not symbolic shapes)
"""
hid_len = hid_shape.prod(-1)
# Use optimized tensor splitting (major performance improvement)
hid_list = _tensor_split(hid, hid_len, dim=0)
# Unflatten each piece
# NOTE: .cpu().numpy() is required for torch.compile compatibility
# .tolist() would fail with symbolic shapes during compilation
result = []
for i, x in enumerate(hid_list):
shape = hid_shape[i]
# Must use .cpu().numpy() for compilation compatibility
# Shape tensors are small, so CPU transfer overhead is minimal
target_shape = list(shape.cpu().numpy())
result.append(x.unflatten(0, target_shape))
return result
def concat(
vid: torch.FloatTensor, # (VL ... c)
txt: torch.FloatTensor, # (TL ... c)
vid_len: torch.LongTensor, # (b)
txt_len: torch.LongTensor, # (b)
) -> torch.FloatTensor: # (L ... c)
"""
Interleave video and text tensors batch-wise.
Splits video and text tensors by batch lengths, then interleaves them:
[vid_0, txt_0, vid_1, txt_1, ..., vid_b, txt_b]
Args:
vid: Video features tensor (VL, c)
txt: Text features tensor (TL, c)
vid_len: Length of each video sequence (b,)
txt_len: Length of each text sequence (b,)
Returns:
Interleaved tensor (L, c) where L = sum(vid_len) + sum(txt_len)
COMPILE OPTIMIZATION: Uses _tensor_split for compile-friendly splitting
"""
# Use tensor-based splitting
vid_splits = _tensor_split(vid, vid_len, dim=0)
txt_splits = _tensor_split(txt, txt_len, dim=0)
# Interleave
interleaved = []
for v, t in zip(vid_splits, txt_splits):
interleaved.extend([v, t])
return torch.cat(interleaved)
def concat_idx(
vid_len: torch.LongTensor, # (b)
txt_len: torch.LongTensor, # (b)
) -> Tuple[
Callable,
Callable,
]:
"""
Create index-based concatenation and un-concatenation functions.
Pre-computes indices for efficient interleaving and de-interleaving operations.
Returns callable functions that can be reused multiple times.
Args:
vid_len: Video sequence lengths (b,)
txt_len: Text sequence lengths (b,)
Returns:
Tuple of (concat_fn, unconcat_fn):
- concat_fn: Lambda that interleaves vid and txt tensors
- unconcat_fn: Lambda that separates interleaved tensor back to vid and txt
COMPILE OPTIMIZATION: Pre-computes all indices using tensor operations
"""
device = vid_len.device
vid_sum = vid_len.sum()
txt_sum = txt_len.sum()
vid_idx = torch.arange(vid_sum, device=device)
txt_idx = torch.arange(vid_sum, vid_sum + txt_sum, device=device)
# Build interleaving indices using compile-friendly _tensor_split
batch_size = len(vid_len)
vid_idx_splits = _tensor_split(vid_idx, vid_len, dim=0)
txt_idx_splits = _tensor_split(txt_idx, txt_len, dim=0)
# Create interleaved target indices
tgt_idx_list = []
for i in range(batch_size):
tgt_idx_list.append(vid_idx_splits[i])
tgt_idx_list.append(txt_idx_splits[i])
tgt_idx = torch.cat(tgt_idx_list)
src_idx = torch.argsort(tgt_idx)
vid_idx_len = len(vid_idx)
return (
lambda vid, txt: torch.index_select(torch.cat([vid, txt]), 0, tgt_idx),
lambda all: torch.index_select(all, 0, src_idx).split([vid_idx_len, len(txt_idx)]),
)
def unconcat(
all: torch.FloatTensor, # (L ... c)
vid_len: torch.LongTensor, # (b)
txt_len: torch.LongTensor, # (b)
) -> Tuple[
torch.FloatTensor, # (VL ... c)
torch.FloatTensor, # (TL ... c)
]:
"""
De-interleave concatenated video and text tensors.
Inverse of concat(). Separates an interleaved tensor back into video and text.
Args:
all: Interleaved tensor (L, c)
vid_len: Video sequence lengths (b,)
txt_len: Text sequence lengths (b,)
Returns:
Tuple of (vid, txt) tensors
COMPILE OPTIMIZATION: Uses tensor operations to build interleave pattern
"""
batch_size = len(vid_len)
# Create interleaved lengths: [vid_0, txt_0, vid_1, txt_1, ...]
interleave_len = torch.stack([vid_len, txt_len], dim=1).flatten()
# Split using compile-friendly operation
all_splits = _tensor_split(all, interleave_len, dim=0)
# Separate even (vid) and odd (txt) indices
vid_parts = [all_splits[i] for i in range(0, len(all_splits), 2)]
txt_parts = [all_splits[i] for i in range(1, len(all_splits), 2)]
vid = torch.cat(vid_parts)
txt = torch.cat(txt_parts)
return vid, txt
def repeat_concat(
vid: torch.FloatTensor, # (VL ... c)
txt: torch.FloatTensor, # (TL ... c)
vid_len: torch.LongTensor, # (n*b)
txt_len: torch.LongTensor, # (b)
txt_repeat: torch.LongTensor, # (n) or (b)
) -> torch.FloatTensor: # (L ... c)
"""
Concatenate video and text with text repetition for window attention.
For windowed attention, text features are repeated and interleaved with
multiple video windows: [vid_0, txt_0, vid_1, txt_0, vid_2, txt_0, ...]
Args:
vid: Video features (VL, c)
txt: Text features (TL, c)
vid_len: Video window lengths (n*b,) where n=num_windows
txt_len: Text sequence lengths (b,)
txt_repeat: Number of times to repeat text (n,) or (b,)
Returns:
Interleaved tensor with repeated text
COMPILE OPTIMIZATION: Uses _tensor_split and tensor-based repetition
"""
# Split using compile-friendly operations
vid_splits = _tensor_split(vid, vid_len, dim=0)
txt_splits = _tensor_split(txt, txt_len, dim=0)
# Handle txt_repeat shape flexibility
if txt_repeat.numel() == len(txt_splits):
repeat_counts = txt_repeat
else:
repeat_counts = txt_repeat.repeat(len(txt_splits))
# Interleave with repetition
result = []
for i, v in enumerate(vid_splits):
result.append(v)
# Get corresponding text sample (cyclic)
batch_idx = i % len(txt_splits) if len(txt_splits) > 0 else 0
if batch_idx < len(txt_splits):
result.append(txt_splits[batch_idx])
return torch.cat(result)
def repeat_concat_idx(
vid_len: torch.LongTensor, # (n*b)
txt_len: torch.LongTensor, # (b)
txt_repeat: torch.LongTensor, # (n) or scalar
) -> Tuple[
Callable,
Callable,
]:
"""
Create index-based repeat-concatenation and un-concatenation with coalescing.
Similar to concat_idx but handles text repetition for window attention.
The unconcat function coalesces (averages) repeated text features.
Args:
vid_len: Video window lengths (n*b,)
txt_len: Text sequence lengths (b,)
txt_repeat: Repetition count (scalar or tensor)
Returns:
Tuple of (concat_fn, unconcat_coalesce_fn):
- concat_fn: Interleaves with text repetition
- unconcat_coalesce_fn: Separates and averages repeated text
Example:
Input: vid=[0,1,2,3,4,5,6,7,8], txt=[9,10], repeat=3
Concat: [0,1,2,9,10, 3,4,5,9,10, 6,7,8,9,10]
Unconcat: vid=[0,1,2,3,4,5,6,7,8], txt=[9,10] (averaged)
COMPILE OPTIMIZATION:
- Eliminates .tolist() calls that caused graph breaks
- Uses pure tensor operations for index building
- Minimizes data-dependent branching
"""
device = vid_len.device
vid_sum = vid_len.sum()
txt_sum = txt_len.sum()
# Create base indices
vid_idx = torch.arange(vid_sum, device=device)
txt_idx = torch.arange(vid_sum, vid_sum + txt_sum, device=device)
# Normalize txt_repeat to tensor
if isinstance(txt_repeat, (int, float)):
txt_repeat = torch.tensor([txt_repeat], dtype=torch.long, device=device)
elif txt_repeat.dim() == 0:
txt_repeat = txt_repeat.unsqueeze(0)
# Calculate repeat pattern - keep as tensor to avoid graph breaks
batch_size = len(txt_len)
if txt_repeat.numel() == 1:
num_repeats_tensor = txt_repeat.reshape(-1) # Ensure 1D tensor
else:
# Use tensor operations for division
num_repeats_tensor = torch.tensor([len(vid_len) // batch_size], dtype=torch.long, device=device)
# Build concatenated indices using compile-friendly _tensor_split
vid_idx_splits = _tensor_split(vid_idx, vid_len, dim=0)
txt_idx_splits = _tensor_split(txt_idx, txt_len, dim=0)
tgt_idx_list = []
for i in range(len(vid_len)):
# Add video window
tgt_idx_list.append(vid_idx_splits[i])
# Add corresponding text (with repeat)
batch_idx = i % batch_size
tgt_idx_list.append(txt_idx_splits[batch_idx])
tgt_idx = torch.cat(tgt_idx_list)
src_idx = torch.argsort(tgt_idx)
txt_idx_len = len(tgt_idx) - len(vid_idx)
# Pre-compute split lengths for coalescing using tensor operations
repeat_txt_len = txt_len * num_repeats_tensor.squeeze()
def unconcat_coalesce(all):
"""
Un-concat vid & txt, and coalesce the repeated txt by averaging.
The text features appear multiple times (once per window) and need
to be averaged to produce a single set of text features.
COMPILE OPTIMIZATION: Uses unflatten with tensor dims (compile-friendly)
"""
vid_out, txt_out = all[src_idx].split([len(vid_idx), txt_idx_len])
# Coalesce repeated text using unflatten and mean
txt_splits = _tensor_split(txt_out, repeat_txt_len, dim=0)
txt_out_coalesced = []
for txt in txt_splits:
# txt has shape (base_len * num_repeats, *other_dims)
# unflatten to (base_len, num_repeats, *other_dims) then average dim 1
txt = txt.unflatten(0, (-1, num_repeats_tensor.squeeze())).mean(1)
txt_out_coalesced.append(txt)
return vid_out, torch.cat(txt_out_coalesced)
# Note: Using direct indexing instead of torch.index_select for backward compatibility
# Direct indexing is deterministic even with repeated indices
return (
lambda vid, txt: torch.cat([vid, txt])[tgt_idx],
lambda all: unconcat_coalesce(all),
)
def rearrange(
hid: torch.FloatTensor, # (L c)
hid_shape: torch.LongTensor, # (b n)
pattern: str,
**kwargs: Dict[str, int],
) -> Tuple[
torch.FloatTensor,
torch.LongTensor,
]:
"""
Rearrange flattened tensor using einops pattern.
Applies einops rearrange to each batch element independently.
Args:
hid: Flattened tensor (L, c)
hid_shape: Shape metadata (b, n)
pattern: Einops rearrange pattern
**kwargs: Additional arguments for einops
Returns:
Tuple of (rearranged tensor, new shape metadata)
"""
unflattened = unflatten(hid, hid_shape)
rearranged = [einops.rearrange(h, pattern, **kwargs) for h in unflattened]
return flatten(rearranged)
def rearrange_idx(
hid_shape: torch.LongTensor, # (b n)
pattern: str,
**kwargs: Dict[str, int],
) -> Tuple[Callable, Callable, torch.LongTensor]:
"""
Create index-based rearrange functions.
Pre-computes indices for efficient rearrangement operations.
Args:
hid_shape: Shape metadata (b, n)
pattern: Einops rearrange pattern
**kwargs: Additional arguments for einops
Returns:
Tuple of (rearrange_fn, reverse_fn, new_shape)
"""
hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1)
tgt_idx, tgt_shape = rearrange(hid_idx, hid_shape, pattern, **kwargs)
tgt_idx = tgt_idx.squeeze(-1)
src_idx = torch.argsort(tgt_idx)
return (
lambda hid: torch.index_select(hid, 0, tgt_idx),
lambda hid: torch.index_select(hid, 0, src_idx),
tgt_shape,
)
def repeat(
hid: torch.FloatTensor, # (L c)
hid_shape: torch.LongTensor, # (b n)
pattern: str,
**kwargs: Dict[str, torch.LongTensor], # (b)
) -> Tuple[
torch.FloatTensor,
torch.LongTensor,
]:
"""
Repeat flattened tensor using einops pattern with per-batch parameters.
Each batch element can have different repeat counts specified in kwargs.
Args:
hid: Flattened tensor (L, c)
hid_shape: Shape metadata (b, n)
pattern: Einops repeat pattern (e.g., "l c -> t l c")
**kwargs: Repeat parameters as tensors (e.g., t=torch.tensor([2,3,4]))
Returns:
Tuple of (repeated tensor, new shape metadata)
COMPILE OPTIMIZATION:
- Minimizes .item() calls
- Only converts to int at the last moment for einops API requirement
"""
unflattened = unflatten(hid, hid_shape)
# Build kwargs for each batch element
repeated = []
for i in range(len(unflattened)):
# Extract values for einops (requires Python int)
batch_kwargs = {}
for k, v in kwargs.items():
if isinstance(v, torch.Tensor):
# Only convert to Python int where required by einops
batch_kwargs[k] = int(v[i].item())
else:
batch_kwargs[k] = v
repeated.append(einops.repeat(unflattened[i], pattern, **batch_kwargs))
return flatten(repeated)
def pack(
samples: List[torch.Tensor], # List of (h w c).
) -> Tuple[
List[torch.Tensor], # groups [(b1 h1 w1 c1), (b2 h2 w2 c2)]
List[List[int]], # reversal indices.
]:
"""
Group samples by shape and return grouped batches with reversal indices.
Useful for batch processing samples with different spatial dimensions.
Args:
samples: List of tensors with potentially different shapes
Returns:
Tuple of (batched_groups, reversal_indices) for unpacking
"""
batches = {}
indices = {}
for i, sample in enumerate(samples):
shape = sample.shape
batches[shape] = batches.get(shape, [])
indices[shape] = indices.get(shape, [])
batches[shape].append(sample)
indices[shape].append(i)
batches = list(map(torch.stack, batches.values()))
indices = list(indices.values())
return batches, indices
def unpack(
batches: List[torch.Tensor],
indices: List[List[int]],
) -> List[torch.Tensor]:
"""
Unpack grouped batches back to original order.
Inverse of pack().
Args:
batches: Grouped batches from pack()
indices: Reversal indices from pack()
Returns:
List of tensors in original order
"""
samples = [None] * (max(chain(*indices)) + 1)
for batch, index in zip(batches, indices):
for sample, i in zip(batch.unbind(), index):
samples[i] = sample
return samples
def window(
hid: torch.FloatTensor, # (L c)
hid_shape: torch.LongTensor, # (b n)
window_fn: Callable[[torch.Tensor], List[torch.Tensor]],
):
"""
Apply windowing function to create non-overlapping windows.
Used for window attention mechanisms where sequences are split into windows.
Args:
hid: Flattened tensor (L, c)
hid_shape: Shape metadata (b, n)
window_fn: Function that splits a tensor into windows
Returns:
Tuple of (windowed_tensor, window_shapes, window_counts)
COMPILE OPTIMIZATION: Uses tensor operation for window count tracking
"""
unflattened = unflatten(hid, hid_shape)
windowed = [window_fn(h) for h in unflattened]
# Track window counts using tensor operations
device = hid_shape.device
hid_windows = torch.tensor([len(w) for w in windowed], dtype=torch.long, device=device)
# Flatten all windows
all_windows = list(chain(*windowed))
hid, hid_shape = flatten(all_windows)
return hid, hid_shape, hid_windows
def window_idx(
hid_shape: torch.LongTensor, # (b n)
window_fn: Callable[[torch.Tensor], List[torch.Tensor]],
):
"""
Create index-based windowing functions.
Pre-computes indices for efficient windowing and reverse operations.
Args:
hid_shape: Shape metadata (b, n)
window_fn: Function that splits a tensor into windows
Returns:
Tuple of (window_fn, reverse_fn, window_shapes, window_counts)
"""
hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1)
tgt_idx, tgt_shape, tgt_windows = window(hid_idx, hid_shape, window_fn)
tgt_idx = tgt_idx.squeeze(-1)
src_idx = torch.argsort(tgt_idx)
return (
lambda hid: torch.index_select(hid, 0, tgt_idx),
lambda hid: torch.index_select(hid, 0, src_idx),
tgt_shape,
tgt_windows,
)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/__init__.py
+26
View File
@@ -0,0 +1,26 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from .mmsr_block import NaMMSRTransformerBlock
nadit_blocks = {
"mmdit_sr": NaMMSRTransformerBlock,
}
def get_nablock(block_type: str):
if block_type in nadit_blocks:
return nadit_blocks[block_type]
raise NotImplementedError(f"{block_type} is not supported")
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/attention/__init__.py
+25
View File
@@ -0,0 +1,25 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from .mmattn import NaMMAttention
attns = {
"mm_full": NaMMAttention,
}
def get_attn(attn_type: str):
if attn_type in attns:
return attns[attn_type]
raise NotImplementedError(f"{attn_type} is not supported")
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/attention/mmattn.py
+271
View File
@@ -0,0 +1,271 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Optional, Tuple, Union
import torch
from einops import rearrange
from torch import nn
from torch.nn import functional as F
from torch.nn.modules.utils import _triple
from .....common.cache import Cache
from .....common.distributed.ops import gather_heads_scatter_seq, gather_seq_scatter_heads_qkv
from .....common.half_precision_fixes import safe_pad_operation
from ... import na
from ...attention import FlashAttentionVarlen
from ...mm import MMArg, MMModule
from ...normalization import norm_layer_type
from ...rope import get_na_rope
from ...window import get_window_op
from itertools import chain
class NaMMAttention(nn.Module):
def __init__(
self,
vid_dim: int,
txt_dim: int,
heads: int,
head_dim: int,
qk_bias: bool,
qk_norm: norm_layer_type,
qk_norm_eps: float,
rope_type: Optional[str],
rope_dim: int,
shared_weights: bool,
attention_mode: str = 'sdpa',
**kwargs,
):
super().__init__()
dim = MMArg(vid_dim, txt_dim)
inner_dim = heads * head_dim
qkv_dim = inner_dim * 3
self.head_dim = head_dim
self.proj_qkv = MMModule(
nn.Linear, dim, qkv_dim, bias=qk_bias, shared_weights=shared_weights
)
self.proj_out = MMModule(nn.Linear, inner_dim, dim, shared_weights=shared_weights)
self.norm_q = MMModule(
qk_norm,
dim=head_dim,
eps=qk_norm_eps,
elementwise_affine=True,
shared_weights=shared_weights,
)
self.norm_k = MMModule(
qk_norm,
dim=head_dim,
eps=qk_norm_eps,
elementwise_affine=True,
shared_weights=shared_weights,
)
self.rope = get_na_rope(rope_type=rope_type, dim=rope_dim)
self.attn = FlashAttentionVarlen(attention_mode=attention_mode)
def forward(
self,
vid: torch.FloatTensor, # l c
txt: torch.FloatTensor, # l c
vid_shape: torch.LongTensor, # b 3
txt_shape: torch.LongTensor, # b 1
cache: Cache,
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
]:
vid_qkv, txt_qkv = self.proj_qkv(vid, txt)
vid_qkv = gather_seq_scatter_heads_qkv(
vid_qkv,
seq_dim=0,
qkv_shape=vid_shape,
cache=cache.namespace("vid"),
)
txt_qkv = gather_seq_scatter_heads_qkv(
txt_qkv,
seq_dim=0,
qkv_shape=txt_shape,
cache=cache.namespace("txt"),
)
vid_qkv = rearrange(vid_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim)
txt_qkv = rearrange(txt_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim)
vid_q, vid_k, vid_v = vid_qkv.unbind(1)
txt_q, txt_k, txt_v = txt_qkv.unbind(1)
vid_q, txt_q = self.norm_q(vid_q, txt_q)
vid_k, txt_k = self.norm_k(vid_k, txt_k)
if self.rope:
if self.rope.mm:
vid_q, vid_k, txt_q, txt_k = self.rope(
vid_q, vid_k, vid_shape, txt_q, txt_k, txt_shape, cache
)
else:
vid_q, vid_k = self.rope(vid_q, vid_k, vid_shape, cache)
vid_len = cache("vid_len", lambda: vid_shape.prod(-1))
txt_len = cache("txt_len", lambda: txt_shape.prod(-1))
all_len = cache("all_len", lambda: vid_len + txt_len)
concat, unconcat = cache("mm_pnp", lambda: na.concat_idx(vid_len, txt_len))
# Attention handles dtype conversion internally using pipeline compute_dtype
attn = self.attn(
q=concat(vid_q, txt_q),
k=concat(vid_k, txt_k),
v=concat(vid_v, txt_v),
cu_seqlens_q=cache("mm_seqlens", lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)).int()),
cu_seqlens_k=cache("mm_seqlens", lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)).int()),
max_seqlen_q=cache("mm_maxlen", lambda: all_len.max()),
max_seqlen_k=cache("mm_maxlen", lambda: all_len.max()),
).type_as(vid_q)
attn = rearrange(attn, "l h d -> l (h d)")
vid_out, txt_out = unconcat(attn)
vid_out = gather_heads_scatter_seq(vid_out, head_dim=1, seq_dim=0)
txt_out = gather_heads_scatter_seq(txt_out, head_dim=1, seq_dim=0)
vid_out, txt_out = self.proj_out(vid_out, txt_out)
return vid_out, txt_out
class NaSwinAttention(NaMMAttention):
def __init__(
self,
*args,
window: Union[int, Tuple[int, int, int]],
window_method: str,
attention_mode: str = 'sdpa',
**kwargs,
):
super().__init__(*args, attention_mode=attention_mode, **kwargs)
self.window = _triple(window)
self.window_method = window_method
assert all(map(lambda v: isinstance(v, int) and v >= 0, self.window))
self.window_op = get_window_op(window_method)
def forward(
self,
vid: torch.FloatTensor, # l c
txt: torch.FloatTensor, # l c
vid_shape: torch.LongTensor, # b 3
txt_shape: torch.LongTensor, # b 1
cache: Cache,
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
]:
vid_qkv, txt_qkv = self.proj_qkv(vid, txt)
vid_qkv = gather_seq_scatter_heads_qkv(
vid_qkv,
seq_dim=0,
qkv_shape=vid_shape,
cache=cache.namespace("vid"),
)
txt_qkv = gather_seq_scatter_heads_qkv(
txt_qkv,
seq_dim=0,
qkv_shape=txt_shape,
cache=cache.namespace("txt"),
)
# re-org the input seq for window attn
cache_win = cache.namespace(f"{self.window_method}_{self.window}_sd3")
def make_window(x: torch.Tensor):
t, h, w, _ = x.shape
window_slices = self.window_op((t, h, w), self.window)
return [x[st, sh, sw] for (st, sh, sw) in window_slices]
window_partition, window_reverse, window_shape, window_count = cache_win(
"win_transform",
lambda: na.window_idx(vid_shape, make_window),
)
vid_qkv_win = window_partition(vid_qkv)
vid_qkv_win = rearrange(vid_qkv_win, "l (o h d) -> l o h d", o=3, d=self.head_dim)
txt_qkv = rearrange(txt_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim)
vid_q, vid_k, vid_v = vid_qkv_win.unbind(1)
txt_q, txt_k, txt_v = txt_qkv.unbind(1)
vid_q, txt_q = self.norm_q(vid_q, txt_q)
vid_k, txt_k = self.norm_k(vid_k, txt_k)
txt_len = cache("txt_len", lambda: txt_shape.prod(-1))
vid_len_win = cache_win("vid_len", lambda: window_shape.prod(-1))
txt_len_win = cache_win("txt_len", lambda: txt_len.repeat_interleave(window_count))
all_len_win = cache_win("all_len", lambda: vid_len_win + txt_len_win)
concat_win, unconcat_win = cache_win(
"mm_pnp", lambda: na.repeat_concat_idx(vid_len_win, txt_len, window_count)
)
# window rope
if self.rope:
if self.rope.mm:
# repeat text q and k for window mmrope
_, num_h, _ = txt_q.shape
txt_q_repeat = rearrange(txt_q, "l h d -> l (h d)")
txt_q_repeat = na.unflatten(txt_q_repeat, txt_shape)
txt_q_repeat = [[x] * n for x, n in zip(txt_q_repeat, window_count)]
txt_q_repeat = list(chain(*txt_q_repeat))
txt_q_repeat, txt_shape_repeat = na.flatten(txt_q_repeat)
txt_q_repeat = rearrange(txt_q_repeat, "l (h d) -> l h d", h=num_h)
txt_k_repeat = rearrange(txt_k, "l h d -> l (h d)")
txt_k_repeat = na.unflatten(txt_k_repeat, txt_shape)
txt_k_repeat = [[x] * n for x, n in zip(txt_k_repeat, window_count)]
txt_k_repeat = list(chain(*txt_k_repeat))
txt_k_repeat, _ = na.flatten(txt_k_repeat)
txt_k_repeat = rearrange(txt_k_repeat, "l (h d) -> l h d", h=num_h)
vid_q, vid_k, txt_q, txt_k = self.rope(
vid_q, vid_k, window_shape, txt_q_repeat, txt_k_repeat, txt_shape_repeat, cache_win
)
else:
vid_q, vid_k = self.rope(vid_q, vid_k, window_shape, cache_win)
# Attention handles dtype conversion internally using pipeline compute_dtype
out = self.attn(
q=concat_win(vid_q, txt_q),
k=concat_win(vid_k, txt_k),
v=concat_win(vid_v, txt_v),
cu_seqlens_q=cache_win(
"vid_seqlens_q", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int()
),
cu_seqlens_k=cache_win(
"vid_seqlens_k", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int()
),
max_seqlen_q=cache_win("vid_max_seqlen_q", lambda: all_len_win.max()),
max_seqlen_k=cache_win("vid_max_seqlen_k", lambda: all_len_win.max()),
).type_as(vid_q)
# text pooling
vid_out, txt_out = unconcat_win(out)
vid_out = rearrange(vid_out, "l h d -> l (h d)")
txt_out = rearrange(txt_out, "l h d -> l (h d)")
vid_out = window_reverse(vid_out)
vid_out = gather_heads_scatter_seq(vid_out, head_dim=1, seq_dim=0)
txt_out = gather_heads_scatter_seq(txt_out, head_dim=1, seq_dim=0)
vid_out, txt_out = self.proj_out(vid_out, txt_out)
return vid_out, txt_out
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/mmsr_block.py
+128
View File
@@ -0,0 +1,128 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Tuple
import torch
import torch.nn as nn
# from ..cache import Cache
from ....common.cache import Cache
from .attention.mmattn import NaSwinAttention
from ..mm import MMArg
from ..modulation import ada_layer_type
from ..normalization import norm_layer_type
from ..mm import MMArg, MMModule
from ..mlp import get_mlp
class NaMMSRTransformerBlock(nn.Module):
def __init__(
self,
*,
vid_dim: int,
txt_dim: int,
emb_dim: int,
heads: int,
head_dim: int,
expand_ratio: int,
norm: norm_layer_type,
norm_eps: float,
ada: ada_layer_type,
qk_bias: bool,
qk_norm: norm_layer_type,
mlp_type: str,
shared_weights: bool,
rope_type: str,
rope_dim: int,
is_last_layer: bool,
attention_mode: str = 'sdpa',
**kwargs,
):
super().__init__()
dim = MMArg(vid_dim, txt_dim)
self.attn_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False, shared_weights=shared_weights,)
self.attn = NaSwinAttention(
vid_dim=vid_dim,
txt_dim=txt_dim,
heads=heads,
head_dim=head_dim,
qk_bias=qk_bias,
qk_norm=qk_norm,
qk_norm_eps=norm_eps,
rope_type=rope_type,
rope_dim=rope_dim,
shared_weights=shared_weights,
attention_mode=attention_mode,
window=kwargs.pop("window", None),
window_method=kwargs.pop("window_method", None),
)
self.mlp_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False, shared_weights=shared_weights, vid_only=is_last_layer)
self.mlp = MMModule(
get_mlp(mlp_type),
dim=dim,
expand_ratio=expand_ratio,
shared_weights=shared_weights,
vid_only=is_last_layer
)
self.ada = MMModule(ada, dim=dim, emb_dim=emb_dim, layers=["attn", "mlp"], shared_weights=shared_weights, vid_only=is_last_layer)
self.is_last_layer = is_last_layer
def forward(
self,
vid: torch.FloatTensor, # l c
txt: torch.FloatTensor, # l c
vid_shape: torch.LongTensor, # b 3
txt_shape: torch.LongTensor, # b 1
emb: torch.FloatTensor,
cache: Cache,
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
torch.LongTensor,
torch.LongTensor,
]:
hid_len = MMArg(
cache("vid_len", lambda: vid_shape.prod(-1)),
cache("txt_len", lambda: txt_shape.prod(-1)),
)
ada_kwargs = {
"emb": emb,
"hid_len": hid_len,
"cache": cache,
"branch_tag": MMArg("vid", "txt"),
}
vid_attn, txt_attn = self.attn_norm(vid, txt)
vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="in", **ada_kwargs)
vid_attn, txt_attn = self.attn(vid_attn, txt_attn, vid_shape, txt_shape, cache)
vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="out", **ada_kwargs)
vid_attn, txt_attn = (vid_attn + vid), (txt_attn + txt)
vid_mlp, txt_mlp = self.mlp_norm(vid_attn, txt_attn)
# ADD BY NUMZ
if vid_mlp.dtype != vid_attn.dtype:
vid_mlp = vid_mlp.to(vid_attn.dtype)
if txt_mlp.dtype != txt_attn.dtype:
txt_mlp = txt_mlp.to(txt_attn.dtype)
# END BY NUMZ
vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="in", **ada_kwargs)
vid_mlp, txt_mlp = self.mlp(vid_mlp, txt_mlp)
vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="out", **ada_kwargs)
vid_mlp, txt_mlp = (vid_mlp + vid_attn), (txt_mlp + txt_attn)
return vid_mlp, txt_mlp, vid_shape, txt_shape
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nadit.py
+248
View File
@@ -0,0 +1,248 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union, Callable
import torch
from torch import nn
from ...common.cache import Cache
from ...common.distributed.ops import slice_inputs
from . import na
from .embedding import TimeEmbedding
from .modulation import get_ada_layer
from .nablocks import get_nablock
from .normalization import get_norm_layer
from .patch import get_na_patch_layers
# Fake func, no checkpointing is required for inference
def gradient_checkpointing(module: Union[Callable, nn.Module], *args, enabled: bool, **kwargs):
return module(*args, **kwargs)
@dataclass
class NaDiTOutput:
vid_sample: torch.Tensor
class NaDiT(nn.Module):
"""
Native Resolution Diffusion Transformer (NaDiT)
"""
gradient_checkpointing = False
def __init__(
self,
vid_in_channels: int,
vid_out_channels: int,
vid_dim: int,
txt_in_dim: Union[int, List[int]],
txt_dim: Optional[int],
emb_dim: int,
heads: int,
head_dim: int,
expand_ratio: int,
norm: Optional[str],
norm_eps: float,
ada: str,
qk_bias: bool,
qk_norm: Optional[str],
patch_size: Union[int, Tuple[int, int, int]],
num_layers: int,
block_type: Union[str, Tuple[str]],
mm_layers: Union[int, Tuple[bool]],
mlp_type: str = "normal",
patch_type: str = "v1",
rope_type: Optional[str] = "rope3d",
rope_dim: Optional[int] = None,
window: Optional[Tuple] = None,
window_method: Optional[Tuple[str]] = None,
msa_type: Optional[Tuple[str]] = None,
mca_type: Optional[Tuple[str]] = None,
txt_in_norm: Optional[str] = None,
txt_in_norm_scale_factor: int = 0.01,
txt_proj_type: Optional[str] = "linear",
vid_out_norm: Optional[str] = None,
attention_mode: str = 'sdpa',
**kwargs,
):
ada = get_ada_layer(ada)
norm = get_norm_layer(norm)
qk_norm = get_norm_layer(qk_norm)
rope_dim = rope_dim if rope_dim is not None else head_dim // 2
if isinstance(block_type, str):
block_type = [block_type] * num_layers
elif len(block_type) != num_layers:
raise ValueError("The ``block_type`` list should equal to ``num_layers``.")
super().__init__()
NaPatchIn, NaPatchOut = get_na_patch_layers(patch_type)
self.vid_in = NaPatchIn(
in_channels=vid_in_channels,
patch_size=patch_size,
dim=vid_dim,
)
if not isinstance(txt_in_dim, int):
self.txt_in = nn.ModuleList([])
for in_dim in txt_in_dim:
txt_norm_layer = get_norm_layer(txt_in_norm)(txt_dim, norm_eps, True)
if txt_proj_type == "linear":
txt_proj_layer = nn.Linear(in_dim, txt_dim)
else:
txt_proj_layer = nn.Sequential(
nn.Linear(in_dim, in_dim), nn.GELU("tanh"), nn.Linear(in_dim, txt_dim)
)
torch.nn.init.constant_(txt_norm_layer.weight, txt_in_norm_scale_factor)
self.txt_in.append(
nn.Sequential(
txt_proj_layer,
txt_norm_layer,
)
)
else:
self.txt_in = (
nn.Linear(txt_in_dim, txt_dim)
if txt_in_dim and txt_in_dim != txt_dim
else nn.Identity()
)
self.emb_in = TimeEmbedding(
sinusoidal_dim=256,
hidden_dim=max(vid_dim, txt_dim),
output_dim=emb_dim,
)
if window is None or isinstance(window[0], int):
window = [window] * num_layers
if window_method is None or isinstance(window_method, str):
window_method = [window_method] * num_layers
if msa_type is None or isinstance(msa_type, str):
msa_type = [msa_type] * num_layers
if mca_type is None or isinstance(mca_type, str):
mca_type = [mca_type] * num_layers
self.blocks = nn.ModuleList(
[
get_nablock(block_type[i])(
vid_dim=vid_dim,
txt_dim=txt_dim,
emb_dim=emb_dim,
heads=heads,
head_dim=head_dim,
expand_ratio=expand_ratio,
norm=norm,
norm_eps=norm_eps,
ada=ada,
qk_bias=qk_bias,
qk_norm=qk_norm,
shared_weights=not (
(i < mm_layers) if isinstance(mm_layers, int) else mm_layers[i]
),
mlp_type=mlp_type,
window=window[i],
window_method=window_method[i],
msa_type=msa_type[i],
mca_type=mca_type[i],
rope_type=rope_type,
rope_dim=rope_dim,
is_last_layer=(i == num_layers - 1),
attention_mode=attention_mode,
**kwargs,
)
for i in range(num_layers)
]
)
self.vid_out_norm = None
if vid_out_norm is not None:
self.vid_out_norm = get_norm_layer(vid_out_norm)(
dim=vid_dim,
eps=norm_eps,
elementwise_affine=True,
)
self.vid_out_ada = ada(
dim=vid_dim,
emb_dim=emb_dim,
layers=["out"],
modes=["in"],
)
self.vid_out = NaPatchOut(
out_channels=vid_out_channels,
patch_size=patch_size,
dim=vid_dim,
)
def set_gradient_checkpointing(self, enable: bool):
self.gradient_checkpointing = enable
def forward(
self,
vid: torch.FloatTensor, # l c
txt: Union[torch.FloatTensor, List[torch.FloatTensor]], # l c
vid_shape: torch.LongTensor, # b 3
txt_shape: Union[torch.LongTensor, List[torch.LongTensor]], # b 1
timestep: Union[int, float, torch.IntTensor, torch.FloatTensor], # b
disable_cache: bool = False, # for test
):
cache = Cache(disable=disable_cache)
# slice vid after patching in when using sequence parallelism
if isinstance(txt, list):
assert isinstance(self.txt_in, nn.ModuleList)
txt = [
na.unflatten(fc(i), s) for fc, i, s in zip(self.txt_in, txt, txt_shape)
] # B L D
txt, txt_shape = na.flatten([torch.cat(t, dim=0) for t in zip(*txt)])
txt = slice_inputs(txt, dim=0)
else:
txt = slice_inputs(txt, dim=0)
txt = self.txt_in(txt)
# Video input.
# Sequence parallel slicing is done inside patching class.
vid, vid_shape = self.vid_in(vid, vid_shape, cache)
# Embedding input.
emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype)
# Body
for i, block in enumerate(self.blocks):
vid, txt, vid_shape, txt_shape = gradient_checkpointing(
enabled=(self.gradient_checkpointing and self.training),
module=block,
vid=vid,
txt=txt,
vid_shape=vid_shape,
txt_shape=txt_shape,
emb=emb,
cache=cache,
)
# Video output norm.
if self.vid_out_norm:
vid = self.vid_out_norm(vid)
vid = self.vid_out_ada(
vid,
emb=emb,
layer="out",
mode="in",
hid_len=cache("vid_len", lambda: vid_shape.prod(-1)),
cache=cache,
branch_tag="vid",
)
# Video output.
vid, vid_shape = self.vid_out(vid, vid_shape, cache)
return NaDiTOutput(vid_sample=vid)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/normalization.py
+150
View File
@@ -0,0 +1,150 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Callable, Optional
from diffusers.models.normalization import RMSNorm
from torch import nn
import torch
import torch.nn.functional as F
import numbers
from torch.nn.parameter import Parameter
from torch.nn import init
# (dim: int, eps: float, elementwise_affine: bool)
norm_layer_type = Callable[[int, float, bool], nn.Module]
class CustomLayerNorm(nn.Module):
"""
Custom LayerNorm implementation to replace Apex FusedLayerNorm
"""
def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
super(CustomLayerNorm, self).__init__()
if isinstance(normalized_shape, numbers.Integral):
normalized_shape = (normalized_shape,)
self.normalized_shape = torch.Size(normalized_shape)
self.eps = eps
self.elementwise_affine = elementwise_affine
if self.elementwise_affine:
self.weight = Parameter(torch.Tensor(*normalized_shape))
self.bias = Parameter(torch.Tensor(*normalized_shape))
else:
self.register_parameter('weight', None)
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
if self.elementwise_affine:
init.ones_(self.weight)
init.zeros_(self.bias)
def forward(self, input):
# 🚀 FP8 COMPATIBILITY: Convert parameters to match input dtype
# This prevents "Promotion for Float8 Types is not supported" errors
weight = self.weight
bias = self.bias
if self.elementwise_affine and weight is not None:
if weight.dtype != input.dtype:
weight = weight.to(input.dtype)
if bias is not None and bias.dtype != input.dtype:
bias = bias.to(input.dtype)
return F.layer_norm(
input, self.normalized_shape, weight, bias, self.eps)
class CustomRMSNorm(nn.Module):
"""
Custom RMSNorm implementation to replace Apex FusedRMSNorm
"""
def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
super(CustomRMSNorm, self).__init__()
if isinstance(normalized_shape, numbers.Integral):
normalized_shape = (normalized_shape,)
self.normalized_shape = torch.Size(normalized_shape)
self.eps = eps
self.elementwise_affine = elementwise_affine
if self.elementwise_affine:
self.weight = Parameter(torch.ones(*normalized_shape))
else:
self.register_parameter('weight', None)
def forward(self, input):
# RMS normalization: x / sqrt(mean(x^2) + eps) * weight
dims = tuple(range(-len(self.normalized_shape), 0))
# Calculate RMS: sqrt(mean(x^2))
variance = input.pow(2).mean(dim=dims, keepdim=True)
rms = torch.sqrt(variance + self.eps)
# Normalize
normalized = input / rms
if self.elementwise_affine:
# Convert FP8 weight to match input dtype for arithmetic operations
if hasattr(torch, 'float8_e4m3fn'):
fp8_types = (torch.float8_e4m3fn, torch.float8_e5m2)
if self.weight.dtype in fp8_types:
# Use input dtype as target (respects pipeline precision)
weight = self.weight.to(input.dtype)
return normalized * weight
return normalized * self.weight
return normalized
def get_norm_layer(norm_type: Optional[str]) -> norm_layer_type:
def _norm_layer(dim: int, eps: float, elementwise_affine: bool):
if norm_type is None:
return nn.Identity()
if norm_type == "layer":
return nn.LayerNorm(
normalized_shape=dim,
eps=eps,
elementwise_affine=elementwise_affine,
)
if norm_type == "rms":
return RMSNorm(
dim=dim,
eps=eps,
elementwise_affine=elementwise_affine,
)
if norm_type == "fusedln":
# Use custom LayerNorm instead of Apex FusedLayerNorm
return CustomLayerNorm(
normalized_shape=dim,
elementwise_affine=elementwise_affine,
eps=eps,
)
if norm_type == "fusedrms":
# Use custom RMSNorm instead of Apex FusedRMSNorm
return CustomRMSNorm(
normalized_shape=dim,
elementwise_affine=elementwise_affine,
eps=eps,
)
raise NotImplementedError(f"{norm_type} is not supported")
return _norm_layer
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/patch/__init__.py
+19
View File
@@ -0,0 +1,19 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
def get_na_patch_layers(patch_type="v1"):
assert patch_type in ["v1"]
if patch_type == "v1":
from .patch_v1 import NaPatchIn, NaPatchOut
return NaPatchIn, NaPatchOut
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/patch/patch_v1.py
+127
View File
@@ -0,0 +1,127 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Tuple, Union
import torch
from einops import rearrange
from torch import nn
from torch.nn.modules.utils import _triple
from ....common.cache import Cache
from ....common.distributed.ops import gather_outputs, slice_inputs
from .. import na
class PatchIn(nn.Module):
def __init__(
self,
in_channels: int,
patch_size: Union[int, Tuple[int, int, int]],
dim: int,
):
super().__init__()
t, h, w = _triple(patch_size)
self.patch_size = t, h, w
self.proj = nn.Linear(in_channels * t * h * w, dim)
def forward(
self,
vid: torch.Tensor,
) -> torch.Tensor:
t, h, w = self.patch_size
if t > 1:
assert vid.size(2) % t == 1
vid = torch.cat([vid[:, :, :1]] * (t - 1) + [vid], dim=2)
vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w)
vid = self.proj(vid)
return vid
class PatchOut(nn.Module):
def __init__(
self,
out_channels: int,
patch_size: Union[int, Tuple[int, int, int]],
dim: int,
):
super().__init__()
t, h, w = _triple(patch_size)
self.patch_size = t, h, w
self.proj = nn.Linear(dim, out_channels * t * h * w)
def forward(
self,
vid: torch.Tensor,
) -> torch.Tensor:
t, h, w = self.patch_size
vid = self.proj(vid)
vid = rearrange(vid, "b T H W (t h w c) -> b c (T t) (H h) (W w)", t=t, h=h, w=w)
if t > 1:
vid = vid[:, :, (t - 1) :]
return vid
class NaPatchIn(PatchIn):
def forward(
self,
vid: torch.Tensor, # l c
vid_shape: torch.LongTensor,
cache: Cache = Cache(disable=True), # for test
) -> torch.Tensor:
cache = cache.namespace("patch")
vid_shape_before_patchify = cache("vid_shape_before_patchify", lambda: vid_shape)
t, h, w = self.patch_size
if not (t == h == w == 1):
vid = na.unflatten(vid, vid_shape)
for i in range(len(vid)):
if t > 1 and vid_shape_before_patchify[i, 0] % t != 0:
vid[i] = torch.cat([vid[i][:1]] * (t - vid[i].size(0) % t) + [vid[i]], dim=0)
vid[i] = rearrange(vid[i], "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w)
vid, vid_shape = na.flatten(vid)
# slice vid after patching in when using sequence parallelism
vid = slice_inputs(vid, dim=0)
vid = self.proj(vid)
return vid, vid_shape
class NaPatchOut(PatchOut):
def forward(
self,
vid: torch.FloatTensor, # l c
vid_shape: torch.LongTensor,
cache: Cache = Cache(disable=True), # for test
) -> Tuple[
torch.FloatTensor,
torch.LongTensor,
]:
cache = cache.namespace("patch")
vid_shape_before_patchify = cache.get("vid_shape_before_patchify")
t, h, w = self.patch_size
vid = self.proj(vid)
# gather vid before patching out when enabling sequence parallelism
vid = gather_outputs(
vid, gather_dim=0, padding_dim=0, unpad_shape=vid_shape, cache=cache.namespace("vid")
)
if not (t == h == w == 1):
vid = na.unflatten(vid, vid_shape)
for i in range(len(vid)):
vid[i] = rearrange(vid[i], "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w)
if t > 1 and vid_shape_before_patchify[i, 0] % t != 0:
vid[i] = vid[i][(t - vid_shape_before_patchify[i, 0] % t) :]
vid, vid_shape = na.flatten(vid)
return vid, vid_shape
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/rope.py
+184
View File
@@ -0,0 +1,184 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from functools import lru_cache
from typing import Optional, Tuple
import torch
from einops import rearrange
from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb
from torch import nn
from ...common.cache import Cache
class RotaryEmbeddingBase(nn.Module):
def __init__(self, dim: int, rope_dim: int):
super().__init__()
self.rope = RotaryEmbedding(
dim=dim // rope_dim,
freqs_for="pixel",
max_freq=256,
)
# 1. Set model.requires_grad_(True) after model creation will make
# the `requires_grad=False` for rope freqs no longer hold.
# 2. Even if we don't set requires_grad_(True) explicitly,
# FSDP is not memory efficient when handling fsdp_wrap
# with mixed requires_grad=True/False.
# With above consideration, it is easier just remove the freqs
# out of nn.Parameters when `learned_freq=False`
freqs = self.rope.freqs
del self.rope.freqs
self.rope.register_buffer("freqs", freqs.data)
@lru_cache(maxsize=128)
def get_axial_freqs(self, *dims):
return self.rope.get_axial_freqs(*dims)
class RotaryEmbedding3d(RotaryEmbeddingBase):
def __init__(self, dim: int):
super().__init__(dim, rope_dim=3)
self.mm = False
def forward(
self,
q: torch.FloatTensor, # b h l d
k: torch.FloatTensor, # b h l d
size: Tuple[int, int, int],
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
]:
T, H, W = size
freqs = self.get_axial_freqs(T, H, W)
q = rearrange(q, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W)
k = rearrange(k, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W)
q = apply_rotary_emb(freqs, q.float()).to(q.dtype)
k = apply_rotary_emb(freqs, k.float()).to(k.dtype)
q = rearrange(q, "b h T H W d -> b h (T H W) d")
k = rearrange(k, "b h T H W d -> b h (T H W) d")
return q, k
class MMRotaryEmbeddingBase(RotaryEmbeddingBase):
def __init__(self, dim: int, rope_dim: int):
super().__init__(dim, rope_dim)
self.rope = RotaryEmbedding(
dim=dim // rope_dim,
freqs_for="lang",
theta=10000,
)
freqs = self.rope.freqs
del self.rope.freqs
self.rope.register_buffer("freqs", freqs.data)
self.mm = True
class NaMMRotaryEmbedding3d(MMRotaryEmbeddingBase):
def __init__(self, dim: int):
super().__init__(dim, rope_dim=3)
def forward(
self,
vid_q: torch.FloatTensor, # L h d
vid_k: torch.FloatTensor, # L h d
vid_shape: torch.LongTensor, # B 3
txt_q: torch.FloatTensor, # L h d
txt_k: torch.FloatTensor, # L h d
txt_shape: torch.LongTensor, # B 1
cache: Cache,
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
torch.FloatTensor,
torch.FloatTensor,
]:
vid_freqs, txt_freqs = cache(
"mmrope_freqs_3d",
lambda: self.get_freqs(vid_shape, txt_shape),
)
target_device = vid_q.device
if vid_freqs.device != target_device:
vid_freqs = vid_freqs.to(target_device)
if txt_freqs.device != target_device:
txt_freqs = txt_freqs.to(target_device)
vid_q = rearrange(vid_q, "L h d -> h L d")
vid_k = rearrange(vid_k, "L h d -> h L d")
vid_q = apply_rotary_emb(vid_freqs, vid_q.float()).to(vid_q.dtype)
vid_k = apply_rotary_emb(vid_freqs, vid_k.float()).to(vid_k.dtype)
vid_q = rearrange(vid_q, "h L d -> L h d")
vid_k = rearrange(vid_k, "h L d -> L h d")
txt_q = rearrange(txt_q, "L h d -> h L d")
txt_k = rearrange(txt_k, "L h d -> h L d")
txt_q = apply_rotary_emb(txt_freqs, txt_q.float()).to(txt_q.dtype)
txt_k = apply_rotary_emb(txt_freqs, txt_k.float()).to(txt_k.dtype)
txt_q = rearrange(txt_q, "h L d -> L h d")
txt_k = rearrange(txt_k, "h L d -> L h d")
return vid_q, vid_k, txt_q, txt_k
@torch._dynamo.disable # Disable compilation: .tolist() is data-dependent and causes graph breaks
def get_freqs(
self,
vid_shape: torch.LongTensor,
txt_shape: torch.LongTensor,
) -> Tuple[
torch.Tensor,
torch.Tensor,
]:
"""
Generate RoPE frequencies for variable batch shapes.
Note: This method uses @torch._dynamo.disable because it requires
data-dependent control flow (shape.tolist()) that cannot be symbolically
traced by torch.compile. The cache() wrapper in the forward pass memoizes
results to reduce recomputation overhead.
"""
# Calculate actual max dimensions needed for this batch
max_temporal = 0
max_height = 0
max_width = 0
max_txt_len = 0
for (f, h, w), l in zip(vid_shape.tolist(), txt_shape[:, 0].tolist()):
max_temporal = max(max_temporal, l + f) # Need up to l+f for temporal
max_height = max(max_height, h)
max_width = max(max_width, w)
max_txt_len = max(max_txt_len, l)
# Compute frequencies for actual max dimensions needed
# Add small buffer to improve cache hits across similar batches
vid_freqs = self.get_axial_freqs(
min(max_temporal + 16, 1024), # Cap at 1024, add small buffer
min(max_height + 4, 128), # Cap at 128, add small buffer
min(max_width + 4, 128) # Cap at 128, add small buffer
)
txt_freqs = self.get_axial_freqs(min(max_txt_len + 16, 1024))
# Now slice as before
vid_freq_list, txt_freq_list = [], []
for (f, h, w), l in zip(vid_shape.tolist(), txt_shape[:, 0].tolist()):
vid_freq = vid_freqs[l : l + f, :h, :w].reshape(-1, vid_freqs.size(-1))
txt_freq = txt_freqs[:l].repeat(1, 3).reshape(-1, vid_freqs.size(-1))
vid_freq_list.append(vid_freq)
txt_freq_list.append(txt_freq)
return torch.cat(vid_freq_list, dim=0), torch.cat(txt_freq_list, dim=0)
def get_na_rope(rope_type: Optional[str], dim: int):
if rope_type is None:
return None
if rope_type == "mmrope3d":
return NaMMRotaryEmbedding3d(dim=dim)
raise NotImplementedError(f"{rope_type} is not supported.")
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/window.py
+83
View File
@@ -0,0 +1,83 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from math import ceil
from typing import Tuple
import math
def get_window_op(name: str):
if name == "720pwin_by_size_bysize":
return make_720Pwindows_bysize
if name == "720pswin_by_size_bysize":
return make_shifted_720Pwindows_bysize
raise ValueError(f"Unknown windowing method: {name}")
# -------------------------------- Windowing -------------------------------- #
def make_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]):
t, h, w = size
resized_nt, resized_nh, resized_nw = num_windows
#cal windows under 720p
scale = math.sqrt((45 * 80) / (h * w))
resized_h, resized_w = round(h * scale), round(w * scale)
wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw) # window size.
wt = ceil(min(t, 30) / resized_nt) # window size.
nt, nh, nw = ceil(t / wt), ceil(h / wh), ceil(w / ww) # window size.
return [
(
slice(it * wt, min((it + 1) * wt, t)),
slice(ih * wh, min((ih + 1) * wh, h)),
slice(iw * ww, min((iw + 1) * ww, w)),
)
for iw in range(nw)
if min((iw + 1) * ww, w) > iw * ww
for ih in range(nh)
if min((ih + 1) * wh, h) > ih * wh
for it in range(nt)
if min((it + 1) * wt, t) > it * wt
]
def make_shifted_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]):
t, h, w = size
resized_nt, resized_nh, resized_nw = num_windows
#cal windows under 720p
scale = math.sqrt((45 * 80) / (h * w))
resized_h, resized_w = round(h * scale), round(w * scale)
wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw) # window size.
wt = ceil(min(t, 30) / resized_nt) # window size.
st, sh, sw = ( # shift size.
0.5 if wt < t else 0,
0.5 if wh < h else 0,
0.5 if ww < w else 0,
)
nt, nh, nw = ceil((t - st) / wt), ceil((h - sh) / wh), ceil((w - sw) / ww) # window size.
nt, nh, nw = ( # number of window.
nt + 1 if st > 0 else 1,
nh + 1 if sh > 0 else 1,
nw + 1 if sw > 0 else 1,
)
return [
(
slice(max(int((it - st) * wt), 0), min(int((it - st + 1) * wt), t)),
slice(max(int((ih - sh) * wh), 0), min(int((ih - sh + 1) * wh), h)),
slice(max(int((iw - sw) * ww), 0), min(int((iw - sw + 1) * ww), w)),
)
for iw in range(nw)
if min(int((iw - sw + 1) * ww), w) > max(int((iw - sw) * ww), 0)
for ih in range(nh)
if min(int((ih - sh + 1) * wh), h) > max(int((ih - sh) * wh), 0)
for it in range(nt)
if min(int((it - st + 1) * wt), t) > max(int((it - st) * wt), 0)
]
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/attention.py
+148
View File
@@ -0,0 +1,148 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
import torch
import torch.nn.functional as F
# Import flash/sage attn with automatic fallback from compatibility layer
from ...optimization.compatibility import (
call_flash_attn_2_varlen, call_flash_attn_3_varlen,
call_sage_attn_2_varlen, call_sage_attn_3_varlen
)
from torch import nn
def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=None, max_seqlen_k=None, dropout_p=0.0, softmax_scale=None, causal=False, deterministic=False):
"""
A PyTorch-based implementation of variable-length attention to replace flash_attn_varlen_func.
It processes each sequence in the batch individually.
NOTE: max_seqlen_q and max_seqlen_k are accepted for API compatibility but not used.
PyTorch's scaled_dot_product_attention automatically handles variable sequence lengths.
COMPILE OPTIMIZATION: Uses torch.tensor_split to avoid .item() graph breaks
"""
# Split q, k, v using cumulative sequence lengths
# NOTE: torch.tensor_split requires int64 dtype and CPU device (PyTorch requirements)
q_splits = list(torch.tensor_split(q, cu_seqlens_q[1:-1].long().cpu(), dim=0))
k_splits = list(torch.tensor_split(k, cu_seqlens_k[1:-1].long().cpu(), dim=0))
v_splits = list(torch.tensor_split(v, cu_seqlens_k[1:-1].long().cpu(), dim=0))
# Process each sequence
output_splits = []
for q_i, k_i, v_i in zip(q_splits, k_splits, v_splits):
# Reshape for torch's scaled_dot_product_attention which expects (batch, heads, seq, dim).
# Here, we treat each sequence as a batch of 1.
q_i = q_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_q, head_dim)
k_i = k_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim)
v_i = v_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim)
# Use PyTorch's built-in scaled dot-product attention.
output_i = F.scaled_dot_product_attention(
q_i, k_i, v_i,
dropout_p=dropout_p if not deterministic else 0.0,
is_causal=causal
)
# Reshape the output back to the original format (seq_len, heads, head_dim)
output_i = output_i.squeeze(0).permute(1, 0, 2)
output_splits.append(output_i)
# Concatenate all outputs
return torch.cat(output_splits, dim=0)
class TorchAttention(nn.Module):
def tflops(self, args, kwargs, output) -> float:
assert len(args) == 0 or len(args) > 2, "query, key should both provided by args / kwargs"
q = kwargs.get("query") or args[0]
k = kwargs.get("key") or args[1]
b, h, sq, d = q.shape
b, h, sk, d = k.shape
return b * h * (4 * d * (sq / 1e6) * (sk / 1e6))
def forward(self, *args, **kwargs):
return F.scaled_dot_product_attention(*args, **kwargs)
class FlashAttentionVarlen(nn.Module):
"""
Variable-length attention with configurable backend.
Supported backends:
- sdpa: PyTorch SDPA (fully compilable, always available)
- flash_attn_2: Flash Attention 2 (Ampere+)
- flash_attn_3: Flash Attention 3 (Hopper+)
- sageattn_2: SageAttention 2
- sageattn_3: SageAttention 3 (Blackwell/RTX 50xx)
All non-SDPA backends use @torch._dynamo.disable wrapper (C++ extensions).
"""
def __init__(self, attention_mode: str = 'sdpa', compute_dtype: torch.dtype = None):
"""
Initialize with specified attention backend.
Args:
attention_mode: 'sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3'
compute_dtype: Compute dtype for attention (set by pipeline, defaults to None for auto-detection)
"""
super().__init__()
self.attention_mode = attention_mode
self.compute_dtype = compute_dtype
def tflops(self, args, kwargs, output) -> float:
cu_seqlens_q = kwargs["cu_seqlens_q"]
cu_seqlens_k = kwargs["cu_seqlens_k"]
_, h, d = output.shape
seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]) / 1e6
seqlens_k = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]) / 1e6
return h * (4 * d * (seqlens_q * seqlens_k).sum())
def forward(self, q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
kwargs["deterministic"] = torch.are_deterministic_algorithms_enabled()
# Convert to pipeline compute_dtype if configured (handles FP8 → fp16/bf16)
if self.compute_dtype is not None and q.dtype != self.compute_dtype:
q = q.to(self.compute_dtype)
k = k.to(self.compute_dtype)
v = v.to(self.compute_dtype)
if self.attention_mode == 'flash_attn_3':
return call_flash_attn_3_varlen(
q, k, v, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, **kwargs
)
elif self.attention_mode == 'flash_attn_2':
return call_flash_attn_2_varlen(
q, k, v, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, **kwargs
)
elif self.attention_mode == 'sageattn_3':
return call_sage_attn_3_varlen(
q, k, v, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, **kwargs
)
elif self.attention_mode == 'sageattn_2':
return call_sage_attn_2_varlen(
q, k, v, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, **kwargs
)
else:
# PyTorch SDPA
return pytorch_varlen_attention(
q, k, v, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, **kwargs
)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/blocks/__init__.py
+25
View File
@@ -0,0 +1,25 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from .mmdit_window_block import MMWindowTransformerBlock
dit_blocks = {
"mmdit_window": MMWindowTransformerBlock,
}
def get_block(block_type: str):
if block_type in dit_blocks:
return dit_blocks[block_type]
raise NotImplementedError(f"{block_type} is not supported")
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/blocks/mmdit_window_block.py
+233
View File
@@ -0,0 +1,233 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Tuple, Union
import torch
from einops import rearrange
from torch import nn
from torch.nn import functional as F
from torch.nn.modules.utils import _triple
from ....common.half_precision_fixes import safe_pad_operation
from ....common.distributed.ops import (
gather_heads,
gather_heads_scatter_seq,
gather_seq_scatter_heads_qkv,
scatter_heads,
)
from ..attention import TorchAttention
from ..mlp import get_mlp
from ..mm import MMArg, MMModule
from ..modulation import ada_layer_type
from ..normalization import norm_layer_type
from ..rope import RotaryEmbedding3d
class MMWindowAttention(nn.Module):
def __init__(
self,
vid_dim: int,
txt_dim: int,
heads: int,
head_dim: int,
qk_bias: bool,
qk_rope: bool,
qk_norm: norm_layer_type,
qk_norm_eps: float,
window: Union[int, Tuple[int, int, int]],
window_method: str,
shared_qkv: bool,
):
super().__init__()
dim = MMArg(vid_dim, txt_dim)
inner_dim = heads * head_dim
qkv_dim = inner_dim * 3
self.window = _triple(window)
self.window_method = window_method
assert all(map(lambda v: isinstance(v, int) and v >= 0, self.window))
self.head_dim = head_dim
self.proj_qkv = MMModule(nn.Linear, dim, qkv_dim, bias=qk_bias, shared_weights=shared_qkv)
self.proj_out = MMModule(nn.Linear, inner_dim, dim, shared_weights=shared_qkv)
self.norm_q = MMModule(qk_norm, dim=head_dim, eps=qk_norm_eps, elementwise_affine=True)
self.norm_k = MMModule(qk_norm, dim=head_dim, eps=qk_norm_eps, elementwise_affine=True)
self.rope = RotaryEmbedding3d(dim=head_dim // 2) if qk_rope else None
self.attn = TorchAttention()
def forward(
self,
vid: torch.FloatTensor, # b T H W c
txt: torch.FloatTensor, # b L c
txt_mask: torch.BoolTensor, # b L
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
]:
# Project q, k, v.
vid_qkv, txt_qkv = self.proj_qkv(vid, txt)
vid_qkv = gather_seq_scatter_heads_qkv(vid_qkv, seq_dim=2)
_, T, H, W, _ = vid_qkv.shape
_, L, _ = txt.shape
if self.window_method == "win":
nt, nh, nw = self.window
tt, hh, ww = T // nt, H // nh, W // nw
elif self.window_method == "win_by_size":
tt, hh, ww = self.window
tt, hh, ww = (
tt if tt > 0 else T,
hh if hh > 0 else H,
ww if ww > 0 else W,
)
nt, nh, nw = T // tt, H // hh, W // ww
else:
raise NotImplementedError
vid_qkv = rearrange(vid_qkv, "b T H W (o h d) -> o b h (T H W) d", o=3, d=self.head_dim)
txt_qkv = rearrange(txt_qkv, "b L (o h d) -> o b h L d", o=3, d=self.head_dim)
txt_qkv = scatter_heads(txt_qkv, dim=2)
vid_q, vid_k, vid_v = vid_qkv.unbind()
txt_q, txt_k, txt_v = txt_qkv.unbind()
vid_q, txt_q = self.norm_q(vid_q, txt_q)
vid_k, txt_k = self.norm_k(vid_k, txt_k)
if self.rope:
vid_q, vid_k = self.rope(vid_q, vid_k, (T, H, W))
def vid_window(v):
return rearrange(
v,
"b h (nt tt nh hh nw ww) d -> b h (nt nh nw) (tt hh ww) d",
hh=hh,
ww=ww,
tt=tt,
nh=nh,
nw=nw,
nt=nt,
)
def txt_window(t):
return rearrange(t, "b h L d -> b h 1 L d").expand(-1, -1, nt * nh * nw, -1, -1)
# Process video attention.
vid_msk = safe_pad_operation(txt_mask, (tt * hh * ww, 0), value=True)
vid_msk = rearrange(vid_msk, "b l -> b 1 1 1 l").expand(-1, 1, 1, tt * hh * ww, -1)
vid_out = self.attn(
vid_window(vid_q),
torch.cat([vid_window(vid_k), txt_window(txt_k)], dim=-2),
torch.cat([vid_window(vid_v), txt_window(txt_v)], dim=-2),
vid_msk,
)
vid_out = rearrange(
vid_out,
"b h (nt nh nw) (tt hh ww) d -> b (nt tt) (nh hh) (nw ww) (h d)",
hh=hh,
ww=ww,
tt=tt,
nh=nh,
nw=nw,
)
vid_out = gather_heads_scatter_seq(vid_out, head_dim=4, seq_dim=2)
# Process text attention.
txt_msk = safe_pad_operation(txt_mask, (T * H * W, 0), value=True)
txt_msk = rearrange(txt_msk, "b l -> b 1 1 l").expand(-1, 1, L, -1)
txt_out = self.attn(
txt_q,
torch.cat([vid_k, txt_k], dim=-2),
torch.cat([vid_v, txt_v], dim=-2),
txt_msk,
)
txt_out = rearrange(txt_out, "b h L d -> b L (h d)")
txt_out = gather_heads(txt_out, dim=2)
# Project output.
vid_out, txt_out = self.proj_out(vid_out, txt_out)
return vid_out, txt_out
class MMWindowTransformerBlock(nn.Module):
def __init__(
self,
*,
vid_dim: int,
txt_dim: int,
emb_dim: int,
heads: int,
head_dim: int,
expand_ratio: int,
norm: norm_layer_type,
norm_eps: float,
ada: ada_layer_type,
qk_bias: bool,
qk_rope: bool,
qk_norm: norm_layer_type,
window: Union[int, Tuple[int, int, int]],
window_method: str,
shared_qkv: bool,
shared_mlp: bool,
mlp_type: str,
**kwargs,
):
super().__init__()
dim = MMArg(vid_dim, txt_dim)
self.attn_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False)
self.attn = MMWindowAttention(
vid_dim=vid_dim,
txt_dim=txt_dim,
heads=heads,
head_dim=head_dim,
qk_bias=qk_bias,
qk_rope=qk_rope,
qk_norm=qk_norm,
qk_norm_eps=norm_eps,
window=window,
window_method=window_method,
shared_qkv=shared_qkv,
)
self.mlp_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False)
self.mlp = MMModule(
get_mlp(mlp_type),
dim=dim,
expand_ratio=expand_ratio,
shared_weights=shared_mlp,
)
self.ada = MMModule(ada, dim=dim, emb_dim=emb_dim, layers=["attn", "mlp"])
def forward(
self,
vid: torch.FloatTensor,
txt: torch.FloatTensor,
txt_mask: torch.BoolTensor,
emb: torch.FloatTensor,
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
]:
vid_attn, txt_attn = self.attn_norm(vid, txt)
vid_attn, txt_attn = self.ada(vid_attn, txt_attn, emb=emb, layer="attn", mode="in")
vid_attn, txt_attn = self.attn(vid_attn, txt_attn, txt_mask=txt_mask)
vid_attn, txt_attn = self.ada(vid_attn, txt_attn, emb=emb, layer="attn", mode="out")
vid_attn, txt_attn = (vid_attn + vid), (txt_attn + txt)
vid_mlp, txt_mlp = self.mlp_norm(vid_attn, txt_attn)
vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, emb=emb, layer="mlp", mode="in")
vid_mlp, txt_mlp = self.mlp(vid_mlp, txt_mlp)
vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, emb=emb, layer="mlp", mode="out")
vid_mlp, txt_mlp = (vid_mlp + vid_attn), (txt_mlp + txt_attn)
return vid_mlp, txt_mlp
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/embedding.py
+62
View File
@@ -0,0 +1,62 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Optional, Union
import torch
from diffusers.models.embeddings import get_timestep_embedding
from torch import nn
def emb_add(emb1: torch.Tensor, emb2: Optional[torch.Tensor]):
return emb1 if emb2 is None else emb1 + emb2
class TimeEmbedding(nn.Module):
def __init__(
self,
sinusoidal_dim: int,
hidden_dim: int,
output_dim: int,
):
super().__init__()
self.sinusoidal_dim = sinusoidal_dim
self.proj_in = nn.Linear(sinusoidal_dim, hidden_dim)
self.proj_hid = nn.Linear(hidden_dim, hidden_dim)
self.proj_out = nn.Linear(hidden_dim, output_dim)
self.act = nn.SiLU()
def forward(
self,
timestep: Union[int, float, torch.IntTensor, torch.FloatTensor],
device: torch.device,
dtype: torch.dtype,
) -> torch.FloatTensor:
if not torch.is_tensor(timestep):
timestep = torch.tensor([timestep], device=device, dtype=dtype)
if timestep.ndim == 0:
timestep = timestep[None]
emb = get_timestep_embedding(
timesteps=timestep,
embedding_dim=self.sinusoidal_dim,
flip_sin_to_cos=False,
downscale_freq_shift=0,
)
emb = emb.to(dtype)
emb = self.proj_in(emb)
emb = self.act(emb)
emb = self.proj_hid(emb)
emb = self.act(emb)
emb = self.proj_out(emb)
return emb
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/mlp.py
+62
View File
@@ -0,0 +1,62 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Optional
import torch
import torch.nn.functional as F
from torch import nn
def get_mlp(mlp_type: Optional[str] = "normal"):
if mlp_type == "normal":
return MLP
elif mlp_type == "swiglu":
return SwiGLUMLP
class MLP(nn.Module):
def __init__(
self,
dim: int,
expand_ratio: int,
):
super().__init__()
self.proj_in = nn.Linear(dim, dim * expand_ratio)
self.act = nn.GELU("tanh")
self.proj_out = nn.Linear(dim * expand_ratio, dim)
def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
x = self.proj_in(x)
x = self.act(x)
x = self.proj_out(x)
return x
class SwiGLUMLP(nn.Module):
def __init__(
self,
dim: int,
expand_ratio: int,
multiple_of: int = 256,
):
super().__init__()
hidden_dim = int(2 * dim * expand_ratio / 3)
hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
self.proj_in_gate = nn.Linear(dim, hidden_dim, bias=False)
self.proj_out = nn.Linear(hidden_dim, dim, bias=False)
self.proj_in = nn.Linear(dim, hidden_dim, bias=False)
def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
x = self.proj_out(F.silu(self.proj_in_gate(x)) * self.proj_in(x))
return x
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/mm.py
+67
View File
@@ -0,0 +1,67 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Tuple
import torch
from torch import nn
@dataclass
class MMArg:
vid: Any
txt: Any
def get_args(key: str, args: List[Any]) -> List[Any]:
return [getattr(v, key) if isinstance(v, MMArg) else v for v in args]
def get_kwargs(key: str, kwargs: Dict[str, Any]) -> Dict[str, Any]:
return {k: getattr(v, key) if isinstance(v, MMArg) else v for k, v in kwargs.items()}
class MMModule(nn.Module):
def __init__(
self,
module: Callable[..., nn.Module],
*args,
shared_weights: bool = False,
**kwargs,
):
super().__init__()
self.shared_weights = shared_weights
if self.shared_weights:
assert get_args("vid", args) == get_args("txt", args)
assert get_kwargs("vid", kwargs) == get_kwargs("txt", kwargs)
self.all = module(*get_args("vid", args), **get_kwargs("vid", kwargs))
else:
self.vid = module(*get_args("vid", args), **get_kwargs("vid", kwargs))
self.txt = module(*get_args("txt", args), **get_kwargs("txt", kwargs))
def forward(
self,
vid: torch.FloatTensor,
txt: torch.FloatTensor,
*args,
**kwargs,
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
]:
vid_module = self.vid if not self.shared_weights else self.all
txt_module = self.txt if not self.shared_weights else self.all
vid = vid_module(vid, *get_args("vid", args), **get_kwargs("vid", kwargs))
txt = txt_module(txt, *get_args("txt", args), **get_kwargs("txt", kwargs))
return vid, txt
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/modulation.py
+112
View File
@@ -0,0 +1,112 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Callable, List, Optional
import torch
from einops import rearrange
from torch import nn
from ...common.cache import Cache
from ...common.distributed.ops import slice_inputs
# (dim: int, emb_dim: int)
ada_layer_type = Callable[[int, int], nn.Module]
def get_ada_layer(ada_layer: str) -> ada_layer_type:
if ada_layer == "single":
return AdaSingle
raise NotImplementedError(f"{ada_layer} is not supported")
def expand_dims(x: torch.Tensor, dim: int, ndim: int):
"""
Expand tensor "x" to "ndim" by adding empty dims at "dim".
Example: x is (b d), target ndim is 5, add dim at 1, return (b 1 1 1 d).
"""
shape = x.shape
shape = shape[:dim] + (1,) * (ndim - len(shape)) + shape[dim:]
return x.reshape(shape)
class AdaSingle(nn.Module):
def __init__(
self,
dim: int,
emb_dim: int,
layers: List[str],
):
assert emb_dim == 6 * dim, "AdaSingle requires emb_dim == 6 * dim"
super().__init__()
self.dim = dim
self.emb_dim = emb_dim
self.layers = layers
for l in layers:
self.register_parameter(f"{l}_shift", nn.Parameter(torch.randn(dim) / dim**0.5))
self.register_parameter(f"{l}_scale", nn.Parameter(torch.randn(dim) / dim**0.5 + 1))
self.register_parameter(f"{l}_gate", nn.Parameter(torch.randn(dim) / dim**0.5))
def forward(
self,
hid: torch.FloatTensor, # b ... c
emb: torch.FloatTensor, # b d
layer: str,
mode: str,
cache: Cache = Cache(disable=True),
branch_tag: str = "",
hid_len: Optional[torch.LongTensor] = None, # b
) -> torch.FloatTensor:
idx = self.layers.index(layer)
emb = rearrange(emb, "b (d l g) -> b d l g", l=len(self.layers), g=3)[..., idx, :]
emb = expand_dims(emb, 1, hid.ndim + 1)
if hid_len is not None:
emb = cache(
f"emb_repeat_{idx}_{branch_tag}",
lambda: slice_inputs(
torch.repeat_interleave(emb, hid_len, dim=0),
dim=0,
),
)
shiftA, scaleA, gateA = emb.unbind(-1)
shiftB, scaleB, gateB = (
getattr(self, f"{layer}_shift"),
getattr(self, f"{layer}_scale"),
getattr(self, f"{layer}_gate"),
)
# Handle potential FP8 parameters - convert to input computation dtype
if hasattr(torch, 'float8_e4m3fn'):
fp8_types = (torch.float8_e4m3fn, torch.float8_e5m2)
# Use input tensor's dtype as target (respects pipeline precision)
target_dtype = hid.dtype
# Convert FP8 parameters to match input dtype for arithmetic operations
if shiftB is not None and shiftB.dtype in fp8_types:
shiftB = shiftB.to(target_dtype)
if scaleB is not None and scaleB.dtype in fp8_types:
scaleB = scaleB.to(target_dtype)
if gateB is not None and gateB.dtype in fp8_types:
gateB = gateB.to(target_dtype)
if mode == "in":
return hid.mul_(scaleA + scaleB).add_(shiftA + shiftB)
if mode == "out":
return hid.mul_(gateA + gateB)
raise NotImplementedError
def extra_repr(self) -> str:
return f"dim={self.dim}, emb_dim={self.emb_dim}, layers={self.layers}"
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/na.py
+641
View File
@@ -0,0 +1,641 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
"""
Native Resolution Transformer (NaDiT) tensor manipulation utilities.
TORCH.COMPILE OPTIMIZED VERSION
================================
This module has been optimized for torch.compile compatibility by eliminating
data-dependent operations that cause graph breaks:
Key Changes from Original:
- Replaced all .tolist() calls with pure tensor operations
- Minimized .item() calls (only used where required by einops API)
- Replaced list comprehensions with tensor-based splitting
- Added _tensor_split helper for compile-friendly splitting
- Proper device management to ensure tensors stay on correct devices
"""
from itertools import chain
from typing import Callable, Dict, List, Tuple
import einops
import torch
def _tensor_split(tensor: torch.Tensor, lengths: torch.LongTensor, dim: int = 0) -> List[torch.Tensor]:
"""
Optimized compile-friendly split using torch.tensor_split.
Uses PyTorch's native C++ implementation for fast eager mode execution
while remaining fully compatible with torch.compile symbolic tracing.
Args:
tensor: Input tensor to split
lengths: Tensor of split lengths (1D)
dim: Dimension along which to split (default: 0)
Returns:
List of split tensors
"""
if lengths.numel() == 0:
return []
if lengths.numel() == 1:
return [tensor]
# Calculate split indices: torch.tensor_split splits BEFORE each index
# So we need cumsum[:-1] to get the split points
# NOTE: torch.tensor_split requires indices on CPU (PyTorch requirement)
split_indices = lengths.cumsum(0)[:-1].cpu()
# Use torch.tensor_split - native C++ implementation
# Compilable: accepts tensor indices (symbolic shapes work)
# Fast: uses optimized CUDA/CPU kernels instead of Python loops
return list(torch.tensor_split(tensor, split_indices, dim=dim))
def flatten(
hid: List[torch.FloatTensor], # List of (*** c)
) -> Tuple[
torch.FloatTensor, # (L c)
torch.LongTensor, # (b n)
]:
"""
Flatten a list of tensors into a single tensor and track their shapes.
Converts a list of tensors with potentially different spatial shapes but
same feature dimension into a flattened tensor and shape metadata.
Args:
hid: List of tensors, each with shape (d1, d2, ..., dn, c)
Returns:
Tuple of:
- Flattened tensor of shape (L, c) where L = sum of all spatial dimensions
- Shape tensor of shape (b, n) tracking original spatial dimensions
COMPILE OPTIMIZATION: Uses tensor operations on correct device from input
"""
assert len(hid) > 0
device = hid[0].device
# Stack shape metadata - ensure tensors are created on correct device
shapes = []
for x in hid:
shape_tensor = torch.tensor(x.shape[:-1], dtype=torch.long, device=device)
shapes.append(shape_tensor)
shape = torch.stack(shapes)
# Flatten and concatenate
hid = torch.cat([x.flatten(0, -2) for x in hid])
return hid, shape
def unflatten(
hid: torch.FloatTensor, # (L c) or (L ... c)
hid_shape: torch.LongTensor, # (b n)
) -> List[torch.Tensor]: # List of (*** c) or (*** ... c)
"""
Unflatten a tensor back to a list using shape metadata.
Inverse operation of flatten(), reconstructing original tensor shapes.
Args:
hid: Flattened tensor of shape (L, c) or (L, ..., c)
hid_shape: Shape metadata tensor of shape (b, n)
Returns:
List of unflattened tensors with original spatial dimensions
COMPILE OPTIMIZATION:
- Uses optimized _tensor_split with torch.tensor_split (fast)
- .cpu().numpy() conversion needed for torch.compile compatibility
(unflatten() requires concrete Python ints, not symbolic shapes)
"""
hid_len = hid_shape.prod(-1)
# Use optimized tensor splitting (major performance improvement)
hid_list = _tensor_split(hid, hid_len, dim=0)
# Unflatten each piece
# NOTE: .cpu().numpy() is required for torch.compile compatibility
# .tolist() would fail with symbolic shapes during compilation
result = []
for i, x in enumerate(hid_list):
shape = hid_shape[i]
# Must use .cpu().numpy() for compilation compatibility
# Shape tensors are small, so CPU transfer overhead is minimal
target_shape = list(shape.cpu().numpy())
result.append(x.unflatten(0, target_shape))
return result
def concat(
vid: torch.FloatTensor, # (VL ... c)
txt: torch.FloatTensor, # (TL ... c)
vid_len: torch.LongTensor, # (b)
txt_len: torch.LongTensor, # (b)
) -> torch.FloatTensor: # (L ... c)
"""
Interleave video and text tensors batch-wise.
Splits video and text tensors by batch lengths, then interleaves them:
[vid_0, txt_0, vid_1, txt_1, ..., vid_b, txt_b]
Args:
vid: Video features tensor (VL, c)
txt: Text features tensor (TL, c)
vid_len: Length of each video sequence (b,)
txt_len: Length of each text sequence (b,)
Returns:
Interleaved tensor (L, c) where L = sum(vid_len) + sum(txt_len)
COMPILE OPTIMIZATION: Uses _tensor_split for compile-friendly splitting
"""
# Use tensor-based splitting
vid_splits = _tensor_split(vid, vid_len, dim=0)
txt_splits = _tensor_split(txt, txt_len, dim=0)
# Interleave
interleaved = []
for v, t in zip(vid_splits, txt_splits):
interleaved.extend([v, t])
return torch.cat(interleaved)
def concat_idx(
vid_len: torch.LongTensor, # (b)
txt_len: torch.LongTensor, # (b)
) -> Tuple[
Callable,
Callable,
]:
"""
Create index-based concatenation and un-concatenation functions.
Pre-computes indices for efficient interleaving and de-interleaving operations.
Returns callable functions that can be reused multiple times.
Args:
vid_len: Video sequence lengths (b,)
txt_len: Text sequence lengths (b,)
Returns:
Tuple of (concat_fn, unconcat_fn):
- concat_fn: Lambda that interleaves vid and txt tensors
- unconcat_fn: Lambda that separates interleaved tensor back to vid and txt
COMPILE OPTIMIZATION: Pre-computes all indices using tensor operations
"""
device = vid_len.device
vid_sum = vid_len.sum()
txt_sum = txt_len.sum()
vid_idx = torch.arange(vid_sum, device=device)
txt_idx = torch.arange(vid_sum, vid_sum + txt_sum, device=device)
# Build interleaving indices using compile-friendly _tensor_split
batch_size = len(vid_len)
vid_idx_splits = _tensor_split(vid_idx, vid_len, dim=0)
txt_idx_splits = _tensor_split(txt_idx, txt_len, dim=0)
# Create interleaved target indices
tgt_idx_list = []
for i in range(batch_size):
tgt_idx_list.append(vid_idx_splits[i])
tgt_idx_list.append(txt_idx_splits[i])
tgt_idx = torch.cat(tgt_idx_list)
src_idx = torch.argsort(tgt_idx)
vid_idx_len = len(vid_idx)
return (
lambda vid, txt: torch.index_select(torch.cat([vid, txt]), 0, tgt_idx),
lambda all: torch.index_select(all, 0, src_idx).split([vid_idx_len, len(txt_idx)]),
)
def unconcat(
all: torch.FloatTensor, # (L ... c)
vid_len: torch.LongTensor, # (b)
txt_len: torch.LongTensor, # (b)
) -> Tuple[
torch.FloatTensor, # (VL ... c)
torch.FloatTensor, # (TL ... c)
]:
"""
De-interleave concatenated video and text tensors.
Inverse of concat(). Separates an interleaved tensor back into video and text.
Args:
all: Interleaved tensor (L, c)
vid_len: Video sequence lengths (b,)
txt_len: Text sequence lengths (b,)
Returns:
Tuple of (vid, txt) tensors
COMPILE OPTIMIZATION: Uses tensor operations to build interleave pattern
"""
batch_size = len(vid_len)
# Create interleaved lengths: [vid_0, txt_0, vid_1, txt_1, ...]
interleave_len = torch.stack([vid_len, txt_len], dim=1).flatten()
# Split using compile-friendly operation
all_splits = _tensor_split(all, interleave_len, dim=0)
# Separate even (vid) and odd (txt) indices
vid_parts = [all_splits[i] for i in range(0, len(all_splits), 2)]
txt_parts = [all_splits[i] for i in range(1, len(all_splits), 2)]
vid = torch.cat(vid_parts)
txt = torch.cat(txt_parts)
return vid, txt
def repeat_concat(
vid: torch.FloatTensor, # (VL ... c)
txt: torch.FloatTensor, # (TL ... c)
vid_len: torch.LongTensor, # (n*b)
txt_len: torch.LongTensor, # (b)
txt_repeat: torch.LongTensor, # (n) or (b)
) -> torch.FloatTensor: # (L ... c)
"""
Concatenate video and text with text repetition for window attention.
For windowed attention, text features are repeated and interleaved with
multiple video windows: [vid_0, txt_0, vid_1, txt_0, vid_2, txt_0, ...]
Args:
vid: Video features (VL, c)
txt: Text features (TL, c)
vid_len: Video window lengths (n*b,) where n=num_windows
txt_len: Text sequence lengths (b,)
txt_repeat: Number of times to repeat text (n,) or (b,)
Returns:
Interleaved tensor with repeated text
COMPILE OPTIMIZATION: Uses _tensor_split and tensor-based repetition
"""
# Split using compile-friendly operations
vid_splits = _tensor_split(vid, vid_len, dim=0)
txt_splits = _tensor_split(txt, txt_len, dim=0)
# Handle txt_repeat shape flexibility
if txt_repeat.numel() == len(txt_splits):
repeat_counts = txt_repeat
else:
repeat_counts = txt_repeat.repeat(len(txt_splits))
# Interleave with repetition
result = []
for i, v in enumerate(vid_splits):
result.append(v)
# Get corresponding text sample (cyclic)
batch_idx = i % len(txt_splits) if len(txt_splits) > 0 else 0
if batch_idx < len(txt_splits):
result.append(txt_splits[batch_idx])
return torch.cat(result)
def repeat_concat_idx(
vid_len: torch.LongTensor, # (n*b)
txt_len: torch.LongTensor, # (b)
txt_repeat: torch.LongTensor, # (n) or scalar
) -> Tuple[
Callable,
Callable,
]:
"""
Create index-based repeat-concatenation and un-concatenation with coalescing.
Similar to concat_idx but handles text repetition for window attention.
The unconcat function coalesces (averages) repeated text features.
Args:
vid_len: Video window lengths (n*b,)
txt_len: Text sequence lengths (b,)
txt_repeat: Repetition count (scalar or tensor)
Returns:
Tuple of (concat_fn, unconcat_coalesce_fn):
- concat_fn: Interleaves with text repetition
- unconcat_coalesce_fn: Separates and averages repeated text
Example:
Input: vid=[0,1,2,3,4,5,6,7,8], txt=[9,10], repeat=3
Concat: [0,1,2,9,10, 3,4,5,9,10, 6,7,8,9,10]
Unconcat: vid=[0,1,2,3,4,5,6,7,8], txt=[9,10] (averaged)
COMPILE OPTIMIZATION:
- Eliminates .tolist() calls that caused graph breaks
- Uses pure tensor operations for index building
- Minimizes data-dependent branching
"""
device = vid_len.device
vid_sum = vid_len.sum()
txt_sum = txt_len.sum()
# Create base indices
vid_idx = torch.arange(vid_sum, device=device)
txt_idx = torch.arange(vid_sum, vid_sum + txt_sum, device=device)
# Normalize txt_repeat to tensor
if isinstance(txt_repeat, (int, float)):
txt_repeat = torch.tensor([txt_repeat], dtype=torch.long, device=device)
elif txt_repeat.dim() == 0:
txt_repeat = txt_repeat.unsqueeze(0)
# Calculate repeat pattern - keep as tensor to avoid graph breaks
batch_size = len(txt_len)
if txt_repeat.numel() == 1:
num_repeats_tensor = txt_repeat.reshape(-1) # Ensure 1D tensor
else:
# Use tensor operations for division
num_repeats_tensor = torch.tensor([len(vid_len) // batch_size], dtype=torch.long, device=device)
# Build concatenated indices using compile-friendly _tensor_split
vid_idx_splits = _tensor_split(vid_idx, vid_len, dim=0)
txt_idx_splits = _tensor_split(txt_idx, txt_len, dim=0)
tgt_idx_list = []
for i in range(len(vid_len)):
# Add video window
tgt_idx_list.append(vid_idx_splits[i])
# Add corresponding text (with repeat)
batch_idx = i % batch_size
tgt_idx_list.append(txt_idx_splits[batch_idx])
tgt_idx = torch.cat(tgt_idx_list)
src_idx = torch.argsort(tgt_idx)
txt_idx_len = len(tgt_idx) - len(vid_idx)
# Pre-compute split lengths for coalescing using tensor operations
repeat_txt_len = txt_len * num_repeats_tensor.squeeze()
def unconcat_coalesce(all):
"""
Un-concat vid & txt, and coalesce the repeated txt by averaging.
The text features appear multiple times (once per window) and need
to be averaged to produce a single set of text features.
COMPILE OPTIMIZATION: Uses unflatten with tensor dims (compile-friendly)
"""
vid_out, txt_out = all[src_idx].split([len(vid_idx), txt_idx_len])
# Coalesce repeated text using unflatten and mean
txt_splits = _tensor_split(txt_out, repeat_txt_len, dim=0)
txt_out_coalesced = []
for txt in txt_splits:
# txt has shape (base_len * num_repeats, *other_dims)
# unflatten to (base_len, num_repeats, *other_dims) then average dim 1
txt = txt.unflatten(0, (-1, num_repeats_tensor.squeeze())).mean(1)
txt_out_coalesced.append(txt)
return vid_out, torch.cat(txt_out_coalesced)
# Note: Using direct indexing instead of torch.index_select for backward compatibility
# Direct indexing is deterministic even with repeated indices
return (
lambda vid, txt: torch.cat([vid, txt])[tgt_idx],
lambda all: unconcat_coalesce(all),
)
def rearrange(
hid: torch.FloatTensor, # (L c)
hid_shape: torch.LongTensor, # (b n)
pattern: str,
**kwargs: Dict[str, int],
) -> Tuple[
torch.FloatTensor,
torch.LongTensor,
]:
"""
Rearrange flattened tensor using einops pattern.
Applies einops rearrange to each batch element independently.
Args:
hid: Flattened tensor (L, c)
hid_shape: Shape metadata (b, n)
pattern: Einops rearrange pattern
**kwargs: Additional arguments for einops
Returns:
Tuple of (rearranged tensor, new shape metadata)
"""
unflattened = unflatten(hid, hid_shape)
rearranged = [einops.rearrange(h, pattern, **kwargs) for h in unflattened]
return flatten(rearranged)
def rearrange_idx(
hid_shape: torch.LongTensor, # (b n)
pattern: str,
**kwargs: Dict[str, int],
) -> Tuple[Callable, Callable, torch.LongTensor]:
"""
Create index-based rearrange functions.
Pre-computes indices for efficient rearrangement operations.
Args:
hid_shape: Shape metadata (b, n)
pattern: Einops rearrange pattern
**kwargs: Additional arguments for einops
Returns:
Tuple of (rearrange_fn, reverse_fn, new_shape)
"""
hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1)
tgt_idx, tgt_shape = rearrange(hid_idx, hid_shape, pattern, **kwargs)
tgt_idx = tgt_idx.squeeze(-1)
src_idx = torch.argsort(tgt_idx)
return (
lambda hid: torch.index_select(hid, 0, tgt_idx),
lambda hid: torch.index_select(hid, 0, src_idx),
tgt_shape,
)
def repeat(
hid: torch.FloatTensor, # (L c)
hid_shape: torch.LongTensor, # (b n)
pattern: str,
**kwargs: Dict[str, torch.LongTensor], # (b)
) -> Tuple[
torch.FloatTensor,
torch.LongTensor,
]:
"""
Repeat flattened tensor using einops pattern with per-batch parameters.
Each batch element can have different repeat counts specified in kwargs.
Args:
hid: Flattened tensor (L, c)
hid_shape: Shape metadata (b, n)
pattern: Einops repeat pattern (e.g., "l c -> t l c")
**kwargs: Repeat parameters as tensors (e.g., t=torch.tensor([2,3,4]))
Returns:
Tuple of (repeated tensor, new shape metadata)
COMPILE OPTIMIZATION:
- Minimizes .item() calls
- Only converts to int at the last moment for einops API requirement
"""
unflattened = unflatten(hid, hid_shape)
# Build kwargs for each batch element
repeated = []
for i in range(len(unflattened)):
# Extract values for einops (requires Python int)
batch_kwargs = {}
for k, v in kwargs.items():
if isinstance(v, torch.Tensor):
# Only convert to Python int where required by einops
batch_kwargs[k] = int(v[i].item())
else:
batch_kwargs[k] = v
repeated.append(einops.repeat(unflattened[i], pattern, **batch_kwargs))
return flatten(repeated)
def pack(
samples: List[torch.Tensor], # List of (h w c).
) -> Tuple[
List[torch.Tensor], # groups [(b1 h1 w1 c1), (b2 h2 w2 c2)]
List[List[int]], # reversal indices.
]:
"""
Group samples by shape and return grouped batches with reversal indices.
Useful for batch processing samples with different spatial dimensions.
Args:
samples: List of tensors with potentially different shapes
Returns:
Tuple of (batched_groups, reversal_indices) for unpacking
"""
batches = {}
indices = {}
for i, sample in enumerate(samples):
shape = sample.shape
batches[shape] = batches.get(shape, [])
indices[shape] = indices.get(shape, [])
batches[shape].append(sample)
indices[shape].append(i)
batches = list(map(torch.stack, batches.values()))
indices = list(indices.values())
return batches, indices
def unpack(
batches: List[torch.Tensor],
indices: List[List[int]],
) -> List[torch.Tensor]:
"""
Unpack grouped batches back to original order.
Inverse of pack().
Args:
batches: Grouped batches from pack()
indices: Reversal indices from pack()
Returns:
List of tensors in original order
"""
samples = [None] * (max(chain(*indices)) + 1)
for batch, index in zip(batches, indices):
for sample, i in zip(batch.unbind(), index):
samples[i] = sample
return samples
def window(
hid: torch.FloatTensor, # (L c)
hid_shape: torch.LongTensor, # (b n)
window_fn: Callable[[torch.Tensor], List[torch.Tensor]],
):
"""
Apply windowing function to create non-overlapping windows.
Used for window attention mechanisms where sequences are split into windows.
Args:
hid: Flattened tensor (L, c)
hid_shape: Shape metadata (b, n)
window_fn: Function that splits a tensor into windows
Returns:
Tuple of (windowed_tensor, window_shapes, window_counts)
COMPILE OPTIMIZATION: Uses tensor operation for window count tracking
"""
unflattened = unflatten(hid, hid_shape)
windowed = [window_fn(h) for h in unflattened]
# Track window counts using tensor operations
device = hid_shape.device
hid_windows = torch.tensor([len(w) for w in windowed], dtype=torch.long, device=device)
# Flatten all windows
all_windows = list(chain(*windowed))
hid, hid_shape = flatten(all_windows)
return hid, hid_shape, hid_windows
def window_idx(
hid_shape: torch.LongTensor, # (b n)
window_fn: Callable[[torch.Tensor], List[torch.Tensor]],
):
"""
Create index-based windowing functions.
Pre-computes indices for efficient windowing and reverse operations.
Args:
hid_shape: Shape metadata (b, n)
window_fn: Function that splits a tensor into windows
Returns:
Tuple of (window_fn, reverse_fn, window_shapes, window_counts)
"""
hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1)
tgt_idx, tgt_shape, tgt_windows = window(hid_idx, hid_shape, window_fn)
tgt_idx = tgt_idx.squeeze(-1)
src_idx = torch.argsort(tgt_idx)
return (
lambda hid: torch.index_select(hid, 0, tgt_idx),
lambda hid: torch.index_select(hid, 0, src_idx),
tgt_shape,
tgt_windows,
)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nablocks/__init__.py
+25
View File
@@ -0,0 +1,25 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from .mmsr_block import NaMMSRTransformerBlock
nadit_blocks = {
"mmdit_sr": NaMMSRTransformerBlock,
}
def get_nablock(block_type: str):
if block_type in nadit_blocks:
return nadit_blocks[block_type]
raise NotImplementedError(f"{block_type} is not supported")
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nablocks/mmsr_block.py
+250
View File
@@ -0,0 +1,250 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Tuple, Union
import torch
from einops import rearrange
from torch.nn import functional as F
# from ..cache import Cache
from ....common.cache import Cache
from ....common.distributed.ops import gather_heads_scatter_seq, gather_seq_scatter_heads_qkv
from .. import na
from ..attention import FlashAttentionVarlen
from ..blocks.mmdit_window_block import MMWindowAttention, MMWindowTransformerBlock
from ..mm import MMArg
from ..modulation import ada_layer_type
from ..normalization import norm_layer_type
from ..rope import NaRotaryEmbedding3d
from ..window import get_window_op
from ....common.half_precision_fixes import safe_pad_operation
class NaSwinAttention(MMWindowAttention):
def __init__(
self,
vid_dim: int,
txt_dim: int,
heads: int,
head_dim: int,
qk_bias: bool,
qk_rope: bool,
qk_norm: norm_layer_type,
qk_norm_eps: float,
window: Union[int, Tuple[int, int, int]],
window_method: str,
shared_qkv: bool,
attention_mode: str = 'sdpa',
**kwargs,
):
super().__init__(
vid_dim=vid_dim,
txt_dim=txt_dim,
heads=heads,
head_dim=head_dim,
qk_bias=qk_bias,
qk_rope=qk_rope,
qk_norm=qk_norm,
qk_norm_eps=qk_norm_eps,
window=window,
window_method=window_method,
shared_qkv=shared_qkv,
)
self.rope = NaRotaryEmbedding3d(dim=head_dim // 2) if qk_rope else None
self.attn = FlashAttentionVarlen(attention_mode=attention_mode)
self.window_op = get_window_op(window_method)
def forward(
self,
vid: torch.FloatTensor, # l c
txt: torch.FloatTensor, # l c
vid_shape: torch.LongTensor, # b 3
txt_shape: torch.LongTensor, # b 1
cache: Cache,
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
]:
vid_qkv, txt_qkv = self.proj_qkv(vid, txt)
vid_qkv = gather_seq_scatter_heads_qkv(
vid_qkv,
seq_dim=0,
qkv_shape=vid_shape,
cache=cache.namespace("vid"),
)
txt_qkv = gather_seq_scatter_heads_qkv(
txt_qkv,
seq_dim=0,
qkv_shape=txt_shape,
cache=cache.namespace("txt"),
)
# re-org the input seq for window attn
cache_win = cache.namespace(f"{self.window_method}_{self.window}_sd3")
def make_window(x: torch.Tensor):
t, h, w, _ = x.shape
window_slices = self.window_op((t, h, w), self.window)
return [x[st, sh, sw] for (st, sh, sw) in window_slices]
window_partition, window_reverse, window_shape, window_count = cache_win(
"win_transform",
lambda: na.window_idx(vid_shape, make_window),
)
vid_qkv_win = window_partition(vid_qkv)
vid_qkv_win = rearrange(vid_qkv_win, "l (o h d) -> l o h d", o=3, d=self.head_dim)
txt_qkv = rearrange(txt_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim)
vid_q, vid_k, vid_v = vid_qkv_win.unbind(1)
txt_q, txt_k, txt_v = txt_qkv.unbind(1)
vid_q, txt_q = self.norm_q(vid_q, txt_q)
vid_k, txt_k = self.norm_k(vid_k, txt_k)
txt_len = cache("txt_len", lambda: txt_shape.prod(-1))
vid_len_win = cache_win("vid_len", lambda: window_shape.prod(-1))
txt_len_win = cache_win("txt_len", lambda: txt_len.repeat_interleave(window_count))
all_len_win = cache_win("all_len", lambda: vid_len_win + txt_len_win)
concat_win, unconcat_win = cache_win(
"mm_pnp", lambda: na.repeat_concat_idx(vid_len_win, txt_len, window_count)
)
# window rope
if self.rope:
vid_q, vid_k = self.rope(vid_q, vid_k, window_shape, cache_win)
# Attention handles dtype conversion internally using pipeline compute_dtype
out = self.attn(
q=concat_win(vid_q, txt_q),
k=concat_win(vid_k, txt_k),
v=concat_win(vid_v, txt_v),
cu_seqlens_q=cache_win(
"vid_seqlens_q", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int()
),
cu_seqlens_k=cache_win(
"vid_seqlens_k", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int()
),
max_seqlen_q=cache_win("vid_max_seqlen_q", lambda: all_len_win.max()),
max_seqlen_k=cache_win("vid_max_seqlen_k", lambda: all_len_win.max()),
).type_as(vid_q)
# text pooling
vid_out, txt_out = unconcat_win(out)
vid_out = rearrange(vid_out, "l h d -> l (h d)")
txt_out = rearrange(txt_out, "l h d -> l (h d)")
vid_out = window_reverse(vid_out)
vid_out = gather_heads_scatter_seq(vid_out, head_dim=1, seq_dim=0)
txt_out = gather_heads_scatter_seq(txt_out, head_dim=1, seq_dim=0)
vid_out, txt_out = self.proj_out(vid_out, txt_out)
return vid_out, txt_out
class NaMMSRTransformerBlock(MMWindowTransformerBlock):
def __init__(
self,
*,
vid_dim: int,
txt_dim: int,
emb_dim: int,
heads: int,
head_dim: int,
expand_ratio: int,
norm: norm_layer_type,
norm_eps: float,
ada: ada_layer_type,
qk_bias: bool,
qk_rope: bool,
qk_norm: norm_layer_type,
shared_qkv: bool,
shared_mlp: bool,
mlp_type: str,
**kwargs,
):
super().__init__(
vid_dim=vid_dim,
txt_dim=txt_dim,
emb_dim=emb_dim,
heads=heads,
head_dim=head_dim,
expand_ratio=expand_ratio,
norm=norm,
norm_eps=norm_eps,
ada=ada,
qk_bias=qk_bias,
qk_rope=qk_rope,
qk_norm=qk_norm,
shared_qkv=shared_qkv,
shared_mlp=shared_mlp,
mlp_type=mlp_type,
**kwargs,
)
self.attn = NaSwinAttention(
vid_dim=vid_dim,
txt_dim=txt_dim,
heads=heads,
head_dim=head_dim,
qk_bias=qk_bias,
qk_rope=qk_rope,
qk_norm=qk_norm,
qk_norm_eps=norm_eps,
shared_qkv=shared_qkv,
**kwargs,
)
def forward(
self,
vid: torch.FloatTensor, # l c
txt: torch.FloatTensor, # l c
vid_shape: torch.LongTensor, # b 3
txt_shape: torch.LongTensor, # b 1
emb: torch.FloatTensor,
cache: Cache,
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
torch.LongTensor,
torch.LongTensor,
]:
hid_len = MMArg(
cache("vid_len", lambda: vid_shape.prod(-1)),
cache("txt_len", lambda: txt_shape.prod(-1)),
)
ada_kwargs = {
"emb": emb,
"hid_len": hid_len,
"cache": cache,
"branch_tag": MMArg("vid", "txt"),
}
vid_attn, txt_attn = self.attn_norm(vid, txt)
vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="in", **ada_kwargs)
vid_attn, txt_attn = self.attn(vid_attn, txt_attn, vid_shape, txt_shape, cache)
vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="out", **ada_kwargs)
vid_attn, txt_attn = (vid_attn + vid), (txt_attn + txt)
vid_mlp, txt_mlp = self.mlp_norm(vid_attn, txt_attn)
vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="in", **ada_kwargs)
vid_mlp, txt_mlp = self.mlp(vid_mlp, txt_mlp)
vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="out", **ada_kwargs)
vid_mlp, txt_mlp = (vid_mlp + vid_attn), (txt_mlp + txt_attn)
return vid_mlp, txt_mlp, vid_shape, txt_shape
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nadit.py
+353
View File
@@ -0,0 +1,353 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from dataclasses import dataclass
from typing import Optional, Tuple, Union, Callable
import torch
from torch import nn
from ...common.cache import Cache
from ...common.distributed.ops import slice_inputs
from . import na
from .embedding import TimeEmbedding
from .modulation import get_ada_layer
from .nablocks import get_nablock
from .normalization import get_norm_layer
from .patch import NaPatchIn, NaPatchOut
# Fake func, no checkpointing is required for inference
def gradient_checkpointing(module: Union[Callable, nn.Module], *args, enabled: bool, **kwargs):
return module(*args, **kwargs)
@dataclass
class NaDiTOutput:
vid_sample: torch.Tensor
class NaDiT(nn.Module):
"""
Native Resolution Diffusion Transformer (NaDiT)
"""
gradient_checkpointing = False
def __init__(
self,
vid_in_channels: int,
vid_out_channels: int,
vid_dim: int,
txt_in_dim: Optional[int],
txt_dim: Optional[int],
emb_dim: int,
heads: int,
head_dim: int,
expand_ratio: int,
norm: Optional[str],
norm_eps: float,
ada: str,
qk_bias: bool,
qk_rope: bool,
qk_norm: Optional[str],
patch_size: Union[int, Tuple[int, int, int]],
num_layers: int,
block_type: Union[str, Tuple[str]],
shared_qkv: bool = False,
shared_mlp: bool = False,
mlp_type: str = "normal",
window: Optional[Tuple] = None,
window_method: Optional[Tuple[str]] = None,
temporal_window_size: int = None,
temporal_shifted: bool = False,
attention_mode: str = 'sdpa',
**kwargs,
):
ada = get_ada_layer(ada)
norm = get_norm_layer(norm)
qk_norm = get_norm_layer(qk_norm)
if isinstance(block_type, str):
block_type = [block_type] * num_layers
elif len(block_type) != num_layers:
raise ValueError("The ``block_type`` list should equal to ``num_layers``.")
super().__init__()
self.vid_in = NaPatchIn(
in_channels=vid_in_channels,
patch_size=patch_size,
dim=vid_dim,
)
self.txt_in = (
nn.Linear(txt_in_dim, txt_dim)
if txt_in_dim and txt_in_dim != txt_dim
else nn.Identity()
)
self.emb_in = TimeEmbedding(
sinusoidal_dim=256,
hidden_dim=max(vid_dim, txt_dim),
output_dim=emb_dim,
)
if window is None or isinstance(window[0], int):
window = [window] * num_layers
if window_method is None or isinstance(window_method, str):
window_method = [window_method] * num_layers
if temporal_window_size is None or isinstance(temporal_window_size, int):
temporal_window_size = [temporal_window_size] * num_layers
if temporal_shifted is None or isinstance(temporal_shifted, bool):
temporal_shifted = [temporal_shifted] * num_layers
self.blocks = nn.ModuleList(
[
get_nablock(block_type[i])(
vid_dim=vid_dim,
txt_dim=txt_dim,
emb_dim=emb_dim,
heads=heads,
head_dim=head_dim,
expand_ratio=expand_ratio,
norm=norm,
norm_eps=norm_eps,
ada=ada,
qk_bias=qk_bias,
qk_rope=qk_rope,
qk_norm=qk_norm,
shared_qkv=shared_qkv,
shared_mlp=shared_mlp,
mlp_type=mlp_type,
window=window[i],
window_method=window_method[i],
temporal_window_size=temporal_window_size[i],
temporal_shifted=temporal_shifted[i],
attention_mode=attention_mode,
**kwargs,
)
for i in range(num_layers)
]
)
self.vid_out = NaPatchOut(
out_channels=vid_out_channels,
patch_size=patch_size,
dim=vid_dim,
)
self.need_txt_repeat = block_type[0] in [
"mmdit_stwin",
"mmdit_stwin_spatial",
"mmdit_stwin_3d_spatial",
]
def set_gradient_checkpointing(self, enable: bool):
self.gradient_checkpointing = enable
def forward(
self,
vid: torch.FloatTensor, # l c
txt: torch.FloatTensor, # l c
vid_shape: torch.LongTensor, # b 3
txt_shape: torch.LongTensor, # b 1
timestep: Union[int, float, torch.IntTensor, torch.FloatTensor], # b
disable_cache: bool = True, # for test
):
# Text input.
if txt_shape.size(-1) == 1 and self.need_txt_repeat:
txt, txt_shape = na.repeat(txt, txt_shape, "l c -> t l c", t=vid_shape[:, 0])
# slice vid after patching in when using sequence parallelism
txt = slice_inputs(txt, dim=0)
txt = self.txt_in(txt)
# Video input.
# Sequence parallel slicing is done inside patching class.
vid, vid_shape = self.vid_in(vid, vid_shape)
# Embedding input.
emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype)
# Body
cache = Cache(disable=disable_cache)
for i, block in enumerate(self.blocks):
vid, txt, vid_shape, txt_shape = gradient_checkpointing(
enabled=(self.gradient_checkpointing and self.training),
module=block,
vid=vid,
txt=txt,
vid_shape=vid_shape,
txt_shape=txt_shape,
emb=emb,
cache=cache,
)
vid, vid_shape = self.vid_out(vid, vid_shape, cache)
return NaDiTOutput(vid_sample=vid)
class NaDiTUpscaler(nn.Module):
"""
Native Resolution Diffusion Transformer (NaDiT)
"""
gradient_checkpointing = False
def __init__(
self,
vid_in_channels: int,
vid_out_channels: int,
vid_dim: int,
txt_in_dim: Optional[int],
txt_dim: Optional[int],
emb_dim: int,
heads: int,
head_dim: int,
expand_ratio: int,
norm: Optional[str],
norm_eps: float,
ada: str,
qk_bias: bool,
qk_rope: bool,
qk_norm: Optional[str],
patch_size: Union[int, Tuple[int, int, int]],
num_layers: int,
block_type: Union[str, Tuple[str]],
shared_qkv: bool = False,
shared_mlp: bool = False,
mlp_type: str = "normal",
window: Optional[Tuple] = None,
window_method: Optional[Tuple[str]] = None,
temporal_window_size: int = None,
temporal_shifted: bool = False,
**kwargs,
):
ada = get_ada_layer(ada)
norm = get_norm_layer(norm)
qk_norm = get_norm_layer(qk_norm)
if isinstance(block_type, str):
block_type = [block_type] * num_layers
elif len(block_type) != num_layers:
raise ValueError("The ``block_type`` list should equal to ``num_layers``.")
super().__init__()
self.vid_in = NaPatchIn(
in_channels=vid_in_channels,
patch_size=patch_size,
dim=vid_dim,
)
self.txt_in = (
nn.Linear(txt_in_dim, txt_dim)
if txt_in_dim and txt_in_dim != txt_dim
else nn.Identity()
)
self.emb_in = TimeEmbedding(
sinusoidal_dim=256,
hidden_dim=max(vid_dim, txt_dim),
output_dim=emb_dim,
)
self.emb_scale = TimeEmbedding(
sinusoidal_dim=256,
hidden_dim=max(vid_dim, txt_dim),
output_dim=emb_dim,
)
if window is None or isinstance(window[0], int):
window = [window] * num_layers
if window_method is None or isinstance(window_method, str):
window_method = [window_method] * num_layers
if temporal_window_size is None or isinstance(temporal_window_size, int):
temporal_window_size = [temporal_window_size] * num_layers
if temporal_shifted is None or isinstance(temporal_shifted, bool):
temporal_shifted = [temporal_shifted] * num_layers
self.blocks = nn.ModuleList(
[
get_nablock(block_type[i])(
vid_dim=vid_dim,
txt_dim=txt_dim,
emb_dim=emb_dim,
heads=heads,
head_dim=head_dim,
expand_ratio=expand_ratio,
norm=norm,
norm_eps=norm_eps,
ada=ada,
qk_bias=qk_bias,
qk_rope=qk_rope,
qk_norm=qk_norm,
shared_qkv=shared_qkv,
shared_mlp=shared_mlp,
mlp_type=mlp_type,
window=window[i],
window_method=window_method[i],
temporal_window_size=temporal_window_size[i],
temporal_shifted=temporal_shifted[i],
attention_mode=attention_mode,
**kwargs,
)
for i in range(num_layers)
]
)
self.vid_out = NaPatchOut(
out_channels=vid_out_channels,
patch_size=patch_size,
dim=vid_dim,
)
self.need_txt_repeat = block_type[0] in [
"mmdit_stwin",
"mmdit_stwin_spatial",
"mmdit_stwin_3d_spatial",
]
def set_gradient_checkpointing(self, enable: bool):
self.gradient_checkpointing = enable
def forward(
self,
vid: torch.FloatTensor, # l c
txt: torch.FloatTensor, # l c
vid_shape: torch.LongTensor, # b 3
txt_shape: torch.LongTensor, # b 1
timestep: Union[int, float, torch.IntTensor, torch.FloatTensor], # b
downscale: Union[int, float, torch.IntTensor, torch.FloatTensor], # b
disable_cache: bool = False, # for test
):
# Text input.
if txt_shape.size(-1) == 1 and self.need_txt_repeat:
txt, txt_shape = na.repeat(txt, txt_shape, "l c -> t l c", t=vid_shape[:, 0])
# slice vid after patching in when using sequence parallelism
txt = slice_inputs(txt, dim=0)
txt = self.txt_in(txt)
# Video input.
# Sequence parallel slicing is done inside patching class.
vid, vid_shape = self.vid_in(vid, vid_shape)
# Embedding input.
emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype)
emb_scale = self.emb_scale(downscale, device=vid.device, dtype=vid.dtype)
emb = emb + emb_scale
# Body
cache = Cache(disable=disable_cache)
for i, block in enumerate(self.blocks):
vid, txt, vid_shape, txt_shape = gradient_checkpointing(
enabled=(self.gradient_checkpointing and self.training),
module=block,
vid=vid,
txt=txt,
vid_shape=vid_shape,
txt_shape=txt_shape,
emb=emb,
cache=cache,
)
vid, vid_shape = self.vid_out(vid, vid_shape, cache)
return NaDiTOutput(vid_sample=vid)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/normalization.py
+139
View File
@@ -0,0 +1,139 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Callable, Optional
from diffusers.models.normalization import RMSNorm
from torch import nn
import torch
import torch.nn.functional as F
import numbers
from torch.nn.parameter import Parameter
from torch.nn import init
# (dim: int, eps: float, elementwise_affine: bool)
norm_layer_type = Callable[[int, float, bool], nn.Module]
class CustomLayerNorm(nn.Module):
"""
Custom LayerNorm implementation to replace Apex FusedLayerNorm
"""
def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
super(CustomLayerNorm, self).__init__()
if isinstance(normalized_shape, numbers.Integral):
normalized_shape = (normalized_shape,)
self.normalized_shape = torch.Size(normalized_shape)
self.eps = eps
self.elementwise_affine = elementwise_affine
if self.elementwise_affine:
self.weight = Parameter(torch.Tensor(*normalized_shape))
self.bias = Parameter(torch.Tensor(*normalized_shape))
else:
self.register_parameter('weight', None)
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
if self.elementwise_affine:
init.ones_(self.weight)
init.zeros_(self.bias)
def forward(self, input):
return F.layer_norm(
input, self.normalized_shape, self.weight, self.bias, self.eps)
class CustomRMSNorm(nn.Module):
"""
Custom RMSNorm implementation to replace Apex FusedRMSNorm
"""
def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
super(CustomRMSNorm, self).__init__()
if isinstance(normalized_shape, numbers.Integral):
normalized_shape = (normalized_shape,)
self.normalized_shape = torch.Size(normalized_shape)
self.eps = eps
self.elementwise_affine = elementwise_affine
if self.elementwise_affine:
self.weight = Parameter(torch.ones(*normalized_shape))
else:
self.register_parameter('weight', None)
def forward(self, input):
# RMS normalization: x / sqrt(mean(x^2) + eps) * weight
dims = tuple(range(-len(self.normalized_shape), 0))
# Calculate RMS: sqrt(mean(x^2))
variance = input.pow(2).mean(dim=dims, keepdim=True)
rms = torch.sqrt(variance + self.eps)
# Normalize
normalized = input / rms
if self.elementwise_affine:
# Convert FP8 weight to match input dtype for arithmetic operations
if hasattr(torch, 'float8_e4m3fn'):
fp8_types = (torch.float8_e4m3fn, torch.float8_e5m2)
if self.weight.dtype in fp8_types:
# Use input dtype as target (respects pipeline precision)
weight = self.weight.to(input.dtype)
return normalized * weight
return normalized * self.weight
return normalized
def get_norm_layer(norm_type: Optional[str]) -> norm_layer_type:
def _norm_layer(dim: int, eps: float, elementwise_affine: bool):
if norm_type is None:
return nn.Identity()
if norm_type == "layer":
return nn.LayerNorm(
normalized_shape=dim,
eps=eps,
elementwise_affine=elementwise_affine,
)
if norm_type == "rms":
return RMSNorm(
dim=dim,
eps=eps,
elementwise_affine=elementwise_affine,
)
if norm_type == "fusedln":
# Use custom LayerNorm instead of Apex FusedLayerNorm
return CustomLayerNorm(
normalized_shape=dim,
elementwise_affine=elementwise_affine,
eps=eps,
)
if norm_type == "fusedrms":
# Use custom RMSNorm instead of Apex FusedRMSNorm
return CustomRMSNorm(
normalized_shape=dim,
elementwise_affine=elementwise_affine,
eps=eps,
)
raise NotImplementedError(f"{norm_type} is not supported")
return _norm_layer
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/patch.py
+112
View File
@@ -0,0 +1,112 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from typing import Tuple, Union
import torch
from einops import rearrange
from torch import nn
from torch.nn.modules.utils import _triple
from ...common.cache import Cache
from ...common.distributed.ops import gather_outputs, slice_inputs
from . import na
class PatchIn(nn.Module):
def __init__(
self,
in_channels: int,
patch_size: Union[int, Tuple[int, int, int]],
dim: int,
):
super().__init__()
t, h, w = _triple(patch_size)
self.patch_size = t, h, w
self.proj = nn.Linear(in_channels * t * h * w, dim)
def forward(
self,
vid: torch.Tensor,
) -> torch.Tensor:
t, h, w = self.patch_size
vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w)
vid = self.proj(vid)
return vid
class PatchOut(nn.Module):
def __init__(
self,
out_channels: int,
patch_size: Union[int, Tuple[int, int, int]],
dim: int,
):
super().__init__()
t, h, w = _triple(patch_size)
self.patch_size = t, h, w
self.proj = nn.Linear(dim, out_channels * t * h * w)
def forward(
self,
vid: torch.Tensor,
) -> torch.Tensor:
t, h, w = self.patch_size
vid = self.proj(vid)
vid = rearrange(vid, "b T H W (t h w c) -> b c (T t) (H h) (W w)", t=t, h=h, w=w)
return vid
class NaPatchIn(PatchIn):
def forward(
self,
vid: torch.Tensor, # l c
vid_shape: torch.LongTensor,
) -> torch.Tensor:
t, h, w = self.patch_size
if not (t == h == w == 1):
vid, vid_shape = na.rearrange(
vid, vid_shape, "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w
)
# slice vid after patching in when using sequence parallelism
vid = slice_inputs(vid, dim=0)
vid = self.proj(vid)
return vid, vid_shape
class NaPatchOut(PatchOut):
def forward(
self,
vid: torch.FloatTensor, # l c
vid_shape: torch.LongTensor,
cache: Cache = Cache(disable=True),
) -> Tuple[
torch.FloatTensor,
torch.LongTensor,
]:
t, h, w = self.patch_size
vid = self.proj(vid)
# gather vid before patching out when enabling sequence parallelism
vid = gather_outputs(
vid,
gather_dim=0,
padding_dim=0,
unpad_shape=vid_shape,
cache=cache.namespace("vid"),
)
if not (t == h == w == 1):
vid, vid_shape = na.rearrange(
vid, vid_shape, "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w
)
return vid, vid_shape
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/rope.py
+111
View File
@@ -0,0 +1,111 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from functools import lru_cache
from typing import Tuple
import torch
from einops import rearrange
from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb
from torch import nn
from ...common.cache import Cache
class RotaryEmbeddingBase(nn.Module):
def __init__(self, dim: int, rope_dim: int):
super().__init__()
self.rope = RotaryEmbedding(
dim=dim // rope_dim,
freqs_for="pixel",
max_freq=256,
)
# 1. Set model.requires_grad_(True) after model creation will make
# the `requires_grad=False` for rope freqs no longer hold.
# 2. Even if we don't set requires_grad_(True) explicitly,
# FSDP is not memory efficient when handling fsdp_wrap
# with mixed requires_grad=True/False.
# With above consideration, it is easier just remove the freqs
# out of nn.Parameters when `learned_freq=False`
freqs = self.rope.freqs
del self.rope.freqs
self.rope.register_buffer("freqs", freqs.data)
@lru_cache(maxsize=128)
def get_axial_freqs(self, *dims):
return self.rope.get_axial_freqs(*dims)
class RotaryEmbedding3d(RotaryEmbeddingBase):
def __init__(self, dim: int):
super().__init__(dim, rope_dim=3)
def forward(
self,
q: torch.FloatTensor, # b h l d
k: torch.FloatTensor, # b h l d
size: Tuple[int, int, int],
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
]:
T, H, W = size
freqs = self.get_axial_freqs(T, H, W)
q = rearrange(q, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W)
k = rearrange(k, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W)
q = apply_rotary_emb(freqs, q)
k = apply_rotary_emb(freqs, k)
q = rearrange(q, "b h T H W d -> b h (T H W) d")
k = rearrange(k, "b h T H W d -> b h (T H W) d")
return q, k
class NaRotaryEmbedding3d(RotaryEmbedding3d):
def forward(
self,
q: torch.FloatTensor, # L h d
k: torch.FloatTensor, # L h d
shape: torch.LongTensor,
cache: Cache,
) -> Tuple[
torch.FloatTensor,
torch.FloatTensor,
]:
freqs = cache("rope_freqs_3d", lambda: self.get_freqs(shape))
freqs = freqs.to(device=q.device, dtype=q.dtype)
q = rearrange(q, "L h d -> h L d")
k = rearrange(k, "L h d -> h L d")
q = apply_rotary_emb(freqs, q.float()).to(q.dtype)
k = apply_rotary_emb(freqs, k.float()).to(k.dtype)
q = rearrange(q, "h L d -> L h d")
k = rearrange(k, "h L d -> L h d")
return q, k
@torch._dynamo.disable # Disable compilation: shape.tolist() is data-dependent and causes graph breaks
def get_freqs(
self,
shape: torch.LongTensor,
) -> torch.Tensor:
"""
Generate RoPE frequencies for video and text with adaptive dimensions.
Note: This method uses @torch._dynamo.disable because it requires
data-dependent control flow (.tolist()) that cannot be symbolically
traced by torch.compile. The cache() wrapper in the forward pass memoizes
results to reduce recomputation overhead.
"""
freq_list = []
for f, h, w in shape.tolist():
freqs = self.get_axial_freqs(f, h, w)
freq_list.append(freqs.view(-1, freqs.size(-1)))
return torch.cat(freq_list, dim=0)
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/window.py
+83
View File
@@ -0,0 +1,83 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
from math import ceil
from typing import Tuple
import math
def get_window_op(name: str):
if name == "720pwin_by_size_bysize":
return make_720Pwindows_bysize
if name == "720pswin_by_size_bysize":
return make_shifted_720Pwindows_bysize
raise ValueError(f"Unknown windowing method: {name}")
# -------------------------------- Windowing -------------------------------- #
def make_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]):
t, h, w = size
resized_nt, resized_nh, resized_nw = num_windows
#cal windows under 720p
scale = math.sqrt((45 * 80) / (h * w))
resized_h, resized_w = round(h * scale), round(w * scale)
wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw) # window size.
wt = ceil(min(t, 30) / resized_nt) # window size.
nt, nh, nw = ceil(t / wt), ceil(h / wh), ceil(w / ww) # window size.
return [
(
slice(it * wt, min((it + 1) * wt, t)),
slice(ih * wh, min((ih + 1) * wh, h)),
slice(iw * ww, min((iw + 1) * ww, w)),
)
for iw in range(nw)
if min((iw + 1) * ww, w) > iw * ww
for ih in range(nh)
if min((ih + 1) * wh, h) > ih * wh
for it in range(nt)
if min((it + 1) * wt, t) > it * wt
]
def make_shifted_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]):
t, h, w = size
resized_nt, resized_nh, resized_nw = num_windows
#cal windows under 720p
scale = math.sqrt((45 * 80) / (h * w))
resized_h, resized_w = round(h * scale), round(w * scale)
wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw) # window size.
wt = ceil(min(t, 30) / resized_nt) # window size.
st, sh, sw = ( # shift size.
0.5 if wt < t else 0,
0.5 if wh < h else 0,
0.5 if ww < w else 0,
)
nt, nh, nw = ceil((t - st) / wt), ceil((h - sh) / wh), ceil((w - sw) / ww) # window size.
nt, nh, nw = ( # number of window.
nt + 1 if st > 0 else 1,
nh + 1 if sh > 0 else 1,
nw + 1 if sw > 0 else 1,
)
return [
(
slice(max(int((it - st) * wt), 0), min(int((it - st + 1) * wt), t)),
slice(max(int((ih - sh) * wh), 0), min(int((ih - sh + 1) * wh), h)),
slice(max(int((iw - sw) * ww), 0), min(int((iw - sw + 1) * ww), w)),
)
for iw in range(nw)
if min(int((iw - sw + 1) * ww), w) > max(int((iw - sw) * ww), 0)
for ih in range(nh)
if min(int((ih - sh + 1) * wh), h) > max(int((ih - sh) * wh), 0)
for it in range(nt)
if min(int((it - st + 1) * wt), t) > max(int((it - st) * wt), 0)
]
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/attn_video_vae.py
+1733
View File
File diff suppressed because it is too large Load Diff
bundled_packs/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/causal_inflation_lib.py
+503
View File
@@ -0,0 +1,503 @@
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // 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
# //
# // http://www.apache.org/licenses/LICENSE-2.0
# //
# // 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.
import math
from contextlib import contextmanager
from typing import List, Optional, Union
import torch
import torch.nn.functional as F
from diffusers.models.normalization import RMSNorm
from einops import rearrange
from torch import Tensor, nn
from torch.nn import Conv3d
from .context_parallel_lib import cache_send_recv, get_cache_size
from .global_config import get_norm_limit
from .types import MemoryState, _inflation_mode_t, _memory_device_t
from ....common.half_precision_fixes import safe_pad_operation
from ....optimization.memory_manager import retry_on_oom
from ....optimization.compatibility import NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND
# Single GPU inference - no distributed processing needed
#print("Warning: Using single GPU inference mode - distributed features disabled in causal_inflation_lib")
# Mock distributed functions for single GPU inference
def get_sequence_parallel_group():
return None
def get_sequence_parallel_rank():
return 0
def get_sequence_parallel_world_size():
return 1
def get_next_sequence_parallel_rank():
return 0
def get_prev_sequence_parallel_rank():
return 0
@contextmanager
def ignore_padding(model):
orig_padding = model.padding
model.padding = (0, 0, 0)
try:
yield
finally:
model.padding = orig_padding
class InflatedCausalConv3d(Conv3d):
def __init__(
self,
*args,
inflation_mode: _inflation_mode_t,
memory_device: _memory_device_t = "same",
**kwargs,
):
self.inflation_mode = inflation_mode
self.memory = None
super().__init__(*args, **kwargs)
self.temporal_padding = self.padding[0]
self.memory_device = memory_device
self.padding = (0, *self.padding[1:]) # Remove temporal pad to keep causal.
self.memory_limit = float("inf")
def set_memory_limit(self, value: float):
self.memory_limit = value
def set_memory_device(self, memory_device: _memory_device_t):
self.memory_device = memory_device
def _conv_forward(self, input, weight, bias, *args, **kwargs):
"""
Override _conv_forward to work around NVIDIA Conv3d memory bug.
Bug: PyTorch 2.9-2.10 with cuDNN >= 91002 uses 3x memory for Conv3d
with fp16/bfloat16 weights due to buggy dispatch layer.
Workaround: Call torch.cudnn_convolution directly to bypass buggy layer.
Status is logged at startup in compatibility.py.
"""
if (NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND and
weight.dtype in (torch.float16, torch.bfloat16) and
hasattr(torch.backends.cudnn, 'is_available') and
torch.backends.cudnn.is_available() and
getattr(torch.backends.cudnn, 'enabled', True)):
try:
# Direct cuDNN call bypasses buggy PyTorch dispatch layer (NVIDIA only)
out = torch.cudnn_convolution(
input, weight, self.padding, self.stride, self.dilation, self.groups,
benchmark=False, deterministic=False, allow_tf32=True
)
if bias is not None:
out += bias.reshape((1, -1) + (1,) * (out.ndim - 2))
return out
except RuntimeError:
# Fallback if direct cuDNN call fails (dev builds, edge cases)
pass
# Use standard path for unaffected configurations or if workaround failed
return super()._conv_forward(input, weight, bias, *args, **kwargs)
def memory_limit_conv(
self,
x,
*,
split_dim=3,
padding=(0, 0, 0, 0, 0, 0),
prev_cache=None,
):
# Compatible with no limit.
if math.isinf(self.memory_limit):
if prev_cache is not None:
x = torch.cat([prev_cache, x], dim=split_dim - 1)
return super().forward(x)
# Compute tensor shape after concat & padding.
shape = torch.tensor(x.size())
if prev_cache is not None:
shape[split_dim - 1] += prev_cache.size(split_dim - 1)
shape[-3:] += torch.tensor(padding).view(3, 2).sum(-1).flip(0)
memory_occupy = shape.prod() * x.element_size() / 1024**3 # GiB
if memory_occupy < self.memory_limit or split_dim == x.ndim:
x_concat = x
if prev_cache is not None:
x_concat = torch.cat([prev_cache, x], dim=split_dim - 1)
def pad_and_forward():
padded = safe_pad_operation(x_concat, padding, mode='constant', value=0.0)
with ignore_padding(self):
return Conv3d.forward(self, padded)
return retry_on_oom(
pad_and_forward,
debug=getattr(self, 'debug', None),
operation_name="InflatedCausalConv3d.pad_and_forward"
)
# Exceed memory limit, splitting tensor
# Split input (& prev_cache).
num_splits = math.ceil(memory_occupy / self.memory_limit)
size_per_split = x.size(split_dim) // num_splits
split_sizes = [size_per_split] * (num_splits - 1)
split_sizes += [x.size(split_dim) - sum(split_sizes)]
x = list(x.split(split_sizes, dim=split_dim))
if prev_cache is not None:
prev_cache = list(prev_cache.split(split_sizes, dim=split_dim))
# Loop Fwd.
cache = None
for idx in range(len(x)):
# Concat prev cache from last dim
if prev_cache is not None:
x[idx] = torch.cat([prev_cache[idx], x[idx]], dim=split_dim - 1)
# Get padding pattern.
lpad_dim = (x[idx].ndim - split_dim - 1) * 2
rpad_dim = lpad_dim + 1
padding = list(padding)
padding[lpad_dim] = self.padding[split_dim - 2] if idx == 0 else 0
padding[rpad_dim] = self.padding[split_dim - 2] if idx == len(x) - 1 else 0
pad_len = padding[lpad_dim] + padding[rpad_dim]
padding = tuple(padding)
# Prepare cache for next slice (this dim).
next_cache = None
cache_len = cache.size(split_dim) if cache is not None else 0
next_catch_size = get_cache_size(
conv_module=self,
input_len=x[idx].size(split_dim) + cache_len,
pad_len=pad_len,
dim=split_dim - 2,
)
if next_catch_size != 0:
assert next_catch_size <= x[idx].size(split_dim)
next_cache = (
x[idx].transpose(0, split_dim)[-next_catch_size:].transpose(0, split_dim)
)
# Recursive.
x[idx] = self.memory_limit_conv(
x[idx],
split_dim=split_dim + 1,
padding=padding,
prev_cache=cache
)
# Update cache.
cache = next_cache
output = retry_on_oom(
torch.cat,
x,
split_dim,
debug=getattr(self, 'debug', None),
operation_name="InflatedCausalConv3d.concat_splits"
)
return output
def forward(
self,
input: Union[Tensor, List[Tensor]],
memory_state: MemoryState = MemoryState.UNSET
) -> Tensor:
assert memory_state != MemoryState.UNSET
if memory_state != MemoryState.ACTIVE:
self.memory = None
if (
math.isinf(self.memory_limit)
and torch.is_tensor(input)
and get_sequence_parallel_group() is None
):
return self.basic_forward(input, memory_state)
return self.slicing_forward(input, memory_state)
def basic_forward(self, input: Tensor, memory_state: MemoryState = MemoryState.UNSET):
mem_size = self.stride[0] - self.kernel_size[0]
if (self.memory is not None) and (memory_state == MemoryState.ACTIVE):
input = extend_head(input, memory=self.memory, times=-1)
else:
input = extend_head(input, times=self.temporal_padding * 2)
memory = (
input[:, :, mem_size:].detach()
if (mem_size != 0 and memory_state != MemoryState.DISABLED)
else None
)
if (
memory_state != MemoryState.DISABLED
and not self.training
and (self.memory_device is not None)
):
self.memory = memory
if self.memory_device == "cpu" and self.memory is not None:
self.memory = self.memory.to("cpu")
return super().forward(input)
def slicing_forward(
self,
input: Union[Tensor, List[Tensor]],
memory_state: MemoryState = MemoryState.UNSET,
) -> Tensor:
squeeze_out = False
if torch.is_tensor(input):
input = [input]
squeeze_out = True
cache_size = self.kernel_size[0] - self.stride[0]
cache = cache_send_recv(
input, cache_size=cache_size, memory=self.memory, times=self.temporal_padding * 2
)
# Single GPU inference - simplified memory management
if (
memory_state in [MemoryState.INITIALIZING, MemoryState.ACTIVE] # use_slicing
and not self.training
and (self.memory_device is not None)
and cache_size != 0
):
if cache_size > input[-1].size(2) and cache is not None and len(input) == 1:
input[0] = torch.cat([cache, input[0]], dim=2)
cache = None
if cache_size <= input[-1].size(2):
self.memory = input[-1][:, :, -cache_size:].detach().contiguous()
if self.memory_device == "cpu" and self.memory is not None:
self.memory = self.memory.to("cpu")
padding = tuple(x for x in reversed(self.padding) for _ in range(2))
for i in range(len(input)):
# Prepare cache for next input slice.
next_cache = None
cache_size = 0
if i < len(input) - 1:
cache_len = cache.size(2) if cache is not None else 0
cache_size = get_cache_size(self, input[i].size(2) + cache_len, pad_len=0)
if cache_size != 0:
if cache_size > input[i].size(2) and cache is not None:
input[i] = torch.cat([cache, input[i]], dim=2)
cache = None
assert cache_size <= input[i].size(2), f"{cache_size} > {input[i].size(2)}"
next_cache = input[i][:, :, -cache_size:]
# Conv forward for this input slice.
input[i] = self.memory_limit_conv(
input[i],
padding=padding,
prev_cache=cache
)
# Update cache.
cache = next_cache
return input[0] if squeeze_out else input
def tflops(self, args, kwargs, output) -> float:
if torch.is_tensor(output):
output_numel = output.numel()
elif isinstance(output, list):
output_numel = sum(o.numel() for o in output)
else:
raise NotImplementedError
return (2 * math.prod(self.kernel_size) * self.in_channels * (output_numel / 1e6)) / 1e6
def _load_from_state_dict(
self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
):
if self.inflation_mode != "none":
state_dict = modify_state_dict(
self,
state_dict,
prefix,
inflate_weight_fn=inflate_weight,
inflate_bias_fn=inflate_bias,
)
super()._load_from_state_dict(
state_dict,
prefix,
local_metadata,
(strict and self.inflation_mode == "none"),
missing_keys,
unexpected_keys,
error_msgs,
)
def init_causal_conv3d(
*args,
inflation_mode: _inflation_mode_t,
**kwargs,
):
"""
Initialize a Causal-3D convolution layer.
Parameters:
inflation_mode: Listed as below. It's compatible with all the 3D-VAE checkpoints we have.
- none: No inflation will be conducted.
The loading logic of state dict will fall back to default.
- tail / replicate: Refer to the definition of `InflatedCausalConv3d`.
"""
return InflatedCausalConv3d(*args, inflation_mode=inflation_mode, **kwargs)
def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor:
if isinstance(norm_layer, (nn.LayerNorm, RMSNorm)):
if x.ndim == 4:
x = rearrange(x, "b c h w -> b h w c")
x = norm_layer(x)
x = rearrange(x, "b h w c -> b c h w")
return x
if x.ndim == 5:
x = rearrange(x, "b c t h w -> b t h w c")
x = norm_layer(x)
x = rearrange(x, "b t h w c -> b c t h w")
return x
if isinstance(norm_layer, (nn.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)):
if x.ndim <= 4:
return norm_layer(x)
if x.ndim == 5:
t = x.size(2)
x = rearrange(x, "b c t h w -> (b t) c h w")
memory_occupy = x.numel() * x.element_size() / 1024**3
if isinstance(norm_layer, nn.GroupNorm) and memory_occupy > get_norm_limit():
num_chunks = min(4 if x.element_size() == 2 else 2, norm_layer.num_groups)
assert norm_layer.num_groups % num_chunks == 0
num_groups_per_chunk = norm_layer.num_groups // num_chunks
x = list(x.chunk(num_chunks, dim=1))
weights = norm_layer.weight.chunk(num_chunks, dim=0)
biases = norm_layer.bias.chunk(num_chunks, dim=0)
for i, (w, b) in enumerate(zip(weights, biases)):
def apply_group_norm():
return F.group_norm(x[i], num_groups_per_chunk, w, b, norm_layer.eps)
x[i] = retry_on_oom(
apply_group_norm,
debug=getattr(norm_layer, 'debug', None),
operation_name=f"GroupNorm.chunk_{i}"
)
x[i] = x[i]
x = retry_on_oom(
torch.cat,
x,
dim=1,
debug=getattr(norm_layer, 'debug', None),
operation_name="GroupNorm.concat_chunks"
)
else:
x = retry_on_oom(
norm_layer,
x,
debug=getattr(norm_layer, 'debug', None),
operation_name="GroupNorm.direct"
)
x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
return x
raise NotImplementedError
def remove_head(tensor: Tensor, times: int = 1) -> Tensor:
"""
Remove duplicated first frame features in the up-sampling process.
"""
# Single GPU inference - always process
if times == 0:
return tensor
return torch.cat(tensors=(tensor[:, :, :1], tensor[:, :, times + 1 :]), dim=2)
def extend_head(tensor: Tensor, times: int = 2, memory: Optional[Tensor] = None) -> Tensor:
"""
When memory is None:
- Duplicate first frame features in the down-sampling process.
When memory is not None:
- Concatenate memory features with the input features to keep temporal consistency.
"""
if memory is not None:
return torch.cat((memory.to(tensor), tensor), dim=2)
assert times >= 0, "Invalid input for function 'extend_head'!"
if times == 0:
return tensor
else:
tile_repeat = [1] * tensor.ndim
tile_repeat[2] = times
return torch.cat(tensors=(torch.tile(tensor[:, :, :1], tile_repeat), tensor), dim=2)
def inflate_weight(weight_2d: torch.Tensor, weight_3d: torch.Tensor, inflation_mode: str):
"""
Inflate a 2D convolution weight matrix to a 3D one.
Parameters:
weight_2d: The weight matrix of 2D conv to be inflated.
weight_3d: The weight matrix of 3D conv to be initialized.
inflation_mode: the mode of inflation
"""
assert inflation_mode in ["tail", "replicate"]
assert weight_3d.shape[:2] == weight_2d.shape[:2]
with torch.no_grad():
if inflation_mode == "replicate":
depth = weight_3d.size(2)
weight_3d.copy_(weight_2d.unsqueeze(2).repeat(1, 1, depth, 1, 1) / depth)
else:
weight_3d.fill_(0.0)
weight_3d[:, :, -1].copy_(weight_2d)
return weight_3d
def inflate_bias(bias_2d: torch.Tensor, bias_3d: torch.Tensor, inflation_mode: str):
"""
Inflate a 2D convolution bias tensor to a 3D one
Parameters:
bias_2d: The bias tensor of 2D conv to be inflated.
bias_3d: The bias tensor of 3D conv to be initialized.
inflation_mode: Placeholder to align `inflate_weight`.
"""
assert bias_3d.shape == bias_2d.shape
with torch.no_grad():
bias_3d.copy_(bias_2d)
return bias_3d
def modify_state_dict(layer, state_dict, prefix, inflate_weight_fn, inflate_bias_fn):
"""
the main function to inflated 2D parameters to 3D.
"""
weight_name = prefix + "weight"
bias_name = prefix + "bias"
if weight_name in state_dict:
weight_2d = state_dict[weight_name]
if weight_2d.dim() == 4:
# Assuming the 2D weights are 4D tensors (out_channels, in_channels, h, w)
weight_3d = inflate_weight_fn(
weight_2d=weight_2d,
weight_3d=layer.weight,
inflation_mode=layer.inflation_mode,
)
state_dict[weight_name] = weight_3d
else:
return state_dict
# It's a 3d state dict, should not do inflation on both bias and weight.
if bias_name in state_dict:
bias_2d = state_dict[bias_name]
if bias_2d.dim() == 1:
# Assuming the 2D biases are 1D tensors (out_channels,)
bias_3d = inflate_bias_fn(
bias_2d=bias_2d,
bias_3d=layer.bias,
inflation_mode=layer.inflation_mode,
)
state_dict[bias_name] = bias_3d
return state_dict

Some files were not shown because too many files have changed in this diff Show More