mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-18 07:41:03 +00:00
new comfyUI
This commit is contained in:
@@ -0,0 +1,160 @@
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/#use-with-ide
|
||||
.pdm.toml
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
@@ -0,0 +1,674 @@
|
||||
GNU GENERAL PUBLIC LICENSE
|
||||
Version 3, 29 June 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 General Public License is a free, copyleft license for
|
||||
software and other kinds of works.
|
||||
|
||||
The licenses for most software and other practical works are designed
|
||||
to take away your freedom to share and change the works. By contrast,
|
||||
the GNU General Public License is 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. We, the Free Software Foundation, use the
|
||||
GNU General Public License for most of our software; it applies also to
|
||||
any other work released this way by its authors. You can apply it to
|
||||
your programs, too.
|
||||
|
||||
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.
|
||||
|
||||
To protect your rights, we need to prevent others from denying you
|
||||
these rights or asking you to surrender the rights. Therefore, you have
|
||||
certain responsibilities if you distribute copies of the software, or if
|
||||
you modify it: responsibilities to respect the freedom of others.
|
||||
|
||||
For example, if you distribute copies of such a program, whether
|
||||
gratis or for a fee, you must pass on to the recipients the same
|
||||
freedoms that you received. You must make sure that they, too, receive
|
||||
or can get the source code. And you must show them these terms so they
|
||||
know their rights.
|
||||
|
||||
Developers that use the GNU GPL protect your rights with two steps:
|
||||
(1) assert copyright on the software, and (2) offer you this License
|
||||
giving you legal permission to copy, distribute and/or modify it.
|
||||
|
||||
For the developers' and authors' protection, the GPL clearly explains
|
||||
that there is no warranty for this free software. For both users' and
|
||||
authors' sake, the GPL requires that modified versions be marked as
|
||||
changed, so that their problems will not be attributed erroneously to
|
||||
authors of previous versions.
|
||||
|
||||
Some devices are designed to deny users access to install or run
|
||||
modified versions of the software inside them, although the manufacturer
|
||||
can do so. This is fundamentally incompatible with the aim of
|
||||
protecting users' freedom to change the software. The systematic
|
||||
pattern of such abuse occurs in the area of products for individuals to
|
||||
use, which is precisely where it is most unacceptable. Therefore, we
|
||||
have designed this version of the GPL to prohibit the practice for those
|
||||
products. If such problems arise substantially in other domains, we
|
||||
stand ready to extend this provision to those domains in future versions
|
||||
of the GPL, as needed to protect the freedom of users.
|
||||
|
||||
Finally, every program is threatened constantly by software patents.
|
||||
States should not allow patents to restrict development and use of
|
||||
software on general-purpose computers, but in those that do, we wish to
|
||||
avoid the special danger that patents applied to a free program could
|
||||
make it effectively proprietary. To prevent this, the GPL assures that
|
||||
patents cannot be used to render the program non-free.
|
||||
|
||||
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 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. Use with the GNU Affero General Public License.
|
||||
|
||||
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 Affero 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 special requirements of the GNU Affero General Public License,
|
||||
section 13, concerning interaction through a network will apply to the
|
||||
combination as such.
|
||||
|
||||
14. Revised Versions of this License.
|
||||
|
||||
The Free Software Foundation may publish revised and/or new versions of
|
||||
the GNU 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 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 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 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 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 General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU 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 the program does terminal interaction, make it output a short
|
||||
notice like this when it starts in an interactive mode:
|
||||
|
||||
<program> Copyright (C) <year> <name of author>
|
||||
This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
|
||||
This is free software, and you are welcome to redistribute it
|
||||
under certain conditions; type `show c' for details.
|
||||
|
||||
The hypothetical commands `show w' and `show c' should show the appropriate
|
||||
parts of the General Public License. Of course, your program's commands
|
||||
might be different; for a GUI interface, you would use an "about box".
|
||||
|
||||
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 GPL, see
|
||||
<https://www.gnu.org/licenses/>.
|
||||
|
||||
The GNU General Public License does not permit incorporating your program
|
||||
into proprietary programs. If your program is a subroutine library, you
|
||||
may consider it more useful to permit linking proprietary applications with
|
||||
the library. If this is what you want to do, use the GNU Lesser General
|
||||
Public License instead of this License. But first, please read
|
||||
<https://www.gnu.org/licenses/why-not-lgpl.html>.
|
||||
@@ -0,0 +1,155 @@
|
||||
# ComfyUI-Advanced-ControlNet
|
||||
Nodes for scheduling ControlNet strength across timesteps and batched latents, as well as applying custom weights and attention masks. The ControlNet nodes here fully support sliding context sampling, like the one used in the [ComfyUI-AnimateDiff-Evolved](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved) nodes. Currently supports ControlNets, T2IAdapters, ControlLoRAs, ControlLLLite, SparseCtrls, and SVD-ControlNets.
|
||||
|
||||
Custom weights allow replication of the "My prompt is more important" feature of Auto1111's sd-webui ControlNet extension.
|
||||
|
||||
ControlNet preprocessors are available through [comfyui_controlnet_aux](https://github.com/Fannovel16/comfyui_controlnet_aux) nodes
|
||||
|
||||
## Features
|
||||
- Timestep and latent strength scheduling
|
||||
- Attention masks
|
||||
- Soft weights to replicate "My prompt is more important" feature from sd-webui ControlNet extension, and also change the scaling.
|
||||
- ControlNet, T2IAdapter, and ControlLoRA support for sliding context windows.
|
||||
- ControlLLLite support (requires model_optional to be passed into and out of Apply Advanced ControlNet node)
|
||||
- SparseCtrl support
|
||||
- SVD-ControlNet support
|
||||
- Stable Video Diffusion ControlNets trained by **CiaraRowles**: [Depth](https://huggingface.co/CiaraRowles/temporal-controlnet-depth-svd-v1/tree/main/controlnet), [Lineart](https://huggingface.co/CiaraRowles/temporal-controlnet-lineart-svd-v1/tree/main/controlnet)
|
||||
|
||||
## Table of Contents:
|
||||
- [Scheduling Explanation](#scheduling-explanation)
|
||||
- [Nodes](#nodes)
|
||||
- [Usage](#usage) (will fill this out soon)
|
||||
|
||||
|
||||
# Scheduling Explanation
|
||||
|
||||
The two core concepts for scheduling are ***Timestep Keyframes*** and ***Latent Keyframes***.
|
||||
|
||||
***Timestep Keyframes*** hold the values that guide the settings for a controlnet, and begin to take effect based on their start_percent, which corresponds to the percentage of the sampling process. They can contain masks for the strengths of each latent, control_net_weights, and latent_keyframes (specific strengths for each latent), all optional.
|
||||
|
||||
***Latent Keyframes*** determine the strength of the controlnet for specific latents - all they contain is the batch_index of the latent, and the strength the controlnet should apply for that latent. As a concept, latent keyframes achieve the same affect as a uniform mask with the chosen strength value.
|
||||
|
||||

|
||||
|
||||
# Nodes
|
||||
|
||||
The ControlNet nodes provided here are the ***Apply Advanced ControlNet*** and ***Load Advanced ControlNet Model*** (or diff) nodes. The vanilla ControlNet nodes are also compatible, and can be used almost interchangeably - the only difference is that **at least one of these nodes must be used** for Advanced versions of ControlNets to be used (important for sliding context sampling, like with AnimateDiff-Evolved).
|
||||
|
||||
Key:
|
||||
- 🟩 - required inputs
|
||||
- 🟨 - optional inputs
|
||||
- 🟦 - start as widgets, can be converted to inputs
|
||||
- 🟥 - optional input/output, but not recommended to use unless needed
|
||||
- 🟪 - output
|
||||
|
||||
## Apply Advanced ControlNet
|
||||

|
||||
|
||||
Same functionality as the vanilla Apply Advanced ControlNet (Advanced) node, except with Advanced ControlNet features added to it. Automatically converts any ControlNet from ControlNet loaders into Advanced versions.
|
||||
|
||||
### Inputs
|
||||
- 🟩***positive***: conditioning (positive).
|
||||
- 🟩***negative***: conditioning (negative).
|
||||
- 🟩***control_net***: loaded controlnet; will be converted to Advanced version automatically by this node, if it's a supported type.
|
||||
- 🟩***image***: images to guide controlnets - if the loaded controlnet requires it, they must preprocessed images. If one image provided, will be used for all latents. If more images provided, will use each image separately for each latent. If not enough images to meet latent count, will repeat the images from the beginning to match vanilla ControlNet functionality.
|
||||
- 🟨***mask_optional***: attention masks to apply to controlnets; basically, decides what part of the image the controlnet to apply to (and the relative strength, if the mask is not binary). Same as image input, if you provide more than one mask, each can apply to a different latent.
|
||||
- 🟨***timestep_kf***: timestep keyframes to guide controlnet effect throughout sampling steps.
|
||||
- 🟨***latent_kf_override***: override for latent keyframes, useful if no other features from timestep keyframes is needed. *NOTE: this latent keyframe will be applied to ALL timesteps, regardless if there are other latent keyframes attached to connected timestep keyframes.*
|
||||
- 🟨***weights_override***: override for weights, useful if no other features from timestep keyframes is needed. *NOTE: this weight will be applied to ALL timesteps, regardless if there are other weights attached to connected timestep keyframes.*
|
||||
- 🟦***strength***: strength of controlnet; 1.0 is full strength, 0.0 is no effect at all.
|
||||
- 🟦***start_percent***: sampling step percentage at which controlnet should start to be applied - no matter what start_percent is set on timestep keyframes, they won't take effect until this start_percent is reached.
|
||||
- 🟦***stop_percent***: sampling step percentage at which controlnet should stop being applied - no matter what start_percent is set on timestep keyframes, they won't take effect once this end_percent is reached.
|
||||
|
||||
### Outputs
|
||||
- 🟪***positive***: conditioning (positive) with applied controlnets
|
||||
- 🟪***negative***: conditioning (negative) with applied controlnets
|
||||
|
||||
## Load Advanced ControlNet Model
|
||||

|
||||
|
||||
Loads a ControlNet model and converts it into an Advanced version that supports all the features in this repo. When used with **Apply Advanced ControlNet** node, there is no reason to use the timestep_keyframe input on this node - use timestep_kf on the Apply node instead.
|
||||
|
||||
### Inputs
|
||||
- 🟥***timestep_keyframe***: optional and likely unnecessary input to have ControlNet use selected timestep_keyframes - should not be used unless you need to. Useful if this node is not attached to **Apply Advanced ControlNet** node, but still want to use Timestep Keyframe, or to use TK_SHORTCUT outputs from ControlWeights in the same scenario. Will be overriden by the timestep_kf input on **Apply Advanced ControlNet** node, if one is provided there.
|
||||
- 🟨***model***: model to plug into the diff version of the node. Some controlnets are designed for receive the model; if you don't know what this does, you probably don't want tot use the diff version of the node.
|
||||
|
||||
### Outputs
|
||||
- 🟪***CONTROL_NET***: loaded Advanced ControlNet
|
||||
|
||||
## Timestep Keyframe
|
||||

|
||||
|
||||
Scheduling node across timesteps (sampling steps) based on the set start_percent. Chaining Timestep Keyframes allows ControlNet scheduling across sampling steps (percentage-wise), through a timestep keyframe schedule.
|
||||
|
||||
### Inputs
|
||||
- 🟨***prev_timestep_kf***: used to chain Timestep Keyframes together to create a schedule. The order does not matter - the Timestep Keyframes sort themselves automatically by their start_percent. *Any Timestep Keyframe contained in the prev_timestep_keyframe that contains the same start_percent as the Timestep Keyframe will be overwritten.*
|
||||
- 🟨***cn_weights***: weights to apply to controlnet while this Timestep Keyframe is in effect. Must be compatible with the loaded controlnet, or will throw an error explaining what weight types are compatible. If inherit_missing is True, if no control_net_weight is passed in, will attempt to reuse the last-used weights in the timestep keyframe schedule. *If Apply Advanced ControlNet node has a weight_override, the weight_override will be used during sampling instead of control_net_weight.*
|
||||
- 🟨***latent_keyframe***: latent keyframes to apply to controlnet while this Timestep Keyframe is in effect. If inherit_missing is True, if no latent_keyframe is passed in, will attempt to reuse the last-used weights in the timestep keyframe schedule. *If Apply Advanced ControlNet node has a latent_kf_override, the latent_lf_override will be used during sampling instead of latent_keyframe.*
|
||||
- 🟨***mask_optional***: attention masks to apply to controlnets; basically, decides what part of the image the controlnet to apply to (and the relative strength, if the mask is not binary). Same as mask_optional on the Apply Advanced ControlNet node, can apply either one maks to all latents, or individual masks for each latent. If inherit_missing is True, if no mask_optional is passed in, will attempt to reuse the last-used mask_optional in the timestep keyframe schedule. It is NOT overriden by mask_optional on the Apply Advanced ControlNet node; will be used together.
|
||||
- 🟦***start_percent***: sampling step percentage at which this Timestep Keyframe qualifies to be used. Acts as the 'key' for the Timestep Keyframe in the timestep keyframe schedule.
|
||||
- 🟦***strength***: strength of the controlnet; multiplies the controlnet by this value, basically, applied alongside the strength on the Apply ControlNet node. If set to 0.0 will not have any effect during the duration of this Timestep Keyframe's effect, and will increase sampling speed by not doing any work.
|
||||
- 🟦***null_latent_kf_strength***: strength to assign to latents that are unaccounted for in the passed in latent_keyframes. Has no effect if no latent_keyframes are passed in, or no batch_indeces are unaccounted in the latent_keyframes for during sampling.
|
||||
- 🟦***inherit_missing***: determines if should reuse values from previous Timestep Keyframes for optional values (control_net_weights, latent_keyframe, and mask_option) that are not included on this TimestepKeyframe. To inherit only specific inputs, use default inputs.
|
||||
- 🟦***guarantee_usage***: when true, even if a Timestep Keyframe's start_percent ahead of this one in the schedule is closer to current sampling percentage, this Timestep Keyframe will still be used for one step before moving on to the next selected Timestep Keyframe in the following step. Whether the Timestep Keyframe is used or not, its inputs will still be accounted for inherit_missing purposes.
|
||||
|
||||
### Outputs
|
||||
- 🟪***TIMESTEP_KF***: the created Timestep Keyframe, that can either be linked to another or into a Timestep Keyframe input.
|
||||
|
||||
## Latent Keyframe
|
||||

|
||||
|
||||
A singular Latent Keyframe, selects the strength for a specific batch_index. If batch_index is not present during sampling, will simply have no effect. Can be chained with any other Latent Keyframe-type node to create a latent keyframe schedule.
|
||||
|
||||
### Inputs
|
||||
- 🟨***prev_latent_kf***: used to chain Latent Keyframes together to create a schedule. *If a Latent Keyframe contained in prev_latent_keyframes have the same batch_index as this Latent Keyframe, they will take priority over this node's value.*
|
||||
- 🟦***batch_index***: index of latent in batch to apply controlnet strength to. Acts as the 'key' for the Latent Keyframe in the latent keyframe schedule.
|
||||
- 🟦***strength***: strength of controlnet to apply to the corresponding latent.
|
||||
|
||||
### Outputs
|
||||
- 🟪***LATENT_KF***: the created Latent Keyframe, that can either be linked to another or into a Latent Keyframe input.
|
||||
|
||||
## Latent Keyframe Group
|
||||

|
||||
|
||||
Allows to create Latent Keyframes via individual indeces or python-style ranges.
|
||||
|
||||
### Inputs
|
||||
- 🟨***prev_latent_kf***: used to chain Latent Keyframes together to create a schedule. *If any Latent Keyframes contained in prev_latent_keyframes have the same batch_index as a this Latent Keyframe, they will take priority over this node's version.*
|
||||
- 🟨***latent_optional***: the latents expected to be passed in for sampling; only required if you wish to use negative indeces (will be automatically converted to real values).
|
||||
- 🟦***index_strengths***: string list of indeces or python-style ranges of indeces to assign strengths to. If latent_optional is passed in, can contain negative indeces or ranges that contain negative numbers, python-style. The different indeces must be comma separated. Individual latents can be specified by ```batch_index=strength```, like ```0=0.9```. Ranges can be specified by ```start_index_inclusive:end_index_exclusive=strength```, like ```0:8=strength```. Negative indeces are possible when latents_optional has an input, with a string such as ```0,-4=0.25```.
|
||||
- 🟦***print_keyframes***: if True, will print the Latent Keyframes generated by this node for debugging purposes.
|
||||
|
||||
### Outputs
|
||||
- 🟪***LATENT_KF***: the created Latent Keyframe, that can either be linked to another or into a Latent Keyframe input.
|
||||
|
||||
## Latent Keyframe Interpolation
|
||||

|
||||
|
||||
Allows to create Latent Keyframes with interpolated values in a range.
|
||||
|
||||
### Inputs
|
||||
- 🟨***prev_latent_kf***: used to chain Latent Keyframes together to create a schedule. *If any Latent Keyframes contained in prev_latent_keyframes have the same batch_index as a this Latent Keyframe, they will take priority over this node's version.*
|
||||
- 🟦***batch_index_from***: starting batch_index of range, included.
|
||||
- 🟦***batch_index_to***: end batch_index of range, excluded (python-style range).
|
||||
- 🟦***strength_from***: starting strength of interpolation.
|
||||
- 🟦***strength_to***: end strength of interpolation.
|
||||
- 🟦***interpolation***: the method of interpolation.
|
||||
- 🟦***print_keyframes***: if True, will print the Latent Keyframes generated by this node for debugging purposes.
|
||||
|
||||
### Outputs
|
||||
- 🟪***LATENT_KF***: the created Latent Keyframe, that can either be linked to another or into a Latent Keyframe input.
|
||||
|
||||
## Latent Keyframe Batched Group
|
||||

|
||||
|
||||
Allows to create Latent Keyframes via a list of floats, such as with Batch Value Schedule from [ComfyUI_FizzNodes](https://github.com/FizzleDorf/ComfyUI_FizzNodes) nodes.
|
||||
|
||||
### Inputs
|
||||
- 🟨***prev_latent_kf***: used to chain Latent Keyframes together to create a schedule. *If any Latent Keyframes contained in prev_latent_keyframes have the same batch_index as a this Latent Keyframe, they will take priority over this node's version.*
|
||||
- 🟦***float_strengths***: a list of floats, that will correspond to the strength of each Latent Keyframe; the batch_index is the index of each float value in the list.
|
||||
- 🟦***print_keyframes***: if True, will print the Latent Keyframes generated by this node for debugging purposes.
|
||||
|
||||
### Outputs
|
||||
- 🟪***LATENT_KF***: the created Latent Keyframe, that can either be linked to another or into a Latent Keyframe input.
|
||||
|
||||
# There are more nodes to document and show usage - will add this soon! TODO
|
||||
@@ -0,0 +1,3 @@
|
||||
from .adv_control.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
+895
@@ -0,0 +1,895 @@
|
||||
from typing import Callable, Union
|
||||
from torch import Tensor
|
||||
import torch
|
||||
import os
|
||||
|
||||
import comfy.utils
|
||||
import comfy.model_management
|
||||
import comfy.model_detection
|
||||
import comfy.controlnet as comfy_cn
|
||||
from comfy.controlnet import ControlBase, ControlNet, ControlLora, T2IAdapter, broadcast_image_to
|
||||
from model_patcher import ModelPatcher
|
||||
|
||||
from .control_sparsectrl import SparseControlNet, SparseCtrlMotionWrapper, SparseMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper
|
||||
from .control_lllite import LLLiteModule, LLLitePatch
|
||||
from .control_svd import svd_unet_config_from_diffusers_unet, SVDControlNet, svd_unet_to_diffusers
|
||||
from .utils import (AdvancedControlBase, TimestepKeyframeGroup, LatentKeyframeGroup, ControlWeightType, ControlWeights, WeightTypeException,
|
||||
manual_cast_clean_groupnorm, disable_weight_init_clean_groupnorm, prepare_mask_batch, get_properly_arranged_t2i_weights, load_torch_file_with_dict_factory)
|
||||
from .logger import logger
|
||||
|
||||
|
||||
class ControlNetAdvanced(ControlNet, AdvancedControlBase):
|
||||
def __init__(self, control_model, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, device=None, load_device=None, manual_cast_dtype=None):
|
||||
super().__init__(control_model=control_model, global_average_pooling=global_average_pooling, device=device, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
|
||||
AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controlnet())
|
||||
|
||||
def get_universal_weights(self) -> ControlWeights:
|
||||
raw_weights = [(self.weights.base_multiplier ** float(12 - i)) for i in range(13)]
|
||||
return ControlWeights.controlnet(raw_weights, self.weights.flip_weights)
|
||||
|
||||
def get_control_advanced(self, x_noisy, t, cond, batched_number):
|
||||
# perform special version of get_control that supports sliding context and masks
|
||||
return self.sliding_get_control(x_noisy, t, cond, batched_number)
|
||||
|
||||
def sliding_get_control(self, x_noisy: Tensor, t, cond, batched_number):
|
||||
control_prev = None
|
||||
if self.previous_controlnet is not None:
|
||||
control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
|
||||
|
||||
if self.timestep_range is not None:
|
||||
if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
|
||||
if control_prev is not None:
|
||||
return control_prev
|
||||
else:
|
||||
return None
|
||||
|
||||
dtype = self.control_model.dtype
|
||||
if self.manual_cast_dtype is not None:
|
||||
dtype = self.manual_cast_dtype
|
||||
|
||||
output_dtype = x_noisy.dtype
|
||||
# make cond_hint appropriate dimensions
|
||||
# TODO: change this to not require cond_hint upscaling every step when self.sub_idxs are present
|
||||
if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
|
||||
if self.cond_hint is not None:
|
||||
del self.cond_hint
|
||||
self.cond_hint = None
|
||||
# if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling
|
||||
if self.sub_idxs is not None and self.cond_hint_original.size(0) >= self.full_latent_length:
|
||||
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
|
||||
else:
|
||||
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
|
||||
if x_noisy.shape[0] != self.cond_hint.shape[0]:
|
||||
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
|
||||
|
||||
# prepare mask_cond_hint
|
||||
self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number, dtype=dtype)
|
||||
|
||||
context = cond.get('crossattn_controlnet', cond['c_crossattn'])
|
||||
# uses 'y' in new ComfyUI update
|
||||
y = cond.get('y', None)
|
||||
if y is None: # TODO: remove this in the future since no longer used by newest ComfyUI
|
||||
y = cond.get('c_adm', None)
|
||||
if y is not None:
|
||||
y = y.to(dtype)
|
||||
timestep = self.model_sampling_current.timestep(t)
|
||||
x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
|
||||
|
||||
control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y)
|
||||
return self.control_merge(None, control, control_prev, output_dtype)
|
||||
|
||||
def copy(self):
|
||||
c = ControlNetAdvanced(self.control_model, self.timestep_keyframes, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
|
||||
self.copy_to(c)
|
||||
self.copy_to_advanced(c)
|
||||
return c
|
||||
|
||||
@staticmethod
|
||||
def from_vanilla(v: ControlNet, timestep_keyframe: TimestepKeyframeGroup=None) -> 'ControlNetAdvanced':
|
||||
return ControlNetAdvanced(control_model=v.control_model, timestep_keyframes=timestep_keyframe,
|
||||
global_average_pooling=v.global_average_pooling, device=v.device, load_device=v.load_device, manual_cast_dtype=v.manual_cast_dtype)
|
||||
|
||||
|
||||
class T2IAdapterAdvanced(T2IAdapter, AdvancedControlBase):
|
||||
def __init__(self, t2i_model, timestep_keyframes: TimestepKeyframeGroup, channels_in, device=None):
|
||||
super().__init__(t2i_model=t2i_model, channels_in=channels_in, device=device)
|
||||
AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.t2iadapter())
|
||||
|
||||
def get_universal_weights(self) -> ControlWeights:
|
||||
raw_weights = [(self.weights.base_multiplier ** float(7 - i)) for i in range(8)]
|
||||
raw_weights = [raw_weights[-8], raw_weights[-3], raw_weights[-2], raw_weights[-1]]
|
||||
raw_weights = get_properly_arranged_t2i_weights(raw_weights)
|
||||
return ControlWeights.t2iadapter(raw_weights, self.weights.flip_weights)
|
||||
|
||||
def get_calc_pow(self, idx: int, layers: int) -> int:
|
||||
# match how T2IAdapterAdvanced deals with universal weights
|
||||
indeces = [7 - i for i in range(8)]
|
||||
indeces = [indeces[-8], indeces[-3], indeces[-2], indeces[-1]]
|
||||
indeces = get_properly_arranged_t2i_weights(indeces)
|
||||
return indeces[idx]
|
||||
|
||||
def get_control_advanced(self, x_noisy, t, cond, batched_number):
|
||||
try:
|
||||
# if sub indexes present, replace original hint with subsection
|
||||
if self.sub_idxs is not None:
|
||||
# cond hints
|
||||
full_cond_hint_original = self.cond_hint_original
|
||||
del self.cond_hint
|
||||
self.cond_hint = None
|
||||
self.cond_hint_original = full_cond_hint_original[self.sub_idxs]
|
||||
# mask hints
|
||||
self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number)
|
||||
return super().get_control(x_noisy, t, cond, batched_number)
|
||||
finally:
|
||||
if self.sub_idxs is not None:
|
||||
# replace original cond hint
|
||||
self.cond_hint_original = full_cond_hint_original
|
||||
del full_cond_hint_original
|
||||
|
||||
def copy(self):
|
||||
c = T2IAdapterAdvanced(self.t2i_model, self.timestep_keyframes, self.channels_in)
|
||||
self.copy_to(c)
|
||||
self.copy_to_advanced(c)
|
||||
return c
|
||||
|
||||
def cleanup(self):
|
||||
super().cleanup()
|
||||
self.cleanup_advanced()
|
||||
|
||||
@staticmethod
|
||||
def from_vanilla(v: T2IAdapter, timestep_keyframe: TimestepKeyframeGroup=None) -> 'T2IAdapterAdvanced':
|
||||
return T2IAdapterAdvanced(t2i_model=v.t2i_model, timestep_keyframes=timestep_keyframe, channels_in=v.channels_in, device=v.device)
|
||||
|
||||
|
||||
class ControlLoraAdvanced(ControlLora, AdvancedControlBase):
|
||||
def __init__(self, control_weights, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, device=None):
|
||||
super().__init__(control_weights=control_weights, global_average_pooling=global_average_pooling, device=device)
|
||||
AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllora())
|
||||
# use some functions from ControlNetAdvanced
|
||||
self.get_control_advanced = ControlNetAdvanced.get_control_advanced.__get__(self, type(self))
|
||||
self.sliding_get_control = ControlNetAdvanced.sliding_get_control.__get__(self, type(self))
|
||||
|
||||
def get_universal_weights(self) -> ControlWeights:
|
||||
raw_weights = [(self.weights.base_multiplier ** float(9 - i)) for i in range(10)]
|
||||
return ControlWeights.controllora(raw_weights, self.weights.flip_weights)
|
||||
|
||||
def copy(self):
|
||||
c = ControlLoraAdvanced(self.control_weights, self.timestep_keyframes, global_average_pooling=self.global_average_pooling)
|
||||
self.copy_to(c)
|
||||
self.copy_to_advanced(c)
|
||||
return c
|
||||
|
||||
def cleanup(self):
|
||||
super().cleanup()
|
||||
self.cleanup_advanced()
|
||||
|
||||
@staticmethod
|
||||
def from_vanilla(v: ControlLora, timestep_keyframe: TimestepKeyframeGroup=None) -> 'ControlLoraAdvanced':
|
||||
return ControlLoraAdvanced(control_weights=v.control_weights, timestep_keyframes=timestep_keyframe,
|
||||
global_average_pooling=v.global_average_pooling, device=v.device)
|
||||
|
||||
|
||||
class SVDControlNetAdvanced(ControlNetAdvanced):
|
||||
def __init__(self, control_model: SVDControlNet, timestep_keyframes: TimestepKeyframeGroup, global_average_pooling=False, device=None, load_device=None, manual_cast_dtype=None):
|
||||
super().__init__(control_model=control_model, timestep_keyframes=timestep_keyframes, global_average_pooling=global_average_pooling, device=device, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
|
||||
|
||||
def set_cond_hint(self, *args, **kwargs):
|
||||
to_return = super().set_cond_hint(*args, **kwargs)
|
||||
# cond hint for SVD-ControlNet needs to be scaled between (-1, 1) instead of (0, 1)
|
||||
self.cond_hint_original = self.cond_hint_original * 2.0 - 1.0
|
||||
return to_return
|
||||
|
||||
def get_control_advanced(self, x_noisy, t, cond, batched_number):
|
||||
control_prev = None
|
||||
if self.previous_controlnet is not None:
|
||||
control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
|
||||
|
||||
if self.timestep_range is not None:
|
||||
if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
|
||||
if control_prev is not None:
|
||||
return control_prev
|
||||
else:
|
||||
return None
|
||||
|
||||
dtype = self.control_model.dtype
|
||||
if self.manual_cast_dtype is not None:
|
||||
dtype = self.manual_cast_dtype
|
||||
|
||||
output_dtype = x_noisy.dtype
|
||||
# make cond_hint appropriate dimensions
|
||||
# TODO: change this to not require cond_hint upscaling every step when self.sub_idxs are present
|
||||
if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
|
||||
if self.cond_hint is not None:
|
||||
del self.cond_hint
|
||||
self.cond_hint = None
|
||||
# if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling
|
||||
if self.sub_idxs is not None and self.cond_hint_original.size(0) >= self.full_latent_length:
|
||||
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
|
||||
else:
|
||||
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
|
||||
if x_noisy.shape[0] != self.cond_hint.shape[0]:
|
||||
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
|
||||
|
||||
# prepare mask_cond_hint
|
||||
self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number, dtype=dtype)
|
||||
|
||||
context = cond.get('crossattn_controlnet', cond['c_crossattn'])
|
||||
# uses 'y' in new ComfyUI update
|
||||
y = cond.get('y', None)
|
||||
if y is not None:
|
||||
y = y.to(dtype)
|
||||
timestep = self.model_sampling_current.timestep(t)
|
||||
x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
|
||||
# concat c_concat if exists (should exist for SVD), doubling channels to 8
|
||||
if cond.get('c_concat', None) is not None:
|
||||
x_noisy = torch.cat([x_noisy] + [cond['c_concat']], dim=1)
|
||||
|
||||
control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y, cond=cond)
|
||||
return self.control_merge(None, control, control_prev, output_dtype)
|
||||
|
||||
def copy(self):
|
||||
c = SVDControlNetAdvanced(self.control_model, self.timestep_keyframes, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
|
||||
self.copy_to(c)
|
||||
self.copy_to_advanced(c)
|
||||
return c
|
||||
|
||||
|
||||
class SparseCtrlAdvanced(ControlNetAdvanced):
|
||||
def __init__(self, control_model, timestep_keyframes: TimestepKeyframeGroup, sparse_settings: SparseSettings=None, global_average_pooling=False, device=None, load_device=None, manual_cast_dtype=None):
|
||||
super().__init__(control_model=control_model, timestep_keyframes=timestep_keyframes, global_average_pooling=global_average_pooling, device=device, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
|
||||
self.add_compatible_weight(ControlWeightType.SPARSECTRL)
|
||||
self.control_model: SparseControlNet = self.control_model # does nothing except help with IDE hints
|
||||
self.sparse_settings = sparse_settings if sparse_settings is not None else SparseSettings.default()
|
||||
self.latent_format = None
|
||||
self.preprocessed = False
|
||||
|
||||
def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int):
|
||||
# normal ControlNet stuff
|
||||
control_prev = None
|
||||
if self.previous_controlnet is not None:
|
||||
control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
|
||||
|
||||
if self.timestep_range is not None:
|
||||
if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
|
||||
if control_prev is not None:
|
||||
return control_prev
|
||||
else:
|
||||
return None
|
||||
|
||||
dtype = self.control_model.dtype
|
||||
if self.manual_cast_dtype is not None:
|
||||
dtype = self.manual_cast_dtype
|
||||
output_dtype = x_noisy.dtype
|
||||
# set actual input length on motion model
|
||||
actual_length = x_noisy.size(0)//batched_number
|
||||
full_length = actual_length if self.sub_idxs is None else self.full_latent_length
|
||||
self.control_model.set_actual_length(actual_length=actual_length, full_length=full_length)
|
||||
# prepare cond_hint, if needed
|
||||
dim_mult = 1 if self.control_model.use_simplified_conditioning_embedding else 8
|
||||
if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2]*dim_mult != self.cond_hint.shape[2] or x_noisy.shape[3]*dim_mult != self.cond_hint.shape[3]:
|
||||
# clear out cond_hint and conditioning_mask
|
||||
if self.cond_hint is not None:
|
||||
del self.cond_hint
|
||||
self.cond_hint = None
|
||||
# first, figure out which cond idxs are relevant, and where they fit in
|
||||
cond_idxs = self.sparse_settings.sparse_method.get_indexes(hint_length=self.cond_hint_original.size(0), full_length=full_length)
|
||||
|
||||
range_idxs = list(range(full_length)) if self.sub_idxs is None else self.sub_idxs
|
||||
hint_idxs = [] # idxs in cond_idxs
|
||||
local_idxs = [] # idx to pun in final cond_hint
|
||||
for i,cond_idx in enumerate(cond_idxs):
|
||||
if cond_idx in range_idxs:
|
||||
hint_idxs.append(i)
|
||||
local_idxs.append(range_idxs.index(cond_idx))
|
||||
# sub_cond_hint now contains the hints relevant to current x_noisy
|
||||
sub_cond_hint = self.cond_hint_original[hint_idxs].to(dtype).to(self.device)
|
||||
|
||||
# scale cond_hints to match noisy input
|
||||
if self.control_model.use_simplified_conditioning_embedding:
|
||||
# RGB SparseCtrl; the inputs are latents - use bilinear to avoid blocky artifacts
|
||||
sub_cond_hint = self.latent_format.process_in(sub_cond_hint) # multiplies by model scale factor
|
||||
sub_cond_hint = comfy.utils.common_upscale(sub_cond_hint, x_noisy.shape[3], x_noisy.shape[2], "nearest-exact", "center").to(dtype).to(self.device)
|
||||
else:
|
||||
# other SparseCtrl; inputs are typical images
|
||||
sub_cond_hint = comfy.utils.common_upscale(sub_cond_hint, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
|
||||
# prepare cond_hint (b, c, h ,w)
|
||||
cond_shape = list(sub_cond_hint.shape)
|
||||
cond_shape[0] = len(range_idxs)
|
||||
self.cond_hint = torch.zeros(cond_shape).to(dtype).to(self.device)
|
||||
self.cond_hint[local_idxs] = sub_cond_hint[:]
|
||||
# prepare cond_mask (b, 1, h, w)
|
||||
cond_shape[1] = 1
|
||||
cond_mask = torch.zeros(cond_shape).to(dtype).to(self.device)
|
||||
cond_mask[local_idxs] = 1.0
|
||||
# combine cond_hint and cond_mask into (b, c+1, h, w)
|
||||
if not self.sparse_settings.merged:
|
||||
self.cond_hint = torch.cat([self.cond_hint, cond_mask], dim=1)
|
||||
del sub_cond_hint
|
||||
del cond_mask
|
||||
# make cond_hint match x_noisy batch
|
||||
if x_noisy.shape[0] != self.cond_hint.shape[0]:
|
||||
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
|
||||
|
||||
# prepare mask_cond_hint
|
||||
self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number, dtype=dtype)
|
||||
|
||||
context = cond['c_crossattn']
|
||||
y = cond.get('y', None)
|
||||
if y is not None:
|
||||
y = y.to(dtype)
|
||||
timestep = self.model_sampling_current.timestep(t)
|
||||
x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
|
||||
|
||||
control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y)
|
||||
return self.control_merge(None, control, control_prev, output_dtype)
|
||||
|
||||
def pre_run_advanced(self, model, percent_to_timestep_function):
|
||||
super().pre_run_advanced(model, percent_to_timestep_function)
|
||||
if type(self.cond_hint_original) == PreprocSparseRGBWrapper:
|
||||
if not self.control_model.use_simplified_conditioning_embedding:
|
||||
raise ValueError("Any model besides RGB SparseCtrl should NOT have its images go through the RGB SparseCtrl preprocessor.")
|
||||
self.cond_hint_original = self.cond_hint_original.condhint
|
||||
self.latent_format = model.latent_format # LatentFormat object, used to process_in latent cond hint
|
||||
if self.control_model.motion_holder is not None:
|
||||
self.control_model.motion_holder.motion_wrapper.reset()
|
||||
self.control_model.motion_holder.motion_wrapper.set_strength(self.sparse_settings.motion_strength)
|
||||
self.control_model.motion_holder.motion_wrapper.set_scale_multiplier(self.sparse_settings.motion_scale)
|
||||
|
||||
def cleanup_advanced(self):
|
||||
super().cleanup_advanced()
|
||||
if self.latent_format is not None:
|
||||
del self.latent_format
|
||||
self.latent_format = None
|
||||
|
||||
def copy(self):
|
||||
c = SparseCtrlAdvanced(self.control_model, self.timestep_keyframes, self.sparse_settings, self.global_average_pooling, self.device, self.load_device, self.manual_cast_dtype)
|
||||
self.copy_to(c)
|
||||
self.copy_to_advanced(c)
|
||||
return c
|
||||
|
||||
|
||||
class ReferenceAdvanced(ControlBase, AdvancedControlBase):
|
||||
def __init__(self, timestep_keyframes: TimestepKeyframeGroup, device=None):
|
||||
super().__init__(device)
|
||||
AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllllite(), require_model=True)
|
||||
# TODO: save attn patches here
|
||||
|
||||
def patch_model(self, model: ModelPatcher):
|
||||
# TODO: do model patching here
|
||||
pass
|
||||
|
||||
def pre_run_advanced(self, *args, **kwargs):
|
||||
AdvancedControlBase.pre_run_advanced(self, *args, **kwargs)
|
||||
# TODO: set control on patches
|
||||
|
||||
def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int):
|
||||
# normal ControlNet stuff
|
||||
control_prev = None
|
||||
if self.previous_controlnet is not None:
|
||||
control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
|
||||
|
||||
if self.timestep_range is not None:
|
||||
if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
|
||||
return control_prev
|
||||
|
||||
dtype = x_noisy.dtype
|
||||
# prepare cond_hint
|
||||
if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
|
||||
if self.cond_hint is not None:
|
||||
del self.cond_hint
|
||||
self.cond_hint = None
|
||||
# if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling
|
||||
if self.sub_idxs is not None and self.cond_hint_original.size(0) >= self.full_latent_length:
|
||||
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
|
||||
else:
|
||||
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
|
||||
if x_noisy.shape[0] != self.cond_hint.shape[0]:
|
||||
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
|
||||
# prepare mask
|
||||
self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number)
|
||||
# done preparing; model patches will take care of everything now.
|
||||
# return normal controlnet stuff
|
||||
return control_prev
|
||||
|
||||
def cleanup_advanced(self):
|
||||
super().cleanup_advanced()
|
||||
# TODO: cleanup patches here
|
||||
|
||||
def copy(self):
|
||||
c = ReferenceAdvanced(self.timestep_keyframes)
|
||||
self.copy_to(c)
|
||||
self.copy_to_advanced(c)
|
||||
return c
|
||||
|
||||
|
||||
class ControlLLLiteAdvanced(ControlBase, AdvancedControlBase):
|
||||
# This ControlNet is more of an attention patch than a traditional controlnet
|
||||
def __init__(self, patch_attn1: LLLitePatch, patch_attn2: LLLitePatch, timestep_keyframes: TimestepKeyframeGroup, device=None):
|
||||
super().__init__(device)
|
||||
AdvancedControlBase.__init__(self, super(), timestep_keyframes=timestep_keyframes, weights_default=ControlWeights.controllllite(), require_model=True)
|
||||
self.patch_attn1 = patch_attn1.set_control(self)
|
||||
self.patch_attn2 = patch_attn2.set_control(self)
|
||||
self.latent_dims_div2 = None
|
||||
self.latent_dims_div4 = None
|
||||
|
||||
def patch_model(self, model: ModelPatcher):
|
||||
model.set_model_attn1_patch(self.patch_attn1)
|
||||
model.set_model_attn2_patch(self.patch_attn2)
|
||||
|
||||
def set_cond_hint(self, *args, **kwargs):
|
||||
to_return = super().set_cond_hint(*args, **kwargs)
|
||||
# cond hint for LLLite needs to be scaled between (-1, 1) instead of (0, 1)
|
||||
self.cond_hint_original = self.cond_hint_original * 2.0 - 1.0
|
||||
return to_return
|
||||
|
||||
def pre_run_advanced(self, *args, **kwargs):
|
||||
AdvancedControlBase.pre_run_advanced(self, *args, **kwargs)
|
||||
#logger.error(f"in cn: {id(self.patch_attn1)},{id(self.patch_attn2)}")
|
||||
self.patch_attn1.set_control(self)
|
||||
self.patch_attn2.set_control(self)
|
||||
#logger.warn(f"in pre_run_advanced: {id(self)}")
|
||||
|
||||
def get_control_advanced(self, x_noisy: Tensor, t, cond, batched_number: int):
|
||||
# normal ControlNet stuff
|
||||
control_prev = None
|
||||
if self.previous_controlnet is not None:
|
||||
control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
|
||||
|
||||
if self.timestep_range is not None:
|
||||
if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
|
||||
return control_prev
|
||||
|
||||
dtype = x_noisy.dtype
|
||||
# prepare cond_hint
|
||||
if self.sub_idxs is not None or self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
|
||||
if self.cond_hint is not None:
|
||||
del self.cond_hint
|
||||
self.cond_hint = None
|
||||
# if self.cond_hint_original length greater or equal to real latent count, subdivide it before scaling
|
||||
if self.sub_idxs is not None and self.cond_hint_original.size(0) >= self.full_latent_length:
|
||||
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original[self.sub_idxs], x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
|
||||
else:
|
||||
self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
|
||||
if x_noisy.shape[0] != self.cond_hint.shape[0]:
|
||||
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
|
||||
# some special logic here compared to other controlnets:
|
||||
# * The cond_emb in attn patches will divide latent dims by 2 or 4, integer
|
||||
# * Due to this loss, the cond_emb will become smaller than x input if latent dims are not divisble by 2 or 4
|
||||
divisible_by_2_h = x_noisy.shape[2]%2==0
|
||||
divisible_by_2_w = x_noisy.shape[3]%2==0
|
||||
if not (divisible_by_2_h and divisible_by_2_w):
|
||||
#logger.warn(f"{x_noisy.shape} not divisible by 2!")
|
||||
new_h = (x_noisy.shape[2]//2)*2
|
||||
new_w = (x_noisy.shape[3]//2)*2
|
||||
if not divisible_by_2_h:
|
||||
new_h += 2
|
||||
if not divisible_by_2_w:
|
||||
new_w += 2
|
||||
self.latent_dims_div2 = (new_h, new_w)
|
||||
divisible_by_4_h = x_noisy.shape[2]%4==0
|
||||
divisible_by_4_w = x_noisy.shape[3]%4==0
|
||||
if not (divisible_by_4_h and divisible_by_4_w):
|
||||
#logger.warn(f"{x_noisy.shape} not divisible by 4!")
|
||||
new_h = (x_noisy.shape[2]//4)*4
|
||||
new_w = (x_noisy.shape[3]//4)*4
|
||||
if not divisible_by_4_h:
|
||||
new_h += 4
|
||||
if not divisible_by_4_w:
|
||||
new_w += 4
|
||||
self.latent_dims_div4 = (new_h, new_w)
|
||||
# prepare mask
|
||||
self.prepare_mask_cond_hint(x_noisy=x_noisy, t=t, cond=cond, batched_number=batched_number)
|
||||
# done preparing; model patches will take care of everything now.
|
||||
# return normal controlnet stuff
|
||||
return control_prev
|
||||
|
||||
def cleanup_advanced(self):
|
||||
super().cleanup_advanced()
|
||||
self.patch_attn1.cleanup()
|
||||
self.patch_attn2.cleanup()
|
||||
self.latent_dims_div2 = None
|
||||
self.latent_dims_div4 = None
|
||||
|
||||
def copy(self):
|
||||
c = ControlLLLiteAdvanced(self.patch_attn1, self.patch_attn2, self.timestep_keyframes)
|
||||
self.copy_to(c)
|
||||
self.copy_to_advanced(c)
|
||||
return c
|
||||
|
||||
# deepcopy needs to properly keep track of objects to work between model.clone calls!
|
||||
# def __deepcopy__(self, *args, **kwargs):
|
||||
# self.cleanup_advanced()
|
||||
# return self
|
||||
|
||||
# def get_models(self):
|
||||
# # get_models is called once at the start of every KSampler run - use to reset already_patched status
|
||||
# out = super().get_models()
|
||||
# logger.error(f"in get_models! {id(self)}")
|
||||
# return out
|
||||
|
||||
|
||||
def load_controlnet(ckpt_path, timestep_keyframe: TimestepKeyframeGroup=None, model=None):
|
||||
controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
|
||||
control = None
|
||||
# check if a non-vanilla ControlNet
|
||||
controlnet_type = ControlWeightType.DEFAULT
|
||||
has_controlnet_key = False
|
||||
has_motion_modules_key = False
|
||||
has_temporal_res_block_key = False
|
||||
for key in controlnet_data:
|
||||
# LLLite check
|
||||
if "lllite" in key:
|
||||
controlnet_type = ControlWeightType.CONTROLLLLITE
|
||||
break
|
||||
# SparseCtrl check
|
||||
elif "motion_modules" in key:
|
||||
has_motion_modules_key = True
|
||||
elif "controlnet" in key:
|
||||
has_controlnet_key = True
|
||||
# SVD-ControlNet check
|
||||
elif "temporal_res_block" in key:
|
||||
has_temporal_res_block_key = True
|
||||
if has_controlnet_key and has_motion_modules_key:
|
||||
controlnet_type = ControlWeightType.SPARSECTRL
|
||||
elif has_controlnet_key and has_temporal_res_block_key:
|
||||
controlnet_type = ControlWeightType.SVD_CONTROLNET
|
||||
|
||||
if controlnet_type != ControlWeightType.DEFAULT:
|
||||
if controlnet_type == ControlWeightType.CONTROLLLLITE:
|
||||
control = load_controllllite(ckpt_path, controlnet_data=controlnet_data, timestep_keyframe=timestep_keyframe)
|
||||
elif controlnet_type == ControlWeightType.SPARSECTRL:
|
||||
control = load_sparsectrl(ckpt_path, controlnet_data=controlnet_data, timestep_keyframe=timestep_keyframe, model=model)
|
||||
elif controlnet_type == ControlWeightType.SVD_CONTROLNET:
|
||||
control = load_svdcontrolnet(ckpt_path, controlnet_data=controlnet_data, timestep_keyframe=timestep_keyframe)
|
||||
#raise Exception(f"SVD-ControlNet is not supported yet!")
|
||||
#control = comfy_cn.load_controlnet(ckpt_path, model=model)
|
||||
# otherwise, load vanilla ControlNet
|
||||
else:
|
||||
try:
|
||||
# hacky way of getting load_torch_file in load_controlnet to use already-present controlnet_data and not redo loading
|
||||
orig_load_torch_file = comfy.utils.load_torch_file
|
||||
comfy.utils.load_torch_file = load_torch_file_with_dict_factory(controlnet_data, orig_load_torch_file)
|
||||
control = comfy_cn.load_controlnet(ckpt_path, model=model)
|
||||
finally:
|
||||
comfy.utils.load_torch_file = orig_load_torch_file
|
||||
return convert_to_advanced(control, timestep_keyframe=timestep_keyframe)
|
||||
|
||||
|
||||
def convert_to_advanced(control, timestep_keyframe: TimestepKeyframeGroup=None):
|
||||
# if already advanced, leave it be
|
||||
if is_advanced_controlnet(control):
|
||||
return control
|
||||
# if exactly ControlNet returned, transform it into ControlNetAdvanced
|
||||
if type(control) == ControlNet:
|
||||
return ControlNetAdvanced.from_vanilla(v=control, timestep_keyframe=timestep_keyframe)
|
||||
# if exactly ControlLora returned, transform it into ControlLoraAdvanced
|
||||
elif type(control) == ControlLora:
|
||||
return ControlLoraAdvanced.from_vanilla(v=control, timestep_keyframe=timestep_keyframe)
|
||||
# if T2IAdapter returned, transform it into T2IAdapterAdvanced
|
||||
elif isinstance(control, T2IAdapter):
|
||||
return T2IAdapterAdvanced.from_vanilla(v=control, timestep_keyframe=timestep_keyframe)
|
||||
# otherwise, leave it be - might be something I am not supporting yet
|
||||
return control
|
||||
|
||||
|
||||
def is_advanced_controlnet(input_object):
|
||||
return hasattr(input_object, "sub_idxs")
|
||||
|
||||
|
||||
def load_sparsectrl(ckpt_path: str, controlnet_data: dict[str, Tensor]=None, timestep_keyframe: TimestepKeyframeGroup=None, sparse_settings=SparseSettings.default(), model=None) -> SparseCtrlAdvanced:
|
||||
if controlnet_data is None:
|
||||
controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
|
||||
# first, separate out motion part from normal controlnet part and attempt to load that portion
|
||||
motion_data = {}
|
||||
for key in list(controlnet_data.keys()):
|
||||
if "temporal" in key:
|
||||
motion_data[key] = controlnet_data.pop(key)
|
||||
if len(motion_data) == 0:
|
||||
raise ValueError(f"No motion-related keys in '{ckpt_path}'; not a valid SparseCtrl model!")
|
||||
motion_wrapper: SparseCtrlMotionWrapper = SparseCtrlMotionWrapper(motion_data).to(comfy.model_management.unet_dtype())
|
||||
missing, unexpected = motion_wrapper.load_state_dict(motion_data)
|
||||
if len(missing) > 0 or len(unexpected) > 0:
|
||||
logger.info(f"SparseCtrlMotionWrapper: {missing}, {unexpected}")
|
||||
|
||||
# now, load as if it was a normal controlnet - mostly copied from comfy load_controlnet function
|
||||
controlnet_config = None
|
||||
is_diffusers = False
|
||||
use_simplified_conditioning_embedding = False
|
||||
if "controlnet_cond_embedding.conv_in.weight" in controlnet_data:
|
||||
is_diffusers = True
|
||||
if "controlnet_cond_embedding.weight" in controlnet_data:
|
||||
is_diffusers = True
|
||||
use_simplified_conditioning_embedding = True
|
||||
if is_diffusers: #diffusers format
|
||||
unet_dtype = comfy.model_management.unet_dtype()
|
||||
controlnet_config = comfy.model_detection.unet_config_from_diffusers_unet(controlnet_data, unet_dtype)
|
||||
diffusers_keys = comfy.utils.unet_to_diffusers(controlnet_config)
|
||||
diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"
|
||||
diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"
|
||||
|
||||
count = 0
|
||||
loop = True
|
||||
while loop:
|
||||
suffix = [".weight", ".bias"]
|
||||
for s in suffix:
|
||||
k_in = "controlnet_down_blocks.{}{}".format(count, s)
|
||||
k_out = "zero_convs.{}.0{}".format(count, s)
|
||||
if k_in not in controlnet_data:
|
||||
loop = False
|
||||
break
|
||||
diffusers_keys[k_in] = k_out
|
||||
count += 1
|
||||
# normal conditioning embedding
|
||||
if not use_simplified_conditioning_embedding:
|
||||
count = 0
|
||||
loop = True
|
||||
while loop:
|
||||
suffix = [".weight", ".bias"]
|
||||
for s in suffix:
|
||||
if count == 0:
|
||||
k_in = "controlnet_cond_embedding.conv_in{}".format(s)
|
||||
else:
|
||||
k_in = "controlnet_cond_embedding.blocks.{}{}".format(count - 1, s)
|
||||
k_out = "input_hint_block.{}{}".format(count * 2, s)
|
||||
if k_in not in controlnet_data:
|
||||
k_in = "controlnet_cond_embedding.conv_out{}".format(s)
|
||||
loop = False
|
||||
diffusers_keys[k_in] = k_out
|
||||
count += 1
|
||||
# simplified conditioning embedding
|
||||
else:
|
||||
count = 0
|
||||
suffix = [".weight", ".bias"]
|
||||
for s in suffix:
|
||||
k_in = "controlnet_cond_embedding{}".format(s)
|
||||
k_out = "input_hint_block.{}{}".format(count, s)
|
||||
diffusers_keys[k_in] = k_out
|
||||
|
||||
new_sd = {}
|
||||
for k in diffusers_keys:
|
||||
if k in controlnet_data:
|
||||
new_sd[diffusers_keys[k]] = controlnet_data.pop(k)
|
||||
|
||||
leftover_keys = controlnet_data.keys()
|
||||
if len(leftover_keys) > 0:
|
||||
logger.info("leftover keys:", leftover_keys)
|
||||
controlnet_data = new_sd
|
||||
|
||||
pth_key = 'control_model.zero_convs.0.0.weight'
|
||||
pth = False
|
||||
key = 'zero_convs.0.0.weight'
|
||||
if pth_key in controlnet_data:
|
||||
pth = True
|
||||
key = pth_key
|
||||
prefix = "control_model."
|
||||
elif key in controlnet_data:
|
||||
prefix = ""
|
||||
else:
|
||||
raise ValueError("The provided model is not a valid SparseCtrl model! [ErrorCode: HORSERADISH]")
|
||||
|
||||
if controlnet_config is None:
|
||||
unet_dtype = comfy.model_management.unet_dtype()
|
||||
controlnet_config = comfy.model_detection.model_config_from_unet(controlnet_data, prefix, unet_dtype, True).unet_config
|
||||
load_device = comfy.model_management.get_torch_device()
|
||||
manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)
|
||||
if manual_cast_dtype is not None:
|
||||
controlnet_config["operations"] = manual_cast_clean_groupnorm
|
||||
else:
|
||||
controlnet_config["operations"] = disable_weight_init_clean_groupnorm
|
||||
controlnet_config.pop("out_channels")
|
||||
# get proper hint channels
|
||||
if use_simplified_conditioning_embedding:
|
||||
controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
|
||||
controlnet_config["use_simplified_conditioning_embedding"] = use_simplified_conditioning_embedding
|
||||
else:
|
||||
controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
|
||||
controlnet_config["use_simplified_conditioning_embedding"] = use_simplified_conditioning_embedding
|
||||
control_model = SparseControlNet(**controlnet_config)
|
||||
|
||||
if pth:
|
||||
if 'difference' in controlnet_data:
|
||||
if model is not None:
|
||||
comfy.model_management.load_models_gpu([model])
|
||||
model_sd = model.model_state_dict()
|
||||
for x in controlnet_data:
|
||||
c_m = "control_model."
|
||||
if x.startswith(c_m):
|
||||
sd_key = "diffusion_model.{}".format(x[len(c_m):])
|
||||
if sd_key in model_sd:
|
||||
cd = controlnet_data[x]
|
||||
cd += model_sd[sd_key].type(cd.dtype).to(cd.device)
|
||||
else:
|
||||
logger.warning("WARNING: Loaded a diff SparseCtrl without a model. It will very likely not work.")
|
||||
|
||||
class WeightsLoader(torch.nn.Module):
|
||||
pass
|
||||
w = WeightsLoader()
|
||||
w.control_model = control_model
|
||||
missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
|
||||
else:
|
||||
missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
|
||||
if len(missing) > 0 or len(unexpected) > 0:
|
||||
logger.info(f"SparseCtrl ControlNet: {missing}, {unexpected}")
|
||||
|
||||
global_average_pooling = False
|
||||
filename = os.path.splitext(ckpt_path)[0]
|
||||
if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling
|
||||
global_average_pooling = True
|
||||
|
||||
# both motion portion and controlnet portions are loaded; bring them together if using motion model
|
||||
if sparse_settings.use_motion:
|
||||
motion_wrapper.inject(control_model)
|
||||
|
||||
control = SparseCtrlAdvanced(control_model, timestep_keyframes=timestep_keyframe, sparse_settings=sparse_settings, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
|
||||
return control
|
||||
|
||||
|
||||
def load_controllllite(ckpt_path: str, controlnet_data: dict[str, Tensor]=None, timestep_keyframe: TimestepKeyframeGroup=None):
|
||||
if controlnet_data is None:
|
||||
controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
|
||||
# adapted from https://github.com/kohya-ss/ControlNet-LLLite-ComfyUI
|
||||
# first, split weights for each module
|
||||
module_weights = {}
|
||||
for key, value in controlnet_data.items():
|
||||
fragments = key.split(".")
|
||||
module_name = fragments[0]
|
||||
weight_name = ".".join(fragments[1:])
|
||||
|
||||
if module_name not in module_weights:
|
||||
module_weights[module_name] = {}
|
||||
module_weights[module_name][weight_name] = value
|
||||
|
||||
# next, load each module
|
||||
modules = {}
|
||||
for module_name, weights in module_weights.items():
|
||||
# kohya planned to do something about how these should be chosen, so I'm not touching this
|
||||
# since I am not familiar with the logic for this
|
||||
if "conditioning1.4.weight" in weights:
|
||||
depth = 3
|
||||
elif weights["conditioning1.2.weight"].shape[-1] == 4:
|
||||
depth = 2
|
||||
else:
|
||||
depth = 1
|
||||
|
||||
module = LLLiteModule(
|
||||
name=module_name,
|
||||
is_conv2d=weights["down.0.weight"].ndim == 4,
|
||||
in_dim=weights["down.0.weight"].shape[1],
|
||||
depth=depth,
|
||||
cond_emb_dim=weights["conditioning1.0.weight"].shape[0] * 2,
|
||||
mlp_dim=weights["down.0.weight"].shape[0],
|
||||
)
|
||||
# load weights into module
|
||||
module.load_state_dict(weights)
|
||||
modules[module_name] = module
|
||||
if len(modules) == 1:
|
||||
module.is_first = True
|
||||
|
||||
#logger.info(f"loaded {ckpt_path} successfully, {len(modules)} modules")
|
||||
|
||||
patch_attn1 = LLLitePatch(modules=modules, patch_type=LLLitePatch.ATTN1)
|
||||
patch_attn2 = LLLitePatch(modules=modules, patch_type=LLLitePatch.ATTN2)
|
||||
control = ControlLLLiteAdvanced(patch_attn1=patch_attn1, patch_attn2=patch_attn2, timestep_keyframes=timestep_keyframe)
|
||||
return control
|
||||
|
||||
|
||||
def load_svdcontrolnet(ckpt_path: str, controlnet_data: dict[str, Tensor]=None, timestep_keyframe: TimestepKeyframeGroup=None, model=None):
|
||||
if controlnet_data is None:
|
||||
controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
|
||||
|
||||
controlnet_config = None
|
||||
if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format
|
||||
unet_dtype = comfy.model_management.unet_dtype()
|
||||
controlnet_config = svd_unet_config_from_diffusers_unet(controlnet_data, unet_dtype)
|
||||
diffusers_keys = svd_unet_to_diffusers(controlnet_config)
|
||||
diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"
|
||||
diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"
|
||||
|
||||
count = 0
|
||||
loop = True
|
||||
while loop:
|
||||
suffix = [".weight", ".bias"]
|
||||
for s in suffix:
|
||||
k_in = "controlnet_down_blocks.{}{}".format(count, s)
|
||||
k_out = "zero_convs.{}.0{}".format(count, s)
|
||||
if k_in not in controlnet_data:
|
||||
loop = False
|
||||
break
|
||||
diffusers_keys[k_in] = k_out
|
||||
count += 1
|
||||
|
||||
count = 0
|
||||
loop = True
|
||||
while loop:
|
||||
suffix = [".weight", ".bias"]
|
||||
for s in suffix:
|
||||
if count == 0:
|
||||
k_in = "controlnet_cond_embedding.conv_in{}".format(s)
|
||||
else:
|
||||
k_in = "controlnet_cond_embedding.blocks.{}{}".format(count - 1, s)
|
||||
k_out = "input_hint_block.{}{}".format(count * 2, s)
|
||||
if k_in not in controlnet_data:
|
||||
k_in = "controlnet_cond_embedding.conv_out{}".format(s)
|
||||
loop = False
|
||||
diffusers_keys[k_in] = k_out
|
||||
count += 1
|
||||
|
||||
new_sd = {}
|
||||
for k in diffusers_keys:
|
||||
if k in controlnet_data:
|
||||
new_sd[diffusers_keys[k]] = controlnet_data.pop(k)
|
||||
|
||||
leftover_keys = controlnet_data.keys()
|
||||
if len(leftover_keys) > 0:
|
||||
spatial_leftover_keys = []
|
||||
temporal_leftover_keys = []
|
||||
other_leftover_keys = []
|
||||
for key in leftover_keys:
|
||||
if "spatial" in key:
|
||||
spatial_leftover_keys.append(key)
|
||||
elif "temporal" in key:
|
||||
temporal_leftover_keys.append(key)
|
||||
else:
|
||||
other_leftover_keys.append(key)
|
||||
logger.warn(f"spatial_leftover_keys ({len(spatial_leftover_keys)}): {spatial_leftover_keys}")
|
||||
logger.warn(f"temporal_leftover_keys ({len(temporal_leftover_keys)}): {temporal_leftover_keys}")
|
||||
logger.warn(f"other_leftover_keys ({len(other_leftover_keys)}): {other_leftover_keys}")
|
||||
#print("leftover keys:", leftover_keys)
|
||||
controlnet_data = new_sd
|
||||
|
||||
pth_key = 'control_model.zero_convs.0.0.weight'
|
||||
pth = False
|
||||
key = 'zero_convs.0.0.weight'
|
||||
if pth_key in controlnet_data:
|
||||
pth = True
|
||||
key = pth_key
|
||||
prefix = "control_model."
|
||||
elif key in controlnet_data:
|
||||
prefix = ""
|
||||
else:
|
||||
raise ValueError("The provided model is not a valid SVD-ControlNet model! [ErrorCode: MUSTARD]")
|
||||
|
||||
if controlnet_config is None:
|
||||
unet_dtype = comfy.model_management.unet_dtype()
|
||||
controlnet_config = comfy.model_detection.model_config_from_unet(controlnet_data, prefix, unet_dtype, True).unet_config
|
||||
load_device = comfy.model_management.get_torch_device()
|
||||
manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)
|
||||
if manual_cast_dtype is not None:
|
||||
controlnet_config["operations"] = comfy.ops.manual_cast
|
||||
controlnet_config.pop("out_channels")
|
||||
controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
|
||||
control_model = SVDControlNet(**controlnet_config)
|
||||
|
||||
if pth:
|
||||
if 'difference' in controlnet_data:
|
||||
if model is not None:
|
||||
comfy.model_management.load_models_gpu([model])
|
||||
model_sd = model.model_state_dict()
|
||||
for x in controlnet_data:
|
||||
c_m = "control_model."
|
||||
if x.startswith(c_m):
|
||||
sd_key = "diffusion_model.{}".format(x[len(c_m):])
|
||||
if sd_key in model_sd:
|
||||
cd = controlnet_data[x]
|
||||
cd += model_sd[sd_key].type(cd.dtype).to(cd.device)
|
||||
else:
|
||||
print("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")
|
||||
|
||||
class WeightsLoader(torch.nn.Module):
|
||||
pass
|
||||
w = WeightsLoader()
|
||||
w.control_model = control_model
|
||||
missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
|
||||
else:
|
||||
missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
|
||||
if len(missing) > 0 or len(unexpected) > 0:
|
||||
logger.info(f"SVD-ControlNet: {missing}, {unexpected}")
|
||||
|
||||
global_average_pooling = False
|
||||
filename = os.path.splitext(ckpt_path)[0]
|
||||
if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling
|
||||
global_average_pooling = True
|
||||
|
||||
control = SVDControlNetAdvanced(control_model, timestep_keyframes=timestep_keyframe, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
|
||||
return control
|
||||
|
||||
+247
@@ -0,0 +1,247 @@
|
||||
# adapted from https://github.com/kohya-ss/ControlNet-LLLite-ComfyUI
|
||||
# basically, all the LLLite core code is from there, which I then combined with
|
||||
# Advanced-ControlNet features and QoL
|
||||
import math
|
||||
from typing import Union
|
||||
from torch import Tensor
|
||||
import torch
|
||||
import os
|
||||
|
||||
import comfy.utils
|
||||
from comfy.controlnet import ControlBase
|
||||
|
||||
from .logger import logger
|
||||
from .utils import AdvancedControlBase, deepcopy_with_sharing, prepare_mask_batch
|
||||
|
||||
|
||||
def extra_options_to_module_prefix(extra_options):
|
||||
# extra_options = {'transformer_index': 2, 'block_index': 8, 'original_shape': [2, 4, 128, 128], 'block': ('input', 7), 'n_heads': 20, 'dim_head': 64}
|
||||
|
||||
# block is: [('input', 4), ('input', 5), ('input', 7), ('input', 8), ('middle', 0),
|
||||
# ('output', 0), ('output', 1), ('output', 2), ('output', 3), ('output', 4), ('output', 5)]
|
||||
# transformer_index is: [0, 1, 2, 3, 4, 5, 6, 7, 8], for each block
|
||||
# block_index is: 0-1 or 0-9, depends on the block
|
||||
# input 7 and 8, middle has 10 blocks
|
||||
|
||||
# make module name from extra_options
|
||||
block = extra_options["block"]
|
||||
block_index = extra_options["block_index"]
|
||||
if block[0] == "input":
|
||||
module_pfx = f"lllite_unet_input_blocks_{block[1]}_1_transformer_blocks_{block_index}"
|
||||
elif block[0] == "middle":
|
||||
module_pfx = f"lllite_unet_middle_block_1_transformer_blocks_{block_index}"
|
||||
elif block[0] == "output":
|
||||
module_pfx = f"lllite_unet_output_blocks_{block[1]}_1_transformer_blocks_{block_index}"
|
||||
else:
|
||||
raise Exception(f"ControlLLLite: invalid block name '{block[0]}'. Expected 'input', 'middle', or 'output'.")
|
||||
return module_pfx
|
||||
|
||||
|
||||
class LLLitePatch:
|
||||
ATTN1 = "attn1"
|
||||
ATTN2 = "attn2"
|
||||
def __init__(self, modules: dict[str, 'LLLiteModule'], patch_type: str, control: Union[AdvancedControlBase, ControlBase]=None):
|
||||
self.modules = modules
|
||||
self.control = control
|
||||
self.patch_type = patch_type
|
||||
#logger.error(f"create LLLitePatch: {id(self)},{control}")
|
||||
|
||||
def __call__(self, q, k, v, extra_options):
|
||||
#logger.error(f"in __call__: {id(self)}")
|
||||
# determine if have anything to run
|
||||
if self.control.timestep_range is not None:
|
||||
# it turns out comparing single-value tensors to floats is extremely slow
|
||||
# a: Tensor = extra_options["sigmas"][0]
|
||||
if self.control.t > self.control.timestep_range[0] or self.control.t < self.control.timestep_range[1]:
|
||||
return q, k, v
|
||||
|
||||
module_pfx = extra_options_to_module_prefix(extra_options)
|
||||
|
||||
is_attn1 = q.shape[-1] == k.shape[-1] # self attention
|
||||
if is_attn1:
|
||||
module_pfx = module_pfx + "_attn1"
|
||||
else:
|
||||
module_pfx = module_pfx + "_attn2"
|
||||
|
||||
module_pfx_to_q = module_pfx + "_to_q"
|
||||
module_pfx_to_k = module_pfx + "_to_k"
|
||||
module_pfx_to_v = module_pfx + "_to_v"
|
||||
|
||||
if module_pfx_to_q in self.modules:
|
||||
q = q + self.modules[module_pfx_to_q](q, self.control)
|
||||
if module_pfx_to_k in self.modules:
|
||||
k = k + self.modules[module_pfx_to_k](k, self.control)
|
||||
if module_pfx_to_v in self.modules:
|
||||
v = v + self.modules[module_pfx_to_v](v, self.control)
|
||||
|
||||
return q, k, v
|
||||
|
||||
def to(self, device):
|
||||
#logger.info(f"to... has control? {self.control}")
|
||||
for d in self.modules.keys():
|
||||
self.modules[d] = self.modules[d].to(device)
|
||||
return self
|
||||
|
||||
def set_control(self, control: Union[AdvancedControlBase, ControlBase]) -> 'LLLitePatch':
|
||||
self.control = control
|
||||
return self
|
||||
#logger.error(f"set control for LLLitePatch: {id(self)}, cn: {id(control)}")
|
||||
|
||||
def clone_with_control(self, control: AdvancedControlBase):
|
||||
#logger.error(f"clone-set control for LLLitePatch: {id(self)},{id(control)}")
|
||||
return LLLitePatch(self.modules, self.patch_type, control)
|
||||
|
||||
def cleanup(self):
|
||||
#total_cleaned = 0
|
||||
for module in self.modules.values():
|
||||
module.cleanup()
|
||||
# total_cleaned += 1
|
||||
#logger.info(f"cleaned modules: {total_cleaned}, {id(self)}")
|
||||
#logger.error(f"cleanup LLLitePatch: {id(self)}")
|
||||
|
||||
# make sure deepcopy does not copy control, and deepcopied LLLitePatch should be assigned to control
|
||||
def __deepcopy__(self, memo):
|
||||
self.cleanup()
|
||||
to_return: LLLitePatch = deepcopy_with_sharing(self, shared_attribute_names = ['control'], memo=memo)
|
||||
#logger.warn(f"patch {id(self)} turned into {id(to_return)}")
|
||||
try:
|
||||
if self.patch_type == self.ATTN1:
|
||||
to_return.control.patch_attn1 = to_return
|
||||
elif self.patch_type == self.ATTN2:
|
||||
to_return.control.patch_attn2 = to_return
|
||||
except Exception:
|
||||
pass
|
||||
return to_return
|
||||
|
||||
|
||||
# TODO: use comfy.ops to support fp8 properly
|
||||
class LLLiteModule(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
is_conv2d: bool,
|
||||
in_dim: int,
|
||||
depth: int,
|
||||
cond_emb_dim: int,
|
||||
mlp_dim: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.name = name
|
||||
self.is_conv2d = is_conv2d
|
||||
self.is_first = False
|
||||
|
||||
modules = []
|
||||
modules.append(torch.nn.Conv2d(3, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0)) # to latent (from VAE) size*2
|
||||
if depth == 1:
|
||||
modules.append(torch.nn.ReLU(inplace=True))
|
||||
modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0))
|
||||
elif depth == 2:
|
||||
modules.append(torch.nn.ReLU(inplace=True))
|
||||
modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=4, stride=4, padding=0))
|
||||
elif depth == 3:
|
||||
# kernel size 8 is too large, so set it to 4
|
||||
modules.append(torch.nn.ReLU(inplace=True))
|
||||
modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0))
|
||||
modules.append(torch.nn.ReLU(inplace=True))
|
||||
modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0))
|
||||
|
||||
self.conditioning1 = torch.nn.Sequential(*modules)
|
||||
|
||||
if self.is_conv2d:
|
||||
self.down = torch.nn.Sequential(
|
||||
torch.nn.Conv2d(in_dim, mlp_dim, kernel_size=1, stride=1, padding=0),
|
||||
torch.nn.ReLU(inplace=True),
|
||||
)
|
||||
self.mid = torch.nn.Sequential(
|
||||
torch.nn.Conv2d(mlp_dim + cond_emb_dim, mlp_dim, kernel_size=1, stride=1, padding=0),
|
||||
torch.nn.ReLU(inplace=True),
|
||||
)
|
||||
self.up = torch.nn.Sequential(
|
||||
torch.nn.Conv2d(mlp_dim, in_dim, kernel_size=1, stride=1, padding=0),
|
||||
)
|
||||
else:
|
||||
self.down = torch.nn.Sequential(
|
||||
torch.nn.Linear(in_dim, mlp_dim),
|
||||
torch.nn.ReLU(inplace=True),
|
||||
)
|
||||
self.mid = torch.nn.Sequential(
|
||||
torch.nn.Linear(mlp_dim + cond_emb_dim, mlp_dim),
|
||||
torch.nn.ReLU(inplace=True),
|
||||
)
|
||||
self.up = torch.nn.Sequential(
|
||||
torch.nn.Linear(mlp_dim, in_dim),
|
||||
)
|
||||
|
||||
self.depth = depth
|
||||
self.cond_emb = None
|
||||
self.cx_shape = None
|
||||
self.prev_batch = 0
|
||||
self.prev_sub_idxs = None
|
||||
|
||||
def cleanup(self):
|
||||
del self.cond_emb
|
||||
self.cond_emb = None
|
||||
self.cx_shape = None
|
||||
self.prev_batch = 0
|
||||
self.prev_sub_idxs = None
|
||||
|
||||
def forward(self, x: Tensor, control: Union[AdvancedControlBase, ControlBase]):
|
||||
mask = None
|
||||
mask_tk = None
|
||||
#logger.info(x.shape)
|
||||
if self.cond_emb is None or control.sub_idxs != self.prev_sub_idxs or x.shape[0] != self.prev_batch:
|
||||
# print(f"cond_emb is None, {self.name}")
|
||||
cond_hint = control.cond_hint.to(x.device, dtype=x.dtype)
|
||||
if control.latent_dims_div2 is not None and x.shape[-1] != 1280:
|
||||
cond_hint = comfy.utils.common_upscale(cond_hint, control.latent_dims_div2[0] * 8, control.latent_dims_div2[1] * 8, 'nearest-exact', "center").to(x.device, dtype=x.dtype)
|
||||
elif control.latent_dims_div4 is not None and x.shape[-1] == 1280:
|
||||
cond_hint = comfy.utils.common_upscale(cond_hint, control.latent_dims_div4[0] * 8, control.latent_dims_div4[1] * 8, 'nearest-exact', "center").to(x.device, dtype=x.dtype)
|
||||
cx = self.conditioning1(cond_hint)
|
||||
self.cx_shape = cx.shape
|
||||
if not self.is_conv2d:
|
||||
# reshape / b,c,h,w -> b,h*w,c
|
||||
n, c, h, w = cx.shape
|
||||
cx = cx.view(n, c, h * w).permute(0, 2, 1)
|
||||
self.cond_emb = cx
|
||||
# save prev values
|
||||
self.prev_batch = x.shape[0]
|
||||
self.prev_sub_idxs = control.sub_idxs
|
||||
|
||||
cx: torch.Tensor = self.cond_emb
|
||||
# print(f"forward {self.name}, {cx.shape}, {x.shape}")
|
||||
|
||||
# TODO: make masks work for conv2d (could not find any ControlLLLites at this time that use them)
|
||||
# create masks
|
||||
if not self.is_conv2d:
|
||||
n, c, h, w = self.cx_shape
|
||||
if control.mask_cond_hint is not None:
|
||||
mask = prepare_mask_batch(control.mask_cond_hint, (1, 1, h, w)).to(cx.dtype)
|
||||
mask = mask.view(mask.shape[0], 1, h * w).permute(0, 2, 1)
|
||||
if control.tk_mask_cond_hint is not None:
|
||||
mask_tk = prepare_mask_batch(control.mask_cond_hint, (1, 1, h, w)).to(cx.dtype)
|
||||
mask_tk = mask_tk.view(mask_tk.shape[0], 1, h * w).permute(0, 2, 1)
|
||||
|
||||
# x in uncond/cond doubles batch size
|
||||
if x.shape[0] != cx.shape[0]:
|
||||
if self.is_conv2d:
|
||||
cx = cx.repeat(x.shape[0] // cx.shape[0], 1, 1, 1)
|
||||
else:
|
||||
# print("x.shape[0] != cx.shape[0]", x.shape[0], cx.shape[0])
|
||||
cx = cx.repeat(x.shape[0] // cx.shape[0], 1, 1)
|
||||
if mask is not None:
|
||||
mask = mask.repeat(x.shape[0] // mask.shape[0], 1, 1)
|
||||
if mask_tk is not None:
|
||||
mask_tk = mask_tk.repeat(x.shape[0] // mask_tk.shape[0], 1, 1)
|
||||
|
||||
if mask is None:
|
||||
mask = 1.0
|
||||
elif mask_tk is not None:
|
||||
mask = mask * mask_tk
|
||||
|
||||
#logger.info(f"cs: {cx.shape}, x: {x.shape}, is_conv2d: {self.is_conv2d}")
|
||||
cx = torch.cat([cx, self.down(x)], dim=1 if self.is_conv2d else 2)
|
||||
cx = self.mid(cx)
|
||||
cx = self.up(cx)
|
||||
if control.latent_keyframes is not None:
|
||||
cx = cx * control.calc_latent_keyframe_mults(x=cx, batched_number=control.batched_number)
|
||||
return cx * mask * control.strength * control.current_timestep_keyframe.strength
|
||||
+907
@@ -0,0 +1,907 @@
|
||||
#taken from: https://github.com/lllyasviel/ControlNet
|
||||
#and modified
|
||||
#and then taken from comfy/cldm/cldm.py and modified again
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
import math
|
||||
import numpy as np
|
||||
from typing import Iterable, Union
|
||||
import torch
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
from einops import rearrange, repeat
|
||||
|
||||
from comfy.ldm.modules.diffusionmodules.util import (
|
||||
zero_module,
|
||||
timestep_embedding,
|
||||
)
|
||||
|
||||
from comfy.cli_args import args
|
||||
from comfy.cldm.cldm import ControlNet as ControlNetCLDM
|
||||
from comfy.ldm.modules.attention import SpatialTransformer
|
||||
from comfy.ldm.modules.attention import attention_basic, attention_pytorch, attention_split, attention_sub_quad, default
|
||||
from comfy.ldm.modules.attention import FeedForward, SpatialTransformer
|
||||
from comfy.ldm.modules.diffusionmodules.openaimodel import TimestepEmbedSequential, ResBlock, Downsample
|
||||
from comfy.controlnet import broadcast_image_to
|
||||
from comfy.utils import repeat_to_batch_size
|
||||
import comfy.ops
|
||||
import comfy.model_management
|
||||
|
||||
from .utils import TimestepKeyframeGroup, disable_weight_init_clean_groupnorm, prepare_mask_batch
|
||||
|
||||
|
||||
# until xformers bug is fixed, do not use xformers for VersatileAttention! TODO: change this when fix is out
|
||||
# logic for choosing optimized_attention method taken from comfy/ldm/modules/attention.py
|
||||
optimized_attention_mm = attention_basic
|
||||
if comfy.model_management.xformers_enabled():
|
||||
pass
|
||||
#optimized_attention_mm = attention_xformers
|
||||
if comfy.model_management.pytorch_attention_enabled():
|
||||
optimized_attention_mm = attention_pytorch
|
||||
else:
|
||||
if args.use_split_cross_attention:
|
||||
optimized_attention_mm = attention_split
|
||||
else:
|
||||
optimized_attention_mm = attention_sub_quad
|
||||
|
||||
|
||||
class SparseControlNet(ControlNetCLDM):
|
||||
def __init__(self, *args,**kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
hint_channels = kwargs.get("hint_channels")
|
||||
operations: disable_weight_init_clean_groupnorm = kwargs.get("operations", disable_weight_init_clean_groupnorm)
|
||||
device = kwargs.get("device", None)
|
||||
self.use_simplified_conditioning_embedding = kwargs.get("use_simplified_conditioning_embedding", False)
|
||||
if self.use_simplified_conditioning_embedding:
|
||||
self.input_hint_block = TimestepEmbedSequential(
|
||||
zero_module(operations.conv_nd(self.dims, hint_channels, self.model_channels, 3, padding=1, dtype=self.dtype, device=device)),
|
||||
)
|
||||
self.motion_holder: MotionWrapperHolder = None
|
||||
|
||||
def set_actual_length(self, actual_length: int, full_length: int):
|
||||
if self.motion_holder is not None:
|
||||
self.motion_holder.motion_wrapper.set_video_length(video_length=actual_length, full_length=full_length)
|
||||
|
||||
def forward(self, x: Tensor, hint: Tensor, timesteps, context, y=None, **kwargs):
|
||||
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
|
||||
emb = self.time_embed(t_emb)
|
||||
|
||||
# SparseCtrl sets noisy input to zeros
|
||||
x = torch.zeros_like(x)
|
||||
guided_hint = self.input_hint_block(hint, emb, context)
|
||||
|
||||
outs = []
|
||||
|
||||
hs = []
|
||||
if self.num_classes is not None:
|
||||
assert y.shape[0] == x.shape[0]
|
||||
emb = emb + self.label_emb(y)
|
||||
|
||||
h = x
|
||||
for module, zero_conv in zip(self.input_blocks, self.zero_convs):
|
||||
if guided_hint is not None:
|
||||
h = module(h, emb, context)
|
||||
h += guided_hint
|
||||
guided_hint = None
|
||||
else:
|
||||
h = module(h, emb, context)
|
||||
outs.append(zero_conv(h, emb, context))
|
||||
|
||||
h = self.middle_block(h, emb, context)
|
||||
outs.append(self.middle_block_out(h, emb, context))
|
||||
|
||||
return outs
|
||||
|
||||
|
||||
class PreprocSparseRGBWrapper:
|
||||
error_msg = "Invalid use of RGB SparseCtrl output. The output of RGB SparseCtrl preprocessor is NOT a usual image, but a latent pretending to be an image - you must connect the output directly to an Apply ControlNet node (advanced or otherwise). It cannot be used for anything else that accepts IMAGE input."
|
||||
def __init__(self, condhint: Tensor):
|
||||
self.condhint = condhint
|
||||
|
||||
def movedim(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
def __getattr__(self, *args, **kwargs):
|
||||
raise AttributeError(self.error_msg)
|
||||
|
||||
def __setattr__(self, name, value):
|
||||
if name != "condhint":
|
||||
raise AttributeError(self.error_msg)
|
||||
super().__setattr__(name, value)
|
||||
|
||||
def __iter__(self, *args, **kwargs):
|
||||
raise AttributeError(self.error_msg)
|
||||
|
||||
def __next__(self, *args, **kwargs):
|
||||
raise AttributeError(self.error_msg)
|
||||
|
||||
def __len__(self, *args, **kwargs):
|
||||
raise AttributeError(self.error_msg)
|
||||
|
||||
def __getitem__(self, *args, **kwargs):
|
||||
raise AttributeError(self.error_msg)
|
||||
|
||||
def __setitem__(self, *args, **kwargs):
|
||||
raise AttributeError(self.error_msg)
|
||||
|
||||
|
||||
class SparseSettings:
|
||||
def __init__(self, sparse_method: 'SparseMethod', use_motion: bool=True, motion_strength=1.0, motion_scale=1.0, merged=False):
|
||||
self.sparse_method = sparse_method
|
||||
self.use_motion = use_motion
|
||||
self.motion_strength = motion_strength
|
||||
self.motion_scale = motion_scale
|
||||
self.merged = merged
|
||||
|
||||
@classmethod
|
||||
def default(cls):
|
||||
return SparseSettings(sparse_method=SparseSpreadMethod(), use_motion=True)
|
||||
|
||||
|
||||
class SparseMethod(ABC):
|
||||
SPREAD = "spread"
|
||||
INDEX = "index"
|
||||
def __init__(self, method: str):
|
||||
self.method = method
|
||||
|
||||
@abstractmethod
|
||||
def get_indexes(self, hint_length: int, full_length: int) -> list[int]:
|
||||
pass
|
||||
|
||||
|
||||
class SparseSpreadMethod(SparseMethod):
|
||||
UNIFORM = "uniform"
|
||||
STARTING = "starting"
|
||||
ENDING = "ending"
|
||||
CENTER = "center"
|
||||
|
||||
LIST = [UNIFORM, STARTING, ENDING, CENTER]
|
||||
|
||||
def __init__(self, spread=UNIFORM):
|
||||
super().__init__(self.SPREAD)
|
||||
self.spread = spread
|
||||
|
||||
def get_indexes(self, hint_length: int, full_length: int) -> list[int]:
|
||||
# if hint_length >= full_length, limit hints to full_length
|
||||
if hint_length >= full_length:
|
||||
return list(range(full_length))
|
||||
# handle special case of 1 hint image
|
||||
if hint_length == 1:
|
||||
if self.spread in [self.UNIFORM, self.STARTING]:
|
||||
return [0]
|
||||
elif self.spread == self.ENDING:
|
||||
return [full_length-1]
|
||||
elif self.spread == self.CENTER:
|
||||
# return second (of three) values as the center
|
||||
return [np.linspace(0, full_length-1, 3, endpoint=True, dtype=int)[1]]
|
||||
else:
|
||||
raise ValueError(f"Unrecognized spread: {self.spread}")
|
||||
# otherwise, handle other cases
|
||||
if self.spread == self.UNIFORM:
|
||||
return list(np.linspace(0, full_length-1, hint_length, endpoint=True, dtype=int))
|
||||
elif self.spread == self.STARTING:
|
||||
# make split 1 larger, remove last element
|
||||
return list(np.linspace(0, full_length-1, hint_length+1, endpoint=True, dtype=int))[:-1]
|
||||
elif self.spread == self.ENDING:
|
||||
# make split 1 larger, remove first element
|
||||
return list(np.linspace(0, full_length-1, hint_length+1, endpoint=True, dtype=int))[1:]
|
||||
elif self.spread == self.CENTER:
|
||||
# if hint length is not 3 greater than full length, do STARTING behavior
|
||||
if full_length-hint_length < 3:
|
||||
return list(np.linspace(0, full_length-1, hint_length+1, endpoint=True, dtype=int))[:-1]
|
||||
# otherwise, get linspace of 2 greater than needed, then cut off first and last
|
||||
return list(np.linspace(0, full_length-1, hint_length+2, endpoint=True, dtype=int))[1:-1]
|
||||
return ValueError(f"Unrecognized spread: {self.spread}")
|
||||
|
||||
|
||||
class SparseIndexMethod(SparseMethod):
|
||||
def __init__(self, idxs: list[int]):
|
||||
super().__init__(self.INDEX)
|
||||
self.idxs = idxs
|
||||
|
||||
def get_indexes(self, hint_length: int, full_length: int) -> list[int]:
|
||||
orig_hint_length = hint_length
|
||||
if hint_length > full_length:
|
||||
hint_length = full_length
|
||||
# if idxs is less than hint_length, throw error
|
||||
if len(self.idxs) < hint_length:
|
||||
err_msg = f"There are not enough indexes ({len(self.idxs)}) provided to fit the usable {hint_length} input images."
|
||||
if orig_hint_length != hint_length:
|
||||
err_msg = f"{err_msg} (original input images: {orig_hint_length})"
|
||||
raise ValueError(err_msg)
|
||||
# cap idxs to hint_length
|
||||
idxs = self.idxs[:hint_length]
|
||||
new_idxs = []
|
||||
real_idxs = set()
|
||||
for idx in idxs:
|
||||
if idx < 0:
|
||||
real_idx = full_length+idx
|
||||
if real_idx in real_idxs:
|
||||
raise ValueError(f"Index '{idx}' maps to '{real_idx}' and is duplicate - indexes in Sparse Index Method must be unique.")
|
||||
else:
|
||||
real_idx = idx
|
||||
if real_idx in real_idxs:
|
||||
raise ValueError(f"Index '{idx}' is duplicate (or a negative index is equivalent) - indexes in Sparse Index Method must be unique.")
|
||||
real_idxs.add(real_idx)
|
||||
new_idxs.append(real_idx)
|
||||
return new_idxs
|
||||
|
||||
|
||||
#########################################
|
||||
# motion-related portion of controlnet
|
||||
class BlockType:
|
||||
UP = "up"
|
||||
DOWN = "down"
|
||||
MID = "mid"
|
||||
|
||||
def get_down_block_max(mm_state_dict: dict[str, Tensor]) -> int:
|
||||
return get_block_max(mm_state_dict, "down_blocks")
|
||||
|
||||
def get_up_block_max(mm_state_dict: dict[str, Tensor]) -> int:
|
||||
return get_block_max(mm_state_dict, "up_blocks")
|
||||
|
||||
def get_block_max(mm_state_dict: dict[str, Tensor], block_name: str) -> int:
|
||||
# keep track of biggest down_block count in module
|
||||
biggest_block = -1
|
||||
for key in mm_state_dict.keys():
|
||||
if block_name in key:
|
||||
try:
|
||||
block_int = key.split(".")[1]
|
||||
block_num = int(block_int)
|
||||
if block_num > biggest_block:
|
||||
biggest_block = block_num
|
||||
except ValueError:
|
||||
pass
|
||||
return biggest_block
|
||||
|
||||
def has_mid_block(mm_state_dict: dict[str, Tensor]):
|
||||
# check if keys contain mid_block
|
||||
for key in mm_state_dict.keys():
|
||||
if key.startswith("mid_block."):
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_position_encoding_max_len(mm_state_dict: dict[str, Tensor], mm_name: str=None) -> int:
|
||||
# use pos_encoder.pe entries to determine max length - [1, {max_length}, {320|640|1280}]
|
||||
for key in mm_state_dict.keys():
|
||||
if key.endswith("pos_encoder.pe"):
|
||||
return mm_state_dict[key].size(1) # get middle dim
|
||||
raise ValueError(f"No pos_encoder.pe found in SparseCtrl state_dict - {mm_name} is not a valid SparseCtrl model!")
|
||||
|
||||
|
||||
class MotionWrapperHolder:
|
||||
def __init__(self, motion_wrapper: 'SparseCtrlMotionWrapper'):
|
||||
self.motion_wrapper = motion_wrapper
|
||||
|
||||
|
||||
class SparseCtrlMotionWrapper(nn.Module):
|
||||
def __init__(self, mm_state_dict: dict[str, Tensor]):
|
||||
super().__init__()
|
||||
self.down_blocks: Iterable[MotionModule] = None
|
||||
self.up_blocks: Iterable[MotionModule] = None
|
||||
self.mid_block: MotionModule = None
|
||||
self.encoding_max_len = get_position_encoding_max_len(mm_state_dict, "")
|
||||
layer_channels = (320, 640, 1280, 1280)
|
||||
if get_down_block_max(mm_state_dict) > -1:
|
||||
self.down_blocks = nn.ModuleList([])
|
||||
for c in layer_channels:
|
||||
self.down_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.DOWN))
|
||||
if get_up_block_max(mm_state_dict) > -1:
|
||||
self.up_blocks = nn.ModuleList([])
|
||||
for c in reversed(layer_channels):
|
||||
self.up_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.UP))
|
||||
if has_mid_block(mm_state_dict):
|
||||
self.mid_block = MotionModule(1280, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.MID)
|
||||
|
||||
def inject(self, unet: SparseControlNet):
|
||||
# inject input (down) blocks
|
||||
self._inject(unet.input_blocks, self.down_blocks)
|
||||
# inject mid block, if present
|
||||
if self.mid_block is not None:
|
||||
self._inject([unet.middle_block], [self.mid_block])
|
||||
unet.motion_holder = MotionWrapperHolder(self)
|
||||
|
||||
def _inject(self, unet_blocks: nn.ModuleList, mm_blocks: nn.ModuleList):
|
||||
# Rules for injection:
|
||||
# For each component list in a unet block:
|
||||
# if SpatialTransformer exists in list, place next block after last occurrence
|
||||
# elif ResBlock exists in list, place next block after first occurrence
|
||||
# else don't place block
|
||||
injection_count = 0
|
||||
unet_idx = 0
|
||||
# details about blocks passed in
|
||||
per_block = len(mm_blocks[0].motion_modules)
|
||||
injection_goal = len(mm_blocks) * per_block
|
||||
# only stop injecting when modules exhausted
|
||||
while injection_count < injection_goal:
|
||||
# figure out which VanillaTemporalModule from mm to inject
|
||||
mm_blk_idx, mm_vtm_idx = injection_count // per_block, injection_count % per_block
|
||||
# figure out layout of unet block components
|
||||
st_idx = -1 # SpatialTransformer index
|
||||
res_idx = -1 # first ResBlock index
|
||||
# first, figure out indeces of relevant blocks
|
||||
for idx, component in enumerate(unet_blocks[unet_idx]):
|
||||
if type(component) == SpatialTransformer:
|
||||
st_idx = idx
|
||||
elif type(component).__name__ == "ResBlock" and res_idx < 0:
|
||||
res_idx = idx
|
||||
# if SpatialTransformer exists, inject right after
|
||||
if st_idx >= 0:
|
||||
unet_blocks[unet_idx].insert(st_idx+1, mm_blocks[mm_blk_idx].motion_modules[mm_vtm_idx])
|
||||
injection_count += 1
|
||||
# otherwise, if only ResBlock exists, inject right after
|
||||
elif res_idx >= 0:
|
||||
unet_blocks[unet_idx].insert(res_idx+1, mm_blocks[mm_blk_idx].motion_modules[mm_vtm_idx])
|
||||
injection_count += 1
|
||||
# increment unet_idx
|
||||
unet_idx += 1
|
||||
|
||||
def eject(self, unet: SparseControlNet):
|
||||
# remove from input blocks (downblocks)
|
||||
self._eject(unet.input_blocks)
|
||||
# remove from middle block (encapsulate in list to make compatible)
|
||||
self._eject([unet.middle_block])
|
||||
del unet.motion_holder
|
||||
unet.motion_holder = None
|
||||
|
||||
def _eject(self, unet_blocks: nn.ModuleList):
|
||||
# eject all VanillaTemporalModule objects from all blocks
|
||||
for block in unet_blocks:
|
||||
idx_to_pop = []
|
||||
for idx, component in enumerate(block):
|
||||
if type(component) == VanillaTemporalModule:
|
||||
idx_to_pop.append(idx)
|
||||
# pop in backwards order, as to not disturb what the indeces refer to
|
||||
for idx in sorted(idx_to_pop, reverse=True):
|
||||
block.pop(idx)
|
||||
|
||||
def set_video_length(self, video_length: int, full_length: int):
|
||||
self.AD_video_length = video_length
|
||||
if self.down_blocks is not None:
|
||||
for block in self.down_blocks:
|
||||
block.set_video_length(video_length, full_length)
|
||||
if self.up_blocks is not None:
|
||||
for block in self.up_blocks:
|
||||
block.set_video_length(video_length, full_length)
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.set_video_length(video_length, full_length)
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
if self.down_blocks is not None:
|
||||
for block in self.down_blocks:
|
||||
block.set_scale_multiplier(multiplier)
|
||||
if self.up_blocks is not None:
|
||||
for block in self.up_blocks:
|
||||
block.set_scale_multiplier(multiplier)
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.set_scale_multiplier(multiplier)
|
||||
|
||||
def set_strength(self, strength: float):
|
||||
if self.down_blocks is not None:
|
||||
for block in self.down_blocks:
|
||||
block.set_strength(strength)
|
||||
if self.up_blocks is not None:
|
||||
for block in self.up_blocks:
|
||||
block.set_strength(strength)
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.set_strength(strength)
|
||||
|
||||
def reset_temp_vars(self):
|
||||
if self.down_blocks is not None:
|
||||
for block in self.down_blocks:
|
||||
block.reset_temp_vars()
|
||||
if self.up_blocks is not None:
|
||||
for block in self.up_blocks:
|
||||
block.reset_temp_vars()
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.reset_temp_vars()
|
||||
|
||||
def reset_scale_multiplier(self):
|
||||
self.set_scale_multiplier(None)
|
||||
|
||||
def reset(self):
|
||||
self.reset_scale_multiplier()
|
||||
self.reset_temp_vars()
|
||||
|
||||
|
||||
class MotionModule(nn.Module):
|
||||
def __init__(self, in_channels, temporal_position_encoding_max_len=24, block_type: str=BlockType.DOWN):
|
||||
super().__init__()
|
||||
if block_type == BlockType.MID:
|
||||
# mid blocks contain only a single VanillaTemporalModule
|
||||
self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList([get_motion_module(in_channels, temporal_position_encoding_max_len)])
|
||||
else:
|
||||
# down blocks contain two VanillaTemporalModules
|
||||
self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList(
|
||||
[
|
||||
get_motion_module(in_channels, temporal_position_encoding_max_len),
|
||||
get_motion_module(in_channels, temporal_position_encoding_max_len)
|
||||
]
|
||||
)
|
||||
# up blocks contain one additional VanillaTemporalModule
|
||||
if block_type == BlockType.UP:
|
||||
self.motion_modules.append(get_motion_module(in_channels, temporal_position_encoding_max_len))
|
||||
|
||||
def set_video_length(self, video_length: int, full_length: int):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_video_length(video_length, full_length)
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_scale_multiplier(multiplier)
|
||||
|
||||
def set_masks(self, masks: Tensor, min_val: float, max_val: float):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_masks(masks, min_val, max_val)
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_sub_idxs(sub_idxs)
|
||||
|
||||
def set_strength(self, strength: float):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_strength(strength)
|
||||
|
||||
def reset_temp_vars(self):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.reset_temp_vars()
|
||||
|
||||
|
||||
def get_motion_module(in_channels, temporal_position_encoding_max_len):
|
||||
# unlike normal AD, there is only one attention block expected in SparseCtrl models
|
||||
return VanillaTemporalModule(in_channels=in_channels, attention_block_types=("Temporal_Self",), temporal_position_encoding_max_len=temporal_position_encoding_max_len)
|
||||
|
||||
|
||||
class VanillaTemporalModule(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
num_attention_heads=8,
|
||||
num_transformer_block=1,
|
||||
attention_block_types=("Temporal_Self", "Temporal_Self"),
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=True,
|
||||
temporal_position_encoding_max_len=24,
|
||||
temporal_attention_dim_div=1,
|
||||
zero_initialize=True,
|
||||
):
|
||||
super().__init__()
|
||||
self.strength = 1.0
|
||||
self.temporal_transformer = TemporalTransformer3DModel(
|
||||
in_channels=in_channels,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=in_channels
|
||||
// num_attention_heads
|
||||
// temporal_attention_dim_div,
|
||||
num_layers=num_transformer_block,
|
||||
attention_block_types=attention_block_types,
|
||||
cross_frame_attention_mode=cross_frame_attention_mode,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
)
|
||||
|
||||
if zero_initialize:
|
||||
self.temporal_transformer.proj_out = zero_module(
|
||||
self.temporal_transformer.proj_out
|
||||
)
|
||||
|
||||
def set_video_length(self, video_length: int, full_length: int):
|
||||
self.temporal_transformer.set_video_length(video_length, full_length)
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
self.temporal_transformer.set_scale_multiplier(multiplier)
|
||||
|
||||
def set_masks(self, masks: Tensor, min_val: float, max_val: float):
|
||||
self.temporal_transformer.set_masks(masks, min_val, max_val)
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
self.temporal_transformer.set_sub_idxs(sub_idxs)
|
||||
|
||||
def set_strength(self, strength: float):
|
||||
self.strength = strength
|
||||
|
||||
def reset_temp_vars(self):
|
||||
self.set_strength(1.0)
|
||||
self.temporal_transformer.reset_temp_vars()
|
||||
|
||||
def forward(self, input_tensor, encoder_hidden_states=None, attention_mask=None):
|
||||
if math.isclose(self.strength, 1.0):
|
||||
return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask)
|
||||
elif math.isclose(self.strength, 0.0):
|
||||
return input_tensor
|
||||
elif self.strength > 1.0:
|
||||
return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask)*self.strength
|
||||
else:
|
||||
return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask)*self.strength + input_tensor*(1.0-self.strength)
|
||||
|
||||
|
||||
class TemporalTransformer3DModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
num_layers,
|
||||
attention_block_types=(
|
||||
"Temporal_Self",
|
||||
"Temporal_Self",
|
||||
),
|
||||
dropout=0.0,
|
||||
norm_num_groups=32,
|
||||
cross_attention_dim=768,
|
||||
activation_fn="geglu",
|
||||
attention_bias=False,
|
||||
upcast_attention=False,
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=24,
|
||||
):
|
||||
super().__init__()
|
||||
self.video_length = 16
|
||||
self.full_length = 16
|
||||
self.scale_min = 1.0
|
||||
self.scale_max = 1.0
|
||||
self.raw_scale_mask: Union[Tensor, None] = None
|
||||
self.temp_scale_mask: Union[Tensor, None] = None
|
||||
self.sub_idxs: Union[list[int], None] = None
|
||||
self.prev_hidden_states_batch = 0
|
||||
|
||||
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
|
||||
self.norm = disable_weight_init_clean_groupnorm.GroupNorm(
|
||||
num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True
|
||||
)
|
||||
self.proj_in = nn.Linear(in_channels, inner_dim)
|
||||
|
||||
self.transformer_blocks: Iterable[TemporalTransformerBlock] = nn.ModuleList(
|
||||
[
|
||||
TemporalTransformerBlock(
|
||||
dim=inner_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
attention_block_types=attention_block_types,
|
||||
dropout=dropout,
|
||||
norm_num_groups=norm_num_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
activation_fn=activation_fn,
|
||||
attention_bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
cross_frame_attention_mode=cross_frame_attention_mode,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
)
|
||||
for d in range(num_layers)
|
||||
]
|
||||
)
|
||||
self.proj_out = nn.Linear(inner_dim, in_channels)
|
||||
|
||||
def set_video_length(self, video_length: int, full_length: int):
|
||||
self.video_length = video_length
|
||||
self.full_length = full_length
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
for block in self.transformer_blocks:
|
||||
block.set_scale_multiplier(multiplier)
|
||||
|
||||
def set_masks(self, masks: Tensor, min_val: float, max_val: float):
|
||||
self.scale_min = min_val
|
||||
self.scale_max = max_val
|
||||
self.raw_scale_mask = masks
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
self.sub_idxs = sub_idxs
|
||||
for block in self.transformer_blocks:
|
||||
block.set_sub_idxs(sub_idxs)
|
||||
|
||||
def reset_temp_vars(self):
|
||||
del self.temp_scale_mask
|
||||
self.temp_scale_mask = None
|
||||
self.prev_hidden_states_batch = 0
|
||||
|
||||
def get_scale_mask(self, hidden_states: Tensor) -> Union[Tensor, None]:
|
||||
# if no raw mask, return None
|
||||
if self.raw_scale_mask is None:
|
||||
return None
|
||||
shape = hidden_states.shape
|
||||
batch, channel, height, width = shape
|
||||
# if temp mask already calculated, return it
|
||||
if self.temp_scale_mask != None:
|
||||
# check if hidden_states batch matches
|
||||
if batch == self.prev_hidden_states_batch:
|
||||
if self.sub_idxs is not None:
|
||||
return self.temp_scale_mask[:, self.sub_idxs, :]
|
||||
return self.temp_scale_mask
|
||||
# if does not match, reset cached temp_scale_mask and recalculate it
|
||||
del self.temp_scale_mask
|
||||
self.temp_scale_mask = None
|
||||
# otherwise, calculate temp mask
|
||||
self.prev_hidden_states_batch = batch
|
||||
mask = prepare_mask_batch(self.raw_scale_mask, shape=(self.full_length, 1, height, width))
|
||||
mask = repeat_to_batch_size(mask, self.full_length)
|
||||
# if mask not the same amount length as full length, make it match
|
||||
if self.full_length != mask.shape[0]:
|
||||
mask = broadcast_image_to(mask, self.full_length, 1)
|
||||
# reshape mask to attention K shape (h*w, latent_count, 1)
|
||||
batch, channel, height, width = mask.shape
|
||||
# first, perform same operations as on hidden_states,
|
||||
# turning (b, c, h, w) -> (b, h*w, c)
|
||||
mask = mask.permute(0, 2, 3, 1).reshape(batch, height*width, channel)
|
||||
# then, make it the same shape as attention's k, (h*w, b, c)
|
||||
mask = mask.permute(1, 0, 2)
|
||||
# make masks match the expected length of h*w
|
||||
batched_number = shape[0] // self.video_length
|
||||
if batched_number > 1:
|
||||
mask = torch.cat([mask] * batched_number, dim=0)
|
||||
# cache mask and set to proper device
|
||||
self.temp_scale_mask = mask
|
||||
# move temp_scale_mask to proper dtype + device
|
||||
self.temp_scale_mask = self.temp_scale_mask.to(dtype=hidden_states.dtype, device=hidden_states.device)
|
||||
# return subset of masks, if needed
|
||||
if self.sub_idxs is not None:
|
||||
return self.temp_scale_mask[:, self.sub_idxs, :]
|
||||
return self.temp_scale_mask
|
||||
|
||||
def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):
|
||||
batch, channel, height, width = hidden_states.shape
|
||||
residual = hidden_states
|
||||
scale_mask = self.get_scale_mask(hidden_states)
|
||||
# add some casts for fp8 purposes - does not affect speed otherwise
|
||||
hidden_states = self.norm(hidden_states).to(hidden_states.dtype)
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
|
||||
batch, height * width, inner_dim
|
||||
)
|
||||
hidden_states = self.proj_in(hidden_states).to(hidden_states.dtype)
|
||||
|
||||
# Transformer Blocks
|
||||
for block in self.transformer_blocks:
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
video_length=self.video_length,
|
||||
scale_mask=scale_mask
|
||||
)
|
||||
|
||||
# output
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
hidden_states = (
|
||||
hidden_states.reshape(batch, height, width, inner_dim)
|
||||
.permute(0, 3, 1, 2)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
output = hidden_states + residual
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class TemporalTransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
attention_block_types=(
|
||||
"Temporal_Self",
|
||||
"Temporal_Self",
|
||||
),
|
||||
dropout=0.0,
|
||||
norm_num_groups=32,
|
||||
cross_attention_dim=768,
|
||||
activation_fn="geglu",
|
||||
attention_bias=False,
|
||||
upcast_attention=False,
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=24,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
attention_blocks = []
|
||||
norms = []
|
||||
|
||||
for block_name in attention_block_types:
|
||||
attention_blocks.append(
|
||||
VersatileAttention(
|
||||
attention_mode=block_name.split("_")[0],
|
||||
context_dim=cross_attention_dim # called context_dim for ComfyUI impl
|
||||
if block_name.endswith("_Cross")
|
||||
else None,
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
#bias=attention_bias, # remove for Comfy CrossAttention
|
||||
#upcast_attention=upcast_attention, # remove for Comfy CrossAttention
|
||||
cross_frame_attention_mode=cross_frame_attention_mode,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
)
|
||||
)
|
||||
norms.append(nn.LayerNorm(dim))
|
||||
|
||||
self.attention_blocks: Iterable[VersatileAttention] = nn.ModuleList(attention_blocks)
|
||||
self.norms = nn.ModuleList(norms)
|
||||
|
||||
self.ff = FeedForward(dim, dropout=dropout, glu=(activation_fn == "geglu"))
|
||||
self.ff_norm = nn.LayerNorm(dim)
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
for block in self.attention_blocks:
|
||||
block.set_scale_multiplier(multiplier)
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
for block in self.attention_blocks:
|
||||
block.set_sub_idxs(sub_idxs)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
video_length=None,
|
||||
scale_mask=None
|
||||
):
|
||||
for attention_block, norm in zip(self.attention_blocks, self.norms):
|
||||
norm_hidden_states = norm(hidden_states).to(hidden_states.dtype)
|
||||
hidden_states = (
|
||||
attention_block(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states
|
||||
if attention_block.is_cross_attention
|
||||
else None,
|
||||
attention_mask=attention_mask,
|
||||
video_length=video_length,
|
||||
scale_mask=scale_mask
|
||||
)
|
||||
+ hidden_states
|
||||
)
|
||||
|
||||
hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states
|
||||
|
||||
output = hidden_states
|
||||
return output
|
||||
|
||||
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(self, d_model, dropout=0.0, max_len=24):
|
||||
super().__init__()
|
||||
self.dropout = nn.Dropout(p=dropout)
|
||||
position = torch.arange(max_len).unsqueeze(1)
|
||||
div_term = torch.exp(
|
||||
torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
|
||||
)
|
||||
pe = torch.zeros(1, max_len, d_model)
|
||||
pe[0, :, 0::2] = torch.sin(position * div_term)
|
||||
pe[0, :, 1::2] = torch.cos(position * div_term)
|
||||
self.register_buffer("pe", pe)
|
||||
self.sub_idxs = None
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
self.sub_idxs = sub_idxs
|
||||
|
||||
def forward(self, x):
|
||||
#if self.sub_idxs is not None:
|
||||
# x = x + self.pe[:, self.sub_idxs]
|
||||
#else:
|
||||
x = x + self.pe[:, : x.size(1)]
|
||||
return self.dropout(x)
|
||||
|
||||
|
||||
class CrossAttentionMM(nn.Module):
|
||||
def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None,
|
||||
operations=comfy.ops.disable_weight_init):
|
||||
super().__init__()
|
||||
inner_dim = dim_head * heads
|
||||
context_dim = default(context_dim, query_dim)
|
||||
|
||||
self.heads = heads
|
||||
self.dim_head = dim_head
|
||||
self.scale = None
|
||||
|
||||
self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device)
|
||||
self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
|
||||
self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
|
||||
|
||||
self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout))
|
||||
|
||||
def forward(self, x, context=None, value=None, mask=None, scale_mask=None):
|
||||
q = self.to_q(x)
|
||||
context = default(context, x)
|
||||
k: Tensor = self.to_k(context)
|
||||
if value is not None:
|
||||
v = self.to_v(value)
|
||||
del value
|
||||
else:
|
||||
v = self.to_v(context)
|
||||
|
||||
# apply custom scale by multiplying k by scale factor
|
||||
if self.scale is not None:
|
||||
k *= self.scale
|
||||
|
||||
# apply scale mask, if present
|
||||
if scale_mask is not None:
|
||||
k *= scale_mask
|
||||
|
||||
out = optimized_attention_mm(q, k, v, self.heads, mask)
|
||||
return self.to_out(out)
|
||||
|
||||
|
||||
class VersatileAttention(CrossAttentionMM):
|
||||
def __init__(
|
||||
self,
|
||||
attention_mode=None,
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=24,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
assert attention_mode == "Temporal"
|
||||
|
||||
self.attention_mode = attention_mode
|
||||
self.is_cross_attention = kwargs["context_dim"] is not None
|
||||
|
||||
self.pos_encoder = (
|
||||
PositionalEncoding(
|
||||
kwargs["query_dim"],
|
||||
dropout=0.0,
|
||||
max_len=temporal_position_encoding_max_len,
|
||||
)
|
||||
if (temporal_position_encoding and attention_mode == "Temporal")
|
||||
else None
|
||||
)
|
||||
|
||||
def extra_repr(self):
|
||||
return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}"
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
if multiplier is None or math.isclose(multiplier, 1.0):
|
||||
self.scale = None
|
||||
else:
|
||||
self.scale = multiplier
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
if self.pos_encoder != None:
|
||||
self.pos_encoder.set_sub_idxs(sub_idxs)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: Tensor,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
video_length=None,
|
||||
scale_mask=None,
|
||||
):
|
||||
if self.attention_mode != "Temporal":
|
||||
raise NotImplementedError
|
||||
|
||||
d = hidden_states.shape[1]
|
||||
hidden_states = rearrange(
|
||||
hidden_states, "(b f) d c -> (b d) f c", f=video_length
|
||||
)
|
||||
|
||||
if self.pos_encoder is not None:
|
||||
hidden_states = self.pos_encoder(hidden_states).to(hidden_states.dtype)
|
||||
|
||||
encoder_hidden_states = (
|
||||
repeat(encoder_hidden_states, "b n c -> (b d) n c", d=d)
|
||||
if encoder_hidden_states is not None
|
||||
else encoder_hidden_states
|
||||
)
|
||||
|
||||
hidden_states = super().forward(
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
value=None,
|
||||
mask=attention_mask,
|
||||
scale_mask=scale_mask,
|
||||
)
|
||||
|
||||
hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
|
||||
|
||||
return hidden_states
|
||||
+517
@@ -0,0 +1,517 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
|
||||
import comfy.model_detection
|
||||
from comfy.utils import UNET_MAP_BASIC, UNET_MAP_RESNET, UNET_MAP_ATTENTIONS, TRANSFORMER_BLOCKS
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
from comfy.ldm.modules.diffusionmodules.util import (
|
||||
zero_module,
|
||||
timestep_embedding,
|
||||
)
|
||||
|
||||
from comfy.ldm.modules.attention import SpatialVideoTransformer
|
||||
from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, VideoResBlock, Downsample
|
||||
from comfy.ldm.util import exists
|
||||
import comfy.ops
|
||||
|
||||
|
||||
class SVDControlNet(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
image_size,
|
||||
in_channels,
|
||||
model_channels,
|
||||
hint_channels,
|
||||
num_res_blocks,
|
||||
dropout=0,
|
||||
channel_mult=(1, 2, 4, 8),
|
||||
conv_resample=True,
|
||||
dims=2,
|
||||
num_classes=None,
|
||||
use_checkpoint=False,
|
||||
dtype=torch.float32,
|
||||
num_heads=-1,
|
||||
num_head_channels=-1,
|
||||
num_heads_upsample=-1,
|
||||
use_scale_shift_norm=False,
|
||||
resblock_updown=False,
|
||||
use_new_attention_order=False,
|
||||
use_spatial_transformer=False, # custom transformer support
|
||||
transformer_depth=1, # custom transformer support
|
||||
context_dim=None, # custom transformer support
|
||||
n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
|
||||
legacy=True,
|
||||
disable_self_attentions=None,
|
||||
num_attention_blocks=None,
|
||||
disable_middle_self_attn=False,
|
||||
use_linear_in_transformer=False,
|
||||
adm_in_channels=None,
|
||||
transformer_depth_middle=None,
|
||||
transformer_depth_output=None,
|
||||
use_spatial_context=False,
|
||||
extra_ff_mix_layer=False,
|
||||
merge_strategy="fixed",
|
||||
merge_factor=0.5,
|
||||
video_kernel_size=3,
|
||||
device=None,
|
||||
operations=comfy.ops.disable_weight_init,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
assert use_spatial_transformer == True, "use_spatial_transformer has to be true"
|
||||
if use_spatial_transformer:
|
||||
assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
|
||||
|
||||
if context_dim is not None:
|
||||
assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
|
||||
# from omegaconf.listconfig import ListConfig
|
||||
# if type(context_dim) == ListConfig:
|
||||
# context_dim = list(context_dim)
|
||||
|
||||
if num_heads_upsample == -1:
|
||||
num_heads_upsample = num_heads
|
||||
|
||||
if num_heads == -1:
|
||||
assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
|
||||
|
||||
if num_head_channels == -1:
|
||||
assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
|
||||
|
||||
self.dims = dims
|
||||
self.image_size = image_size
|
||||
self.in_channels = in_channels
|
||||
self.model_channels = model_channels
|
||||
|
||||
if isinstance(num_res_blocks, int):
|
||||
self.num_res_blocks = len(channel_mult) * [num_res_blocks]
|
||||
else:
|
||||
if len(num_res_blocks) != len(channel_mult):
|
||||
raise ValueError("provide num_res_blocks either as an int (globally constant) or "
|
||||
"as a list/tuple (per-level) with the same length as channel_mult")
|
||||
self.num_res_blocks = num_res_blocks
|
||||
|
||||
if disable_self_attentions is not None:
|
||||
# should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
|
||||
assert len(disable_self_attentions) == len(channel_mult)
|
||||
if num_attention_blocks is not None:
|
||||
assert len(num_attention_blocks) == len(self.num_res_blocks)
|
||||
assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
|
||||
|
||||
transformer_depth = transformer_depth[:]
|
||||
|
||||
self.dropout = dropout
|
||||
self.channel_mult = channel_mult
|
||||
self.conv_resample = conv_resample
|
||||
self.num_classes = num_classes
|
||||
self.use_checkpoint = use_checkpoint
|
||||
self.dtype = dtype
|
||||
self.num_heads = num_heads
|
||||
self.num_head_channels = num_head_channels
|
||||
self.num_heads_upsample = num_heads_upsample
|
||||
self.predict_codebook_ids = n_embed is not None
|
||||
|
||||
time_embed_dim = model_channels * 4
|
||||
self.time_embed = nn.Sequential(
|
||||
operations.Linear(model_channels, time_embed_dim, dtype=self.dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
|
||||
)
|
||||
|
||||
if self.num_classes is not None:
|
||||
if isinstance(self.num_classes, int):
|
||||
self.label_emb = nn.Embedding(num_classes, time_embed_dim)
|
||||
elif self.num_classes == "continuous":
|
||||
print("setting up linear c_adm embedding layer")
|
||||
self.label_emb = nn.Linear(1, time_embed_dim)
|
||||
elif self.num_classes == "sequential":
|
||||
assert adm_in_channels is not None
|
||||
self.label_emb = nn.Sequential(
|
||||
nn.Sequential(
|
||||
operations.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise ValueError()
|
||||
|
||||
self.input_blocks = nn.ModuleList(
|
||||
[
|
||||
TimestepEmbedSequential(
|
||||
operations.conv_nd(dims, in_channels, model_channels, 3, padding=1, dtype=self.dtype, device=device)
|
||||
)
|
||||
]
|
||||
)
|
||||
self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels, operations=operations, dtype=self.dtype, device=device)])
|
||||
|
||||
self.input_hint_block = TimestepEmbedSequential(
|
||||
operations.conv_nd(dims, hint_channels, 16, 3, padding=1, dtype=self.dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.conv_nd(dims, 16, 16, 3, padding=1, dtype=self.dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.conv_nd(dims, 16, 32, 3, padding=1, stride=2, dtype=self.dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.conv_nd(dims, 32, 32, 3, padding=1, dtype=self.dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.conv_nd(dims, 32, 96, 3, padding=1, stride=2, dtype=self.dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.conv_nd(dims, 96, 96, 3, padding=1, dtype=self.dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.conv_nd(dims, 96, 256, 3, padding=1, stride=2, dtype=self.dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.conv_nd(dims, 256, model_channels, 3, padding=1, dtype=self.dtype, device=device)
|
||||
)
|
||||
|
||||
self._feature_size = model_channels
|
||||
input_block_chans = [model_channels]
|
||||
ch = model_channels
|
||||
ds = 1
|
||||
for level, mult in enumerate(channel_mult):
|
||||
for nr in range(self.num_res_blocks[level]):
|
||||
layers = [
|
||||
VideoResBlock(
|
||||
ch,
|
||||
time_embed_dim,
|
||||
dropout,
|
||||
out_channels=mult * model_channels,
|
||||
dims=dims,
|
||||
use_checkpoint=use_checkpoint,
|
||||
use_scale_shift_norm=use_scale_shift_norm,
|
||||
dtype=self.dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
video_kernel_size=video_kernel_size,
|
||||
merge_strategy=merge_strategy, merge_factor=merge_factor,
|
||||
)
|
||||
]
|
||||
ch = mult * model_channels
|
||||
num_transformers = transformer_depth.pop(0)
|
||||
if num_transformers > 0:
|
||||
if num_head_channels == -1:
|
||||
dim_head = ch // num_heads
|
||||
else:
|
||||
num_heads = ch // num_head_channels
|
||||
dim_head = num_head_channels
|
||||
if legacy:
|
||||
#num_heads = 1
|
||||
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
|
||||
if exists(disable_self_attentions):
|
||||
disabled_sa = disable_self_attentions[level]
|
||||
else:
|
||||
disabled_sa = False
|
||||
|
||||
if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
|
||||
layers.append(
|
||||
SpatialVideoTransformer(
|
||||
ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim,
|
||||
disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
|
||||
checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations,
|
||||
use_spatial_context=use_spatial_context, ff_in=extra_ff_mix_layer,
|
||||
merge_strategy=merge_strategy, merge_factor=merge_factor,
|
||||
)
|
||||
)
|
||||
self.input_blocks.append(TimestepEmbedSequential(*layers))
|
||||
self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device))
|
||||
self._feature_size += ch
|
||||
input_block_chans.append(ch)
|
||||
if level != len(channel_mult) - 1:
|
||||
out_ch = ch
|
||||
self.input_blocks.append(
|
||||
TimestepEmbedSequential(
|
||||
VideoResBlock(
|
||||
ch,
|
||||
time_embed_dim,
|
||||
dropout,
|
||||
out_channels=out_ch,
|
||||
dims=dims,
|
||||
use_checkpoint=use_checkpoint,
|
||||
use_scale_shift_norm=use_scale_shift_norm,
|
||||
down=True,
|
||||
dtype=self.dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
video_kernel_size=video_kernel_size,
|
||||
merge_strategy=merge_strategy, merge_factor=merge_factor,
|
||||
)
|
||||
if resblock_updown
|
||||
else Downsample(
|
||||
ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device, operations=operations
|
||||
)
|
||||
)
|
||||
)
|
||||
ch = out_ch
|
||||
input_block_chans.append(ch)
|
||||
self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device))
|
||||
ds *= 2
|
||||
self._feature_size += ch
|
||||
|
||||
if num_head_channels == -1:
|
||||
dim_head = ch // num_heads
|
||||
else:
|
||||
num_heads = ch // num_head_channels
|
||||
dim_head = num_head_channels
|
||||
if legacy:
|
||||
#num_heads = 1
|
||||
dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
|
||||
mid_block = [
|
||||
VideoResBlock(
|
||||
ch,
|
||||
time_embed_dim,
|
||||
dropout,
|
||||
dims=dims,
|
||||
use_checkpoint=use_checkpoint,
|
||||
use_scale_shift_norm=use_scale_shift_norm,
|
||||
dtype=self.dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
video_kernel_size=video_kernel_size,
|
||||
merge_strategy=merge_strategy, merge_factor=merge_factor,
|
||||
)]
|
||||
if transformer_depth_middle >= 0:
|
||||
mid_block += [SpatialVideoTransformer( # always uses a self-attn
|
||||
ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim,
|
||||
disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
|
||||
checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations,
|
||||
use_spatial_context=use_spatial_context, ff_in=extra_ff_mix_layer,
|
||||
merge_strategy=merge_strategy, merge_factor=merge_factor,
|
||||
),
|
||||
VideoResBlock(
|
||||
ch,
|
||||
time_embed_dim,
|
||||
dropout,
|
||||
dims=dims,
|
||||
use_checkpoint=use_checkpoint,
|
||||
use_scale_shift_norm=use_scale_shift_norm,
|
||||
dtype=self.dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
video_kernel_size=video_kernel_size,
|
||||
merge_strategy=merge_strategy, merge_factor=merge_factor,
|
||||
)]
|
||||
self.middle_block = TimestepEmbedSequential(*mid_block)
|
||||
self.middle_block_out = self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device)
|
||||
self._feature_size += ch
|
||||
|
||||
def make_zero_conv(self, channels, operations=None, dtype=None, device=None):
|
||||
return TimestepEmbedSequential(operations.conv_nd(self.dims, channels, channels, 1, padding=0, dtype=dtype, device=device))
|
||||
|
||||
def forward(self, x, hint, timesteps, context, y=None, **kwargs):
|
||||
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
|
||||
emb = self.time_embed(t_emb)
|
||||
|
||||
cond = kwargs["cond"]
|
||||
num_video_frames = cond["num_video_frames"]
|
||||
image_only_indicator = cond["image_only_indicator"]
|
||||
time_context = cond.get("time_context", None)
|
||||
del cond
|
||||
|
||||
guided_hint = self.input_hint_block(hint, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
|
||||
|
||||
outs = []
|
||||
|
||||
hs = []
|
||||
if self.num_classes is not None:
|
||||
assert y.shape[0] == x.shape[0]
|
||||
emb = emb + self.label_emb(y)
|
||||
|
||||
h = x
|
||||
for module, zero_conv in zip(self.input_blocks, self.zero_convs):
|
||||
if guided_hint is not None:
|
||||
h = module(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
|
||||
h += guided_hint
|
||||
guided_hint = None
|
||||
else:
|
||||
h = module(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
|
||||
outs.append(zero_conv(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator))
|
||||
|
||||
h = self.middle_block(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
|
||||
outs.append(self.middle_block_out(h, emb, context, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator))
|
||||
|
||||
return outs
|
||||
|
||||
|
||||
TEMPORAL_TRANSFORMER_BLOCKS = {
|
||||
"norm_in.weight",
|
||||
"norm_in.bias",
|
||||
"ff_in.net.0.proj.weight",
|
||||
"ff_in.net.0.proj.bias",
|
||||
"ff_in.net.2.weight",
|
||||
"ff_in.net.2.bias",
|
||||
}
|
||||
TEMPORAL_TRANSFORMER_BLOCKS.update(TRANSFORMER_BLOCKS)
|
||||
|
||||
|
||||
TEMPORAL_UNET_MAP_ATTENTIONS = {
|
||||
"time_mixer.mix_factor",
|
||||
}
|
||||
TEMPORAL_UNET_MAP_ATTENTIONS.update(UNET_MAP_ATTENTIONS)
|
||||
|
||||
|
||||
TEMPORAL_TRANSFORMER_MAP = {
|
||||
"time_pos_embed.0.weight": "time_pos_embed.linear_1.weight",
|
||||
"time_pos_embed.0.bias": "time_pos_embed.linear_1.bias",
|
||||
"time_pos_embed.2.weight": "time_pos_embed.linear_2.weight",
|
||||
"time_pos_embed.2.bias": "time_pos_embed.linear_2.bias",
|
||||
}
|
||||
|
||||
|
||||
TEMPORAL_RESNET = {
|
||||
"time_mixer.mix_factor",
|
||||
}
|
||||
|
||||
|
||||
def svd_unet_config_from_diffusers_unet(state_dict: dict[str, Tensor], dtype):
|
||||
match = {}
|
||||
transformer_depth = []
|
||||
|
||||
attn_res = 1
|
||||
down_blocks = comfy.model_detection.count_blocks(state_dict, "down_blocks.{}")
|
||||
for i in range(down_blocks):
|
||||
attn_blocks = comfy.model_detection.count_blocks(state_dict, "down_blocks.{}.attentions.".format(i) + '{}')
|
||||
for ab in range(attn_blocks):
|
||||
transformer_count = comfy.model_detection.count_blocks(state_dict, "down_blocks.{}.attentions.{}.transformer_blocks.".format(i, ab) + '{}')
|
||||
transformer_depth.append(transformer_count)
|
||||
if transformer_count > 0:
|
||||
match["context_dim"] = state_dict["down_blocks.{}.attentions.{}.transformer_blocks.0.attn2.to_k.weight".format(i, ab)].shape[1]
|
||||
|
||||
attn_res *= 2
|
||||
if attn_blocks == 0:
|
||||
transformer_depth.append(0)
|
||||
transformer_depth.append(0)
|
||||
|
||||
match["transformer_depth"] = transformer_depth
|
||||
|
||||
match["model_channels"] = state_dict["conv_in.weight"].shape[0]
|
||||
match["in_channels"] = state_dict["conv_in.weight"].shape[1]
|
||||
match["adm_in_channels"] = None
|
||||
if "class_embedding.linear_1.weight" in state_dict:
|
||||
match["adm_in_channels"] = state_dict["class_embedding.linear_1.weight"].shape[1]
|
||||
elif "add_embedding.linear_1.weight" in state_dict:
|
||||
match["adm_in_channels"] = state_dict["add_embedding.linear_1.weight"].shape[1]
|
||||
|
||||
# based on unet_config of SVD
|
||||
SVD = {
|
||||
'use_checkpoint': False,
|
||||
'image_size': 32,
|
||||
'use_spatial_transformer': True,
|
||||
'legacy': False,
|
||||
'num_classes': 'sequential',
|
||||
'adm_in_channels': 768,
|
||||
'dtype': dtype,
|
||||
'in_channels': 8,
|
||||
'out_channels': 4,
|
||||
'model_channels': 320,
|
||||
'num_res_blocks': [2, 2, 2, 2],
|
||||
'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0],
|
||||
'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
|
||||
'channel_mult': [1, 2, 4, 4],
|
||||
'transformer_depth_middle': 1,
|
||||
'use_linear_in_transformer': True,
|
||||
'context_dim': 1024,
|
||||
'extra_ff_mix_layer': True,
|
||||
'use_spatial_context': True,
|
||||
'merge_strategy': 'learned_with_images',
|
||||
'merge_factor': 0.0,
|
||||
'video_kernel_size': [3, 1, 1],
|
||||
'use_temporal_attention': True,
|
||||
'use_temporal_resblock': True,
|
||||
'num_heads': -1,
|
||||
'num_head_channels': 64,
|
||||
}
|
||||
|
||||
supported_models = [SVD]
|
||||
|
||||
for unet_config in supported_models:
|
||||
matches = True
|
||||
for k in match:
|
||||
if match[k] != unet_config[k]:
|
||||
matches = False
|
||||
break
|
||||
if matches:
|
||||
return comfy.model_detection.convert_config(unet_config)
|
||||
return None
|
||||
|
||||
|
||||
def svd_unet_to_diffusers(unet_config):
|
||||
num_res_blocks = unet_config["num_res_blocks"]
|
||||
channel_mult = unet_config["channel_mult"]
|
||||
transformer_depth = unet_config["transformer_depth"][:]
|
||||
transformer_depth_output = unet_config["transformer_depth_output"][:]
|
||||
num_blocks = len(channel_mult)
|
||||
|
||||
transformers_mid = unet_config.get("transformer_depth_middle", None)
|
||||
|
||||
diffusers_unet_map = {}
|
||||
for x in range(num_blocks):
|
||||
n = 1 + (num_res_blocks[x] + 1) * x
|
||||
for i in range(num_res_blocks[x]):
|
||||
for b in TEMPORAL_RESNET:
|
||||
diffusers_unet_map["down_blocks.{}.resnets.{}.{}".format(x, i, b)] = "input_blocks.{}.0.{}".format(n, b)
|
||||
for b in UNET_MAP_RESNET:
|
||||
diffusers_unet_map["down_blocks.{}.resnets.{}.spatial_res_block.{}".format(x, i, UNET_MAP_RESNET[b])] = "input_blocks.{}.0.{}".format(n, b)
|
||||
diffusers_unet_map["down_blocks.{}.resnets.{}.temporal_res_block.{}".format(x, i, UNET_MAP_RESNET[b])] = "input_blocks.{}.0.time_stack.{}".format(n, b)
|
||||
#diffusers_unet_map["down_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])] = "input_blocks.{}.0.{}".format(n, b)
|
||||
num_transformers = transformer_depth.pop(0)
|
||||
if num_transformers > 0:
|
||||
for b in TEMPORAL_UNET_MAP_ATTENTIONS:
|
||||
diffusers_unet_map["down_blocks.{}.attentions.{}.{}".format(x, i, b)] = "input_blocks.{}.1.{}".format(n, b)
|
||||
for b in TEMPORAL_TRANSFORMER_MAP:
|
||||
diffusers_unet_map["down_blocks.{}.attentions.{}.{}".format(x, i, TEMPORAL_TRANSFORMER_MAP[b])] = "input_blocks.{}.1.{}".format(n, b)
|
||||
for t in range(num_transformers):
|
||||
for b in TRANSFORMER_BLOCKS:
|
||||
diffusers_unet_map["down_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(x, i, t, b)] = "input_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b)
|
||||
for b in TEMPORAL_TRANSFORMER_BLOCKS:
|
||||
diffusers_unet_map["down_blocks.{}.attentions.{}.temporal_transformer_blocks.{}.{}".format(x, i, t, b)] = "input_blocks.{}.1.time_stack.{}.{}".format(n, t, b)
|
||||
n += 1
|
||||
for k in ["weight", "bias"]:
|
||||
diffusers_unet_map["down_blocks.{}.downsamplers.0.conv.{}".format(x, k)] = "input_blocks.{}.0.op.{}".format(n, k)
|
||||
|
||||
i = 0
|
||||
for b in TEMPORAL_UNET_MAP_ATTENTIONS:
|
||||
diffusers_unet_map["mid_block.attentions.{}.{}".format(i, b)] = "middle_block.1.{}".format(b)
|
||||
for b in TEMPORAL_TRANSFORMER_MAP:
|
||||
diffusers_unet_map["mid_block.attentions.{}.{}".format(i, TEMPORAL_TRANSFORMER_MAP[b])] = "middle_block.1.{}".format(b)
|
||||
for t in range(transformers_mid):
|
||||
for b in TRANSFORMER_BLOCKS:
|
||||
diffusers_unet_map["mid_block.attentions.{}.transformer_blocks.{}.{}".format(i, t, b)] = "middle_block.1.transformer_blocks.{}.{}".format(t, b)
|
||||
for b in TEMPORAL_TRANSFORMER_BLOCKS:
|
||||
diffusers_unet_map["mid_block.attentions.{}.temporal_transformer_blocks.{}.{}".format(i, t, b)] = "middle_block.1.time_stack.{}.{}".format(t, b)
|
||||
|
||||
for i, n in enumerate([0, 2]):
|
||||
for b in TEMPORAL_RESNET:
|
||||
diffusers_unet_map["mid_block.resnets.{}.{}".format(i, b)] = "middle_block.{}.{}".format(n, b)
|
||||
for b in UNET_MAP_RESNET:
|
||||
diffusers_unet_map["mid_block.resnets.{}.spatial_res_block.{}".format(i, UNET_MAP_RESNET[b])] = "middle_block.{}.{}".format(n, b)
|
||||
diffusers_unet_map["mid_block.resnets.{}.temporal_res_block.{}".format(i, UNET_MAP_RESNET[b])] = "middle_block.{}.time_stack.{}".format(n, b)
|
||||
#diffusers_unet_map["mid_block.resnets.{}.{}".format(i, UNET_MAP_RESNET[b])] = "middle_block.{}.{}".format(n, b)
|
||||
|
||||
num_res_blocks = list(reversed(num_res_blocks))
|
||||
for x in range(num_blocks):
|
||||
n = (num_res_blocks[x] + 1) * x
|
||||
l = num_res_blocks[x] + 1
|
||||
for i in range(l):
|
||||
c = 0
|
||||
for b in UNET_MAP_RESNET:
|
||||
diffusers_unet_map["up_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])] = "output_blocks.{}.0.{}".format(n, b)
|
||||
c += 1
|
||||
num_transformers = transformer_depth_output.pop()
|
||||
if num_transformers > 0:
|
||||
c += 1
|
||||
for b in UNET_MAP_ATTENTIONS:
|
||||
diffusers_unet_map["up_blocks.{}.attentions.{}.{}".format(x, i, b)] = "output_blocks.{}.1.{}".format(n, b)
|
||||
for t in range(num_transformers):
|
||||
for b in TRANSFORMER_BLOCKS:
|
||||
diffusers_unet_map["up_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(x, i, t, b)] = "output_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b)
|
||||
if i == l - 1:
|
||||
for k in ["weight", "bias"]:
|
||||
diffusers_unet_map["up_blocks.{}.upsamplers.0.conv.{}".format(x, k)] = "output_blocks.{}.{}.conv.{}".format(n, c, k)
|
||||
n += 1
|
||||
|
||||
for k in UNET_MAP_BASIC:
|
||||
diffusers_unet_map[k[1]] = k[0]
|
||||
|
||||
return diffusers_unet_map
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
import sys
|
||||
import copy
|
||||
import logging
|
||||
|
||||
|
||||
class ColoredFormatter(logging.Formatter):
|
||||
COLORS = {
|
||||
"DEBUG": "\033[0;36m", # CYAN
|
||||
"INFO": "\033[0;32m", # GREEN
|
||||
"WARNING": "\033[0;33m", # YELLOW
|
||||
"ERROR": "\033[0;31m", # RED
|
||||
"CRITICAL": "\033[0;37;41m", # WHITE ON RED
|
||||
"RESET": "\033[0m", # RESET COLOR
|
||||
}
|
||||
|
||||
def format(self, record):
|
||||
colored_record = copy.copy(record)
|
||||
levelname = colored_record.levelname
|
||||
seq = self.COLORS.get(levelname, self.COLORS["RESET"])
|
||||
colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}"
|
||||
return super().format(colored_record)
|
||||
|
||||
|
||||
# Create a new logger
|
||||
logger = logging.getLogger("Advanced-ControlNet")
|
||||
logger.propagate = False
|
||||
|
||||
# Add handler if we don't have one.
|
||||
if not logger.handlers:
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
handler.setFormatter(ColoredFormatter("[%(name)s] - %(levelname)s - %(message)s"))
|
||||
logger.addHandler(handler)
|
||||
|
||||
# Configure logger
|
||||
loglevel = logging.INFO
|
||||
logger.setLevel(loglevel)
|
||||
+272
@@ -0,0 +1,272 @@
|
||||
import numpy as np
|
||||
from torch import Tensor
|
||||
|
||||
import folder_paths
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from .control import load_controlnet, convert_to_advanced, is_advanced_controlnet
|
||||
from .utils import ControlWeights, ControlWeightType, LatentKeyframeGroup, TimestepKeyframe, TimestepKeyframeGroup
|
||||
from .utils import StrengthInterpolation as SI
|
||||
from .nodes_weight import (DefaultWeights, ScaledSoftMaskedUniversalWeights, ScaledSoftUniversalWeights, SoftControlNetWeights, CustomControlNetWeights,
|
||||
SoftT2IAdapterWeights, CustomT2IAdapterWeights)
|
||||
from .nodes_latent_keyframe import LatentKeyframeGroupNode, LatentKeyframeInterpolationNode, LatentKeyframeBatchedGroupNode, LatentKeyframeNode
|
||||
from .nodes_sparsectrl import SparseCtrlMergedLoaderAdvanced, SparseCtrlLoaderAdvanced, SparseIndexMethodNode, SparseSpreadMethodNode, RgbSparseCtrlPreprocessor
|
||||
from .nodes_loosecontrol import ControlNetLoaderWithLoraAdvanced
|
||||
from .nodes_deprecated import LoadImagesFromDirectory
|
||||
from .logger import logger
|
||||
|
||||
|
||||
class TimestepKeyframeNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}, ),
|
||||
},
|
||||
"optional": {
|
||||
"prev_timestep_kf": ("TIMESTEP_KEYFRAME", ),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"cn_weights": ("CONTROL_NET_WEIGHTS", ),
|
||||
"latent_keyframe": ("LATENT_KEYFRAME", ),
|
||||
"null_latent_kf_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"inherit_missing": ("BOOLEAN", {"default": True}, ),
|
||||
"guarantee_usage": ("BOOLEAN", {"default": True}, ),
|
||||
"mask_optional": ("MASK", ),
|
||||
#"interpolation": ([SI.LINEAR, SI.EASE_IN, SI.EASE_OUT, SI.EASE_IN_OUT, SI.NONE], {"default": SI.NONE}, ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_NAMES = ("TIMESTEP_KF", )
|
||||
RETURN_TYPES = ("TIMESTEP_KEYFRAME", )
|
||||
FUNCTION = "load_keyframe"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/keyframes"
|
||||
|
||||
def load_keyframe(self,
|
||||
start_percent: float,
|
||||
strength: float=1.0,
|
||||
cn_weights: ControlWeights=None, control_net_weights: ControlWeights=None, # old name
|
||||
latent_keyframe: LatentKeyframeGroup=None,
|
||||
prev_timestep_kf: TimestepKeyframeGroup=None, prev_timestep_keyframe: TimestepKeyframeGroup=None, # old name
|
||||
null_latent_kf_strength: float=0.0,
|
||||
inherit_missing=True,
|
||||
guarantee_usage=True,
|
||||
mask_optional=None,
|
||||
interpolation: str=SI.NONE,):
|
||||
control_net_weights = control_net_weights if control_net_weights else cn_weights
|
||||
prev_timestep_keyframe = prev_timestep_keyframe if prev_timestep_keyframe else prev_timestep_kf
|
||||
if not prev_timestep_keyframe:
|
||||
prev_timestep_keyframe = TimestepKeyframeGroup()
|
||||
else:
|
||||
prev_timestep_keyframe = prev_timestep_keyframe.clone()
|
||||
keyframe = TimestepKeyframe(start_percent=start_percent, strength=strength, interpolation=interpolation, null_latent_kf_strength=null_latent_kf_strength,
|
||||
control_weights=control_net_weights, latent_keyframes=latent_keyframe, inherit_missing=inherit_missing, guarantee_usage=guarantee_usage,
|
||||
mask_hint_orig=mask_optional)
|
||||
prev_timestep_keyframe.add(keyframe)
|
||||
return (prev_timestep_keyframe,)
|
||||
|
||||
|
||||
class ControlNetLoaderAdvanced:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"control_net_name": (folder_paths.get_filename_list("controlnet"), ),
|
||||
},
|
||||
"optional": {
|
||||
"timestep_keyframe": ("TIMESTEP_KEYFRAME", ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET", )
|
||||
FUNCTION = "load_controlnet"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝"
|
||||
|
||||
def load_controlnet(self, control_net_name,
|
||||
timestep_keyframe: TimestepKeyframeGroup=None
|
||||
):
|
||||
controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
|
||||
controlnet = load_controlnet(controlnet_path, timestep_keyframe)
|
||||
return (controlnet,)
|
||||
|
||||
|
||||
class DiffControlNetLoaderAdvanced:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"control_net_name": (folder_paths.get_filename_list("controlnet"), )
|
||||
},
|
||||
"optional": {
|
||||
"timestep_keyframe": ("TIMESTEP_KEYFRAME", ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET", )
|
||||
FUNCTION = "load_controlnet"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝"
|
||||
|
||||
def load_controlnet(self, control_net_name, model,
|
||||
timestep_keyframe: TimestepKeyframeGroup=None
|
||||
):
|
||||
controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
|
||||
controlnet = load_controlnet(controlnet_path, timestep_keyframe, model)
|
||||
if is_advanced_controlnet(controlnet):
|
||||
controlnet.verify_all_weights()
|
||||
return (controlnet,)
|
||||
|
||||
|
||||
class AdvancedControlNetApply:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"positive": ("CONDITIONING", ),
|
||||
"negative": ("CONDITIONING", ),
|
||||
"control_net": ("CONTROL_NET", ),
|
||||
"image": ("IMAGE", ),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001})
|
||||
},
|
||||
"optional": {
|
||||
"mask_optional": ("MASK", ),
|
||||
"timestep_kf": ("TIMESTEP_KEYFRAME", ),
|
||||
"latent_kf_override": ("LATENT_KEYFRAME", ),
|
||||
"weights_override": ("CONTROL_NET_WEIGHTS", ),
|
||||
"model_optional": ("MODEL",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING","CONDITIONING","MODEL",)
|
||||
RETURN_NAMES = ("positive", "negative", "model_opt")
|
||||
FUNCTION = "apply_controlnet"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝"
|
||||
|
||||
def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent,
|
||||
mask_optional: Tensor=None, model_optional: ModelPatcher=None,
|
||||
timestep_kf: TimestepKeyframeGroup=None, latent_kf_override: LatentKeyframeGroup=None,
|
||||
weights_override: ControlWeights=None):
|
||||
if strength == 0:
|
||||
return (positive, negative, model_optional)
|
||||
if model_optional:
|
||||
model_optional = model_optional.clone()
|
||||
|
||||
control_hint = image.movedim(-1,1)
|
||||
cnets = {}
|
||||
|
||||
out = []
|
||||
for conditioning in [positive, negative]:
|
||||
c = []
|
||||
for t in conditioning:
|
||||
d = t[1].copy()
|
||||
|
||||
prev_cnet = d.get('control', None)
|
||||
if prev_cnet in cnets:
|
||||
c_net = cnets[prev_cnet]
|
||||
else:
|
||||
# copy, convert to advanced if needed, and set cond
|
||||
c_net = convert_to_advanced(control_net.copy()).set_cond_hint(control_hint, strength, (start_percent, end_percent))
|
||||
if is_advanced_controlnet(c_net):
|
||||
# disarm node check
|
||||
c_net.disarm()
|
||||
# if model required, verify model is passed in, and if so patch it
|
||||
if c_net.require_model:
|
||||
if not model_optional:
|
||||
raise Exception(f"Type '{type(c_net).__name__}' requires model_optional input, but got None.")
|
||||
c_net.patch_model(model=model_optional)
|
||||
# apply optional parameters and overrides, if provided
|
||||
if timestep_kf is not None:
|
||||
c_net.set_timestep_keyframes(timestep_kf)
|
||||
if latent_kf_override is not None:
|
||||
c_net.latent_keyframe_override = latent_kf_override
|
||||
if weights_override is not None:
|
||||
c_net.weights_override = weights_override
|
||||
# verify weights are compatible
|
||||
c_net.verify_all_weights()
|
||||
# set cond hint mask
|
||||
if mask_optional is not None:
|
||||
mask_optional = mask_optional.clone()
|
||||
# if not in the form of a batch, make it so
|
||||
if len(mask_optional.shape) < 3:
|
||||
mask_optional = mask_optional.unsqueeze(0)
|
||||
c_net.set_cond_hint_mask(mask_optional)
|
||||
c_net.set_previous_controlnet(prev_cnet)
|
||||
cnets[prev_cnet] = c_net
|
||||
|
||||
d['control'] = c_net
|
||||
d['control_apply_to_uncond'] = False
|
||||
n = [t[0], d]
|
||||
c.append(n)
|
||||
out.append(c)
|
||||
return (out[0], out[1], model_optional)
|
||||
|
||||
|
||||
# NODE MAPPING
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
# Keyframes
|
||||
"TimestepKeyframe": TimestepKeyframeNode,
|
||||
"LatentKeyframe": LatentKeyframeNode,
|
||||
"LatentKeyframeGroup": LatentKeyframeGroupNode,
|
||||
"LatentKeyframeBatchedGroup": LatentKeyframeBatchedGroupNode,
|
||||
"LatentKeyframeTiming": LatentKeyframeInterpolationNode,
|
||||
# Conditioning
|
||||
"ACN_AdvancedControlNetApply": AdvancedControlNetApply,
|
||||
# Loaders
|
||||
"ControlNetLoaderAdvanced": ControlNetLoaderAdvanced,
|
||||
"DiffControlNetLoaderAdvanced": DiffControlNetLoaderAdvanced,
|
||||
# Weights
|
||||
"ScaledSoftControlNetWeights": ScaledSoftUniversalWeights,
|
||||
"ScaledSoftMaskedUniversalWeights": ScaledSoftMaskedUniversalWeights,
|
||||
"SoftControlNetWeights": SoftControlNetWeights,
|
||||
"CustomControlNetWeights": CustomControlNetWeights,
|
||||
"SoftT2IAdapterWeights": SoftT2IAdapterWeights,
|
||||
"CustomT2IAdapterWeights": CustomT2IAdapterWeights,
|
||||
"ACN_DefaultUniversalWeights": DefaultWeights,
|
||||
# SparseCtrl
|
||||
"ACN_SparseCtrlRGBPreprocessor": RgbSparseCtrlPreprocessor,
|
||||
"ACN_SparseCtrlLoaderAdvanced": SparseCtrlLoaderAdvanced,
|
||||
"ACN_SparseCtrlMergedLoaderAdvanced": SparseCtrlMergedLoaderAdvanced,
|
||||
"ACN_SparseCtrlIndexMethodNode": SparseIndexMethodNode,
|
||||
"ACN_SparseCtrlSpreadMethodNode": SparseSpreadMethodNode,
|
||||
# LOOSEControl
|
||||
#"ACN_ControlNetLoaderWithLoraAdvanced": ControlNetLoaderWithLoraAdvanced,
|
||||
# Deprecated
|
||||
"LoadImagesFromDirectory": LoadImagesFromDirectory,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
# Keyframes
|
||||
"TimestepKeyframe": "Timestep Keyframe 🛂🅐🅒🅝",
|
||||
"LatentKeyframe": "Latent Keyframe 🛂🅐🅒🅝",
|
||||
"LatentKeyframeGroup": "Latent Keyframe Group 🛂🅐🅒🅝",
|
||||
"LatentKeyframeBatchedGroup": "Latent Keyframe Batched Group 🛂🅐🅒🅝",
|
||||
"LatentKeyframeTiming": "Latent Keyframe Interpolation 🛂🅐🅒🅝",
|
||||
# Conditioning
|
||||
"ACN_AdvancedControlNetApply": "Apply Advanced ControlNet 🛂🅐🅒🅝",
|
||||
# Loaders
|
||||
"ControlNetLoaderAdvanced": "Load Advanced ControlNet Model 🛂🅐🅒🅝",
|
||||
"DiffControlNetLoaderAdvanced": "Load Advanced ControlNet Model (diff) 🛂🅐🅒🅝",
|
||||
# Weights
|
||||
"ScaledSoftControlNetWeights": "Scaled Soft Weights 🛂🅐🅒🅝",
|
||||
"ScaledSoftMaskedUniversalWeights": "Scaled Soft Masked Weights 🛂🅐🅒🅝",
|
||||
"SoftControlNetWeights": "ControlNet Soft Weights 🛂🅐🅒🅝",
|
||||
"CustomControlNetWeights": "ControlNet Custom Weights 🛂🅐🅒🅝",
|
||||
"SoftT2IAdapterWeights": "T2IAdapter Soft Weights 🛂🅐🅒🅝",
|
||||
"CustomT2IAdapterWeights": "T2IAdapter Custom Weights 🛂🅐🅒🅝",
|
||||
"ACN_DefaultUniversalWeights": "Force Default Weights 🛂🅐🅒🅝",
|
||||
# SparseCtrl
|
||||
"ACN_SparseCtrlRGBPreprocessor": "RGB SparseCtrl 🛂🅐🅒🅝",
|
||||
"ACN_SparseCtrlLoaderAdvanced": "Load SparseCtrl Model 🛂🅐🅒🅝",
|
||||
"ACN_SparseCtrlMergedLoaderAdvanced": "Load Merged SparseCtrl Model 🛂🅐🅒🅝",
|
||||
"ACN_SparseCtrlIndexMethodNode": "SparseCtrl Index Method 🛂🅐🅒🅝",
|
||||
"ACN_SparseCtrlSpreadMethodNode": "SparseCtrl Spread Method 🛂🅐🅒🅝",
|
||||
# LOOSEControl
|
||||
#"ACN_ControlNetLoaderWithLoraAdvanced": "Load Adv. ControlNet Model w/ LoRA 🛂🅐🅒🅝",
|
||||
# Deprecated
|
||||
"LoadImagesFromDirectory": "Load Images [DEPRECATED] 🛂🅐🅒🅝",
|
||||
}
|
||||
+103
@@ -0,0 +1,103 @@
|
||||
import os
|
||||
|
||||
import torch
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageOps
|
||||
from .utils import ControlWeights, LatentKeyframeGroup, TimestepKeyframeGroup, TimestepKeyframe, BIGMAX
|
||||
from .logger import logger
|
||||
|
||||
|
||||
class LoadImagesFromDirectory:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"directory": ("STRING", {"default": ""}),
|
||||
},
|
||||
"optional": {
|
||||
"image_load_cap": ("INT", {"default": 0, "min": 0, "max": BIGMAX, "step": 1}),
|
||||
"start_index": ("INT", {"default": 0, "min": 0, "max": BIGMAX, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "INT")
|
||||
FUNCTION = "load_images"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/deprecated"
|
||||
|
||||
def load_images(self, directory: str, image_load_cap: int = 0, start_index: int = 0):
|
||||
if not os.path.isdir(directory):
|
||||
raise FileNotFoundError(f"Directory '{directory} cannot be found.'")
|
||||
dir_files = os.listdir(directory)
|
||||
if len(dir_files) == 0:
|
||||
raise FileNotFoundError(f"No files in directory '{directory}'.")
|
||||
|
||||
dir_files = sorted(dir_files)
|
||||
dir_files = [os.path.join(directory, x) for x in dir_files]
|
||||
# start at start_index
|
||||
dir_files = dir_files[start_index:]
|
||||
|
||||
images = []
|
||||
masks = []
|
||||
|
||||
limit_images = False
|
||||
if image_load_cap > 0:
|
||||
limit_images = True
|
||||
image_count = 0
|
||||
|
||||
for image_path in dir_files:
|
||||
if os.path.isdir(image_path):
|
||||
continue
|
||||
if limit_images and image_count >= image_load_cap:
|
||||
break
|
||||
i = Image.open(image_path)
|
||||
i = ImageOps.exif_transpose(i)
|
||||
image = i.convert("RGB")
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = torch.from_numpy(image)[None,]
|
||||
if 'A' in i.getbands():
|
||||
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
||||
mask = 1. - torch.from_numpy(mask)
|
||||
else:
|
||||
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
|
||||
images.append(image)
|
||||
masks.append(mask)
|
||||
image_count += 1
|
||||
|
||||
if len(images) == 0:
|
||||
raise FileNotFoundError(f"No images could be loaded from directory '{directory}'.")
|
||||
|
||||
return (torch.cat(images, dim=0), torch.stack(masks, dim=0), image_count)
|
||||
|
||||
|
||||
class TimestepKeyframeNodeDeprecated:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}, ),
|
||||
},
|
||||
"optional": {
|
||||
"control_net_weights": ("CONTROL_NET_WEIGHTS", ),
|
||||
"t2i_adapter_weights": ("T2I_ADAPTER_WEIGHTS", ),
|
||||
"latent_keyframe": ("LATENT_KEYFRAME", ),
|
||||
"prev_timestep_keyframe": ("TIMESTEP_KEYFRAME", ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("TIMESTEP_KEYFRAME", )
|
||||
FUNCTION = "load_keyframe"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/keyframes"
|
||||
|
||||
def load_keyframe(self,
|
||||
start_percent: float,
|
||||
control_net_weights: ControlWeights=None,
|
||||
latent_keyframe: LatentKeyframeGroup=None,
|
||||
prev_timestep_keyframe: TimestepKeyframeGroup=None):
|
||||
if not prev_timestep_keyframe:
|
||||
prev_timestep_keyframe = TimestepKeyframeGroup()
|
||||
keyframe = TimestepKeyframe(start_percent, control_net_weights, latent_keyframe)
|
||||
prev_timestep_keyframe.add(keyframe)
|
||||
return (prev_timestep_keyframe,)
|
||||
+283
@@ -0,0 +1,283 @@
|
||||
from typing import Union
|
||||
import numpy as np
|
||||
from collections.abc import Iterable
|
||||
|
||||
from .utils import LatentKeyframe, LatentKeyframeGroup, BIGMIN, BIGMAX
|
||||
from .utils import StrengthInterpolation as SI
|
||||
from .logger import logger
|
||||
|
||||
|
||||
class LatentKeyframeNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"batch_index": ("INT", {"default": 0, "min": BIGMIN, "max": BIGMAX, "step": 1}),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
},
|
||||
"optional": {
|
||||
"prev_latent_kf": ("LATENT_KEYFRAME", ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_NAMES = ("LATENT_KF", )
|
||||
RETURN_TYPES = ("LATENT_KEYFRAME", )
|
||||
FUNCTION = "load_keyframe"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/keyframes"
|
||||
|
||||
def load_keyframe(self,
|
||||
batch_index: int,
|
||||
strength: float,
|
||||
prev_latent_kf: LatentKeyframeGroup=None,
|
||||
prev_latent_keyframe: LatentKeyframeGroup=None, # old name
|
||||
):
|
||||
prev_latent_keyframe = prev_latent_keyframe if prev_latent_keyframe else prev_latent_kf
|
||||
if not prev_latent_keyframe:
|
||||
prev_latent_keyframe = LatentKeyframeGroup()
|
||||
else:
|
||||
prev_latent_keyframe = prev_latent_keyframe.clone()
|
||||
keyframe = LatentKeyframe(batch_index, strength)
|
||||
prev_latent_keyframe.add(keyframe)
|
||||
return (prev_latent_keyframe,)
|
||||
|
||||
|
||||
class LatentKeyframeGroupNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"index_strengths": ("STRING", {"multiline": True, "default": ""}),
|
||||
},
|
||||
"optional": {
|
||||
"prev_latent_kf": ("LATENT_KEYFRAME", ),
|
||||
"latent_optional": ("LATENT", ),
|
||||
"print_keyframes": ("BOOLEAN", {"default": False})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_NAMES = ("LATENT_KF", )
|
||||
RETURN_TYPES = ("LATENT_KEYFRAME", )
|
||||
FUNCTION = "load_keyframes"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/keyframes"
|
||||
|
||||
def validate_index(self, index: int, latent_count: int = 0, is_range: bool = False, allow_negative = False) -> int:
|
||||
# if part of range, do nothing
|
||||
if is_range:
|
||||
return index
|
||||
# otherwise, validate index
|
||||
# validate not out of range - only when latent_count is passed in
|
||||
if latent_count > 0 and index > latent_count-1:
|
||||
raise IndexError(f"Index '{index}' out of range for the total {latent_count} latents.")
|
||||
# if negative, validate not out of range
|
||||
if index < 0:
|
||||
if not allow_negative:
|
||||
raise IndexError(f"Negative indeces not allowed, but was {index}.")
|
||||
conv_index = latent_count+index
|
||||
if conv_index < 0:
|
||||
raise IndexError(f"Index '{index}', converted to '{conv_index}' out of range for the total {latent_count} latents.")
|
||||
index = conv_index
|
||||
return index
|
||||
|
||||
def convert_to_index_int(self, raw_index: str, latent_count: int = 0, is_range: bool = False, allow_negative = False) -> int:
|
||||
try:
|
||||
return self.validate_index(int(raw_index), latent_count=latent_count, is_range=is_range, allow_negative=allow_negative)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"index '{raw_index}' must be an integer.", e)
|
||||
|
||||
def convert_to_latent_keyframes(self, latent_indeces: str, latent_count: int) -> set[LatentKeyframe]:
|
||||
if not latent_indeces:
|
||||
return set()
|
||||
int_latent_indeces = [i for i in range(0, latent_count)]
|
||||
allow_negative = latent_count > 0
|
||||
chosen_indeces = set()
|
||||
# parse string - allow positive ints, negative ints, and ranges separated by ':'
|
||||
groups = latent_indeces.split(",")
|
||||
groups = [g.strip() for g in groups]
|
||||
for g in groups:
|
||||
# parse strengths - default to 1.0 if no strength given
|
||||
strength = 1.0
|
||||
if '=' in g:
|
||||
g, strength_str = g.split("=", 1)
|
||||
g = g.strip()
|
||||
try:
|
||||
strength = float(strength_str.strip())
|
||||
except ValueError as e:
|
||||
raise ValueError(f"strength '{strength_str}' must be a float.", e)
|
||||
if strength < 0:
|
||||
raise ValueError(f"Strength '{strength}' cannot be negative.")
|
||||
# parse range of indeces (e.g. 2:16)
|
||||
if ':' in g:
|
||||
index_range = g.split(":", 1)
|
||||
index_range = [r.strip() for r in index_range]
|
||||
start_index = self.convert_to_index_int(index_range[0], latent_count=latent_count, is_range=True, allow_negative=allow_negative)
|
||||
end_index = self.convert_to_index_int(index_range[1], latent_count=latent_count, is_range=True, allow_negative=allow_negative)
|
||||
# if latents were passed in, base indeces on known latent count
|
||||
if len(int_latent_indeces) > 0:
|
||||
for i in int_latent_indeces[start_index:end_index]:
|
||||
chosen_indeces.add(LatentKeyframe(i, strength))
|
||||
# otherwise, assume indeces are valid
|
||||
else:
|
||||
for i in range(start_index, end_index):
|
||||
chosen_indeces.add(LatentKeyframe(i, strength))
|
||||
# parse individual indeces
|
||||
else:
|
||||
chosen_indeces.add(LatentKeyframe(self.convert_to_index_int(g, latent_count=latent_count, allow_negative=allow_negative), strength))
|
||||
return chosen_indeces
|
||||
|
||||
def load_keyframes(self,
|
||||
index_strengths: str,
|
||||
prev_latent_kf: LatentKeyframeGroup=None,
|
||||
prev_latent_keyframe: LatentKeyframeGroup=None, # old name
|
||||
latent_image_opt=None,
|
||||
print_keyframes=False):
|
||||
prev_latent_keyframe = prev_latent_keyframe if prev_latent_keyframe else prev_latent_kf
|
||||
if not prev_latent_keyframe:
|
||||
prev_latent_keyframe = LatentKeyframeGroup()
|
||||
else:
|
||||
prev_latent_keyframe = prev_latent_keyframe.clone()
|
||||
curr_latent_keyframe = LatentKeyframeGroup()
|
||||
|
||||
latent_count = -1
|
||||
if latent_image_opt:
|
||||
latent_count = latent_image_opt['samples'].size()[0]
|
||||
latent_keyframes = self.convert_to_latent_keyframes(index_strengths, latent_count=latent_count)
|
||||
|
||||
for latent_keyframe in latent_keyframes:
|
||||
curr_latent_keyframe.add(latent_keyframe)
|
||||
|
||||
if print_keyframes:
|
||||
for keyframe in curr_latent_keyframe.keyframes:
|
||||
logger.info(f"keyframe {keyframe.batch_index}:{keyframe.strength}")
|
||||
|
||||
# replace values with prev_latent_keyframes
|
||||
for latent_keyframe in prev_latent_keyframe.keyframes:
|
||||
curr_latent_keyframe.add(latent_keyframe)
|
||||
|
||||
return (curr_latent_keyframe,)
|
||||
|
||||
|
||||
class LatentKeyframeInterpolationNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"batch_index_from": ("INT", {"default": 0, "min": BIGMIN, "max": BIGMAX, "step": 1}),
|
||||
"batch_index_to_excl": ("INT", {"default": 0, "min": BIGMIN, "max": BIGMAX, "step": 1}),
|
||||
"strength_from": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"strength_to": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"interpolation": ([SI.LINEAR, SI.EASE_IN, SI.EASE_OUT, SI.EASE_IN_OUT], ),
|
||||
},
|
||||
"optional": {
|
||||
"prev_latent_kf": ("LATENT_KEYFRAME", ),
|
||||
"print_keyframes": ("BOOLEAN", {"default": False})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_NAMES = ("LATENT_KF", )
|
||||
RETURN_TYPES = ("LATENT_KEYFRAME", )
|
||||
FUNCTION = "load_keyframe"
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/keyframes"
|
||||
|
||||
def load_keyframe(self,
|
||||
batch_index_from: int,
|
||||
strength_from: float,
|
||||
batch_index_to_excl: int,
|
||||
strength_to: float,
|
||||
interpolation: str,
|
||||
prev_latent_kf: LatentKeyframeGroup=None,
|
||||
prev_latent_keyframe: LatentKeyframeGroup=None, # old name
|
||||
print_keyframes=False):
|
||||
|
||||
if (batch_index_from > batch_index_to_excl):
|
||||
raise ValueError("batch_index_from must be less than or equal to batch_index_to.")
|
||||
|
||||
if (batch_index_from < 0 and batch_index_to_excl >= 0):
|
||||
raise ValueError("batch_index_from and batch_index_to must be either both positive or both negative.")
|
||||
|
||||
prev_latent_keyframe = prev_latent_keyframe if prev_latent_keyframe else prev_latent_kf
|
||||
if not prev_latent_keyframe:
|
||||
prev_latent_keyframe = LatentKeyframeGroup()
|
||||
else:
|
||||
prev_latent_keyframe = prev_latent_keyframe.clone()
|
||||
curr_latent_keyframe = LatentKeyframeGroup()
|
||||
|
||||
steps = batch_index_to_excl - batch_index_from
|
||||
diff = strength_to - strength_from
|
||||
if interpolation == SI.LINEAR:
|
||||
weights = np.linspace(strength_from, strength_to, steps)
|
||||
elif interpolation == SI.EASE_IN:
|
||||
index = np.linspace(0, 1, steps)
|
||||
weights = diff * np.power(index, 2) + strength_from
|
||||
elif interpolation == SI.EASE_OUT:
|
||||
index = np.linspace(0, 1, steps)
|
||||
weights = diff * (1 - np.power(1 - index, 2)) + strength_from
|
||||
elif interpolation == SI.EASE_IN_OUT:
|
||||
index = np.linspace(0, 1, steps)
|
||||
weights = diff * ((1 - np.cos(index * np.pi)) / 2) + strength_from
|
||||
|
||||
for i in range(steps):
|
||||
keyframe = LatentKeyframe(batch_index_from + i, float(weights[i]))
|
||||
curr_latent_keyframe.add(keyframe)
|
||||
|
||||
if print_keyframes:
|
||||
for keyframe in curr_latent_keyframe.keyframes:
|
||||
logger.info(f"keyframe {keyframe.batch_index}:{keyframe.strength}")
|
||||
|
||||
# replace values with prev_latent_keyframes
|
||||
for latent_keyframe in prev_latent_keyframe.keyframes:
|
||||
curr_latent_keyframe.add(latent_keyframe)
|
||||
|
||||
return (curr_latent_keyframe,)
|
||||
|
||||
|
||||
class LatentKeyframeBatchedGroupNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"float_strengths": ("FLOAT", {"default": -1, "min": -1, "step": 0.001, "forceInput": True}),
|
||||
},
|
||||
"optional": {
|
||||
"prev_latent_kf": ("LATENT_KEYFRAME", ),
|
||||
"print_keyframes": ("BOOLEAN", {"default": False})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_NAMES = ("LATENT_KF", )
|
||||
RETURN_TYPES = ("LATENT_KEYFRAME", )
|
||||
FUNCTION = "load_keyframe"
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/keyframes"
|
||||
|
||||
def load_keyframe(self, float_strengths: Union[float, list[float]],
|
||||
prev_latent_kf: LatentKeyframeGroup=None,
|
||||
prev_latent_keyframe: LatentKeyframeGroup=None, # old name
|
||||
print_keyframes=False):
|
||||
prev_latent_keyframe = prev_latent_keyframe if prev_latent_keyframe else prev_latent_kf
|
||||
if not prev_latent_keyframe:
|
||||
prev_latent_keyframe = LatentKeyframeGroup()
|
||||
else:
|
||||
prev_latent_keyframe = prev_latent_keyframe.clone()
|
||||
curr_latent_keyframe = LatentKeyframeGroup()
|
||||
|
||||
# if received a normal float input, do nothing
|
||||
if type(float_strengths) in (float, int):
|
||||
logger.info("No batched float_strengths passed into Latent Keyframe Batch Group node; will not create any new keyframes.")
|
||||
# if iterable, attempt to create LatentKeyframes with chosen strengths
|
||||
elif isinstance(float_strengths, Iterable):
|
||||
for idx, strength in enumerate(float_strengths):
|
||||
keyframe = LatentKeyframe(idx, strength)
|
||||
curr_latent_keyframe.add(keyframe)
|
||||
else:
|
||||
raise ValueError(f"Expected strengths to be an iterable input, but was {type(float_strengths).__repr__}.")
|
||||
|
||||
if print_keyframes:
|
||||
for keyframe in curr_latent_keyframe.keyframes:
|
||||
logger.info(f"keyframe {keyframe.batch_index}:{keyframe.strength}")
|
||||
|
||||
# replace values with prev_latent_keyframes
|
||||
for latent_keyframe in prev_latent_keyframe.keyframes:
|
||||
curr_latent_keyframe.add(latent_keyframe)
|
||||
|
||||
return (curr_latent_keyframe,)
|
||||
+67
@@ -0,0 +1,67 @@
|
||||
import folder_paths
|
||||
import comfy.utils
|
||||
import comfy.model_detection
|
||||
import comfy.model_management
|
||||
import comfy.lora
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from .utils import TimestepKeyframeGroup
|
||||
from .control import ControlNetAdvanced, load_controlnet
|
||||
|
||||
|
||||
|
||||
|
||||
def convert_cn_lora_from_diffusers(cn_model: ModelPatcher, lora_path: str):
|
||||
lora_data = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||||
unet_dtype = comfy.model_management.unet_dtype()
|
||||
for key, value in lora_data.items():
|
||||
lora_data[key] = value.to(unet_dtype)
|
||||
diffusers_keys = comfy.utils.unet_to_diffusers(cn_model.model.state_dict())
|
||||
|
||||
#lora_data = comfy.model_detection.unet_config_from_diffusers_unet(lora_data, dtype=unet_dtype)
|
||||
|
||||
|
||||
|
||||
#key_map = comfy.lora.model_lora_keys_unet(cn_model.model, key_map)
|
||||
lora_data = comfy.lora.load_lora(lora_data, to_load=diffusers_keys)
|
||||
|
||||
# TODO: detect if diffusers for sure? not sure if needed at this time, since cn loras are
|
||||
# only used currently for LOOSEControl, and those are all in diffusers format
|
||||
#unet_dtype = comfy.model_management.unet_dtype()
|
||||
#lora_data = comfy.model_detection.unet_config_from_diffusers_unet(lora_data, unet_dtype)
|
||||
return lora_data
|
||||
|
||||
|
||||
class ControlNetLoaderWithLoraAdvanced:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"control_net_name": (folder_paths.get_filename_list("controlnet"), ),
|
||||
"cn_lora_name": (folder_paths.get_filename_list("controlnet"), ),
|
||||
"cn_lora_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
},
|
||||
"optional": {
|
||||
"timestep_keyframe": ("TIMESTEP_KEYFRAME", ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET", )
|
||||
FUNCTION = "load_controlnet"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/LOOSEControl"
|
||||
|
||||
def load_controlnet(self, control_net_name, cn_lora_name, cn_lora_strength: float,
|
||||
timestep_keyframe: TimestepKeyframeGroup=None
|
||||
):
|
||||
controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
|
||||
controlnet: ControlNetAdvanced = load_controlnet(controlnet_path, timestep_keyframe)
|
||||
if not isinstance(controlnet, ControlNetAdvanced):
|
||||
raise ValueError("Type {} is not compatible with CN LoRA features at this time.")
|
||||
# now, try to load CN LoRA
|
||||
lora_path = folder_paths.get_full_path("controlnet", cn_lora_name)
|
||||
lora_data = convert_cn_lora_from_diffusers(cn_model=controlnet.control_model_wrapped, lora_path=lora_path)
|
||||
# apply patches to wrapped control_model
|
||||
controlnet.control_model_wrapped.add_patches(lora_data, strength_patch=cn_lora_strength)
|
||||
# all done
|
||||
return (controlnet,)
|
||||
+163
@@ -0,0 +1,163 @@
|
||||
from torch import Tensor
|
||||
|
||||
import folder_paths
|
||||
from nodes import VAEEncode
|
||||
import comfy.utils
|
||||
from comfy.sd import VAE
|
||||
|
||||
from .utils import TimestepKeyframeGroup
|
||||
from .control_sparsectrl import SparseMethod, SparseIndexMethod, SparseSettings, SparseSpreadMethod, PreprocSparseRGBWrapper
|
||||
from .control import load_sparsectrl, load_controlnet, ControlNetAdvanced, SparseCtrlAdvanced
|
||||
|
||||
|
||||
# node for SparseCtrl loading
|
||||
class SparseCtrlLoaderAdvanced:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"sparsectrl_name": (folder_paths.get_filename_list("controlnet"), ),
|
||||
"use_motion": ("BOOLEAN", {"default": True}, ),
|
||||
"motion_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
},
|
||||
"optional": {
|
||||
"sparse_method": ("SPARSE_METHOD", ),
|
||||
"tk_optional": ("TIMESTEP_KEYFRAME", ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET", )
|
||||
FUNCTION = "load_controlnet"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/SparseCtrl"
|
||||
|
||||
def load_controlnet(self, sparsectrl_name: str, use_motion: bool, motion_strength: float, motion_scale: float, sparse_method: SparseMethod=SparseSpreadMethod(), tk_optional: TimestepKeyframeGroup=None):
|
||||
sparsectrl_path = folder_paths.get_full_path("controlnet", sparsectrl_name)
|
||||
sparse_settings = SparseSettings(sparse_method=sparse_method, use_motion=use_motion, motion_strength=motion_strength, motion_scale=motion_scale)
|
||||
sparsectrl = load_sparsectrl(sparsectrl_path, timestep_keyframe=tk_optional, sparse_settings=sparse_settings)
|
||||
return (sparsectrl,)
|
||||
|
||||
|
||||
class SparseCtrlMergedLoaderAdvanced:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"sparsectrl_name": (folder_paths.get_filename_list("controlnet"), ),
|
||||
"control_net_name": (folder_paths.get_filename_list("controlnet"), ),
|
||||
"use_motion": ("BOOLEAN", {"default": True}, ),
|
||||
"motion_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
},
|
||||
"optional": {
|
||||
"sparse_method": ("SPARSE_METHOD", ),
|
||||
"tk_optional": ("TIMESTEP_KEYFRAME", ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET", )
|
||||
FUNCTION = "load_controlnet"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/SparseCtrl/experimental"
|
||||
|
||||
def load_controlnet(self, sparsectrl_name: str, control_net_name: str, use_motion: bool, motion_strength: float, motion_scale: float, sparse_method: SparseMethod=SparseSpreadMethod(), tk_optional: TimestepKeyframeGroup=None):
|
||||
sparsectrl_path = folder_paths.get_full_path("controlnet", sparsectrl_name)
|
||||
controlnet_path = folder_paths.get_full_path("controlnet", control_net_name)
|
||||
sparse_settings = SparseSettings(sparse_method=sparse_method, use_motion=use_motion, motion_strength=motion_strength, motion_scale=motion_scale, merged=True)
|
||||
# first, load normal controlnet
|
||||
controlnet = load_controlnet(controlnet_path, timestep_keyframe=tk_optional)
|
||||
# confirm that controlnet is ControlNetAdvanced
|
||||
if controlnet is None or type(controlnet) != ControlNetAdvanced:
|
||||
raise ValueError(f"controlnet_path must point to a normal ControlNet, but instead: {type(controlnet).__name__}")
|
||||
# next, load sparsectrl, making sure to load motion portion
|
||||
sparsectrl = load_sparsectrl(sparsectrl_path, timestep_keyframe=tk_optional, sparse_settings=SparseSettings.default())
|
||||
# now, combine state dicts
|
||||
new_state_dict = controlnet.control_model.state_dict()
|
||||
for key, value in sparsectrl.control_model.motion_holder.motion_wrapper.state_dict().items():
|
||||
new_state_dict[key] = value
|
||||
# now, reload sparsectrl with real settings
|
||||
sparsectrl = load_sparsectrl(sparsectrl_path, controlnet_data=new_state_dict, timestep_keyframe=tk_optional, sparse_settings=sparse_settings)
|
||||
return (sparsectrl,)
|
||||
|
||||
|
||||
class SparseIndexMethodNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"indexes": ("STRING", {"default": "0"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SPARSE_METHOD",)
|
||||
FUNCTION = "get_method"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/SparseCtrl"
|
||||
|
||||
def get_method(self, indexes: str):
|
||||
idxs = []
|
||||
unique_idxs = set()
|
||||
# get indeces from string
|
||||
str_idxs = [x.strip() for x in indexes.strip().split(",")]
|
||||
for str_idx in str_idxs:
|
||||
try:
|
||||
idx = int(str_idx)
|
||||
if idx in unique_idxs:
|
||||
raise ValueError(f"'{idx}' is duplicated; indexes must be unique.")
|
||||
idxs.append(idx)
|
||||
unique_idxs.add(idx)
|
||||
except ValueError:
|
||||
raise ValueError(f"'{str_idx}' is not a valid integer index.")
|
||||
if len(idxs) == 0:
|
||||
raise ValueError(f"No indexes were listed in Sparse Index Method.")
|
||||
return (SparseIndexMethod(idxs),)
|
||||
|
||||
|
||||
class SparseSpreadMethodNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"spread": (SparseSpreadMethod.LIST,),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SPARSE_METHOD",)
|
||||
FUNCTION = "get_method"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/SparseCtrl"
|
||||
|
||||
def get_method(self, spread: str):
|
||||
return (SparseSpreadMethod(spread=spread),)
|
||||
|
||||
|
||||
class RgbSparseCtrlPreprocessor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", ),
|
||||
"vae": ("VAE", ),
|
||||
"latent_size": ("LATENT", ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("proc_IMAGE",)
|
||||
FUNCTION = "preprocess_images"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/SparseCtrl/preprocess"
|
||||
|
||||
def preprocess_images(self, vae: VAE, image: Tensor, latent_size: Tensor):
|
||||
# first, resize image to match latents
|
||||
image = image.movedim(-1,1)
|
||||
image = comfy.utils.common_upscale(image, latent_size["samples"].shape[3] * 8, latent_size["samples"].shape[2] * 8, 'nearest-exact', "center")
|
||||
image = image.movedim(1,-1)
|
||||
# then, vae encode
|
||||
try:
|
||||
image = vae.vae_encode_crop_pixels(image)
|
||||
except Exception:
|
||||
image = VAEEncode.vae_encode_crop_pixels(image)
|
||||
encoded = vae.encode(image[:,:,:,:3])
|
||||
return (PreprocSparseRGBWrapper(condhint=encoded),)
|
||||
+201
@@ -0,0 +1,201 @@
|
||||
from torch import Tensor
|
||||
import torch
|
||||
from .utils import TimestepKeyframe, TimestepKeyframeGroup, ControlWeights, get_properly_arranged_t2i_weights, linear_conversion
|
||||
from .logger import logger
|
||||
|
||||
|
||||
WEIGHTS_RETURN_NAMES = ("CN_WEIGHTS", "TK_SHORTCUT")
|
||||
|
||||
|
||||
class DefaultWeights:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",)
|
||||
RETURN_NAMES = WEIGHTS_RETURN_NAMES
|
||||
FUNCTION = "load_weights"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/weights"
|
||||
|
||||
def load_weights(self):
|
||||
weights = ControlWeights.default()
|
||||
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))
|
||||
|
||||
|
||||
class ScaledSoftMaskedUniversalWeights:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK", ),
|
||||
"min_base_multiplier": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}, ),
|
||||
"max_base_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}, ),
|
||||
#"lock_min": ("BOOLEAN", {"default": False}, ),
|
||||
#"lock_max": ("BOOLEAN", {"default": False}, ),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",)
|
||||
RETURN_NAMES = WEIGHTS_RETURN_NAMES
|
||||
FUNCTION = "load_weights"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/weights"
|
||||
|
||||
def load_weights(self, mask: Tensor, min_base_multiplier: float, max_base_multiplier: float, lock_min=False, lock_max=False):
|
||||
# normalize mask
|
||||
mask = mask.clone()
|
||||
x_min = 0.0 if lock_min else mask.min()
|
||||
x_max = 1.0 if lock_max else mask.max()
|
||||
if x_min == x_max:
|
||||
mask = torch.ones_like(mask) * max_base_multiplier
|
||||
else:
|
||||
mask = linear_conversion(mask, x_min, x_max, min_base_multiplier, max_base_multiplier)
|
||||
weights = ControlWeights.universal_mask(weight_mask=mask)
|
||||
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))
|
||||
|
||||
|
||||
class ScaledSoftUniversalWeights:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"base_multiplier": ("FLOAT", {"default": 0.825, "min": 0.0, "max": 1.0, "step": 0.001}, ),
|
||||
"flip_weights": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",)
|
||||
RETURN_NAMES = WEIGHTS_RETURN_NAMES
|
||||
FUNCTION = "load_weights"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/weights"
|
||||
|
||||
def load_weights(self, base_multiplier, flip_weights):
|
||||
weights = ControlWeights.universal(base_multiplier=base_multiplier, flip_weights=flip_weights)
|
||||
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))
|
||||
|
||||
|
||||
class SoftControlNetWeights:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"weight_00": ("FLOAT", {"default": 0.09941396206337118, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_01": ("FLOAT", {"default": 0.12050177219802567, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_02": ("FLOAT", {"default": 0.14606275417942507, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_03": ("FLOAT", {"default": 0.17704576264172736, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_04": ("FLOAT", {"default": 0.214600924414215, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_05": ("FLOAT", {"default": 0.26012233262329093, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_06": ("FLOAT", {"default": 0.3152997971191405, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_07": ("FLOAT", {"default": 0.3821815722656249, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_08": ("FLOAT", {"default": 0.4632503906249999, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_09": ("FLOAT", {"default": 0.561515625, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_10": ("FLOAT", {"default": 0.6806249999999999, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_11": ("FLOAT", {"default": 0.825, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_12": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"flip_weights": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",)
|
||||
RETURN_NAMES = WEIGHTS_RETURN_NAMES
|
||||
FUNCTION = "load_weights"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/weights/ControlNet"
|
||||
|
||||
def load_weights(self, weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06,
|
||||
weight_07, weight_08, weight_09, weight_10, weight_11, weight_12, flip_weights):
|
||||
weights = [weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06,
|
||||
weight_07, weight_08, weight_09, weight_10, weight_11, weight_12]
|
||||
weights = ControlWeights.controlnet(weights, flip_weights=flip_weights)
|
||||
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))
|
||||
|
||||
|
||||
class CustomControlNetWeights:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"weight_00": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_01": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_02": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_03": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_04": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_05": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_06": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_07": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_08": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_09": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_10": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_11": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_12": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"flip_weights": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",)
|
||||
RETURN_NAMES = WEIGHTS_RETURN_NAMES
|
||||
FUNCTION = "load_weights"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/weights/ControlNet"
|
||||
|
||||
def load_weights(self, weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06,
|
||||
weight_07, weight_08, weight_09, weight_10, weight_11, weight_12, flip_weights):
|
||||
weights = [weight_00, weight_01, weight_02, weight_03, weight_04, weight_05, weight_06,
|
||||
weight_07, weight_08, weight_09, weight_10, weight_11, weight_12]
|
||||
weights = ControlWeights.controlnet(weights, flip_weights=flip_weights)
|
||||
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))
|
||||
|
||||
|
||||
class SoftT2IAdapterWeights:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"weight_00": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_01": ("FLOAT", {"default": 0.62, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_02": ("FLOAT", {"default": 0.825, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_03": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"flip_weights": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",)
|
||||
RETURN_NAMES = WEIGHTS_RETURN_NAMES
|
||||
FUNCTION = "load_weights"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/weights/T2IAdapter"
|
||||
|
||||
def load_weights(self, weight_00, weight_01, weight_02, weight_03, flip_weights):
|
||||
weights = [weight_00, weight_01, weight_02, weight_03]
|
||||
weights = get_properly_arranged_t2i_weights(weights)
|
||||
weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights)
|
||||
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))
|
||||
|
||||
|
||||
class CustomT2IAdapterWeights:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"weight_00": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_01": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_02": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"weight_03": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}, ),
|
||||
"flip_weights": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTROL_NET_WEIGHTS", "TIMESTEP_KEYFRAME",)
|
||||
RETURN_NAMES = WEIGHTS_RETURN_NAMES
|
||||
FUNCTION = "load_weights"
|
||||
|
||||
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/weights/T2IAdapter"
|
||||
|
||||
def load_weights(self, weight_00, weight_01, weight_02, weight_03, flip_weights):
|
||||
weights = [weight_00, weight_01, weight_02, weight_03]
|
||||
weights = get_properly_arranged_t2i_weights(weights)
|
||||
weights = ControlWeights.t2iadapter(weights, flip_weights=flip_weights)
|
||||
return (weights, TimestepKeyframeGroup.default(TimestepKeyframe(control_weights=weights)))
|
||||
+677
@@ -0,0 +1,677 @@
|
||||
from copy import deepcopy
|
||||
from typing import Callable, Union
|
||||
import torch
|
||||
from torch import Tensor
|
||||
import torch.nn.functional as F
|
||||
|
||||
import comfy.ops
|
||||
import comfy.utils
|
||||
from comfy.controlnet import ControlBase, broadcast_image_to
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from .logger import logger
|
||||
|
||||
BIGMIN = -(2**63-1)
|
||||
BIGMAX = (2**63-1)
|
||||
|
||||
def load_torch_file_with_dict_factory(controlnet_data: dict[str, Tensor], orig_load_torch_file: Callable):
|
||||
def load_torch_file_with_dict(*args, **kwargs):
|
||||
# immediately restore load_torch_file to original version
|
||||
comfy.utils.load_torch_file = orig_load_torch_file
|
||||
return controlnet_data
|
||||
return load_torch_file_with_dict
|
||||
|
||||
|
||||
def get_properly_arranged_t2i_weights(initial_weights: list[float]):
|
||||
new_weights = []
|
||||
new_weights.extend([initial_weights[0]]*3)
|
||||
new_weights.extend([initial_weights[1]]*3)
|
||||
new_weights.extend([initial_weights[2]]*3)
|
||||
new_weights.extend([initial_weights[3]]*3)
|
||||
return new_weights
|
||||
|
||||
|
||||
class ControlWeightType:
|
||||
DEFAULT = "default"
|
||||
UNIVERSAL = "universal"
|
||||
T2IADAPTER = "t2iadapter"
|
||||
CONTROLNET = "controlnet"
|
||||
CONTROLLORA = "controllora"
|
||||
CONTROLLLLITE = "controllllite"
|
||||
SVD_CONTROLNET = "svd_controlnet"
|
||||
SPARSECTRL = "sparsectrl"
|
||||
|
||||
|
||||
class ControlWeights:
|
||||
def __init__(self, weight_type: str, base_multiplier: float=1.0, flip_weights: bool=False, weights: list[float]=None, weight_mask: Tensor=None):
|
||||
self.weight_type = weight_type
|
||||
self.base_multiplier = base_multiplier
|
||||
self.flip_weights = flip_weights
|
||||
self.weights = weights
|
||||
if self.weights is not None and self.flip_weights:
|
||||
self.weights.reverse()
|
||||
self.weight_mask = weight_mask
|
||||
|
||||
def get(self, idx: int) -> Union[float, Tensor]:
|
||||
# if weights is not none, return index
|
||||
if self.weights is not None:
|
||||
return self.weights[idx]
|
||||
return 1.0
|
||||
|
||||
@classmethod
|
||||
def default(cls):
|
||||
return cls(ControlWeightType.DEFAULT)
|
||||
|
||||
@classmethod
|
||||
def universal(cls, base_multiplier: float, flip_weights: bool=False):
|
||||
return cls(ControlWeightType.UNIVERSAL, base_multiplier=base_multiplier, flip_weights=flip_weights)
|
||||
|
||||
@classmethod
|
||||
def universal_mask(cls, weight_mask: Tensor):
|
||||
return cls(ControlWeightType.UNIVERSAL, weight_mask=weight_mask)
|
||||
|
||||
@classmethod
|
||||
def t2iadapter(cls, weights: list[float]=None, flip_weights: bool=False):
|
||||
if weights is None:
|
||||
weights = [1.0]*12
|
||||
return cls(ControlWeightType.T2IADAPTER, weights=weights,flip_weights=flip_weights)
|
||||
|
||||
@classmethod
|
||||
def controlnet(cls, weights: list[float]=None, flip_weights: bool=False):
|
||||
if weights is None:
|
||||
weights = [1.0]*13
|
||||
return cls(ControlWeightType.CONTROLNET, weights=weights, flip_weights=flip_weights)
|
||||
|
||||
@classmethod
|
||||
def controllora(cls, weights: list[float]=None, flip_weights: bool=False):
|
||||
if weights is None:
|
||||
weights = [1.0]*10
|
||||
return cls(ControlWeightType.CONTROLLORA, weights=weights, flip_weights=flip_weights)
|
||||
|
||||
@classmethod
|
||||
def controllllite(cls, weights: list[float]=None, flip_weights: bool=False):
|
||||
if weights is None:
|
||||
# TODO: make this have a real value
|
||||
weights = [1.0]*200
|
||||
return cls(ControlWeightType.CONTROLLLLITE, weights=weights, flip_weights=flip_weights)
|
||||
|
||||
|
||||
class StrengthInterpolation:
|
||||
LINEAR = "linear"
|
||||
EASE_IN = "ease-in"
|
||||
EASE_OUT = "ease-out"
|
||||
EASE_IN_OUT = "ease-in-out"
|
||||
NONE = "none"
|
||||
|
||||
|
||||
class LatentKeyframe:
|
||||
def __init__(self, batch_index: int, strength: float) -> None:
|
||||
self.batch_index = batch_index
|
||||
self.strength = strength
|
||||
|
||||
|
||||
# always maintain sorted state (by batch_index of LatentKeyframe)
|
||||
class LatentKeyframeGroup:
|
||||
def __init__(self) -> None:
|
||||
self.keyframes: list[LatentKeyframe] = []
|
||||
|
||||
def add(self, keyframe: LatentKeyframe) -> None:
|
||||
added = False
|
||||
# replace existing keyframe if same batch_index
|
||||
for i in range(len(self.keyframes)):
|
||||
if self.keyframes[i].batch_index == keyframe.batch_index:
|
||||
self.keyframes[i] = keyframe
|
||||
added = True
|
||||
break
|
||||
if not added:
|
||||
self.keyframes.append(keyframe)
|
||||
self.keyframes.sort(key=lambda k: k.batch_index)
|
||||
|
||||
def get_index(self, index: int) -> Union[LatentKeyframe, None]:
|
||||
try:
|
||||
return self.keyframes[index]
|
||||
except IndexError:
|
||||
return None
|
||||
|
||||
def __getitem__(self, index) -> LatentKeyframe:
|
||||
return self.keyframes[index]
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return len(self.keyframes) == 0
|
||||
|
||||
def clone(self) -> 'LatentKeyframeGroup':
|
||||
cloned = LatentKeyframeGroup()
|
||||
for tk in self.keyframes:
|
||||
cloned.add(tk)
|
||||
return cloned
|
||||
|
||||
|
||||
class TimestepKeyframe:
|
||||
def __init__(self,
|
||||
start_percent: float = 0.0,
|
||||
strength: float = 1.0,
|
||||
interpolation: str = StrengthInterpolation.NONE,
|
||||
control_weights: ControlWeights = None,
|
||||
latent_keyframes: LatentKeyframeGroup = None,
|
||||
null_latent_kf_strength: float = 0.0,
|
||||
inherit_missing: bool = True,
|
||||
guarantee_usage: bool = True,
|
||||
mask_hint_orig: Tensor = None) -> None:
|
||||
self.start_percent = start_percent
|
||||
self.start_t = 999999999.9
|
||||
self.strength = strength
|
||||
self.interpolation = interpolation
|
||||
self.control_weights = control_weights
|
||||
self.latent_keyframes = latent_keyframes
|
||||
self.null_latent_kf_strength = null_latent_kf_strength
|
||||
self.inherit_missing = inherit_missing
|
||||
self.guarantee_usage = guarantee_usage
|
||||
self.mask_hint_orig = mask_hint_orig
|
||||
|
||||
def has_control_weights(self):
|
||||
return self.control_weights is not None
|
||||
|
||||
def has_latent_keyframes(self):
|
||||
return self.latent_keyframes is not None
|
||||
|
||||
def has_mask_hint(self):
|
||||
return self.mask_hint_orig is not None
|
||||
|
||||
|
||||
@classmethod
|
||||
def default(cls) -> 'TimestepKeyframe':
|
||||
return cls(0.0)
|
||||
|
||||
|
||||
# always maintain sorted state (by start_percent of TimestepKeyFrame)
|
||||
class TimestepKeyframeGroup:
|
||||
def __init__(self) -> None:
|
||||
self.keyframes: list[TimestepKeyframe] = []
|
||||
self.keyframes.append(TimestepKeyframe.default())
|
||||
|
||||
def add(self, keyframe: TimestepKeyframe) -> None:
|
||||
added = False
|
||||
# replace existing keyframe if same start_percent
|
||||
for i in range(len(self.keyframes)):
|
||||
if self.keyframes[i].start_percent == keyframe.start_percent:
|
||||
self.keyframes[i] = keyframe
|
||||
added = True
|
||||
break
|
||||
if not added:
|
||||
self.keyframes.append(keyframe)
|
||||
self.keyframes.sort(key=lambda k: k.start_percent)
|
||||
|
||||
def get_index(self, index: int) -> Union[TimestepKeyframe, None]:
|
||||
try:
|
||||
return self.keyframes[index]
|
||||
except IndexError:
|
||||
return None
|
||||
|
||||
def has_index(self, index: int) -> int:
|
||||
return index >=0 and index < len(self.keyframes)
|
||||
|
||||
def __getitem__(self, index) -> TimestepKeyframe:
|
||||
return self.keyframes[index]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.keyframes)
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return len(self.keyframes) == 0
|
||||
|
||||
def clone(self) -> 'TimestepKeyframeGroup':
|
||||
cloned = TimestepKeyframeGroup()
|
||||
for tk in self.keyframes:
|
||||
cloned.add(tk)
|
||||
return cloned
|
||||
|
||||
@classmethod
|
||||
def default(cls, keyframe: TimestepKeyframe) -> 'TimestepKeyframeGroup':
|
||||
group = cls()
|
||||
group.keyframes[0] = keyframe
|
||||
return group
|
||||
|
||||
|
||||
# depending on model, AnimateDiff may inject into GroupNorm, so make sure GroupNorm will be clean
|
||||
class disable_weight_init_clean_groupnorm(comfy.ops.disable_weight_init):
|
||||
class GroupNorm(comfy.ops.disable_weight_init.GroupNorm):
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return F.group_norm(
|
||||
input, self.num_groups, self.weight, self.bias, self.eps)
|
||||
class manual_cast_clean_groupnorm(comfy.ops.manual_cast):
|
||||
class GroupNorm(comfy.ops.manual_cast.GroupNorm):
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return F.group_norm(
|
||||
input, self.num_groups, self.weight, self.bias, self.eps)
|
||||
|
||||
|
||||
# adapted from comfy/sample.py
|
||||
def prepare_mask_batch(mask: Tensor, shape: Tensor, multiplier: int=1, match_dim1=False):
|
||||
mask = mask.clone()
|
||||
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[2]*multiplier, shape[3]*multiplier), mode="bilinear")
|
||||
if match_dim1:
|
||||
mask = torch.cat([mask] * shape[1], dim=1)
|
||||
return mask
|
||||
|
||||
|
||||
# applies min-max normalization, from:
|
||||
# https://stackoverflow.com/questions/68791508/min-max-normalization-of-a-tensor-in-pytorch
|
||||
def normalize_min_max(x: Tensor, new_min = 0.0, new_max = 1.0):
|
||||
x_min, x_max = x.min(), x.max()
|
||||
return (((x - x_min)/(x_max - x_min)) * (new_max - new_min)) + new_min
|
||||
|
||||
def linear_conversion(x, x_min=0.0, x_max=1.0, new_min=0.0, new_max=1.0):
|
||||
return (((x - x_min)/(x_max - x_min)) * (new_max - new_min)) + new_min
|
||||
|
||||
|
||||
# from https://stackoverflow.com/a/24621200
|
||||
def deepcopy_with_sharing(obj, shared_attribute_names, memo=None):
|
||||
'''
|
||||
Deepcopy an object, except for a given list of attributes, which should
|
||||
be shared between the original object and its copy.
|
||||
|
||||
obj is some object
|
||||
shared_attribute_names: A list of strings identifying the attributes that
|
||||
should be shared between the original and its copy.
|
||||
memo is the dictionary passed into __deepcopy__. Ignore this argument if
|
||||
not calling from within __deepcopy__.
|
||||
'''
|
||||
assert isinstance(shared_attribute_names, (list, tuple))
|
||||
|
||||
shared_attributes = {k: getattr(obj, k) for k in shared_attribute_names}
|
||||
|
||||
if hasattr(obj, '__deepcopy__'):
|
||||
# Do hack to prevent infinite recursion in call to deepcopy
|
||||
deepcopy_method = obj.__deepcopy__
|
||||
obj.__deepcopy__ = None
|
||||
|
||||
for attr in shared_attribute_names:
|
||||
del obj.__dict__[attr]
|
||||
|
||||
clone = deepcopy(obj)
|
||||
|
||||
for attr, val in shared_attributes.items():
|
||||
setattr(obj, attr, val)
|
||||
setattr(clone, attr, val)
|
||||
|
||||
if hasattr(obj, '__deepcopy__'):
|
||||
# Undo hack
|
||||
obj.__deepcopy__ = deepcopy_method
|
||||
del clone.__deepcopy__
|
||||
|
||||
return clone
|
||||
|
||||
|
||||
class WeightTypeException(TypeError):
|
||||
"Raised when weight not compatible with AdvancedControlBase object"
|
||||
pass
|
||||
|
||||
|
||||
class AdvancedControlBase:
|
||||
def __init__(self, base: ControlBase, timestep_keyframes: TimestepKeyframeGroup, weights_default: ControlWeights, require_model=False):
|
||||
self.base = base
|
||||
self.compatible_weights = [ControlWeightType.UNIVERSAL]
|
||||
self.add_compatible_weight(weights_default.weight_type)
|
||||
# mask for which parts of controlnet output to keep
|
||||
self.mask_cond_hint_original = None
|
||||
self.mask_cond_hint = None
|
||||
self.tk_mask_cond_hint_original = None
|
||||
self.tk_mask_cond_hint = None
|
||||
self.weight_mask_cond_hint = None
|
||||
# actual index values
|
||||
self.sub_idxs = None
|
||||
self.full_latent_length = 0
|
||||
self.context_length = 0
|
||||
# timesteps
|
||||
self.t: Tensor = None
|
||||
self.batched_number: int = None
|
||||
# weights + override
|
||||
self.weights: ControlWeights = None
|
||||
self.weights_default: ControlWeights = weights_default
|
||||
self.weights_override: ControlWeights = None
|
||||
# latent keyframe + override
|
||||
self.latent_keyframes: LatentKeyframeGroup = None
|
||||
self.latent_keyframe_override: LatentKeyframeGroup = None
|
||||
# initialize timestep_keyframes
|
||||
self.set_timestep_keyframes(timestep_keyframes)
|
||||
# override some functions
|
||||
self.get_control = self.get_control_inject
|
||||
self.control_merge = self.control_merge_inject#.__get__(self, type(self))
|
||||
self.pre_run = self.pre_run_inject
|
||||
self.cleanup = self.cleanup_inject
|
||||
self.set_previous_controlnet = self.set_previous_controlnet_inject
|
||||
# require model to be passed into Apply Advanced ControlNet 🛂🅐🅒🅝 node
|
||||
self.require_model = require_model
|
||||
# disarm - when set to False, used to force usage of Apply Advanced ControlNet 🛂🅐🅒🅝 node (which will set it to True)
|
||||
self.disarmed = not require_model
|
||||
|
||||
def patch_model(self, model: ModelPatcher):
|
||||
pass
|
||||
|
||||
def add_compatible_weight(self, control_weight_type: str):
|
||||
self.compatible_weights.append(control_weight_type)
|
||||
|
||||
def verify_all_weights(self, throw_error=True):
|
||||
# first, check if override exists - if so, only need to check the override
|
||||
if self.weights_override is not None:
|
||||
if self.weights_override.weight_type not in self.compatible_weights:
|
||||
msg = f"Weight override is type {self.weights_override.weight_type}, but loaded {type(self).__name__}" + \
|
||||
f"only supports {self.compatible_weights} weights."
|
||||
raise WeightTypeException(msg)
|
||||
# otherwise, check all timestep keyframe weights
|
||||
else:
|
||||
for tk in self.timestep_keyframes.keyframes:
|
||||
if tk.has_control_weights() and tk.control_weights.weight_type not in self.compatible_weights:
|
||||
msg = f"Weight on Timestep Keyframe with start_percent={tk.start_percent} is type" + \
|
||||
f"{tk.control_weights.weight_type}, but loaded {type(self).__name__} only supports {self.compatible_weights} weights."
|
||||
raise WeightTypeException(msg)
|
||||
|
||||
def set_timestep_keyframes(self, timestep_keyframes: TimestepKeyframeGroup):
|
||||
self.timestep_keyframes = timestep_keyframes if timestep_keyframes else TimestepKeyframeGroup()
|
||||
# prepare first timestep_keyframe related stuff
|
||||
self.current_timestep_keyframe = None
|
||||
self.current_timestep_index = -1
|
||||
self.next_timestep_keyframe = None
|
||||
self.weights = None
|
||||
self.latent_keyframes = None
|
||||
|
||||
def prepare_current_timestep(self, t: Tensor, batched_number: int):
|
||||
self.t = float(t[0])
|
||||
self.batched_number = batched_number
|
||||
# get current step percent
|
||||
curr_t: float = self.t
|
||||
prev_index = self.current_timestep_index
|
||||
# if has next index, loop through and see if need to switch
|
||||
if self.timestep_keyframes.has_index(self.current_timestep_index+1):
|
||||
for i in range(self.current_timestep_index+1, len(self.timestep_keyframes)):
|
||||
eval_tk = self.timestep_keyframes[i]
|
||||
# check if start percent is less or equal to curr_t
|
||||
if eval_tk.start_t >= curr_t:
|
||||
self.current_timestep_index = i
|
||||
self.current_timestep_keyframe = eval_tk
|
||||
# keep track of control weights, latent keyframes, and masks,
|
||||
# accounting for inherit_missing
|
||||
if self.current_timestep_keyframe.has_control_weights():
|
||||
self.weights = self.current_timestep_keyframe.control_weights
|
||||
elif not self.current_timestep_keyframe.inherit_missing:
|
||||
self.weights = self.weights_default
|
||||
if self.current_timestep_keyframe.has_latent_keyframes():
|
||||
self.latent_keyframes = self.current_timestep_keyframe.latent_keyframes
|
||||
elif not self.current_timestep_keyframe.inherit_missing:
|
||||
self.latent_keyframes = None
|
||||
if self.current_timestep_keyframe.has_mask_hint():
|
||||
self.tk_mask_cond_hint_original = self.current_timestep_keyframe.mask_hint_orig
|
||||
elif not self.current_timestep_keyframe.inherit_missing:
|
||||
del self.tk_mask_cond_hint_original
|
||||
self.tk_mask_cond_hint_original = None
|
||||
# if guarantee_usage, stop searching for other TKs
|
||||
if self.current_timestep_keyframe.guarantee_usage:
|
||||
break
|
||||
# if eval_tk is outside of percent range, stop looking further
|
||||
else:
|
||||
break
|
||||
|
||||
# if index changed, apply overrides
|
||||
if prev_index != self.current_timestep_index:
|
||||
if self.weights_override is not None:
|
||||
self.weights = self.weights_override
|
||||
if self.latent_keyframe_override is not None:
|
||||
self.latent_keyframes = self.latent_keyframe_override
|
||||
|
||||
# make sure weights and latent_keyframes are in a workable state
|
||||
# Note: each AdvancedControlBase should create their own get_universal_weights class
|
||||
self.prepare_weights()
|
||||
|
||||
def prepare_weights(self):
|
||||
if self.weights is None or self.weights.weight_type == ControlWeightType.DEFAULT:
|
||||
self.weights = self.weights_default
|
||||
elif self.weights.weight_type == ControlWeightType.UNIVERSAL:
|
||||
# if universal and weight_mask present, no need to convert
|
||||
if self.weights.weight_mask is not None:
|
||||
return
|
||||
self.weights = self.get_universal_weights()
|
||||
|
||||
def get_universal_weights(self) -> ControlWeights:
|
||||
return self.weights
|
||||
|
||||
def set_cond_hint_mask(self, mask_hint):
|
||||
self.mask_cond_hint_original = mask_hint
|
||||
return self
|
||||
|
||||
def pre_run_inject(self, model, percent_to_timestep_function):
|
||||
self.base.pre_run(model, percent_to_timestep_function)
|
||||
self.pre_run_advanced(model, percent_to_timestep_function)
|
||||
|
||||
def pre_run_advanced(self, model, percent_to_timestep_function):
|
||||
# for each timestep keyframe, calculate the start_t
|
||||
for tk in self.timestep_keyframes.keyframes:
|
||||
tk.start_t = percent_to_timestep_function(tk.start_percent)
|
||||
# clear variables
|
||||
self.cleanup_advanced()
|
||||
|
||||
def set_previous_controlnet_inject(self, *args, **kwargs):
|
||||
to_return = self.base.set_previous_controlnet(*args, **kwargs)
|
||||
if not self.disarmed:
|
||||
raise Exception(f"Type '{type(self).__name__}' must be used with Apply Advanced ControlNet 🛂🅐🅒🅝 node (with model_optional passed in); otherwise, it will not work.")
|
||||
return to_return
|
||||
|
||||
def disarm(self):
|
||||
self.disarmed = True
|
||||
|
||||
def get_control_inject(self, x_noisy, t, cond, batched_number):
|
||||
# prepare timestep and everything related
|
||||
self.prepare_current_timestep(t=t, batched_number=batched_number)
|
||||
# if should not perform any actions for the controlnet, exit without doing any work
|
||||
if self.strength == 0.0 or self.current_timestep_keyframe.strength == 0.0:
|
||||
return self.default_control_actions(x_noisy, t, cond, batched_number)
|
||||
# otherwise, perform normal function
|
||||
return self.get_control_advanced(x_noisy, t, cond, batched_number)
|
||||
|
||||
def get_control_advanced(self, x_noisy, t, cond, batched_number):
|
||||
return self.default_control_actions(x_noisy, t, cond, batched_number)
|
||||
|
||||
def default_control_actions(self, x_noisy, t, cond, batched_number):
|
||||
control_prev = None
|
||||
if self.previous_controlnet is not None:
|
||||
control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
|
||||
return control_prev
|
||||
|
||||
def calc_weight(self, idx: int, x: Tensor, layers: int) -> Union[float, Tensor]:
|
||||
if self.weights.weight_mask is not None:
|
||||
# prepare weight mask
|
||||
self.prepare_weight_mask_cond_hint(x, self.batched_number)
|
||||
# adjust mask for current layer and return
|
||||
return torch.pow(self.weight_mask_cond_hint, self.get_calc_pow(idx=idx, layers=layers))
|
||||
return self.weights.get(idx=idx)
|
||||
|
||||
def get_calc_pow(self, idx: int, layers: int) -> int:
|
||||
return (layers-1)-idx
|
||||
|
||||
def calc_latent_keyframe_mults(self, x: Tensor, batched_number: int) -> Tensor:
|
||||
# apply strengths, and get batch indeces to null out
|
||||
# AKA latents that should not be influenced by ControlNet
|
||||
final_mults = [1.0] * x.shape[0]
|
||||
if self.latent_keyframes:
|
||||
latent_count = x.shape[0] // batched_number
|
||||
indeces_to_null = set(range(latent_count))
|
||||
mapped_indeces = None
|
||||
# if expecting subdivision, will need to translate between subset and actual idx values
|
||||
if self.sub_idxs:
|
||||
mapped_indeces = {}
|
||||
for i, actual in enumerate(self.sub_idxs):
|
||||
mapped_indeces[actual] = i
|
||||
for keyframe in self.latent_keyframes:
|
||||
real_index = keyframe.batch_index
|
||||
# if negative, count from end
|
||||
if real_index < 0:
|
||||
real_index += latent_count if self.sub_idxs is None else self.full_latent_length
|
||||
|
||||
# if not mapping indeces, what you see is what you get
|
||||
if mapped_indeces is None:
|
||||
if real_index in indeces_to_null:
|
||||
indeces_to_null.remove(real_index)
|
||||
# otherwise, see if batch_index is even included in this set of latents
|
||||
else:
|
||||
real_index = mapped_indeces.get(real_index, None)
|
||||
if real_index is None:
|
||||
continue
|
||||
indeces_to_null.remove(real_index)
|
||||
|
||||
# if real_index is outside the bounds of latents, don't apply
|
||||
if real_index >= latent_count or real_index < 0:
|
||||
continue
|
||||
|
||||
# apply strength for each batched cond/uncond
|
||||
for b in range(batched_number):
|
||||
final_mults[(latent_count*b)+real_index] = keyframe.strength
|
||||
# null them out by multiplying by null_latent_kf_strength
|
||||
for batch_index in indeces_to_null:
|
||||
# apply null for each batched cond/uncond
|
||||
for b in range(batched_number):
|
||||
final_mults[(latent_count*b)+batch_index] = self.current_timestep_keyframe.null_latent_kf_strength
|
||||
# convert final_mults into tensor and match expected dimension count
|
||||
final_tensor = torch.tensor(final_mults, dtype=x.dtype, device=x.device)
|
||||
while len(final_tensor.shape) < len(x.shape):
|
||||
final_tensor = final_tensor.unsqueeze(-1)
|
||||
return final_tensor
|
||||
|
||||
def apply_advanced_strengths_and_masks(self, x: Tensor, batched_number: int):
|
||||
if self.latent_keyframes is not None:
|
||||
x[:] = x[:] * self.calc_latent_keyframe_mults(x=x, batched_number=batched_number)
|
||||
# apply masks, resizing mask to required dims
|
||||
if self.mask_cond_hint is not None:
|
||||
masks = prepare_mask_batch(self.mask_cond_hint, x.shape)
|
||||
x[:] = x[:] * masks
|
||||
if self.tk_mask_cond_hint is not None:
|
||||
masks = prepare_mask_batch(self.tk_mask_cond_hint, x.shape)
|
||||
x[:] = x[:] * masks
|
||||
# apply timestep keyframe strengths
|
||||
if self.current_timestep_keyframe.strength != 1.0:
|
||||
x[:] *= self.current_timestep_keyframe.strength
|
||||
|
||||
def control_merge_inject(self: 'AdvancedControlBase', control_input, control_output, control_prev, output_dtype):
|
||||
out = {'input':[], 'middle':[], 'output': []}
|
||||
|
||||
if control_input is not None:
|
||||
for i in range(len(control_input)):
|
||||
key = 'input'
|
||||
x = control_input[i]
|
||||
if x is not None:
|
||||
self.apply_advanced_strengths_and_masks(x, self.batched_number)
|
||||
|
||||
x *= self.strength * self.calc_weight(i, x, len(control_input))
|
||||
if x.dtype != output_dtype:
|
||||
x = x.to(output_dtype)
|
||||
out[key].insert(0, x)
|
||||
|
||||
if control_output is not None:
|
||||
for i in range(len(control_output)):
|
||||
if i == (len(control_output) - 1):
|
||||
key = 'middle'
|
||||
index = 0
|
||||
else:
|
||||
key = 'output'
|
||||
index = i
|
||||
x = control_output[i]
|
||||
if x is not None:
|
||||
self.apply_advanced_strengths_and_masks(x, self.batched_number)
|
||||
|
||||
if self.global_average_pooling:
|
||||
x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3])
|
||||
|
||||
x *= self.strength * self.calc_weight(i, x, len(control_output))
|
||||
if x.dtype != output_dtype:
|
||||
x = x.to(output_dtype)
|
||||
|
||||
out[key].append(x)
|
||||
if control_prev is not None:
|
||||
for x in ['input', 'middle', 'output']:
|
||||
o = out[x]
|
||||
for i in range(len(control_prev[x])):
|
||||
prev_val = control_prev[x][i]
|
||||
if i >= len(o):
|
||||
o.append(prev_val)
|
||||
elif prev_val is not None:
|
||||
if o[i] is None:
|
||||
o[i] = prev_val
|
||||
else:
|
||||
if o[i].shape[0] < prev_val.shape[0]:
|
||||
o[i] = prev_val + o[i]
|
||||
else:
|
||||
o[i] += prev_val
|
||||
return out
|
||||
|
||||
def prepare_mask_cond_hint(self, x_noisy: Tensor, t, cond, batched_number, dtype=None):
|
||||
self._prepare_mask("mask_cond_hint", self.mask_cond_hint_original, x_noisy, t, cond, batched_number, dtype)
|
||||
self.prepare_tk_mask_cond_hint(x_noisy, t, cond, batched_number, dtype)
|
||||
|
||||
def prepare_tk_mask_cond_hint(self, x_noisy: Tensor, t, cond, batched_number, dtype=None):
|
||||
return self._prepare_mask("tk_mask_cond_hint", self.current_timestep_keyframe.mask_hint_orig, x_noisy, t, cond, batched_number, dtype)
|
||||
|
||||
def prepare_weight_mask_cond_hint(self, x_noisy: Tensor, batched_number, dtype=None):
|
||||
return self._prepare_mask("weight_mask_cond_hint", self.weights.weight_mask, x_noisy, t=None, cond=None, batched_number=batched_number, dtype=dtype, direct_attn=True)
|
||||
|
||||
def _prepare_mask(self, attr_name, orig_mask: Tensor, x_noisy: Tensor, t, cond, batched_number, dtype=None, direct_attn=False):
|
||||
# make mask appropriate dimensions, if present
|
||||
if orig_mask is not None:
|
||||
out_mask = getattr(self, attr_name)
|
||||
if self.sub_idxs is not None or out_mask is None or x_noisy.shape[2] * 8 != out_mask.shape[1] or x_noisy.shape[3] * 8 != out_mask.shape[2]:
|
||||
self._reset_attr(attr_name)
|
||||
del out_mask
|
||||
# TODO: perform upscale on only the sub_idxs masks at a time instead of all to conserve RAM
|
||||
# resize mask and match batch count
|
||||
multiplier = 1 if direct_attn else 8
|
||||
out_mask = prepare_mask_batch(orig_mask, x_noisy.shape, multiplier=multiplier)
|
||||
actual_latent_length = x_noisy.shape[0] // batched_number
|
||||
out_mask = comfy.utils.repeat_to_batch_size(out_mask, actual_latent_length if self.sub_idxs is None else self.full_latent_length)
|
||||
if self.sub_idxs is not None:
|
||||
out_mask = out_mask[self.sub_idxs]
|
||||
# make cond_hint_mask length match x_noise
|
||||
if x_noisy.shape[0] != out_mask.shape[0]:
|
||||
out_mask = broadcast_image_to(out_mask, x_noisy.shape[0], batched_number)
|
||||
# default dtype to be same as x_noisy
|
||||
if dtype is None:
|
||||
dtype = x_noisy.dtype
|
||||
setattr(self, attr_name, out_mask.to(dtype=dtype).to(self.device))
|
||||
del out_mask
|
||||
|
||||
def _reset_attr(self, attr_name, new_value=None):
|
||||
if hasattr(self, attr_name):
|
||||
delattr(self, attr_name)
|
||||
setattr(self, attr_name, new_value)
|
||||
|
||||
def cleanup_inject(self):
|
||||
self.base.cleanup()
|
||||
self.cleanup_advanced()
|
||||
|
||||
def cleanup_advanced(self):
|
||||
self.sub_idxs = None
|
||||
self.full_latent_length = 0
|
||||
self.context_length = 0
|
||||
self.t = None
|
||||
self.batched_number = None
|
||||
self.weights = None
|
||||
self.latent_keyframes = None
|
||||
# timestep stuff
|
||||
self.current_timestep_keyframe = None
|
||||
self.next_timestep_keyframe = None
|
||||
self.current_timestep_index = -1
|
||||
# clear mask hints
|
||||
if self.mask_cond_hint is not None:
|
||||
del self.mask_cond_hint
|
||||
self.mask_cond_hint = None
|
||||
if self.tk_mask_cond_hint_original is not None:
|
||||
del self.tk_mask_cond_hint_original
|
||||
self.tk_mask_cond_hint_original = None
|
||||
if self.tk_mask_cond_hint is not None:
|
||||
del self.tk_mask_cond_hint
|
||||
self.tk_mask_cond_hint = None
|
||||
if self.weight_mask_cond_hint is not None:
|
||||
del self.weight_mask_cond_hint
|
||||
self.weight_mask_cond_hint = None
|
||||
|
||||
def copy_to_advanced(self, copied: 'AdvancedControlBase'):
|
||||
copied.mask_cond_hint_original = self.mask_cond_hint_original
|
||||
copied.weights_override = self.weights_override
|
||||
copied.latent_keyframe_override = self.latent_keyframe_override
|
||||
copied.disarmed = self.disarmed
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
# These are supported funding model platforms
|
||||
|
||||
github: [Kosinkadink]
|
||||
patreon: Kosinkadink
|
||||
open_collective: # Replace with a single Open Collective username
|
||||
ko_fi: # Replace with a single Ko-fi username
|
||||
tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel
|
||||
community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry
|
||||
liberapay: # Replace with a single Liberapay username
|
||||
issuehunt: # Replace with a single IssueHunt username
|
||||
otechie: # Replace with a single Otechie username
|
||||
lfx_crowdfunding: # Replace with a single LFX Crowdfunding project-name e.g., cloud-foundry
|
||||
custom: # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2']
|
||||
@@ -0,0 +1,8 @@
|
||||
__pycache__
|
||||
/venv
|
||||
.vscode
|
||||
*.ckpt
|
||||
*.safetensors
|
||||
*.pth
|
||||
manual_testing.py
|
||||
|
||||
@@ -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 2023 Jedrzej Kosinski
|
||||
|
||||
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.
|
||||
@@ -0,0 +1,495 @@
|
||||
# AnimateDiff for ComfyUI
|
||||
|
||||
Improved [AnimateDiff](https://github.com/guoyww/AnimateDiff/) integration for ComfyUI, as well as advanced sampling options dubbed Evolved Sampling usable outside of AnimateDiff. Please read the AnimateDiff repo README and Wiki for more information about how it works at its core.
|
||||
|
||||
AnimateDiff workflows will often make use of these helpful node packs:
|
||||
- [ComfyUI_FizzNodes](https://github.com/FizzleDorf/ComfyUI_FizzNodes) for prompt-travel functionality with the BatchPromptSchedule node. Maintained by FizzleDorf.
|
||||
- [ComfyUI-Advanced-ControlNet](https://github.com/Kosinkadink/ComfyUI-Advanced-ControlNet) for making ControlNets work with Context Options and controlling which latents should be affected by the ControlNet inputs. Includes SparseCtrl support. Maintained by me.
|
||||
- [ComfyUI-VideoHelperSuite](https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite) for loading videos, combining images into videos, and doing various image/latent operations like appending, splitting, duplicating, selecting, or counting. Actively maintained by AustinMroz and I.
|
||||
- [comfyui_controlnet_aux](https://github.com/Fannovel16/comfyui_controlnet_aux) for ControlNet preprocessors not present in vanilla ComfyUI. Maintained by Fannovel16.
|
||||
- [ComfyUI_IPAdapter_plus](https://github.com/cubiq/ComfyUI_IPAdapter_plus) for IPAdapter support. Maintained by cubiq (matt3o).
|
||||
|
||||
# Installation
|
||||
|
||||
## If using ComfyUI Manager:
|
||||
|
||||
1. Look for ```AnimateDiff Evolved```, and be sure the author is ```Kosinkadink```. Install it.
|
||||

|
||||
|
||||
|
||||
## If installing manually:
|
||||
1. Clone this repo into `custom_nodes` folder.
|
||||
|
||||
# Model Setup:
|
||||
1. Download motion modules. You will need at least 1. Different modules produce different results.
|
||||
- Original models ```mm_sd_v14```, ```mm_sd_v15```, ```mm_sd_v15_v2```, ```v3_sd15_mm```: [HuggingFace](https://huggingface.co/guoyww/animatediff/tree/cd71ae134a27ec6008b968d6419952b0c0494cf2) | [Google Drive](https://drive.google.com/drive/folders/1EqLC65eR1-W-sGD0Im7fkED6c8GkiNFI) | [CivitAI](https://civitai.com/models/108836)
|
||||
- Stabilized finetunes of mm_sd_v14, ```mm-Stabilized_mid``` and ```mm-Stabilized_high```, by **manshoety**: [HuggingFace](https://huggingface.co/manshoety/AD_Stabilized_Motion/tree/main)
|
||||
- Finetunes of mm_sd_v15_v2, ```mm-p_0.5.pth``` and ```mm-p_0.75.pth```, by **manshoety**: [HuggingFace](https://huggingface.co/manshoety/beta_testing_models/tree/main)
|
||||
- Higher resolution finetune,```temporaldiff-v1-animatediff``` by **CiaraRowles**: [HuggingFace](https://huggingface.co/CiaraRowles/TemporalDiff/tree/main)
|
||||
- FP16/safetensor versions of vanilla motion models, hosted by **continue-revolution** (takes up less storage space, but uses up the same amount of VRAM as ComfyUI loads models in fp16 by default): [HuffingFace](https://huggingface.co/conrevo/AnimateDiff-A1111/tree/main)
|
||||
2. Place models in one of these locations (you can rename models if you wish):
|
||||
- ```ComfyUI/custom_nodes/ComfyUI-AnimateDiff-Evolved/models```
|
||||
- ```ComfyUI/models/animatediff_models```
|
||||
3. Optionally, you can use Motion LoRAs to influence movement of v2-based motion models like mm_sd_v15_v2.
|
||||
- [Google Drive](https://drive.google.com/drive/folders/1EqLC65eR1-W-sGD0Im7fkED6c8GkiNFI?usp=sharing) | [HuggingFace](https://huggingface.co/guoyww/animatediff) | [CivitAI](https://civitai.com/models/108836/animatediff-motion-modules)
|
||||
- Place Motion LoRAs in one of these locations (you can rename Motion LoRAs if you wish):
|
||||
- ```ComfyUI/custom_nodes/ComfyUI-AnimateDiff-Evolved/motion_lora```
|
||||
- ```ComfyUI/models/animatediff_motion_lora```
|
||||
4. Get creative! If it works for normal image generation, it (probably) will work for AnimateDiff generations. Latent upscales? Go for it. ControlNets, one or more stacked? You betcha. Masking the conditioning of ControlNets to only affect part of the animation? Sure. Try stuff and you will be surprised by what you can do. Samples with workflows are included below.
|
||||
|
||||
NOTE: you can also use custom locations for models/motion loras by making use of the ComfyUI ```extra_model_paths.yaml``` file. The id for motion model folder is ```animatediff_models``` and the id for motion lora folder is ```animatediff_motion_lora```.
|
||||
|
||||
|
||||
# Features
|
||||
- Compatible with almost any vanilla or custom KSampler node.
|
||||
- ControlNet, SparseCtrl, and IPAdapter support
|
||||
- Infinite animation length support via sliding context windows across whole unet (Context Options) and/or within motion module (View Options)
|
||||
- Scheduling Context Options to change across different points in the sampling process
|
||||
- FreeInit and FreeNoise support (FreeInit is under iteration opts, FreeNoise is in SampleSettings' noise_type dropdown)
|
||||
- Mixable Motion LoRAs from [original AnimateDiff repository](https://github.com/guoyww/animatediff/) implemented. Caveat: the original loras really only work on v2-based motion models like ```mm_sd_v15_v2```, ```mm-p_0.5.pth```, and ```mm-p_0.75.pth```.
|
||||
- UPDATE: New motion LoRAs without the v2 limitation can now be trained via the [AnimateDiff-MotionDirector repo](https://github.com/ExponentialML/AnimateDiff-MotionDirector). Shoutout to ExponentialML for implementing MotionDirector for AnimateDiff purposes!
|
||||
- Prompt travel using BatchPromptSchedule node from [ComfyUI_FizzNodes](https://github.com/FizzleDorf/ComfyUI_FizzNodes)
|
||||
- Scale and Effect multival inputs to control motion amount and motion model influence on generation.
|
||||
- Can be float, list of floats, or masks
|
||||
- Custom noise scheduling via Noise Types, Noise Layers, and seed_override/seed_offset/batch_offset in Sample Settings and related nodes
|
||||
- AnimateDiff model v1/v2/v3 support
|
||||
- Using multiple motion models at once via Gen2 nodes (each supporting
|
||||
- [HotshotXL](https://huggingface.co/hotshotco/Hotshot-XL/tree/main) support (an SDXL motion module arch), ```hsxl_temporal_layers.safetensors```.
|
||||
- NOTE: You will need to use ```autoselect``` or ```linear (HotshotXL/default)``` beta_schedule, the sweetspot for context_length or total frames (when not using context) is 8 frames, and you will need to use an SDXL checkpoint.
|
||||
- AnimateDiff-SDXL support, with corresponding model. Currently, a beta version is out, which you can find info about at [AnimateDiff](https://github.com/guoyww/AnimateDiff/).
|
||||
- NOTE: You will need to use ```autoselect``` or ```linear (AnimateDiff-SDXL)``` beta_schedule. Other than that, same rules of thumb apply to AnimateDiff-SDXL as AnimateDiff.
|
||||
- [AnimateLCM](https://github.com/G-U-N/AnimateLCM) support
|
||||
- NOTE: You will need to use ```autoselect``` or ```lcm``` or ```lcm[100_ots]``` beta_schedule. To use fully with LCM, be sure to use appropriate LCM lora, use the ```lcm``` sampler_name in KSampler nodes, and lower cfg to somewhere around 1.0 to 2.0. Don't forget to decrease steps (minimum = ~4 steps), since LCM converges faster (less steps). Increase step count to increase detail as desired.
|
||||
- AnimateDiff Keyframes to change Scale and Effect at different points in the sampling process.
|
||||
- fp8 support; requires newest ComfyUI and torch >= 2.1 (decreases VRAM usage, but changes outputs)
|
||||
- Mac M1/M2/M3 support
|
||||
- Usage of Context Options and Sample Settings outside of AnimateDiff via Gen2 Use Evolved Sampling node
|
||||
|
||||
## Upcoming Features
|
||||
- Maskable Motion LoRA
|
||||
- Maskable SD LoRA (and perhaps maskable SD Models as well)
|
||||
- [PIA](https://github.com/open-mmlab/PIA) support
|
||||
- Anything else AnimateDiff-related that comes out
|
||||
|
||||
|
||||
# Basic Usage And Nodes
|
||||
|
||||
There are two families of nodes that can be used to use AnimateDiff/Evolved Sampling - **Gen1** and **Gen2**. Other than nodes marked specifically for Gen1/Gen2, all other nodes can be used for both Gen1 and Gen2.
|
||||
|
||||
Gen1 and Gen2 produce the exact same results (the backend code is identical), the only difference is in how the modes are used. Overall, Gen1 is the simplest way to use basic AnimateDiff features, while Gen2 separates model loading and application from the Evolved Sampling features. This means in practice, Gen2's Use Evolved Sampling node can be used without a model model, letting Context Options and Sample Settings be used without AnimateDiff.
|
||||
|
||||
In the following documentation, inputs/outputs will be color coded as follows:
|
||||
- 🟩 - required inputs
|
||||
- 🟨 - optional inputs
|
||||
- 🟦 - start as widgets, can be converted to inputs
|
||||
- 🟪 - output
|
||||
|
||||
## Gen1/Gen2 Nodes
|
||||
|
||||
| ① Gen1 ① | ② Gen2 ② |
|
||||
|---|---|
|
||||
| - All-in-One node<br/> - If same model is loaded by multiple Gen1 nodes, duplicates RAM usage. | - Separates model loading from application and Evolved Sampling<br/> - Enables no motion model usage while preserving Evolved Sampling features<br/> - Enables multiple motion model usage with Apply AnimateDiff Model (Adv.) Node|
|
||||
| | |
|
||||
| |  |
|
||||
|
||||
|
||||
### Inputs
|
||||
- 🟩*model*: StableDiffusion (SD) Model input.
|
||||
- 🟦*model_name*: AnimateDiff (AD) model to load and/or apply during the sampling process. Certain motion models work with SD1.5, while others work with SDXL.
|
||||
- 🟦*beta_schedule*: Applies selected beta_schedule to SD model; ```autoselect``` will automatically select the recommended beta_schedule for selected motion models - or will use_existing if no motion model selected for Gen2.
|
||||
- 🟨*context_options*: Context Options node from the context_opts submenu - should be used when needing to go back the sweetspot of an AnimateDiff model. Works with no motion models as well (Gen2 only).
|
||||
- 🟨*sample_settings*: Sample Settings node input - used to apply custom sampling options such as FreeNoise (noise_type), FreeInit (iter_opts), custom seeds, Noise Layers, etc. Works with no motion models as well (Gen2 only).
|
||||
- 🟨*motion_lora*: For v2-based models, Motion LoRA will influence the generated movement. Only a few official motion LoRAs were released - soon, I will be working with some community members to create training code to create (and test) new Motion LoRAs that might work with non-v2 models.
|
||||
- 🟨*ad_settings*: Modifies motion models during loading process, allowing the Positional Encoders (PEs) to be adjusted to extend a model's sweetspot or modify overall motion.
|
||||
- 🟨*ad_keyframes*: Allows scheduling of ```scale_multival``` and ```effect_multival``` inputs across sampling timesteps.
|
||||
- 🟨*scale_multival*: Uses a ```Multival``` input (defaults to ```1.0```). Previously called motion_scale, it directly influences the amount of motion generated by the model. With the Multival nodes, it can accept a float, list of floats, and/or mask inputs, allowing different scale to be applied to not only different frames, but different areas of frames (including per-frame).
|
||||
- 🟨*effect_multival*: Uses a ```Multival``` input (defaults to ```1.0```). Determines the influence of the motion models on the sampling process. Value of ```0.0``` is equivalent to normal SD output with no AnimateDiff influence. With the Multival nodes, it can accept a float, list of floats, and/or mask inputs, allowing different effect amount to be applied to not only different frames, but different areas of frames (including per-frame).
|
||||
|
||||
#### Gen2-Only Inputs
|
||||
- 🟨*motion_model*: Input for loaded motion_model.
|
||||
- 🟨*m_models*: One (or more) motion models outputted from Apply AnimateDiff Model nodes.
|
||||
|
||||
#### Gen2 Adv.-Only Inputs
|
||||
- 🟨*prev_m_models*: Previous applied motion models to use alongside this one.
|
||||
- 🟨*start_percent*: Determines when connected motion_model should take effect (supercedes any ad_keyframes).
|
||||
- 🟨*end_percent*: Determines when connected motion_model should stop taking effect (supercedes any ad_keyframes).
|
||||
|
||||
#### Gen1 (Legacy) Inputs
|
||||
- 🟦*motion_scale*: legacy version of ```scale_multival```, can only be a float.
|
||||
- 🟦*apply_v2_models_properly*: backwards compatible toggle for months-old workflows that used code that did not turn off groupnorm hack for v2 models. **Only affects v2 models, nothing else.** All nodes default this value to ```True``` now.
|
||||
|
||||
### Outputs
|
||||
- 🟪*MODEL*: Injected SD model with Evolved Sampling/AnimateDiff.
|
||||
|
||||
#### Gen2-Only Outputs
|
||||
- 🟪*MOTION_MODEL*: Loaded motion model.
|
||||
- 🟪*M_MODELS*: One (or more) applied motion models, to be either plugged into Use Evolved Sampling or another Apply AnimateDiff Model (Adv.) node.
|
||||
|
||||
|
||||
## Multival Nodes
|
||||
|
||||
For Multival inputs, these nodes allow the use of floats, list of floats, and/or masks to use as input. Scaled Mask node allows customization of dark/light areas of masks in terms of what the values correspond to.
|
||||
|
||||
| Node | Inputs |
|
||||
|---|---|
|
||||
|  | 🟨*mask_optional*: Mask for float values - black means 0.0, white means 1.0 (multiplied by float_val). <br/> 🟦*float_val*: Float multiplier.|
|
||||
|  | 🟩*mask*: Mask for float values. <br/> 🟦*min_float_val*: Minimum value. <br/>🟦*max_float_val*: Maximum value. <br/> 🟦*scaling*: When ```absolute```, black means min_float_val, white means max_float_val. When ```relative```, darkest area in masks (total) means min_float_val, lighest area in massk (total) means max_float_val. |
|
||||
|
||||
|
||||
## AnimateDiff Keyframe
|
||||
|
||||
Allows scheduling (in terms of timesteps) for scale_multival and effect_multival.
|
||||
|
||||
The two settings to determine schedule are ***start_percent*** and ***guarantee_steps***. When multiple keyframes have the same start_percent, they will be executed in the order they are connected, and run for guarantee_steps before moving on to the next node.
|
||||
|
||||
| Node |
|
||||
|---|
|
||||
|  |
|
||||
|
||||
### Inputs
|
||||
- 🟨*prev_ad_keyframes*: Chained keyframes to create schedule.
|
||||
- 🟨*scale_multival*: Value of scale to use for this keyframe.
|
||||
- 🟨*effect_multival*: Value of effect to use for this keyframe.
|
||||
- 🟦*start_percent*: Percent of timesteps to start usage of this keyframe. If multiple keyframes have same start_percent, order of execution is determined by their chained order, and will last for guarantee_steps timesteps.
|
||||
- 🟦*guarantee_steps*: Minimum amount of steps the keyframe will be used - when set to 0, this keyframe will only be used when no other keyframes are better matches for current timestep.
|
||||
- 🟦*inherit_missing*: When set to ```True```, any missing scale_multival or effect_multival inputs will inherit the previous keyframe's values - if the previous keyframe also inherits missing, the last inherited value will be used.
|
||||
|
||||
|
||||
## Context Options and View Options
|
||||
|
||||
These nodes provide techniques used to extend the lengths of animations to get around the sweetspot limitations of AnimateDiff models (typically 16 frames) and HotshotXL model (8 frames).
|
||||
|
||||
Context Options works by diffusing portions of the animation at a time, including main SD diffusion, ControlNets, IPAdapters, etc., effectively limiting VRAM usage to be equivalent to be context_length latents.
|
||||
|
||||
View Options, in contrast, work by portioning the latents seen by the motion model. This does NOT decrease VRAM usage, but in general is more stable and faster than Context Options, since the latents don't have to go through the whole SD unet.
|
||||
|
||||
Context Options and View Options can be combined to get the best of both worlds - longer context_length can be used to gain more stable output, at the cost of using more VRAM (since context_length determines how much SD sampling is done at the same time on the GPU). Provided you have the VRAM, you could also use Views Only Context Options to use only View Options (and automatically make context_length equivalent to full latents) to get a speed boost in return for the higher VRAM usage.
|
||||
|
||||
There are two types of Context/View Options: ***Standard*** and ***Looped***. ***Standard*** options do not cause looping in the output. ***Looped*** options, as the name implies, causes looping in the output (from end to beginning). Prior to the code rework, the only context available was the looping kind.
|
||||
|
||||
***I recommend using Standard Static at first when not wanting looped outputs.***
|
||||
|
||||
In the below animations, ***green*** shows the Contexts, and ***red*** shows the Views. TL;DR green is the amount of latents that are loaded into VRAM (and sampled), while red is the amount of latents that get passed into the motion model at a time.
|
||||
|
||||
### Context Options◆Standard Static
|
||||
| Behavior |
|
||||
|---|
|
||||
|  <br/> (latent count: 64, context_length: 16, context_overlap: 4, total steps: 20)|
|
||||
|
||||
| Node | Inputs |
|
||||
|---|---|
|
||||
|  | 🟦*context_length*: Amount of latents to diffuse at once.<br/> 🟦*context_overlap*: Minimum common latents between adjacent windows.<br/> 🟦*fuse_method*: Method for averaging results of windows.<br/> 🟦*use_on_equal_length*: When True, allows context to be used when latent count matches context_length.<br/> 🟦*start_percent*: When multiple Context Options are chained, allows scheduling.<br/> 🟦*guarantee_steps*: When scheduling contexts, determines the *minimum* amount of sampling steps context should be used.<br/> 🟦*context_length*: Amount of latents to diffuse at once.<br/> 🟨*prev_context*: Allows chaining of contexts.<br/> 🟨*view_options*: When context_length > view_length (unless otherwise specified), allows view_options to be used within each context window.|
|
||||
|
||||
### Context Options◆Standard Uniform
|
||||
| Behavior |
|
||||
|---|
|
||||
|  <br/> (latent count: 64, context_length: 16, context_overlap: 4, context_stride: 1, total steps: 20) |
|
||||
|  <br/> (latent count: 64, context_length: 16, context_overlap: 4, context_stride: 2, total steps: 20) |
|
||||
|
||||
| Node | Inputs |
|
||||
|---|---|
|
||||
|  | 🟦*context_length*: Amount of latents to diffuse at once.<br/> 🟦*context_overlap*: Minimum common latents between adjacent windows.<br/> 🟦*context_stride*: Maximum 2^(stride-1) distance between adjacent latents.<br/> 🟦*fuse_method*: Method for averaging results of windows.<br/> 🟦*use_on_equal_length*: When True, allows context to be used when latent count matches context_length.<br/> 🟦*start_percent*: When multiple Context Options are chained, allows scheduling.<br/> 🟦*guarantee_steps*: When scheduling contexts, determines the *minimum* amount of sampling steps context should be used.<br/> 🟦*context_length*: Amount of latents to diffuse at once.<br/> 🟨*prev_context*: Allows chaining of contexts.<br/> 🟨*view_options*: When context_length > view_length (unless otherwise specified), allows view_options to be used within each context window.|
|
||||
|
||||
### Context Options◆Looped Uniform
|
||||
| Behavior |
|
||||
|---|
|
||||
|  <br/> (latent count: 64, context_length: 16, context_overlap: 4, context_stride: 1, closed_loop: False, total steps: 20) |
|
||||
|  <br/> (latent count: 64, context_length: 16, context_overlap: 4, context_stride: 1, closed_loop: True, total steps: 20) |
|
||||
|
||||
| Node | Inputs |
|
||||
|---|---|
|
||||
|  | 🟦*context_length*: Amount of latents to diffuse at once.<br/> 🟦*context_overlap*: Minimum common latents between adjacent windows.<br/> 🟦*context_stride*: Maximum 2^(stride-1) distance between adjacent latents.<br/> 🟦*closed_loop*: When True, adds additional windows to enhance looping.<br/> 🟦*fuse_method*: Method for averaging results of windows.<br/> 🟦*use_on_equal_length*: When True, allows context to be used when latent count matches context_length - allows loops to be made when latent count == context_length.<br/> 🟦*start_percent*: When multiple Context Options are chained, allows scheduling.<br/> 🟦*guarantee_steps*: When scheduling contexts, determines the *minimum* amount of sampling steps context should be used.<br/> 🟦*context_length*: Amount of latents to diffuse at once.<br/> 🟨*prev_context*: Allows chaining of contexts.<br/> 🟨*view_options*: When context_length > view_length (unless otherwise specified), allows view_options to be used within each context window.|
|
||||
|
||||
### Context Options◆Views Only [VRAM⇈]
|
||||
| Behavior |
|
||||
|---|
|
||||
|  <br/> (latent count: 64, view_length: 16, view_overlap: 4, View Options◆Standard Static, total steps: 20) |
|
||||
|
||||
| Node | Inputs |
|
||||
|---|---|
|
||||
|  | 🟩*view_opts_req*: View_options to be used across all latents. <br/> 🟨*prev_context*: Allows chaining of contexts.<br/> |
|
||||
|
||||
|
||||
There are View Options equivalent of these schedules:
|
||||
|
||||
### View Options◆Standard Static
|
||||
| Behavior |
|
||||
|---|
|
||||
|  <br/> (latent count: 64, view_length: 16, view_overlap: 4, Context Options◆Standard Static, context_length: 32, context_overlap: 8, total steps: 20) |
|
||||
|
||||
| Node | Inputs |
|
||||
|---|---|
|
||||
|  | 🟦*view_length*: Amount of latents in context to pass into motion model at a time.<br/> 🟦*view_overlap*: Minimum common latents between adjacent windows.<br/> 🟦*fuse_method*: Method for averaging results of windows.<br/> |
|
||||
|
||||
### View Options◆Standard Uniform
|
||||
| Behavior |
|
||||
|---|
|
||||
|  <br/> (latent count: 64, view_length: 16, view_overlap: 4, view_stride: 1, Context Options◆Standard Static, context_length: 32, context_overlap: 8, total steps: 20) |
|
||||
|
||||
| Node | Inputs |
|
||||
|---|---|
|
||||
|  | 🟦*view_length*: Amount of latents in context to pass into motion model at a time.<br/> 🟦*view_overlap*: Minimum common latents between adjacent windows.<br/> 🟦*view_stride*: Maximum 2^(stride-1) distance between adjacent latents.<br/> 🟦*fuse_method*: Method for averaging results of windows.<br/> |
|
||||
|
||||
### View Options◆Looped Uniform
|
||||
| Behavior |
|
||||
|---|
|
||||
|  <br/> (latent count: 64, view_length: 16, view_overlap: 4, view_stride: 1, closed_loop: False, Context Options◆Standard Static, context_length: 32, context_overlap: 8, total steps: 20) |
|
||||
| NOTE: this one is probably not going to come out looking well unless you are using this for a very specific reason. |
|
||||
|
||||
| Node | Inputs |
|
||||
|---|---|
|
||||
|  | 🟦*view_length*: Amount of latents in context to pass into motion model at a time.<br/> 🟦*view_overlap*: Minimum common latents between adjacent windows.<br/> 🟦*view_stride*: Maximum 2^(stride-1) distance between adjacent latents.<br/> 🟦*closed_loop*: When True, adds additional windows to enhance looping.<br/> 🟦*use_on_equal_length*: When True, allows context to be used when latent count matches context_length - allows loops to be made when latent count == context_length.<br/> 🟦*fuse_method*: Method for averaging results of windows.<br/> |
|
||||
|
||||
## Sample Settings
|
||||
|
||||
The Sample Settings node allows customization of the sampling process beyond what is exposed on most KSampler nodes. With its default values, it will NOT have any effect, and can safely be attached without changing any behavior.
|
||||
|
||||
TL;DR To use FreeNoise, select ```FreeNoise``` from the noise_type dropdown. FreeNoise does not decrease performance in any way. To use FreeInit, attach the FreeInit Iteration Options to the iteration_opts input. NOTE: FreeInit, despite it's name, works by resampling the latents ```iterations``` amount of times - this means if you use iteration=2, total sampling time will be exactly twice as slow since it will be performing the sampling twice.
|
||||
|
||||
Noise Layers with the inputs of the same name (or very close to same name) have same intended behavior as the ones for Sample Settings - refer to the inputs below.
|
||||
|
||||
| Node |
|
||||
|---|
|
||||
|  |
|
||||
|
||||
### Inputs
|
||||
- 🟨*noise_layers*: Customizable, stackable noise to add to/modify initial noise.
|
||||
- 🟨*iteration_opts*: Options for determining if (and how) sampling should be repeated consecutively; if you want to check out FreeInit, this is how to use it.
|
||||
- 🟨*seed_override*: Accepts a single int to use a seed instead of the seed passed into the KSampler, or a list of ints (like via FizzNodes' BatchedValueSchedule) to assign individual seeds to each latent in the batch.
|
||||
- 🟦*seed_offset*: When not set to 0, adds value to current seed, predictably changing it, whatever the original seed may have been.
|
||||
- 🟦*batch_offset*: When not set to 0, will 'offset' the noise as if the first latent was actually the batch_offset-nth latent, shifting all the noises over.
|
||||
- 🟦*noise_type*: Selects type of noise to be generated. Values include:
|
||||
- **default**: generates different noise for all latents as usual.
|
||||
- **constant**: generates exact same noise for all latents (based on seed).
|
||||
- **empty**: generates no noise for all latents (as if noise was turned off).
|
||||
- **repeated_context**: repeats noise every context_length (or view_length) amount of latents; stabilizes longer generations, but has very obvious repetition.
|
||||
- **FreeNoise**: repeats noise such that it is repeated every context_length (or view_length), but the overlapped noise between contexts/views is shuffled to make repetition less prevelant while still achieving stabilization.
|
||||
- 🟦*seed_gen*: Allows choosing between ComfyUI and Auto1111 methods of noise generation. One is not better than the other (noise distributions are the same), they are just different methods.
|
||||
- **comfy**: Noise is generated for the entire latent batch tensor at once based on the provided seed.
|
||||
- **auto1111**: Noise is generated individually for each latent, with each latent receiving an increasing +1 seed offset (first latent uses seed, second latent uses seed+1, etc.).
|
||||
- 🟦*adapt_denoise_steps*: When True, KSamplers with a 'denoise' input will automatically scale down the total steps to run like the default options in Auto1111.
|
||||
- **True**: Steps will decrease with lower denoise, i.e. 20 steps with 0.5 denoise will be 10 total steps executed, but sigmas will be selected that still achieve 0.5 denoise. Trades speed for quality (since less steps are sampled).
|
||||
- **False**: Default behavior; 20 steps with 0.5 denoise will execute 20 steps.
|
||||
|
||||
|
||||
## Iteration Options
|
||||
|
||||
These options allow KSamplers to re-sample the same latents without needing to chain multiple KSamplers together, and also allows specialized iteration behavior to implement features such as FreeInit.
|
||||
|
||||
### Default Iteration Options
|
||||
|
||||
Simply re-runs the KSampler, plugging in the output of the previous iteration into the next one. At the dafault iterations=1, it is no different than not having this node plugged in at all.
|
||||
|
||||
| Node | Inputs |
|
||||
|---|---|
|
||||
|  | 🟦*iterations*: Total amount of times KSampler should run back-to-back. <br/> 🟦*iter_batch_offset*: batch_offset to apply on each subsequent iteration. <br/> 🟦*iter_seed_offset*: seed_offset to apply on each subsequent iteration. |
|
||||
|
||||
### FreeInit Iteration Options
|
||||
|
||||
Implements [FreeInit](https://github.com/TianxingWu/FreeInit), which is the idea that AnimateDiff was trained on latents of existing videos (images with temporal coherence between them) that were then noised rather than from random initial noise, and that when noising existing latents, low-frequency data still remains in the noised latents. It combines the low-frequency noise from existing videos (or, as is the default behavior, the previous iteration) with the high-frequency noise in randomly generated noise to run the subsequent iterations. ***Each iteration is a full sample - 2 iterations means it will take twice as long to run as compared to having 1 iteration/no iteration_opts connected.***
|
||||
|
||||
When apply_to_1st_iter is False, the noising/low-freq/high-freq combination will not occur on the first iteration, with the assumption that there are no useful latents passed in to do the noise combining in the first place, thus requiring at least 2 iterations for FreeInit to take effect.
|
||||
|
||||
If you have an existing set of latents to use to get low-freq noise from, you may set apply_to_1st_iter to True, and then even if you set iterations=1, FreeInit will still take effect.
|
||||
|
||||
| Node |
|
||||
|---|
|
||||
|  |
|
||||
|
||||
#### Inputs
|
||||
- 🟦*iterations*: Total amount of times KSampler should run back-to-back. Refer to explanation above why it is 2 by default (and when it can be set to 1 instead).
|
||||
- 🟦*init_type*: Code implementation for applying FreeInit.
|
||||
- ***FreeInit [sampler sigma]***: likely closest to intended implementation, and gets the sigma for noising from the sampler instead of the model (when possible).
|
||||
- ***FreeInit [model sigma]***: gets sigma for noising from the model; when using Custom KSampler, this is the method that will be used for both FreeInit options.
|
||||
- ***DinkInit_v1***: my initial, flawed implementation of FreeInit before I figured out how to exactly copy the noising behavior. By sheer luck and trial and error, I managed to have it actually sort of work with this method. Mainly for backwards compatibility now, but might produce useful results too.
|
||||
|
||||
- 🟦*apply_to_1st_iter*: When set to True, will do FreeInit low-freq/high-freq combo work even on the 1st iteration it runs Refer to explanation in the above FreeInit Iteration Options section for when this can be set to True.
|
||||
- 🟦*init_type*: Code implementation for applying FreeInit.
|
||||
- 🟦*iter_batch_offset*: batch_offset to apply on each subsequent iteration.
|
||||
- 🟦*iter_seed_offset*: seed_offset to apply on each subsequent iteration. Defaults to 1 so that new random noise is used for each iteration.
|
||||
|
||||
- 🟦*filter*: Determines low-freq filter to apply to noise. Very technical, look into code/online resources to figure out how the individual filters act.
|
||||
- 🟦*d_s*: Spatial parameter of filter (within latents, I think); very technical. Look into code/online resources if you wish to know what exactly it does.
|
||||
- 🟦*d_t*: Temporal parameter of filter (across latents, I think); very technical. Look into code/online resources if you wish to know what exactly it does.
|
||||
- 🟦*n_butterworth*: Only applies to ```butterworth``` filter; very technical. Look into code/online resources if you wish to know what exactly it does.
|
||||
- 🟦*sigma_step*: Noising step to use/emulate when noising latents to then get low-freq noise out of. 999 actually means last (-1), and any number under 999 will mean the distance away from last. Leave at 999 unless you know what you're trying to do with it.
|
||||
|
||||
|
||||
## Noise Layers
|
||||
|
||||
These nodes allow initial noise to be added onto, weighted, or replaced. In near future, I will add the ability for masks to 'move' the noise relative to the masks' movement instead of just 'cutting and pasting' the noise.
|
||||
|
||||
The inputs that are shared with Sample Settings have the same exact effect - only new option is in seed_gen_override, which by default will use same seed_gen as Sample Settings (use existing). You can make a noise layer use a different seed_gen strategy at will, or use a different seed/set of seeds, etc.
|
||||
|
||||
The ```mask_optional``` parameter determines where on the initial noise the noise layer should be applied.
|
||||
|
||||
| Node | Behavior + Inputs |
|
||||
|---|---|
|
||||
|  | [Add]; Adds noise directly on top. <br/> 🟦*noise_weight*: Multiplier for noise layer before being added on top. |
|
||||
|  | [Add Weighted]; Adds noise, but takes a weighted average between what is already there and itself. <br/> 🟦*noise_weight*: Weight of new noise in the weighted average with existing noise. <br/> 🟦*balance_multipler*: Scale for how much noise_weight should affect existing noise; 1.0 means normal weighted average, and below 1.0 will lessen the weighted reduction by that amount (i.e. if balance_multiplier is set to 0.5 and noise_weight is 0.25, existing noise will only be reduced by 0.125 instead of 0.25, but new noise will be added with the unmodified 0.25 weight). |
|
||||
|  | [Replace]; Directly replaces existing noise from layers underneath with itself. |
|
||||
|
||||
|
||||
# Samples (download or drag images of the workflows into ComfyUI to instantly load the corresponding workflows!)
|
||||
|
||||
NOTE: I've scaled down the gifs to 0.75x size to make them take up less space on the README.
|
||||
|
||||
### txt2img
|
||||
|
||||
| Result |
|
||||
|---|
|
||||
|  |
|
||||
| Workflow |
|
||||
|  |
|
||||
|
||||
### txt2img - (prompt travel)
|
||||
|
||||
| Result |
|
||||
|---|
|
||||
|  |
|
||||
| Workflow |
|
||||
|  |
|
||||
|
||||
|
||||
|
||||
### txt2img - 48 frame animation with 16 context_length (Context Options◆Standard Static) + FreeNoise
|
||||
|
||||
| Result |
|
||||
|---|
|
||||
|  |
|
||||
| Workflow |
|
||||
|  |
|
||||
|
||||
|
||||
# Old Samples (TODO: update all of these + add new ones when I get sleep)
|
||||
|
||||
### txt2img - 32 frame animation with 16 context_length (uniform) - PanLeft and ZoomOut Motion LoRAs
|
||||
|
||||

|
||||
|
||||

|
||||
|
||||
[aaa_readme_00095_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/d730ae2e-188c-4a61-8a6d-bd48f60a2d07)
|
||||
|
||||
|
||||
### txt2img w/ latent upscale (partial denoise on upscale)
|
||||
|
||||

|
||||
|
||||

|
||||
|
||||
[aaa_readme_up_00002_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/2f44342f-3fd8-4863-8e3d-360377d608b7)
|
||||
|
||||
|
||||
|
||||
### txt2img w/ latent upscale (partial denoise on upscale) - PanLeft and ZoomOut Motion LoRAs
|
||||
|
||||

|
||||
|
||||

|
||||
|
||||
[aaa_readme_up_00024_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/414c16d8-231c-422f-8dfc-a93d4b68ffcc)
|
||||
|
||||
|
||||
|
||||
### txt2img w/ latent upscale (partial denoise on upscale) - 48 frame animation with 16 context_length (uniform)
|
||||
|
||||

|
||||
|
||||
[aaa_readme_up_00009_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/f7a45f81-e700-4bfe-9fdd-fbcaa4fa8a4e)
|
||||
|
||||
|
||||
|
||||
### txt2img w/ latent upscale (full denoise on upscale)
|
||||
|
||||

|
||||
|
||||

|
||||
|
||||
[aaa_readme_up_00012_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/3eb575cf-92dd-434a-b3db-1a2064ff0033)
|
||||
|
||||
|
||||
|
||||
### txt2img w/ latent upscale (full denoise on upscale) - 48 frame animation with 16 context_length (uniform)
|
||||
|
||||

|
||||
|
||||
[aaa_readme_up_00014_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/034aff4c-f814-4b87-b5d1-407b1089af0d)
|
||||
|
||||
|
||||
|
||||
### txt2img w/ ControlNet-stabilized latent-upscale (partial denoise on upscale, Scaled Soft ControlNet Weights)
|
||||
|
||||

|
||||
|
||||

|
||||
|
||||
[aaa_readme_up_00019_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/b562251d-a4fb-4141-94dd-9f8bca9f3ce8)
|
||||
|
||||
|
||||
|
||||
### txt2img w/ ControlNet-stabilized latent-upscale (partial denoise on upscale, Scaled Soft ControlNet Weights) 48 frame animation with 16 context_length (uniform)
|
||||
|
||||

|
||||
|
||||
[aaa_readme_up_00003_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/0f57c949-0af3-4da4-b7c4-5c1fb1549927)
|
||||
|
||||
|
||||
|
||||
### txt2img w/ Initial ControlNet input (using Normal LineArt preprocessor on first txt2img as an example)
|
||||
|
||||

|
||||
|
||||

|
||||
|
||||
[aaa_readme_cn_00003_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/9c9d425d-2378-4af0-8464-2c6c0d1a68bf)
|
||||
|
||||
|
||||
|
||||
### txt2img w/ Initial ControlNet input (using Normal LineArt preprocessor on first txt2img 48 frame as an example) 48 frame animation with 16 context_length (uniform)
|
||||
|
||||

|
||||
|
||||

|
||||
|
||||
[aaa_readme_cn_00006_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/5d3665b7-c2da-46a1-88d8-ab43ba8eb0c6)
|
||||
|
||||
|
||||
|
||||
### txt2img w/ Initial ControlNet input (using OpenPose images) + latent upscale w/ full denoise
|
||||
|
||||

|
||||
|
||||
(open_pose images provided courtesy of toyxyz)
|
||||
|
||||

|
||||
|
||||

|
||||
|
||||
[aaa_readme_cn_00033_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/c5df09a5-8c64-4811-9ecf-57ac73d82377)
|
||||
|
||||
|
||||
|
||||
### txt2img w/ Initial ControlNet input (using OpenPose images) + latent upscale w/ full denoise, 48 frame animation with 16 context_length (uniform)
|
||||
|
||||

|
||||
|
||||
(open_pose images provided courtesy of toyxyz)
|
||||
|
||||

|
||||
|
||||
[aaa_readme_cn_00024_.webm](https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/assets/7365912/8f4c840c-06a2-4c64-b97e-568dd5ff6f46)
|
||||
|
||||
|
||||
|
||||
### img2img
|
||||
|
||||
TODO: fill this out with a few useful ways, some using control net tile. I'm sorry there is nothing here right now, I have a lot of code to write. I'll try to fill this section out + Advance ControlNet use piece by piece.
|
||||
|
||||
|
||||
|
||||
## Known Issues
|
||||
|
||||
### Some motion models have visible watermark on resulting images (especially when using mm_sd_v15)
|
||||
|
||||
Training data used by the authors of the AnimateDiff paper contained Shutterstock watermarks. Since mm_sd_v15 was finetuned on finer, less drastic movement, the motion module attempts to replicate the transparency of that watermark and does not get blurred away like mm_sd_v14. Using other motion modules, or combinations of them using Advanced KSamplers should alleviate watermark issues.
|
||||
@@ -0,0 +1,10 @@
|
||||
import folder_paths
|
||||
from .animatediff.logger import logger
|
||||
from .animatediff.utils_model import get_available_motion_models, Folders
|
||||
from .animatediff.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
if len(get_available_motion_models()) == 0:
|
||||
logger.error(f"No motion models found. Please download one and place in: {folder_paths.get_folder_paths(Folders.ANIMATEDIFF_MODELS)}")
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
+143
@@ -0,0 +1,143 @@
|
||||
from torch import Tensor
|
||||
|
||||
from .utils_motion import normalize_min_max
|
||||
|
||||
|
||||
class AnimateDiffSettings:
|
||||
def __init__(self,
|
||||
adjust_pe: 'AdjustPEGroup'=None,
|
||||
pe_strength: float=1.0,
|
||||
attn_strength: float=1.0,
|
||||
attn_q_strength: float=1.0,
|
||||
attn_k_strength: float=1.0,
|
||||
attn_v_strength: float=1.0,
|
||||
attn_out_weight_strength: float=1.0,
|
||||
attn_out_bias_strength: float=1.0,
|
||||
other_strength: float=1.0,
|
||||
attn_scale: float=1.0,
|
||||
mask_attn_scale: Tensor=None,
|
||||
mask_attn_scale_min: float=1.0,
|
||||
mask_attn_scale_max: float=1.0,
|
||||
):
|
||||
# PE-interpolation settings
|
||||
self.adjust_pe = adjust_pe if adjust_pe is not None else AdjustPEGroup()
|
||||
# general strengths
|
||||
self.pe_strength = pe_strength
|
||||
self.attn_strength = attn_strength
|
||||
self.other_strength = other_strength
|
||||
# specific attn strengths
|
||||
self.attn_q_strength = attn_q_strength
|
||||
self.attn_k_strength = attn_k_strength
|
||||
self.attn_v_strength = attn_v_strength
|
||||
self.attn_out_weight_strength = attn_out_weight_strength
|
||||
self.attn_out_bias_strength = attn_out_bias_strength
|
||||
# attention scale settings - DEPRECATED
|
||||
self.attn_scale = attn_scale
|
||||
# attention scale mask settings - DEPRECATED
|
||||
self.mask_attn_scale = mask_attn_scale.clone() if mask_attn_scale is not None else mask_attn_scale
|
||||
self.mask_attn_scale_min = mask_attn_scale_min
|
||||
self.mask_attn_scale_max = mask_attn_scale_max
|
||||
self._prepare_mask_attn_scale()
|
||||
|
||||
def _prepare_mask_attn_scale(self):
|
||||
if self.mask_attn_scale is not None:
|
||||
self.mask_attn_scale = normalize_min_max(self.mask_attn_scale, self.mask_attn_scale_min, self.mask_attn_scale_max)
|
||||
|
||||
def has_mask_attn_scale(self) -> bool:
|
||||
return self.mask_attn_scale is not None
|
||||
|
||||
def has_pe_strength(self) -> bool:
|
||||
return self.pe_strength != 1.0
|
||||
|
||||
def has_attn_strength(self) -> bool:
|
||||
return self.attn_strength != 1.0
|
||||
|
||||
def has_other_strength(self) -> bool:
|
||||
return self.other_strength != 1.0
|
||||
|
||||
def has_anything_to_apply(self) -> bool:
|
||||
return self.adjust_pe.has_anything_to_apply() \
|
||||
or self.has_pe_strength() \
|
||||
or self.has_attn_strength() \
|
||||
or self.has_other_strength() \
|
||||
or self.has_any_attn_sub_strength()
|
||||
|
||||
def has_any_attn_sub_strength(self) -> bool:
|
||||
return self.has_attn_q_strength() \
|
||||
or self.has_attn_k_strength() \
|
||||
or self.has_attn_v_strength() \
|
||||
or self.has_attn_out_weight_strength() \
|
||||
or self.has_attn_out_bias_strength()
|
||||
|
||||
def has_attn_q_strength(self) -> bool:
|
||||
return self.attn_q_strength != 1.0
|
||||
|
||||
def has_attn_k_strength(self) -> bool:
|
||||
return self.attn_k_strength != 1.0
|
||||
|
||||
def has_attn_v_strength(self) -> bool:
|
||||
return self.attn_v_strength != 1.0
|
||||
|
||||
def has_attn_out_weight_strength(self) -> bool:
|
||||
return self.attn_out_weight_strength != 1.0
|
||||
|
||||
def has_attn_out_bias_strength(self) -> bool:
|
||||
return self.attn_out_bias_strength != 1.0
|
||||
|
||||
|
||||
class AdjustPE:
|
||||
def __init__(self,
|
||||
cap_initial_pe_length: int=0, interpolate_pe_to_length: int=0,
|
||||
initial_pe_idx_offset: int=0, final_pe_idx_offset: int=0,
|
||||
motion_pe_stretch: int=0, print_adjustment=False):
|
||||
# PE-interpolation settings
|
||||
self.cap_initial_pe_length = cap_initial_pe_length
|
||||
self.interpolate_pe_to_length = interpolate_pe_to_length
|
||||
self.initial_pe_idx_offset = initial_pe_idx_offset
|
||||
self.final_pe_idx_offset = final_pe_idx_offset
|
||||
self.motion_pe_stretch = motion_pe_stretch
|
||||
self.print_adjustment = print_adjustment
|
||||
|
||||
def has_cap_initial_pe_length(self) -> bool:
|
||||
return self.cap_initial_pe_length > 0
|
||||
|
||||
def has_interpolate_pe_to_length(self) -> bool:
|
||||
return self.interpolate_pe_to_length > 0
|
||||
|
||||
def has_initial_pe_idx_offset(self) -> bool:
|
||||
return self.initial_pe_idx_offset > 0
|
||||
|
||||
def has_final_pe_idx_offset(self) -> bool:
|
||||
return self.final_pe_idx_offset > 0
|
||||
|
||||
def has_motion_pe_stretch(self) -> bool:
|
||||
return self.motion_pe_stretch > 0
|
||||
|
||||
def has_anything_to_apply(self) -> bool:
|
||||
return self.has_cap_initial_pe_length() \
|
||||
or self.has_interpolate_pe_to_length() \
|
||||
or self.has_initial_pe_idx_offset() \
|
||||
or self.has_final_pe_idx_offset() \
|
||||
or self.has_motion_pe_stretch()
|
||||
|
||||
|
||||
class AdjustPEGroup:
|
||||
def __init__(self, initial: AdjustPE=None):
|
||||
self.adjusts: list[AdjustPE] = []
|
||||
if initial is not None:
|
||||
self.add(initial)
|
||||
|
||||
def add(self, adjust_pe: AdjustPE):
|
||||
self.adjusts.append(adjust_pe)
|
||||
|
||||
def has_anything_to_apply(self):
|
||||
for adjust in self.adjusts:
|
||||
if adjust.has_anything_to_apply():
|
||||
return True
|
||||
return False
|
||||
|
||||
def clone(self):
|
||||
new_group = AdjustPEGroup()
|
||||
for adjust in self.adjusts:
|
||||
new_group.add(adjust)
|
||||
return new_group
|
||||
+389
@@ -0,0 +1,389 @@
|
||||
from typing import Callable, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
from torch import Tensor
|
||||
|
||||
from comfy.model_base import BaseModel
|
||||
|
||||
from .utils_motion import get_sorted_list_via_attr
|
||||
|
||||
class ContextFuseMethod:
|
||||
FLAT = "flat"
|
||||
PYRAMID = "pyramid"
|
||||
RELATIVE = "relative"
|
||||
|
||||
LIST = [PYRAMID, FLAT]
|
||||
LIST_STATIC = [PYRAMID, RELATIVE, FLAT]
|
||||
|
||||
|
||||
class ContextType:
|
||||
UNIFORM_WINDOW = "uniform window"
|
||||
|
||||
|
||||
class ContextOptions:
|
||||
def __init__(self, context_length: int=None, context_stride: int=None, context_overlap: int=None,
|
||||
context_schedule: str=None, closed_loop: bool=False, fuse_method: str=ContextFuseMethod.FLAT,
|
||||
use_on_equal_length: bool=False, view_options: 'ContextOptions'=None,
|
||||
start_percent=0.0, guarantee_steps=1):
|
||||
# permanent settings
|
||||
self.context_length = context_length
|
||||
self.context_stride = context_stride
|
||||
self.context_overlap = context_overlap
|
||||
self.context_schedule = context_schedule
|
||||
self.closed_loop = closed_loop
|
||||
self.fuse_method = fuse_method
|
||||
self.sync_context_to_pe = False # this feature is likely bad and stay unused, so I might remove this
|
||||
self.use_on_equal_length = use_on_equal_length
|
||||
self.view_options = view_options.clone() if view_options else view_options
|
||||
# scheduling
|
||||
self.start_percent = float(start_percent)
|
||||
self.start_t = 999999999.9
|
||||
self.guarantee_steps = guarantee_steps
|
||||
# temporary vars
|
||||
self._step: int = 0
|
||||
|
||||
@property
|
||||
def step(self):
|
||||
return self._step
|
||||
@step.setter
|
||||
def step(self, value: int):
|
||||
self._step = value
|
||||
if self.view_options:
|
||||
self.view_options.step = value
|
||||
|
||||
def clone(self):
|
||||
n = ContextOptions(context_length=self.context_length, context_stride=self.context_stride,
|
||||
context_overlap=self.context_overlap, context_schedule=self.context_schedule,
|
||||
closed_loop=self.closed_loop, fuse_method=self.fuse_method,
|
||||
use_on_equal_length=self.use_on_equal_length, view_options=self.view_options,
|
||||
start_percent=self.start_percent, guarantee_steps=self.guarantee_steps)
|
||||
n.start_t = self.start_t
|
||||
return n
|
||||
|
||||
|
||||
class ContextOptionsGroup:
|
||||
def __init__(self):
|
||||
self.contexts: list[ContextOptions] = []
|
||||
self._current_context: ContextOptions = None
|
||||
self._current_used_steps: int = 0
|
||||
self._current_index: int = 0
|
||||
self.step = 0
|
||||
|
||||
def reset(self):
|
||||
self._current_context = None
|
||||
self._current_used_steps = 0
|
||||
self._current_index = 0
|
||||
self.step = 0
|
||||
self._set_first_as_current()
|
||||
|
||||
@classmethod
|
||||
def default(cls):
|
||||
def_context = ContextOptions()
|
||||
new_group = ContextOptionsGroup()
|
||||
new_group.add(def_context)
|
||||
return new_group
|
||||
|
||||
def add(self, context: ContextOptions):
|
||||
# add to end of list, then sort
|
||||
self.contexts.append(context)
|
||||
self.contexts = get_sorted_list_via_attr(self.contexts, "start_percent")
|
||||
self._set_first_as_current()
|
||||
|
||||
def add_to_start(self, context: ContextOptions):
|
||||
# add to start of list, then sort
|
||||
self.contexts.insert(0, context)
|
||||
self.contexts = get_sorted_list_via_attr(self.contexts, "start_percent")
|
||||
self._set_first_as_current()
|
||||
|
||||
def has_index(self, index: int) -> int:
|
||||
return index >=0 and index < len(self.contexts)
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return len(self.contexts) == 0
|
||||
|
||||
def clone(self):
|
||||
cloned = ContextOptionsGroup()
|
||||
for context in self.contexts:
|
||||
cloned.contexts.append(context)
|
||||
cloned._set_first_as_current()
|
||||
return cloned
|
||||
|
||||
def initialize_timesteps(self, model: BaseModel):
|
||||
for context in self.contexts:
|
||||
context.start_t = model.model_sampling.percent_to_sigma(context.start_percent)
|
||||
|
||||
def prepare_current_context(self, t: Tensor):
|
||||
curr_t: float = t[0]
|
||||
prev_index = self._current_index
|
||||
# if met guaranteed steps, look for next context in case need to switch
|
||||
if self._current_used_steps >= self._current_context.guarantee_steps:
|
||||
# if has next index, loop through and see if need to switch
|
||||
if self.has_index(self._current_index+1):
|
||||
for i in range(self._current_index+1, len(self.contexts)):
|
||||
eval_c = self.contexts[i]
|
||||
# check if start_t is greater or equal to curr_t
|
||||
# NOTE: t is in terms of sigmas, not percent, so bigger number = earlier step in sampling
|
||||
if eval_c.start_t >= curr_t:
|
||||
self._current_index = i
|
||||
self._current_context = eval_c
|
||||
self._current_used_steps = 0
|
||||
# if guarantee_steps greater than zero, stop searching for other keyframes
|
||||
if self._current_context.guarantee_steps > 0:
|
||||
break
|
||||
# if eval_c is outside the percent range, stop looking further
|
||||
else:
|
||||
break
|
||||
# update steps current context is used
|
||||
self._current_used_steps += 1
|
||||
|
||||
def _set_first_as_current(self):
|
||||
if len(self.contexts) > 0:
|
||||
self._current_context = self.contexts[0]
|
||||
|
||||
# properties shadow those of ContextOptions
|
||||
@property
|
||||
def context_length(self):
|
||||
return self._current_context.context_length
|
||||
|
||||
@property
|
||||
def context_overlap(self):
|
||||
return self._current_context.context_overlap
|
||||
|
||||
@property
|
||||
def context_stride(self):
|
||||
return self._current_context.context_stride
|
||||
|
||||
@property
|
||||
def context_schedule(self):
|
||||
return self._current_context.context_schedule
|
||||
|
||||
@property
|
||||
def closed_loop(self):
|
||||
return self._current_context.closed_loop
|
||||
|
||||
@property
|
||||
def fuse_method(self):
|
||||
return self._current_context.fuse_method
|
||||
|
||||
@property
|
||||
def use_on_equal_length(self):
|
||||
return self._current_context.use_on_equal_length
|
||||
|
||||
@property
|
||||
def view_options(self):
|
||||
return self._current_context.view_options
|
||||
|
||||
|
||||
class ContextSchedules:
|
||||
UNIFORM_LOOPED = "looped_uniform"
|
||||
UNIFORM_STANDARD = "standard_uniform"
|
||||
STATIC_STANDARD = "standard_static"
|
||||
BATCHED = "batched"
|
||||
VIEW_AS_CONTEXT = "view_as_context"
|
||||
|
||||
LEGACY_UNIFORM_LOOPED = "uniform"
|
||||
LEGACY_UNIFORM_SCHEDULE_LIST = [LEGACY_UNIFORM_LOOPED]
|
||||
|
||||
|
||||
# from https://github.com/neggles/animatediff-cli/blob/main/src/animatediff/pipelines/context.py
|
||||
def create_windows_uniform_looped(num_frames: int, opts: Union[ContextOptionsGroup, ContextOptions]):
|
||||
windows = []
|
||||
if num_frames < opts.context_length:
|
||||
windows.append(list(range(num_frames)))
|
||||
return windows
|
||||
|
||||
context_stride = min(opts.context_stride, int(np.ceil(np.log2(num_frames / opts.context_length))) + 1)
|
||||
# obtain uniform windows as normal, looping and all
|
||||
for context_step in 1 << np.arange(context_stride):
|
||||
pad = int(round(num_frames * ordered_halving(opts.step)))
|
||||
for j in range(
|
||||
int(ordered_halving(opts.step) * context_step) + pad,
|
||||
num_frames + pad + (0 if opts.closed_loop else -opts.context_overlap),
|
||||
(opts.context_length * context_step - opts.context_overlap),
|
||||
):
|
||||
windows.append([e % num_frames for e in range(j, j + opts.context_length * context_step, context_step)])
|
||||
|
||||
return windows
|
||||
|
||||
|
||||
def create_windows_uniform_standard(num_frames: int, opts: Union[ContextOptionsGroup, ContextOptions]):
|
||||
# unlike looped, uniform_straight does NOT allow windows that loop back to the beginning;
|
||||
# instead, they get shifted to the corresponding end of the frames.
|
||||
# in the case that a window (shifted or not) is identical to the previous one, it gets skipped.
|
||||
windows = []
|
||||
if num_frames <= opts.context_length:
|
||||
windows.append(list(range(num_frames)))
|
||||
return windows
|
||||
|
||||
context_stride = min(opts.context_stride, int(np.ceil(np.log2(num_frames / opts.context_length))) + 1)
|
||||
# first, obtain uniform windows as normal, looping and all
|
||||
for context_step in 1 << np.arange(context_stride):
|
||||
pad = int(round(num_frames * ordered_halving(opts.step)))
|
||||
for j in range(
|
||||
int(ordered_halving(opts.step) * context_step) + pad,
|
||||
num_frames + pad + (-opts.context_overlap),
|
||||
(opts.context_length * context_step - opts.context_overlap),
|
||||
):
|
||||
windows.append([e % num_frames for e in range(j, j + opts.context_length * context_step, context_step)])
|
||||
|
||||
# now that windows are created, shift any windows that loop, and delete duplicate windows
|
||||
delete_idxs = []
|
||||
win_i = 0
|
||||
while win_i < len(windows):
|
||||
# if window is rolls over itself, need to shift it
|
||||
is_roll, roll_idx = does_window_roll_over(windows[win_i], num_frames)
|
||||
if is_roll:
|
||||
roll_val = windows[win_i][roll_idx] # roll_val might not be 0 for windows of higher strides
|
||||
shift_window_to_end(windows[win_i], num_frames=num_frames)
|
||||
# check if next window (cyclical) is missing roll_val
|
||||
if roll_val not in windows[(win_i+1) % len(windows)]:
|
||||
# need to insert new window here - just insert window starting at roll_val
|
||||
windows.insert(win_i+1, list(range(roll_val, roll_val + opts.context_length)))
|
||||
# delete window if it's not unique
|
||||
for pre_i in range(0, win_i):
|
||||
if windows[win_i] == windows[pre_i]:
|
||||
delete_idxs.append(win_i)
|
||||
break
|
||||
win_i += 1
|
||||
|
||||
# reverse delete_idxs so that they will be deleted in an order that doesn't break idx correlation
|
||||
delete_idxs.reverse()
|
||||
for i in delete_idxs:
|
||||
windows.pop(i)
|
||||
|
||||
return windows
|
||||
|
||||
|
||||
def create_windows_static_standard(num_frames: int, opts: Union[ContextOptionsGroup, ContextOptions]):
|
||||
windows = []
|
||||
if num_frames <= opts.context_length:
|
||||
windows.append(list(range(num_frames)))
|
||||
return windows
|
||||
# always return the same set of windows
|
||||
delta = opts.context_length - opts.context_overlap
|
||||
for start_idx in range(0, num_frames, delta):
|
||||
# if past the end of frames, move start_idx back to allow same context_length
|
||||
ending = start_idx + opts.context_length
|
||||
if ending >= num_frames:
|
||||
final_delta = ending - num_frames
|
||||
final_start_idx = start_idx - final_delta
|
||||
windows.append(list(range(final_start_idx, final_start_idx + opts.context_length)))
|
||||
break
|
||||
windows.append(list(range(start_idx, start_idx + opts.context_length)))
|
||||
return windows
|
||||
|
||||
|
||||
def create_windows_batched(num_frames: int, opts: Union[ContextOptionsGroup, ContextOptions]):
|
||||
windows = []
|
||||
if num_frames <= opts.context_length:
|
||||
windows.append(list(range(num_frames)))
|
||||
return windows
|
||||
# always return the same set of windows;
|
||||
# no overlap, just cut up based on context_length;
|
||||
# last window size will be different if num_frames % opts.context_length != 0
|
||||
for start_idx in range(0, num_frames, opts.context_length):
|
||||
windows.append(list(range(start_idx, min(start_idx + opts.context_length, num_frames))))
|
||||
return windows
|
||||
|
||||
|
||||
def create_windows_default(num_frames: int, opts: Union[ContextOptionsGroup, ContextOptions]):
|
||||
return [list(range(num_frames))]
|
||||
|
||||
|
||||
def get_context_windows(num_frames: int, opts: Union[ContextOptionsGroup, ContextOptions]):
|
||||
context_func = CONTEXT_MAPPING.get(opts.context_schedule, None)
|
||||
if not context_func:
|
||||
raise ValueError(f"Unknown context_schedule '{opts.context_schedule}'.")
|
||||
return context_func(num_frames, opts)
|
||||
|
||||
|
||||
CONTEXT_MAPPING = {
|
||||
ContextSchedules.UNIFORM_LOOPED: create_windows_uniform_looped,
|
||||
ContextSchedules.UNIFORM_STANDARD: create_windows_uniform_standard,
|
||||
ContextSchedules.STATIC_STANDARD: create_windows_static_standard,
|
||||
ContextSchedules.BATCHED: create_windows_batched,
|
||||
ContextSchedules.VIEW_AS_CONTEXT: create_windows_default, # just return all to allow Views to do all the work
|
||||
}
|
||||
|
||||
|
||||
def get_context_weights(num_frames: int, fuse_method: str):
|
||||
weights_func = FUSE_MAPPING.get(fuse_method, None)
|
||||
if not weights_func:
|
||||
raise ValueError(f"Unknown fuse_method '{fuse_method}'.")
|
||||
return weights_func(num_frames)
|
||||
|
||||
|
||||
def create_weights_flat(length: int, **kwargs) -> list[float]:
|
||||
# weight is the same for all
|
||||
return [1.0] * length
|
||||
|
||||
|
||||
def create_weights_pyramid(length: int, **kwargs) -> list[float]:
|
||||
# weight is based on the distance away from the edge of the context window;
|
||||
# based on weighted average concept in FreeNoise paper
|
||||
if length % 2 == 0:
|
||||
max_weight = length // 2
|
||||
weight_sequence = list(range(1, max_weight + 1, 1)) + list(range(max_weight, 0, -1))
|
||||
else:
|
||||
max_weight = (length + 1) // 2
|
||||
weight_sequence = list(range(1, max_weight, 1)) + [max_weight] + list(range(max_weight - 1, 0, -1))
|
||||
return weight_sequence
|
||||
|
||||
|
||||
FUSE_MAPPING = {
|
||||
ContextFuseMethod.FLAT: create_weights_flat,
|
||||
ContextFuseMethod.PYRAMID: create_weights_pyramid,
|
||||
ContextFuseMethod.RELATIVE: create_weights_pyramid,
|
||||
}
|
||||
|
||||
|
||||
# Returns fraction that has denominator that is a power of 2
|
||||
def ordered_halving(val):
|
||||
# get binary value, padded with 0s for 64 bits
|
||||
bin_str = f"{val:064b}"
|
||||
# flip binary value, padding included
|
||||
bin_flip = bin_str[::-1]
|
||||
# convert binary to int
|
||||
as_int = int(bin_flip, 2)
|
||||
# divide by 1 << 64, equivalent to 2**64, or 18446744073709551616,
|
||||
# or b10000000000000000000000000000000000000000000000000000000000000000 (1 with 64 zero's)
|
||||
return as_int / (1 << 64)
|
||||
|
||||
|
||||
def get_missing_indexes(windows: list[list[int]], num_frames: int) -> list[int]:
|
||||
all_indexes = list(range(num_frames))
|
||||
for w in windows:
|
||||
for val in w:
|
||||
try:
|
||||
all_indexes.remove(val)
|
||||
except ValueError:
|
||||
pass
|
||||
return all_indexes
|
||||
|
||||
|
||||
def does_window_roll_over(window: list[int], num_frames: int) -> tuple[bool, int]:
|
||||
prev_val = -1
|
||||
for i, val in enumerate(window):
|
||||
val = val % num_frames
|
||||
if val < prev_val:
|
||||
return True, i
|
||||
prev_val = val
|
||||
return False, -1
|
||||
|
||||
|
||||
def shift_window_to_start(window: list[int], num_frames: int):
|
||||
start_val = window[0]
|
||||
for i in range(len(window)):
|
||||
# 1) subtract each element by start_val to move vals relative to the start of all frames
|
||||
# 2) add num_frames and take modulus to get adjusted vals
|
||||
window[i] = ((window[i] - start_val) + num_frames) % num_frames
|
||||
|
||||
|
||||
def shift_window_to_end(window: list[int], num_frames: int):
|
||||
# 1) shift window to start
|
||||
shift_window_to_start(window, num_frames)
|
||||
end_val = window[-1]
|
||||
end_delta = num_frames - end_val - 1
|
||||
for i in range(len(window)):
|
||||
# 2) add end_delta to each val to slide windows to end
|
||||
window[i] = window[i] + end_delta
|
||||
+162
@@ -0,0 +1,162 @@
|
||||
# S-Lab License 1.0
|
||||
|
||||
# Copyright 2023 S-Lab
|
||||
# Redistribution and use for non-commercial purpose in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
|
||||
# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
# 4. In the event that redistribution and/or use for commercial purpose in source or binary forms, with or without modification is required, please contact the contributor(s) of the work.
|
||||
|
||||
# Code has been modified from https://github.com/TianxingWu/FreeInit
|
||||
|
||||
import torch
|
||||
import torch.fft as fft
|
||||
import math
|
||||
|
||||
|
||||
class FreeInitFilter:
|
||||
GAUSSIAN = "gaussian"
|
||||
IDEAL = "ideal"
|
||||
BOX = "box"
|
||||
BUTTERWORTH = "butterworth"
|
||||
|
||||
LIST = [GAUSSIAN, BUTTERWORTH, IDEAL, BOX]
|
||||
|
||||
|
||||
def freq_mix_3d(x, noise, LPF):
|
||||
"""
|
||||
Noise reinitialization.
|
||||
|
||||
Args:
|
||||
x: diffused latent
|
||||
noise: randomly sampled noise
|
||||
LPF: low pass filter
|
||||
"""
|
||||
# FFT
|
||||
x_freq = fft.fftn(x, dim=(-4, -2, -1))
|
||||
x_freq = fft.fftshift(x_freq, dim=(-4, -2, -1))
|
||||
noise_freq = fft.fftn(noise, dim=(-4, -2, -1))
|
||||
noise_freq = fft.fftshift(noise_freq, dim=(-4, -2, -1))
|
||||
|
||||
# frequency mix
|
||||
HPF = 1 - LPF
|
||||
x_freq_low = x_freq * LPF
|
||||
noise_freq_high = noise_freq * HPF
|
||||
x_freq_mixed = x_freq_low + noise_freq_high # mix in freq domain
|
||||
|
||||
# IFFT
|
||||
x_freq_mixed = fft.ifftshift(x_freq_mixed, dim=(-4, -2, -1))
|
||||
x_mixed = fft.ifftn(x_freq_mixed, dim=(-4, -2, -1)).real
|
||||
|
||||
return x_mixed
|
||||
|
||||
|
||||
def get_freq_filter(shape, device, filter_type, n, d_s, d_t):
|
||||
"""
|
||||
Form the frequency filter for noise reinitialization.
|
||||
|
||||
Args:
|
||||
shape: shape of latent (T, C, H, W)
|
||||
filter_type: type of the freq filter
|
||||
n: (only for butterworth) order of the filter, larger n ~ ideal, smaller n ~ gaussian
|
||||
d_s: normalized stop frequency for spatial dimensions (0.0-1.0)
|
||||
d_t: normalized stop frequency for temporal dimension (0.0-1.0)
|
||||
"""
|
||||
if filter_type == FreeInitFilter.GAUSSIAN:
|
||||
return gaussian_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device)
|
||||
elif filter_type == FreeInitFilter.IDEAL:
|
||||
return ideal_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device)
|
||||
elif filter_type == FreeInitFilter.BOX:
|
||||
return box_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device)
|
||||
elif filter_type == FreeInitFilter.BUTTERWORTH:
|
||||
return butterworth_low_pass_filter(shape=shape, n=n, d_s=d_s, d_t=d_t).to(device)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
def gaussian_low_pass_filter(shape, d_s=0.25, d_t=0.25):
|
||||
"""
|
||||
Compute the gaussian low pass filter mask.
|
||||
|
||||
Args:
|
||||
shape: shape of the filter (volume)
|
||||
d_s: normalized stop frequency for spatial dimensions (0.0-1.0)
|
||||
d_t: normalized stop frequency for temporal dimension (0.0-1.0)
|
||||
"""
|
||||
T, H, W = shape[-4], shape[-2], shape[-1]
|
||||
mask = torch.zeros(shape)
|
||||
if d_s==0 or d_t==0:
|
||||
return mask
|
||||
for t in range(T):
|
||||
for h in range(H):
|
||||
for w in range(W):
|
||||
d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2)
|
||||
mask[t, ..., h,w] = math.exp(-1/(2*d_s**2) * d_square)
|
||||
return mask
|
||||
|
||||
|
||||
def butterworth_low_pass_filter(shape, n=4, d_s=0.25, d_t=0.25):
|
||||
"""
|
||||
Compute the butterworth low pass filter mask.
|
||||
|
||||
Args:
|
||||
shape: shape of the filter (volume)
|
||||
n: order of the filter, larger n ~ ideal, smaller n ~ gaussian
|
||||
d_s: normalized stop frequency for spatial dimensions (0.0-1.0)
|
||||
d_t: normalized stop frequency for temporal dimension (0.0-1.0)
|
||||
"""
|
||||
T, H, W = shape[-4], shape[-2], shape[-1]
|
||||
mask = torch.zeros(shape)
|
||||
if d_s==0 or d_t==0:
|
||||
return mask
|
||||
for t in range(T):
|
||||
for h in range(H):
|
||||
for w in range(W):
|
||||
d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2)
|
||||
mask[t, ..., h,w] = 1 / (1 + (d_square / d_s**2)**n)
|
||||
return mask
|
||||
|
||||
|
||||
def ideal_low_pass_filter(shape, d_s=0.25, d_t=0.25):
|
||||
"""
|
||||
Compute the ideal low pass filter mask.
|
||||
|
||||
Args:
|
||||
shape: shape of the filter (volume)
|
||||
d_s: normalized stop frequency for spatial dimensions (0.0-1.0)
|
||||
d_t: normalized stop frequency for temporal dimension (0.0-1.0)
|
||||
"""
|
||||
T, H, W = shape[-4], shape[-2], shape[-1]
|
||||
mask = torch.zeros(shape)
|
||||
if d_s==0 or d_t==0:
|
||||
return mask
|
||||
for t in range(T):
|
||||
for h in range(H):
|
||||
for w in range(W):
|
||||
d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2)
|
||||
mask[t, ...,h,w] = 1 if d_square <= d_s*2 else 0
|
||||
return mask
|
||||
|
||||
|
||||
def box_low_pass_filter(shape, d_s=0.25, d_t=0.25):
|
||||
"""
|
||||
Compute the ideal low pass filter mask (approximated version).
|
||||
|
||||
Args:
|
||||
shape: shape of the filter (volume)
|
||||
d_s: normalized stop frequency for spatial dimensions (0.0-1.0)
|
||||
d_t: normalized stop frequency for temporal dimension (0.0-1.0)
|
||||
"""
|
||||
T, H, W = shape[-4], shape[-2], shape[-1]
|
||||
mask = torch.zeros(shape)
|
||||
if d_s==0 or d_t==0:
|
||||
return mask
|
||||
|
||||
threshold_s = round(int(H // 2) * d_s)
|
||||
threshold_t = round(T // 2 * d_t)
|
||||
|
||||
cframe, crow, ccol = T // 2, H // 2, W //2
|
||||
mask[cframe - threshold_t:cframe + threshold_t, ..., crow - threshold_s:crow + threshold_s, ccol - threshold_s:ccol + threshold_s] = 1.0
|
||||
|
||||
return mask
|
||||
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
import copy
|
||||
import logging
|
||||
import sys
|
||||
|
||||
|
||||
class ColoredFormatter(logging.Formatter):
|
||||
COLORS = {
|
||||
"DEBUG": "\033[0;36m", # CYAN
|
||||
"INFO": "\033[0;32m", # GREEN
|
||||
"WARNING": "\033[0;33m", # YELLOW
|
||||
"ERROR": "\033[0;31m", # RED
|
||||
"CRITICAL": "\033[0;37;41m", # WHITE ON RED
|
||||
"RESET": "\033[0m", # RESET COLOR
|
||||
}
|
||||
|
||||
def format(self, record):
|
||||
colored_record = copy.copy(record)
|
||||
levelname = colored_record.levelname
|
||||
seq = self.COLORS.get(levelname, self.COLORS["RESET"])
|
||||
colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}"
|
||||
return super().format(colored_record)
|
||||
|
||||
|
||||
# Create a new logger
|
||||
logger = logging.getLogger("AnimateDiffEvo")
|
||||
logger.propagate = False
|
||||
|
||||
# Add handler if we don't have one.
|
||||
if not logger.handlers:
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
handler.setFormatter(ColoredFormatter("[%(name)s] - %(levelname)s - %(message)s"))
|
||||
logger.addHandler(handler)
|
||||
|
||||
# Configure logger
|
||||
loglevel = logging.INFO
|
||||
logger.setLevel(loglevel)
|
||||
+581
@@ -0,0 +1,581 @@
|
||||
import copy
|
||||
from typing import Union
|
||||
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.utils
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
from comfy.model_base import BaseModel
|
||||
|
||||
from .ad_settings import AnimateDiffSettings
|
||||
from .context import ContextOptions, ContextOptions, ContextOptionsGroup
|
||||
from .motion_module_ad import AnimateDiffModel, AnimateDiffFormat, has_mid_block, normalize_ad_state_dict
|
||||
from .logger import logger
|
||||
from .utils_motion import ADKeyframe, ADKeyframeGroup, MotionCompatibilityError, get_combined_multival, normalize_min_max
|
||||
from .motion_lora import MotionLoraInfo, MotionLoraList
|
||||
from .utils_model import get_motion_lora_path, get_motion_model_path, get_sd_model_type
|
||||
from .sample_settings import SampleSettings, SeedNoiseGeneration
|
||||
|
||||
|
||||
# some motion_model casts here might fail if model becomes metatensor or is not castable;
|
||||
# should not really matter if it fails, so ignore raised Exceptions
|
||||
class ModelPatcherAndInjector(ModelPatcher):
|
||||
def __init__(self, m: ModelPatcher):
|
||||
# replicate ModelPatcher.clone() to initialize ModelPatcherAndInjector
|
||||
super().__init__(m.model, m.load_device, m.offload_device, m.size, m.current_device, weight_inplace_update=m.weight_inplace_update)
|
||||
self.patches = {}
|
||||
for k in m.patches:
|
||||
self.patches[k] = m.patches[k][:]
|
||||
|
||||
self.object_patches = m.object_patches.copy()
|
||||
self.model_options = copy.deepcopy(m.model_options)
|
||||
self.model_keys = m.model_keys
|
||||
|
||||
# injection stuff
|
||||
self.motion_injection_params: InjectionParams = None
|
||||
self.sample_settings: SampleSettings = SampleSettings()
|
||||
self.motion_models: MotionModelGroup = None
|
||||
|
||||
def model_patches_to(self, device):
|
||||
super().model_patches_to(device)
|
||||
if self.motion_models is not None:
|
||||
for motion_model in self.motion_models.models:
|
||||
try:
|
||||
motion_model.model.to(device)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def patch_model(self, device_to=None):
|
||||
# first, perform model patching
|
||||
patched_model = super().patch_model(device_to)
|
||||
# finally, perform motion model injection
|
||||
self.inject_model(device_to=device_to)
|
||||
return patched_model
|
||||
|
||||
def unpatch_model(self, device_to=None):
|
||||
# first, eject motion model from unet
|
||||
self.eject_model(device_to=device_to)
|
||||
# finally, do normal model unpatching
|
||||
return super().unpatch_model(device_to)
|
||||
|
||||
def inject_model(self, device_to=None):
|
||||
if self.motion_models is not None:
|
||||
for motion_model in self.motion_models.models:
|
||||
motion_model.model.inject(self)
|
||||
try:
|
||||
motion_model.model.to(device_to)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def eject_model(self, device_to=None):
|
||||
if self.motion_models is not None:
|
||||
for motion_model in self.motion_models.models:
|
||||
motion_model.model.eject(self)
|
||||
try:
|
||||
motion_model.model.to(device_to)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def clone(self):
|
||||
cloned = ModelPatcherAndInjector(self)
|
||||
cloned.motion_models = self.motion_models.clone() if self.motion_models else self.motion_models
|
||||
cloned.sample_settings = self.sample_settings
|
||||
cloned.motion_injection_params = self.motion_injection_params.clone() if self.motion_injection_params else self.motion_injection_params
|
||||
return cloned
|
||||
|
||||
|
||||
class MotionModelPatcher(ModelPatcher):
|
||||
# Mostly here so that type hints work in IDEs
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.model: AnimateDiffModel = self.model
|
||||
self.timestep_percent_range = (0.0, 1.0)
|
||||
self.timestep_range: tuple[float, float] = None
|
||||
self.keyframes: ADKeyframeGroup = ADKeyframeGroup()
|
||||
|
||||
self.scale_multival = None
|
||||
self.effect_multival = None
|
||||
# temporary variables
|
||||
self.current_used_steps = 0
|
||||
self.current_keyframe: ADKeyframe = None
|
||||
self.current_index = -1
|
||||
self.current_scale: Union[float, Tensor] = None
|
||||
self.current_effect: Union[float, Tensor] = None
|
||||
self.combined_scale: Union[float, Tensor] = None
|
||||
self.combined_effect: Union[float, Tensor] = None
|
||||
self.was_within_range = False
|
||||
|
||||
def patch_model(self, *args, **kwargs):
|
||||
# patch as normal, but prepare_weights so that lowvram meta device works properly
|
||||
patched_model = super().patch_model(*args, **kwargs)
|
||||
self.prepare_weights()
|
||||
return patched_model
|
||||
|
||||
def prepare_weights(self):
|
||||
# in case lowvram is active and meta device is used, need to convert weights
|
||||
# otherwise, will get exceptions thrown related to meta device
|
||||
# TODO: with new comfy lowvram system, this is unnecessary
|
||||
state_dict = self.model.state_dict()
|
||||
for key in state_dict:
|
||||
weight = comfy.model_management.resolve_lowvram_weight(state_dict[key], self.model, key)
|
||||
try:
|
||||
comfy.utils.set_attr(self.model, key, weight)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def pre_run(self, model: ModelPatcherAndInjector):
|
||||
self.cleanup()
|
||||
self.model.reset()
|
||||
# just in case, prepare_weights before every run
|
||||
self.prepare_weights()
|
||||
self.model.set_scale(self.scale_multival)
|
||||
self.model.set_effect(self.effect_multival)
|
||||
|
||||
def initialize_timesteps(self, model: BaseModel):
|
||||
self.timestep_range = (model.model_sampling.percent_to_sigma(self.timestep_percent_range[0]),
|
||||
model.model_sampling.percent_to_sigma(self.timestep_percent_range[1]))
|
||||
if self.keyframes is not None:
|
||||
for keyframe in self.keyframes.keyframes:
|
||||
keyframe.start_t = model.model_sampling.percent_to_sigma(keyframe.start_percent)
|
||||
|
||||
def prepare_current_keyframe(self, t: Tensor):
|
||||
curr_t: float = t[0]
|
||||
prev_index = self.current_index
|
||||
# if met guaranteed steps, look for next keyframe in case need to switch
|
||||
if self.current_keyframe is None or self.current_used_steps >= self.current_keyframe.guarantee_steps:
|
||||
# if has next index, loop through and see if need to switch
|
||||
if self.keyframes.has_index(self.current_index+1):
|
||||
for i in range(self.current_index+1, len(self.keyframes)):
|
||||
eval_kf = self.keyframes[i]
|
||||
# check if start_t is greater or equal to curr_t
|
||||
# NOTE: t is in terms of sigmas, not percent, so bigger number = earlier step in sampling
|
||||
if eval_kf.start_t >= curr_t:
|
||||
self.current_index = i
|
||||
self.current_keyframe = eval_kf
|
||||
self.current_used_steps = 0
|
||||
# keep track of scale and effect multivals, accounting for inherit_missing
|
||||
if self.current_keyframe.has_scale():
|
||||
self.current_scale = self.current_keyframe.scale_multival
|
||||
elif not self.current_keyframe.inherit_missing:
|
||||
self.current_scale = None
|
||||
if self.current_keyframe.has_effect():
|
||||
self.current_effect = self.current_keyframe.effect_multival
|
||||
elif not self.current_keyframe.inherit_missing:
|
||||
self.current_effect = None
|
||||
# if guarantee_steps greater than zero, stop searching for other keyframes
|
||||
if self.current_keyframe.guarantee_steps > 0:
|
||||
break
|
||||
# if eval_kf is outside the percent range, stop looking further
|
||||
else:
|
||||
break
|
||||
# if index changed, apply new combined values
|
||||
if prev_index != self.current_index:
|
||||
# combine model's scale and effect with keyframe's scale and effect
|
||||
self.combined_scale = get_combined_multival(self.scale_multival, self.current_scale)
|
||||
self.combined_effect = get_combined_multival(self.effect_multival, self.current_effect)
|
||||
# apply scale and effect
|
||||
self.model.set_scale(self.combined_scale)
|
||||
self.model.set_effect(self.combined_effect)
|
||||
# apply effect - if not within range, set effect to 0, effectively turning model off
|
||||
if curr_t > self.timestep_range[0] or curr_t < self.timestep_range[1]:
|
||||
self.model.set_effect(0.0)
|
||||
self.was_within_range = False
|
||||
else:
|
||||
# if was not in range last step, apply effect to toggle AD status
|
||||
if not self.was_within_range:
|
||||
self.model.set_effect(self.combined_effect)
|
||||
self.was_within_range = True
|
||||
# update steps current keyframe is used
|
||||
self.current_used_steps += 1
|
||||
|
||||
def cleanup(self):
|
||||
if self.model is not None:
|
||||
self.model.cleanup()
|
||||
self.current_used_steps = 0
|
||||
self.current_keyframe = None
|
||||
self.current_index = -1
|
||||
self.current_scale = None
|
||||
self.current_effect = None
|
||||
self.combined_scale = None
|
||||
self.combined_effect = None
|
||||
self.was_within_range = False
|
||||
|
||||
def clone(self):
|
||||
# normal ModelPatcher clone actions
|
||||
n = MotionModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device, weight_inplace_update=self.weight_inplace_update)
|
||||
n.patches = {}
|
||||
for k in self.patches:
|
||||
n.patches[k] = self.patches[k][:]
|
||||
|
||||
n.object_patches = self.object_patches.copy()
|
||||
n.model_options = copy.deepcopy(self.model_options)
|
||||
n.model_keys = self.model_keys
|
||||
# extra cloned params
|
||||
n.timestep_percent_range = self.timestep_percent_range
|
||||
n.timestep_range = self.timestep_range
|
||||
n.keyframes = self.keyframes.clone()
|
||||
n.scale_multival = self.scale_multival
|
||||
n.effect_multival = self.effect_multival
|
||||
return n
|
||||
|
||||
|
||||
class MotionModelGroup:
|
||||
def __init__(self, init_motion_model: MotionModelPatcher=None):
|
||||
self.models: list[MotionModelPatcher] = []
|
||||
if init_motion_model is not None:
|
||||
self.add(init_motion_model)
|
||||
|
||||
def add(self, mm: MotionModelPatcher):
|
||||
# add to end of list
|
||||
self.models.append(mm)
|
||||
|
||||
def add_to_start(self, mm: MotionModelPatcher):
|
||||
self.models.insert(0, mm)
|
||||
|
||||
def __getitem__(self, index) -> MotionModelPatcher:
|
||||
return self.models[index]
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return len(self.models) == 0
|
||||
|
||||
def clone(self) -> 'MotionModelGroup':
|
||||
cloned = MotionModelGroup()
|
||||
for mm in self.models:
|
||||
cloned.add(mm)
|
||||
return cloned
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
for motion_model in self.models:
|
||||
motion_model.model.set_sub_idxs(sub_idxs=sub_idxs)
|
||||
|
||||
def set_view_options(self, view_options: ContextOptions):
|
||||
for motion_model in self.models:
|
||||
motion_model.model.set_view_options(view_options)
|
||||
|
||||
def set_video_length(self, video_length: int, full_length: int):
|
||||
for motion_model in self.models:
|
||||
motion_model.model.set_video_length(video_length=video_length, full_length=full_length)
|
||||
|
||||
def initialize_timesteps(self, model: BaseModel):
|
||||
for motion_model in self.models:
|
||||
motion_model.initialize_timesteps(model)
|
||||
|
||||
def pre_run(self, model: ModelPatcherAndInjector):
|
||||
for motion_model in self.models:
|
||||
motion_model.pre_run(model)
|
||||
|
||||
def prepare_current_keyframe(self, t: Tensor):
|
||||
for motion_model in self.models:
|
||||
motion_model.prepare_current_keyframe(t=t)
|
||||
|
||||
def get_name_string(self, show_version=False):
|
||||
identifiers = []
|
||||
for motion_model in self.models:
|
||||
id = motion_model.model.mm_info.mm_name
|
||||
if show_version:
|
||||
id += f":{motion_model.model.mm_info.mm_version}"
|
||||
identifiers.append(id)
|
||||
return ", ".join(identifiers)
|
||||
|
||||
|
||||
def get_vanilla_model_patcher(m: ModelPatcher) -> ModelPatcher:
|
||||
model = ModelPatcher(m.model, m.load_device, m.offload_device, m.size, m.current_device, weight_inplace_update=m.weight_inplace_update)
|
||||
model.patches = {}
|
||||
for k in m.patches:
|
||||
model.patches[k] = m.patches[k][:]
|
||||
|
||||
model.object_patches = m.object_patches.copy()
|
||||
model.model_options = copy.deepcopy(m.model_options)
|
||||
model.model_keys = m.model_keys
|
||||
return model
|
||||
|
||||
# adapted from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/utils/convert_lora_safetensor_to_diffusers.py
|
||||
# Example LoRA keys:
|
||||
# down_blocks.0.motion_modules.0.temporal_transformer.transformer_blocks.0.attention_blocks.0.processor.to_q_lora.down.weight
|
||||
# down_blocks.0.motion_modules.0.temporal_transformer.transformer_blocks.0.attention_blocks.0.processor.to_q_lora.up.weight
|
||||
#
|
||||
# Example model keys:
|
||||
# down_blocks.0.motion_modules.0.temporal_transformer.transformer_blocks.0.attention_blocks.0.to_q.weight
|
||||
#
|
||||
def load_motion_lora_as_patches(motion_model: MotionModelPatcher, lora: MotionLoraInfo) -> None:
|
||||
def get_version(has_midblock: bool):
|
||||
return "v2" if has_midblock else "v1"
|
||||
|
||||
lora_path = get_motion_lora_path(lora.name)
|
||||
logger.info(f"Loading motion LoRA {lora.name}")
|
||||
state_dict = comfy.utils.load_torch_file(lora_path)
|
||||
|
||||
# remove all non-temporal keys (in case model has extra stuff in it)
|
||||
for key in list(state_dict.keys()):
|
||||
if "temporal" not in key:
|
||||
del state_dict[key]
|
||||
if len(state_dict) == 0:
|
||||
raise ValueError(f"'{lora.name}' contains no temporal keys; it is not a valid motion LoRA!")
|
||||
|
||||
model_has_midblock = motion_model.model.mid_block != None
|
||||
lora_has_midblock = has_mid_block(state_dict)
|
||||
logger.info(f"Applying a {get_version(lora_has_midblock)} LoRA ({lora.name}) to a { motion_model.model.mm_info.mm_version} motion model.")
|
||||
|
||||
patches = {}
|
||||
# convert lora state dict to one that matches motion_module keys and tensors
|
||||
for key in state_dict:
|
||||
# if motion_module doesn't have a midblock, skip mid_block entries
|
||||
if not model_has_midblock:
|
||||
if "mid_block" in key: continue
|
||||
# only process lora down key (we will process up at the same time as down)
|
||||
if "up." in key: continue
|
||||
|
||||
# get up key version of down key
|
||||
up_key = key.replace(".down.", ".up.")
|
||||
|
||||
# adapt key to match motion_module key format - remove 'processor.', '_lora', 'down.', and 'up.'
|
||||
model_key = key.replace("processor.", "").replace("_lora", "").replace("down.", "").replace("up.", "")
|
||||
# motion_module keys have a '0.' after all 'to_out.' weight keys
|
||||
model_key = model_key.replace("to_out.", "to_out.0.")
|
||||
|
||||
weight_down = state_dict[key]
|
||||
weight_up = state_dict[up_key]
|
||||
# actual weights obtained by matrix multiplication of up and down weights
|
||||
# save as a tuple, so that (Motion)ModelPatcher's calculate_weight function detects len==1, applying it correctly
|
||||
patches[model_key] = (torch.mm(weight_up, weight_down),)
|
||||
del state_dict
|
||||
# add patches to motion ModelPatcher
|
||||
motion_model.add_patches(patches=patches, strength_patch=lora.strength)
|
||||
|
||||
|
||||
def load_motion_module_gen1(model_name: str, model: ModelPatcher, motion_lora: MotionLoraList = None, motion_model_settings: AnimateDiffSettings = None) -> MotionModelPatcher:
|
||||
model_path = get_motion_model_path(model_name)
|
||||
logger.info(f"Loading motion module {model_name}")
|
||||
mm_state_dict = comfy.utils.load_torch_file(model_path, safe_load=True)
|
||||
# TODO: check for empty state dict?
|
||||
# get normalized state_dict and motion model info
|
||||
mm_state_dict, mm_info = normalize_ad_state_dict(mm_state_dict=mm_state_dict, mm_name=model_name)
|
||||
# check that motion model is compatible with sd model
|
||||
model_sd_type = get_sd_model_type(model)
|
||||
if model_sd_type != mm_info.sd_type:
|
||||
raise MotionCompatibilityError(f"Motion module '{mm_info.mm_name}' is intended for {mm_info.sd_type} models, " \
|
||||
+ f"but the provided model is type {model_sd_type}.")
|
||||
# apply motion model settings
|
||||
mm_state_dict = apply_mm_settings(model_dict=mm_state_dict, mm_settings=motion_model_settings)
|
||||
# initialize AnimateDiffModelWrapper
|
||||
ad_wrapper = AnimateDiffModel(mm_state_dict=mm_state_dict, mm_info=mm_info)
|
||||
ad_wrapper.to(model.model_dtype())
|
||||
ad_wrapper.to(model.offload_device)
|
||||
is_animatelcm = mm_info.mm_format==AnimateDiffFormat.ANIMATELCM
|
||||
load_result = ad_wrapper.load_state_dict(mm_state_dict, strict=not is_animatelcm)
|
||||
# TODO: report load_result of motion_module loading?
|
||||
# wrap motion_module into a ModelPatcher, to allow motion lora patches
|
||||
motion_model = MotionModelPatcher(model=ad_wrapper, load_device=model.load_device, offload_device=model.offload_device)
|
||||
# load motion_lora, if present
|
||||
if motion_lora is not None:
|
||||
for lora in motion_lora.loras:
|
||||
load_motion_lora_as_patches(motion_model, lora)
|
||||
return motion_model
|
||||
|
||||
|
||||
def load_motion_module_gen2(model_name: str, motion_model_settings: AnimateDiffSettings = None) -> MotionModelPatcher:
|
||||
model_path = get_motion_model_path(model_name)
|
||||
logger.info(f"Loading motion module {model_name} via Gen2")
|
||||
mm_state_dict = comfy.utils.load_torch_file(model_path, safe_load=True)
|
||||
# TODO: check for empty state dict?
|
||||
# get normalized state_dict and motion model info (converts alternate AD models like HotshotXL into AD keys)
|
||||
mm_state_dict, mm_info = normalize_ad_state_dict(mm_state_dict=mm_state_dict, mm_name=model_name)
|
||||
# apply motion model settings
|
||||
mm_state_dict = apply_mm_settings(model_dict=mm_state_dict, mm_settings=motion_model_settings)
|
||||
# initialize AnimateDiffModelWrapper
|
||||
ad_wrapper = AnimateDiffModel(mm_state_dict=mm_state_dict, mm_info=mm_info)
|
||||
ad_wrapper.to(comfy.model_management.unet_dtype())
|
||||
ad_wrapper.to(comfy.model_management.unet_offload_device())
|
||||
is_animatelcm = mm_info.mm_format==AnimateDiffFormat.ANIMATELCM
|
||||
load_result = ad_wrapper.load_state_dict(mm_state_dict, strict=not is_animatelcm)
|
||||
# TODO: manually check load_results for AnimateLCM models
|
||||
if is_animatelcm:
|
||||
pass
|
||||
# TODO: report load_result of motion_module loading?
|
||||
# wrap motion_module into a ModelPatcher, to allow motion lora patches
|
||||
motion_model = MotionModelPatcher(model=ad_wrapper, load_device=comfy.model_management.get_torch_device(),
|
||||
offload_device=comfy.model_management.unet_offload_device())
|
||||
return motion_model
|
||||
|
||||
|
||||
def create_fresh_motion_module(motion_model: MotionModelPatcher) -> MotionModelPatcher:
|
||||
ad_wrapper = AnimateDiffModel(mm_state_dict=motion_model.model.state_dict(), mm_info=motion_model.model.mm_info)
|
||||
ad_wrapper.to(comfy.model_management.unet_dtype())
|
||||
ad_wrapper.to(comfy.model_management.unet_offload_device())
|
||||
ad_wrapper.load_state_dict(motion_model.model.state_dict())
|
||||
return MotionModelPatcher(model=ad_wrapper, load_device=comfy.model_management.get_torch_device(),
|
||||
offload_device=comfy.model_management.unet_offload_device())
|
||||
|
||||
|
||||
def validate_model_compatibility_gen2(model: ModelPatcher, motion_model: MotionModelPatcher):
|
||||
# check that motion model is compatible with sd model
|
||||
model_sd_type = get_sd_model_type(model)
|
||||
mm_info = motion_model.model.mm_info
|
||||
if model_sd_type != mm_info.sd_type:
|
||||
raise MotionCompatibilityError(f"Motion module '{mm_info.mm_name}' is intended for {mm_info.sd_type} models, " \
|
||||
+ f"but the provided model is type {model_sd_type}.")
|
||||
|
||||
|
||||
def interpolate_pe_to_length(model_dict: dict[str, Tensor], key: str, new_length: int):
|
||||
pe_shape = model_dict[key].shape
|
||||
temp_pe = rearrange(model_dict[key], "(t b) f d -> t b f d", t=1)
|
||||
temp_pe = F.interpolate(temp_pe, size=(new_length, pe_shape[-1]), mode="bilinear")
|
||||
temp_pe = rearrange(temp_pe, "t b f d -> (t b) f d", t=1)
|
||||
model_dict[key] = temp_pe
|
||||
del temp_pe
|
||||
|
||||
|
||||
def interpolate_pe_to_length_diffs(model_dict: dict[str, Tensor], key: str, new_length: int):
|
||||
# TODO: fill out and try out
|
||||
pe_shape = model_dict[key].shape
|
||||
temp_pe = rearrange(model_dict[key], "(t b) f d -> t b f d", t=1)
|
||||
temp_pe = F.interpolate(temp_pe, size=(new_length, pe_shape[-1]), mode="bilinear")
|
||||
temp_pe = rearrange(temp_pe, "t b f d -> (t b) f d", t=1)
|
||||
model_dict[key] = temp_pe
|
||||
del temp_pe
|
||||
|
||||
|
||||
def interpolate_pe_to_length_pingpong(model_dict: dict[str, Tensor], key: str, new_length: int):
|
||||
if model_dict[key].shape[1] < new_length:
|
||||
temp_pe = model_dict[key]
|
||||
flipped_temp_pe = torch.flip(temp_pe[:, 1:-1, :], [1])
|
||||
use_flipped = True
|
||||
preview_pe = None
|
||||
while model_dict[key].shape[1] < new_length:
|
||||
preview_pe = model_dict[key]
|
||||
model_dict[key] = torch.cat([model_dict[key], flipped_temp_pe if use_flipped else temp_pe], dim=1)
|
||||
use_flipped = not use_flipped
|
||||
del temp_pe
|
||||
del flipped_temp_pe
|
||||
del preview_pe
|
||||
model_dict[key] = model_dict[key][:, :new_length]
|
||||
|
||||
|
||||
def freeze_mask_of_pe(model_dict: dict[str, Tensor], key: str):
|
||||
pe_portion = model_dict[key].shape[2] // 64
|
||||
first_pe = model_dict[key][:,:1,:]
|
||||
model_dict[key][:,:,pe_portion:] = first_pe[:,:,pe_portion:]
|
||||
del first_pe
|
||||
|
||||
|
||||
def freeze_mask_of_attn(model_dict: dict[str, Tensor], key: str):
|
||||
attn_portion = model_dict[key].shape[0] // 2
|
||||
model_dict[key][:attn_portion,:attn_portion] *= 1.5
|
||||
|
||||
|
||||
def apply_mm_settings(model_dict: dict[str, Tensor], mm_settings: AnimateDiffSettings) -> dict[str, Tensor]:
|
||||
if mm_settings is None:
|
||||
return model_dict
|
||||
if not mm_settings.has_anything_to_apply():
|
||||
return model_dict
|
||||
# first, handle PE Adjustments
|
||||
for adjust in mm_settings.adjust_pe.adjusts:
|
||||
if adjust.has_anything_to_apply():
|
||||
already_printed = False
|
||||
for key in model_dict:
|
||||
if "attention_blocks" in key and "pos_encoder" in key:
|
||||
# apply simple motion pe stretch, if needed
|
||||
if adjust.has_motion_pe_stretch():
|
||||
original_length = model_dict[key].shape[1]
|
||||
new_pe_length = original_length + adjust.motion_pe_stretch
|
||||
interpolate_pe_to_length(model_dict, key, new_length=new_pe_length)
|
||||
if adjust.print_adjustment and not already_printed:
|
||||
logger.info(f"[Adjust PE]: PE Stretch from {original_length} to {new_pe_length}.")
|
||||
# apply pe_idx_offset, if needed
|
||||
if adjust.has_initial_pe_idx_offset():
|
||||
original_length = model_dict[key].shape[1]
|
||||
model_dict[key] = model_dict[key][:, adjust.initial_pe_idx_offset:]
|
||||
if adjust.print_adjustment and not already_printed:
|
||||
logger.info(f"[Adjust PE]: Offsetting PEs by {adjust.initial_pe_idx_offset}; PE length to shortens from {original_length} to {model_dict[key].shape[1]}.")
|
||||
# apply has_cap_initial_pe_length, if needed
|
||||
if adjust.has_cap_initial_pe_length():
|
||||
original_length = model_dict[key].shape[1]
|
||||
model_dict[key] = model_dict[key][:, :adjust.cap_initial_pe_length]
|
||||
if adjust.print_adjustment and not already_printed:
|
||||
logger.info(f"[Adjust PE]: Capping PEs (initial) from {original_length} to {model_dict[key].shape[1]}.")
|
||||
# apply interpolate_pe_to_length, if needed
|
||||
if adjust.has_interpolate_pe_to_length():
|
||||
original_length = model_dict[key].shape[1]
|
||||
interpolate_pe_to_length(model_dict, key, new_length=adjust.interpolate_pe_to_length)
|
||||
if adjust.print_adjustment and not already_printed:
|
||||
logger.info(f"[Adjust PE]: Interpolating PE length from {original_length} to {model_dict[key].shape[1]}.")
|
||||
# apply final_pe_idx_offset, if needed
|
||||
if adjust.has_final_pe_idx_offset():
|
||||
original_length = model_dict[key].shape[1]
|
||||
model_dict[key] = model_dict[key][:, adjust.final_pe_idx_offset:]
|
||||
if adjust.print_adjustment and not already_printed:
|
||||
logger.info(f"[Adjust PE]: Capping PEs (final) from {original_length} to {model_dict[key].shape[1]}.")
|
||||
already_printed = True
|
||||
# finally, apply any weight changes
|
||||
for key in model_dict:
|
||||
if "attention_blocks" in key:
|
||||
if "pos_encoder" in key and mm_settings.adjust_pe.has_anything_to_apply():
|
||||
# apply pe_strength, if needed
|
||||
if mm_settings.has_pe_strength():
|
||||
model_dict[key] *= mm_settings.pe_strength
|
||||
else:
|
||||
# apply attn_strenth, if needed
|
||||
if mm_settings.has_attn_strength():
|
||||
model_dict[key] *= mm_settings.attn_strength
|
||||
# apply specific attn_strengths, if needed
|
||||
if mm_settings.has_any_attn_sub_strength():
|
||||
if "to_q" in key and mm_settings.has_attn_q_strength():
|
||||
model_dict[key] *= mm_settings.attn_q_strength
|
||||
elif "to_k" in key and mm_settings.has_attn_k_strength():
|
||||
model_dict[key] *= mm_settings.attn_k_strength
|
||||
elif "to_v" in key and mm_settings.has_attn_v_strength():
|
||||
model_dict[key] *= mm_settings.attn_v_strength
|
||||
elif "to_out" in key:
|
||||
if key.strip().endswith("weight") and mm_settings.has_attn_out_weight_strength():
|
||||
model_dict[key] *= mm_settings.attn_out_weight_strength
|
||||
elif key.strip().endswith("bias") and mm_settings.has_attn_out_bias_strength():
|
||||
model_dict[key] *= mm_settings.attn_out_bias_strength
|
||||
# apply other strength, if needed
|
||||
elif mm_settings.has_other_strength():
|
||||
model_dict[key] *= mm_settings.other_strength
|
||||
return model_dict
|
||||
|
||||
|
||||
class InjectionParams:
|
||||
def __init__(self, unlimited_area_hack: bool=False, apply_mm_groupnorm_hack: bool=True, model_name: str="",
|
||||
apply_v2_properly: bool=True) -> None:
|
||||
self.full_length = None
|
||||
self.unlimited_area_hack = unlimited_area_hack
|
||||
self.apply_mm_groupnorm_hack = apply_mm_groupnorm_hack
|
||||
self.model_name = model_name
|
||||
self.apply_v2_properly = apply_v2_properly
|
||||
self.context_options: ContextOptionsGroup = ContextOptionsGroup.default()
|
||||
self.motion_model_settings = AnimateDiffSettings() # Gen1
|
||||
self.sub_idxs = None # value should NOT be included in clone, so it will auto reset
|
||||
|
||||
def set_noise_extra_args(self, noise_extra_args: dict):
|
||||
noise_extra_args["context_options"] = self.context_options.clone()
|
||||
|
||||
def set_context(self, context_options: ContextOptionsGroup):
|
||||
self.context_options = context_options.clone() if context_options else ContextOptionsGroup.default()
|
||||
|
||||
def is_using_sliding_context(self) -> bool:
|
||||
return self.context_options.context_length is not None
|
||||
|
||||
def set_motion_model_settings(self, motion_model_settings: AnimateDiffSettings): # Gen1
|
||||
if motion_model_settings is None:
|
||||
self.motion_model_settings = AnimateDiffSettings()
|
||||
else:
|
||||
self.motion_model_settings = motion_model_settings
|
||||
|
||||
def reset_context(self):
|
||||
self.context_options = ContextOptionsGroup.default()
|
||||
|
||||
def clone(self) -> 'InjectionParams':
|
||||
new_params = InjectionParams(
|
||||
self.unlimited_area_hack, self.apply_mm_groupnorm_hack,
|
||||
self.model_name, apply_v2_properly=self.apply_v2_properly,
|
||||
)
|
||||
new_params.full_length = self.full_length
|
||||
new_params.set_context(self.context_options)
|
||||
new_params.set_motion_model_settings(self.motion_model_settings) # Gen1
|
||||
return new_params
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
class MotionLoraInfo:
|
||||
def __init__(self, name: str, strength: float = 1.0, hash: str=""):
|
||||
self.name = name
|
||||
self.strength = strength
|
||||
self.hash = ""
|
||||
|
||||
def set_hash(self, hash: str):
|
||||
self.hash = hash
|
||||
|
||||
def clone(self):
|
||||
return MotionLoraInfo(self.name, self.strength, self.hash)
|
||||
|
||||
|
||||
class MotionLoraList:
|
||||
def __init__(self):
|
||||
self.loras: list[MotionLoraInfo] = []
|
||||
|
||||
def add_lora(self, lora: MotionLoraInfo):
|
||||
self.loras.append(lora)
|
||||
|
||||
def clone(self):
|
||||
new_list = MotionLoraList()
|
||||
for lora in self.loras:
|
||||
new_list.add_lora(lora.clone())
|
||||
return new_list
|
||||
+971
@@ -0,0 +1,971 @@
|
||||
import math
|
||||
from typing import Iterable, Tuple, Union
|
||||
import re
|
||||
|
||||
import torch
|
||||
from einops import rearrange, repeat
|
||||
from torch import Tensor, nn
|
||||
|
||||
from comfy.ldm.modules.attention import FeedForward, SpatialTransformer
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
from comfy.ldm.modules.diffusionmodules import openaimodel
|
||||
from comfy.ldm.modules.diffusionmodules.openaimodel import SpatialTransformer
|
||||
from comfy.controlnet import broadcast_image_to
|
||||
from comfy.utils import repeat_to_batch_size
|
||||
import comfy.ops
|
||||
import comfy.model_management
|
||||
|
||||
from .context import ContextFuseMethod, ContextOptions, get_context_weights, get_context_windows
|
||||
from .utils_motion import CrossAttentionMM, MotionCompatibilityError, extend_to_batch_size, prepare_mask_batch
|
||||
from .utils_model import BetaSchedules, ModelTypeSD
|
||||
from .logger import logger
|
||||
|
||||
|
||||
def zero_module(module):
|
||||
# Zero out the parameters of a module and return it.
|
||||
for p in module.parameters():
|
||||
p.detach().zero_()
|
||||
return module
|
||||
|
||||
|
||||
class AnimateDiffFormat:
|
||||
ANIMATEDIFF = "AnimateDiff"
|
||||
HOTSHOTXL = "HotshotXL"
|
||||
ANIMATELCM = "AnimateLCM"
|
||||
|
||||
|
||||
class AnimateDiffVersion:
|
||||
V1 = "v1"
|
||||
V2 = "v2"
|
||||
V3 = "v3"
|
||||
|
||||
|
||||
class AnimateDiffInfo:
|
||||
def __init__(self, sd_type: str, mm_format: str, mm_version: str, mm_name: str):
|
||||
self.sd_type = sd_type
|
||||
self.mm_format = mm_format
|
||||
self.mm_version = mm_version
|
||||
self.mm_name = mm_name
|
||||
|
||||
def get_string(self):
|
||||
return f"{self.mm_name}:{self.mm_version}:{self.mm_format}:{self.sd_type}"
|
||||
|
||||
|
||||
def is_hotshotxl(mm_state_dict: dict[str, Tensor]) -> bool:
|
||||
# use pos_encoder naming to determine if hotshotxl model
|
||||
for key in mm_state_dict.keys():
|
||||
if key.endswith("pos_encoder.positional_encoding"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_animatelcm(mm_state_dict: dict[str, Tensor]) -> bool:
|
||||
# use lack of ANY pos_encoder keys to determine if animatelcm model
|
||||
for key in mm_state_dict.keys():
|
||||
if "pos_encoder" in key:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_down_block_max(mm_state_dict: dict[str, Tensor]) -> int:
|
||||
# keep track of biggest down_block count in module
|
||||
biggest_block = 0
|
||||
for key in mm_state_dict.keys():
|
||||
if "down_blocks" in key:
|
||||
try:
|
||||
block_int = key.split(".")[1]
|
||||
block_num = int(block_int)
|
||||
if block_num > biggest_block:
|
||||
biggest_block = block_num
|
||||
except ValueError:
|
||||
pass
|
||||
return biggest_block
|
||||
|
||||
|
||||
def has_mid_block(mm_state_dict: dict[str, Tensor]):
|
||||
# check if keys contain mid_block
|
||||
for key in mm_state_dict.keys():
|
||||
if key.startswith("mid_block."):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def get_position_encoding_max_len(mm_state_dict: dict[str, Tensor], mm_name: str, mm_format: str) -> Union[int, None]:
|
||||
# use pos_encoder.pe entries to determine max length - [1, {max_length}, {320|640|1280}]
|
||||
for key in mm_state_dict.keys():
|
||||
if key.endswith("pos_encoder.pe"):
|
||||
return mm_state_dict[key].size(1) # get middle dim
|
||||
# AnimateLCM models should have no pos_encoder entries, and assumed to be 64
|
||||
if mm_format == AnimateDiffFormat.ANIMATELCM:
|
||||
return 64
|
||||
raise MotionCompatibilityError(f"No pos_encoder.pe found in mm_state_dict - {mm_name} is not a valid AnimateDiff motion module!")
|
||||
|
||||
|
||||
_regex_hotshotxl_module_num = re.compile(r'temporal_attentions\.(\d+)\.')
|
||||
def find_hotshot_module_num(key: str) -> Union[int, None]:
|
||||
found = _regex_hotshotxl_module_num.search(key)
|
||||
if found:
|
||||
return int(found.group(1))
|
||||
return None
|
||||
|
||||
|
||||
def normalize_ad_state_dict(mm_state_dict: dict[str, Tensor], mm_name: str) -> Tuple[dict[str, Tensor], AnimateDiffInfo]:
|
||||
# from pathlib import Path
|
||||
# with open(Path(__file__).parent.parent.parent / f"keys_{mm_name}.txt", "w") as afile:
|
||||
# for key, value in mm_state_dict.items():
|
||||
# afile.write(f"{key}:\t{value.shape}\n")
|
||||
|
||||
# remove all non-temporal keys (in case model has extra stuff in it)
|
||||
for key in list(mm_state_dict.keys()):
|
||||
if "temporal" not in key:
|
||||
del mm_state_dict[key]
|
||||
# determine what SD model the motion module is intended for
|
||||
sd_type: str = None
|
||||
down_block_max = get_down_block_max(mm_state_dict)
|
||||
if down_block_max == 3:
|
||||
sd_type = ModelTypeSD.SD1_5
|
||||
elif down_block_max == 2:
|
||||
sd_type = ModelTypeSD.SDXL
|
||||
else:
|
||||
raise ValueError(f"'{mm_name}' is not a valid SD1.5 nor SDXL motion module - contained {down_block_max} downblocks.")
|
||||
# determine the model's format
|
||||
mm_format = AnimateDiffFormat.ANIMATEDIFF
|
||||
if is_hotshotxl(mm_state_dict):
|
||||
mm_format = AnimateDiffFormat.HOTSHOTXL
|
||||
if is_animatelcm(mm_state_dict):
|
||||
mm_format = AnimateDiffFormat.ANIMATELCM
|
||||
# determine the model's version
|
||||
mm_version = AnimateDiffVersion.V1
|
||||
if has_mid_block(mm_state_dict):
|
||||
mm_version = AnimateDiffVersion.V2
|
||||
elif sd_type==ModelTypeSD.SD1_5 and get_position_encoding_max_len(mm_state_dict, mm_name, mm_format)==32:
|
||||
mm_version = AnimateDiffVersion.V3
|
||||
info = AnimateDiffInfo(sd_type=sd_type, mm_format=mm_format, mm_version=mm_version, mm_name=mm_name)
|
||||
# convert to AnimateDiff format, if needed
|
||||
if mm_format == AnimateDiffFormat.HOTSHOTXL:
|
||||
# HotshotXL is AD-based architecture applied to SDXL instead of SD1.5
|
||||
# By renaming the keys, no code needs to be adapted at all
|
||||
#
|
||||
# reformat temporal_attentions:
|
||||
# HSXL: temporal_attentions.#.
|
||||
# AD: motion_modules.#.temporal_transformer.
|
||||
# HSXL: pos_encoder.positional_encoding
|
||||
# AD: pos_encoder.pe
|
||||
for key in list(mm_state_dict.keys()):
|
||||
module_num = find_hotshot_module_num(key)
|
||||
if module_num is not None:
|
||||
new_key = key.replace(f"temporal_attentions.{module_num}",
|
||||
f"motion_modules.{module_num}.temporal_transformer", 1)
|
||||
new_key = new_key.replace("pos_encoder.positional_encoding", "pos_encoder.pe")
|
||||
mm_state_dict[new_key] = mm_state_dict[key]
|
||||
del mm_state_dict[key]
|
||||
# return adjusted mm_state_dict and info
|
||||
return mm_state_dict, info
|
||||
|
||||
|
||||
class BlockType:
|
||||
UP = "up"
|
||||
DOWN = "down"
|
||||
MID = "mid"
|
||||
|
||||
|
||||
class AnimateDiffModel(nn.Module):
|
||||
def __init__(self, mm_state_dict: dict[str, Tensor], mm_info: AnimateDiffInfo):
|
||||
super().__init__()
|
||||
self.mm_info = mm_info
|
||||
self.down_blocks: Iterable[MotionModule] = nn.ModuleList([])
|
||||
self.up_blocks: Iterable[MotionModule] = nn.ModuleList([])
|
||||
self.mid_block: Union[MotionModule, None] = None
|
||||
self.encoding_max_len = get_position_encoding_max_len(mm_state_dict, mm_info.mm_name, mm_info.mm_format)
|
||||
self.has_position_encoding = self.encoding_max_len is not None
|
||||
# determine ops to use (to support fp8 properly)
|
||||
if comfy.model_management.unet_manual_cast(comfy.model_management.unet_dtype(), comfy.model_management.get_torch_device()) is None:
|
||||
ops = comfy.ops.disable_weight_init
|
||||
else:
|
||||
ops = comfy.ops.manual_cast
|
||||
# SDXL has 3 up/down blocks, SD1.5 has 4 up/down blocks
|
||||
if mm_info.sd_type == ModelTypeSD.SDXL:
|
||||
layer_channels = (320, 640, 1280)
|
||||
else:
|
||||
layer_channels = (320, 640, 1280, 1280)
|
||||
# fill out down/up blocks and middle block, if present
|
||||
for c in layer_channels:
|
||||
self.down_blocks.append(MotionModule(c, temporal_position_encoding=self.has_position_encoding,
|
||||
temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.DOWN, ops=ops))
|
||||
for c in reversed(layer_channels):
|
||||
self.up_blocks.append(MotionModule(c, temporal_position_encoding=self.has_position_encoding,
|
||||
temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.UP, ops=ops))
|
||||
if has_mid_block(mm_state_dict):
|
||||
self.mid_block = MotionModule(1280, temporal_position_encoding=self.has_position_encoding,
|
||||
temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.MID, ops=ops)
|
||||
self.AD_video_length: int = 24
|
||||
|
||||
def get_device_debug(self):
|
||||
return self.down_blocks[0].motion_modules[0].temporal_transformer.proj_in.weight.device
|
||||
|
||||
def is_length_valid_for_encoding_max_len(self, length: int):
|
||||
if self.encoding_max_len is None:
|
||||
return True
|
||||
return length <= self.encoding_max_len
|
||||
|
||||
def get_best_beta_schedule(self, log=False) -> str:
|
||||
to_return = None
|
||||
if self.mm_info.sd_type == ModelTypeSD.SD1_5:
|
||||
if self.mm_info.mm_format == AnimateDiffFormat.ANIMATELCM:
|
||||
to_return = BetaSchedules.LCM # while LCM_100 is the intended schedule, I find LCM to have much less flicker
|
||||
else:
|
||||
to_return = BetaSchedules.SQRT_LINEAR
|
||||
elif self.mm_info.sd_type == ModelTypeSD.SDXL:
|
||||
if self.mm_info.mm_format == AnimateDiffFormat.HOTSHOTXL:
|
||||
to_return = BetaSchedules.LINEAR
|
||||
else:
|
||||
to_return = BetaSchedules.LINEAR_ADXL
|
||||
if to_return is not None:
|
||||
if log: logger.info(f"[Autoselect]: '{to_return}' beta_schedule for {self.mm_info.get_string()}")
|
||||
else:
|
||||
to_return = BetaSchedules.USE_EXISTING
|
||||
if log: logger.info(f"[Autoselect]: could not find beta_schedule for {self.mm_info.get_string()}, defaulting to '{to_return}'")
|
||||
return to_return
|
||||
|
||||
def cleanup(self):
|
||||
pass
|
||||
|
||||
def inject(self, model: ModelPatcher):
|
||||
unet: openaimodel.UNetModel = model.model.diffusion_model
|
||||
# inject input (down) blocks
|
||||
# SD15 mm contains 4 downblocks, each with 2 TemporalTransformers - 8 in total
|
||||
# SDXL mm contains 3 downblocks, each with 2 TemporalTransformers - 6 in total
|
||||
self._inject(unet.input_blocks, self.down_blocks)
|
||||
# inject output (up) blocks
|
||||
# SD15 mm contains 4 upblocks, each with 3 TemporalTransformers - 12 in total
|
||||
# SDXL mm contains 3 upblocks, each with 3 TemporalTransformers - 9 in total
|
||||
self._inject(unet.output_blocks, self.up_blocks)
|
||||
# inject mid block, if needed (encapsulate in list to make structure compatible)
|
||||
if self.mid_block is not None:
|
||||
self._inject([unet.middle_block], [self.mid_block])
|
||||
del unet
|
||||
|
||||
def _inject(self, unet_blocks: nn.ModuleList, mm_blocks: nn.ModuleList):
|
||||
# Rules for injection:
|
||||
# For each component list in a unet block:
|
||||
# if SpatialTransformer exists in list, place next block after last occurrence
|
||||
# elif ResBlock exists in list, place next block after first occurrence
|
||||
# else don't place block
|
||||
injection_count = 0
|
||||
unet_idx = 0
|
||||
# details about blocks passed in
|
||||
per_block = len(mm_blocks[0].motion_modules)
|
||||
injection_goal = len(mm_blocks) * per_block
|
||||
# only stop injecting when modules exhausted
|
||||
while injection_count < injection_goal:
|
||||
# figure out which VanillaTemporalModule from mm to inject
|
||||
mm_blk_idx, mm_vtm_idx = injection_count // per_block, injection_count % per_block
|
||||
# figure out layout of unet block components
|
||||
st_idx = -1 # SpatialTransformer index
|
||||
res_idx = -1 # first ResBlock index
|
||||
# first, figure out indeces of relevant blocks
|
||||
for idx, component in enumerate(unet_blocks[unet_idx]):
|
||||
if type(component) == SpatialTransformer:
|
||||
st_idx = idx
|
||||
elif type(component).__name__ == "ResBlock" and res_idx < 0:
|
||||
res_idx = idx
|
||||
# if SpatialTransformer exists, inject right after
|
||||
if st_idx >= 0:
|
||||
#logger.info(f"AD: injecting after ST({st_idx})")
|
||||
unet_blocks[unet_idx].insert(st_idx+1, mm_blocks[mm_blk_idx].motion_modules[mm_vtm_idx])
|
||||
injection_count += 1
|
||||
# otherwise, if only ResBlock exists, inject right after
|
||||
elif res_idx >= 0:
|
||||
#logger.info(f"AD: injecting after Res({res_idx})")
|
||||
unet_blocks[unet_idx].insert(res_idx+1, mm_blocks[mm_blk_idx].motion_modules[mm_vtm_idx])
|
||||
injection_count += 1
|
||||
# increment unet_idx
|
||||
unet_idx += 1
|
||||
|
||||
def eject(self, model: ModelPatcher):
|
||||
unet: openaimodel.UNetModel = model.model.diffusion_model
|
||||
# remove from input blocks (downblocks)
|
||||
self._eject(unet.input_blocks)
|
||||
# remove from output blocks (upblocks)
|
||||
self._eject(unet.output_blocks)
|
||||
# remove from middle block (encapsulate in list to make compatible)
|
||||
self._eject([unet.middle_block])
|
||||
del unet
|
||||
|
||||
def _eject(self, unet_blocks: nn.ModuleList):
|
||||
# eject all VanillaTemporalModule objects from all blocks
|
||||
for block in unet_blocks:
|
||||
idx_to_pop = []
|
||||
for idx, component in enumerate(block):
|
||||
if type(component) == VanillaTemporalModule:
|
||||
idx_to_pop.append(idx)
|
||||
# pop in backwards order, as to not disturb what the indeces refer to
|
||||
for idx in sorted(idx_to_pop, reverse=True):
|
||||
block.pop(idx)
|
||||
|
||||
def set_video_length(self, video_length: int, full_length: int):
|
||||
self.AD_video_length = video_length
|
||||
for block in self.down_blocks:
|
||||
block.set_video_length(video_length, full_length)
|
||||
for block in self.up_blocks:
|
||||
block.set_video_length(video_length, full_length)
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.set_video_length(video_length, full_length)
|
||||
|
||||
def set_scale(self, multival: Union[float, Tensor]):
|
||||
if multival is None:
|
||||
multival = 1.0
|
||||
if type(multival) == Tensor:
|
||||
self._set_scale_multiplier(1.0)
|
||||
self._set_scale_mask(multival)
|
||||
else:
|
||||
self._set_scale_multiplier(multival)
|
||||
self._set_scale_mask(None)
|
||||
|
||||
def set_effect(self, multival: Union[float, Tensor]):
|
||||
for block in self.down_blocks:
|
||||
block.set_effect(multival)
|
||||
for block in self.up_blocks:
|
||||
block.set_effect(multival)
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.set_effect(multival)
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
for block in self.down_blocks:
|
||||
block.set_sub_idxs(sub_idxs)
|
||||
for block in self.up_blocks:
|
||||
block.set_sub_idxs(sub_idxs)
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.set_sub_idxs(sub_idxs)
|
||||
|
||||
def set_view_options(self, view_options: ContextOptions):
|
||||
for block in self.down_blocks:
|
||||
block.set_view_options(view_options)
|
||||
for block in self.up_blocks:
|
||||
block.set_view_options(view_options)
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.set_view_options(view_options)
|
||||
|
||||
def reset(self):
|
||||
self._reset_sub_idxs()
|
||||
self._reset_scale_multiplier()
|
||||
self._reset_temp_vars()
|
||||
|
||||
def _set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
for block in self.down_blocks:
|
||||
block.set_scale_multiplier(multiplier)
|
||||
for block in self.up_blocks:
|
||||
block.set_scale_multiplier(multiplier)
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.set_scale_multiplier(multiplier)
|
||||
|
||||
def _set_scale_mask(self, mask: Tensor):
|
||||
for block in self.down_blocks:
|
||||
block.set_scale_mask(mask)
|
||||
for block in self.up_blocks:
|
||||
block.set_scale_mask(mask)
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.set_scale_mask(mask)
|
||||
|
||||
def _reset_temp_vars(self):
|
||||
for block in self.down_blocks:
|
||||
block.reset_temp_vars()
|
||||
for block in self.up_blocks:
|
||||
block.reset_temp_vars()
|
||||
if self.mid_block is not None:
|
||||
self.mid_block.reset_temp_vars()
|
||||
|
||||
def _reset_scale_multiplier(self):
|
||||
self._set_scale_multiplier(None)
|
||||
|
||||
def _reset_sub_idxs(self):
|
||||
self.set_sub_idxs(None)
|
||||
|
||||
|
||||
class MotionModule(nn.Module):
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
temporal_position_encoding=True,
|
||||
temporal_position_encoding_max_len=24,
|
||||
block_type: str=BlockType.DOWN,
|
||||
ops=comfy.ops.disable_weight_init
|
||||
):
|
||||
super().__init__()
|
||||
if block_type == BlockType.MID:
|
||||
# mid blocks contain only a single VanillaTemporalModule
|
||||
self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList([get_motion_module(in_channels, temporal_position_encoding, temporal_position_encoding_max_len, ops=ops)])
|
||||
else:
|
||||
# down blocks contain two VanillaTemporalModules
|
||||
self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList(
|
||||
[
|
||||
get_motion_module(in_channels, temporal_position_encoding, temporal_position_encoding_max_len, ops=ops),
|
||||
get_motion_module(in_channels, temporal_position_encoding, temporal_position_encoding_max_len, ops=ops)
|
||||
]
|
||||
)
|
||||
# up blocks contain one additional VanillaTemporalModule
|
||||
if block_type == BlockType.UP:
|
||||
self.motion_modules.append(get_motion_module(in_channels, temporal_position_encoding, temporal_position_encoding_max_len, ops=ops))
|
||||
|
||||
def set_video_length(self, video_length: int, full_length: int):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_video_length(video_length, full_length)
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_scale_multiplier(multiplier)
|
||||
|
||||
def set_scale_mask(self, mask: Tensor):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_scale_mask(mask)
|
||||
|
||||
def set_effect(self, multival: Union[float, Tensor]):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_effect(multival)
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_sub_idxs(sub_idxs)
|
||||
|
||||
def set_view_options(self, view_options: ContextOptions):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.set_view_options(view_options=view_options)
|
||||
|
||||
def reset_temp_vars(self):
|
||||
for motion_module in self.motion_modules:
|
||||
motion_module.reset_temp_vars()
|
||||
|
||||
|
||||
def get_motion_module(in_channels, temporal_position_encoding, temporal_position_encoding_max_len, ops=comfy.ops.disable_weight_init):
|
||||
return VanillaTemporalModule(in_channels=in_channels, temporal_position_encoding=temporal_position_encoding, temporal_position_encoding_max_len=temporal_position_encoding_max_len, ops=ops)
|
||||
|
||||
|
||||
class VanillaTemporalModule(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
num_attention_heads=8,
|
||||
num_transformer_block=1,
|
||||
attention_block_types=("Temporal_Self", "Temporal_Self"),
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=True,
|
||||
temporal_position_encoding_max_len=24,
|
||||
temporal_attention_dim_div=1,
|
||||
zero_initialize=True,
|
||||
ops=comfy.ops.disable_weight_init,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.video_length = 16
|
||||
self.full_length = 16
|
||||
self.sub_idxs = None
|
||||
self.view_options = None
|
||||
|
||||
self.effect = None
|
||||
self.temp_effect_mask: Tensor = None
|
||||
self.prev_input_tensor_batch = 0
|
||||
|
||||
self.temporal_transformer = TemporalTransformer3DModel(
|
||||
in_channels=in_channels,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=in_channels
|
||||
// num_attention_heads
|
||||
// temporal_attention_dim_div,
|
||||
num_layers=num_transformer_block,
|
||||
attention_block_types=attention_block_types,
|
||||
cross_frame_attention_mode=cross_frame_attention_mode,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
ops=ops
|
||||
)
|
||||
|
||||
if zero_initialize:
|
||||
self.temporal_transformer.proj_out = zero_module(
|
||||
self.temporal_transformer.proj_out
|
||||
)
|
||||
|
||||
def set_video_length(self, video_length: int, full_length: int):
|
||||
self.video_length = video_length
|
||||
self.full_length = full_length
|
||||
self.temporal_transformer.set_video_length(video_length, full_length)
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
self.temporal_transformer.set_scale_multiplier(multiplier)
|
||||
|
||||
def set_scale_mask(self, mask: Tensor):
|
||||
self.temporal_transformer.set_scale_mask(mask)
|
||||
|
||||
def set_effect(self, multival: Union[float, Tensor]):
|
||||
if type(multival) == Tensor:
|
||||
self.effect = multival
|
||||
elif multival is not None and math.isclose(multival, 1.0):
|
||||
self.effect = None
|
||||
else:
|
||||
self.effect = multival
|
||||
self.temp_effect_mask = None
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
self.sub_idxs = sub_idxs
|
||||
self.temporal_transformer.set_sub_idxs(sub_idxs)
|
||||
|
||||
def set_view_options(self, view_options: ContextOptions):
|
||||
self.view_options = view_options
|
||||
|
||||
def reset_temp_vars(self):
|
||||
self.set_effect(None)
|
||||
self.set_view_options(None)
|
||||
self.temporal_transformer.reset_temp_vars()
|
||||
|
||||
def get_effect_mask(self, input_tensor: Tensor):
|
||||
batch, channel, height, width = input_tensor.shape
|
||||
batched_number = batch // self.video_length
|
||||
full_batched_idxs = list(range(self.video_length))*batched_number
|
||||
# if there is a cached temp_effect_mask and it is valid for current input, return it
|
||||
if batch == self.prev_input_tensor_batch and self.temp_effect_mask is not None:
|
||||
if self.sub_idxs is not None:
|
||||
return self.temp_effect_mask[self.sub_idxs*batched_number]
|
||||
return self.temp_effect_mask[full_batched_idxs]
|
||||
# clear any existing mask
|
||||
del self.temp_effect_mask
|
||||
self.temp_effect_mask = None
|
||||
# recalculate temp mask
|
||||
self.prev_input_tensor_batch = batch
|
||||
# make sure mask matches expected dimensions
|
||||
mask = prepare_mask_batch(self.effect, shape=(self.full_length, 1, height, width))
|
||||
# make sure mask is as long as full_length - clone last element of list if too short
|
||||
self.temp_effect_mask = extend_to_batch_size(mask, self.full_length).to(
|
||||
dtype=input_tensor.dtype, device=input_tensor.device)
|
||||
# return finalized mask
|
||||
if self.sub_idxs is not None:
|
||||
return self.temp_effect_mask[self.sub_idxs*batched_number]
|
||||
return self.temp_effect_mask[full_batched_idxs]
|
||||
|
||||
def forward(self, input_tensor: Tensor, encoder_hidden_states=None, attention_mask=None):
|
||||
if self.effect is None:
|
||||
return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask, self.view_options)
|
||||
# return weighted average of input_tensor and AD output
|
||||
if type(self.effect) != Tensor:
|
||||
effect = self.effect
|
||||
# do nothing if effect is 0
|
||||
if math.isclose(effect, 0.0):
|
||||
return input_tensor
|
||||
else:
|
||||
effect = self.get_effect_mask(input_tensor)
|
||||
return input_tensor*(1.0-effect) + self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask, self.view_options)*effect
|
||||
|
||||
|
||||
class TemporalTransformer3DModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
num_layers,
|
||||
attention_block_types=(
|
||||
"Temporal_Self",
|
||||
"Temporal_Self",
|
||||
),
|
||||
dropout=0.0,
|
||||
norm_num_groups=32,
|
||||
cross_attention_dim=768,
|
||||
activation_fn="geglu",
|
||||
attention_bias=False,
|
||||
upcast_attention=False,
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=24,
|
||||
ops=comfy.ops.disable_weight_init,
|
||||
):
|
||||
super().__init__()
|
||||
self.video_length = 16
|
||||
self.full_length = 16
|
||||
self.raw_scale_mask: Union[Tensor, None] = None
|
||||
self.temp_scale_mask: Union[Tensor, None] = None
|
||||
self.sub_idxs: Union[list[int], None] = None
|
||||
self.prev_hidden_states_batch = 0
|
||||
|
||||
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
|
||||
self.norm = ops.GroupNorm(
|
||||
num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True
|
||||
)
|
||||
self.proj_in = ops.Linear(in_channels, inner_dim)
|
||||
|
||||
self.transformer_blocks: Iterable[TemporalTransformerBlock] = nn.ModuleList(
|
||||
[
|
||||
TemporalTransformerBlock(
|
||||
dim=inner_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
attention_block_types=attention_block_types,
|
||||
dropout=dropout,
|
||||
norm_num_groups=norm_num_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
activation_fn=activation_fn,
|
||||
attention_bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
cross_frame_attention_mode=cross_frame_attention_mode,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
ops=ops,
|
||||
)
|
||||
for d in range(num_layers)
|
||||
]
|
||||
)
|
||||
self.proj_out = ops.Linear(inner_dim, in_channels)
|
||||
|
||||
def set_video_length(self, video_length: int, full_length: int):
|
||||
self.video_length = video_length
|
||||
self.full_length = full_length
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
for block in self.transformer_blocks:
|
||||
block.set_scale_multiplier(multiplier)
|
||||
|
||||
def set_scale_mask(self, mask: Tensor):
|
||||
self.raw_scale_mask = mask
|
||||
self.temp_scale_mask = None
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
self.sub_idxs = sub_idxs
|
||||
for block in self.transformer_blocks:
|
||||
block.set_sub_idxs(sub_idxs)
|
||||
|
||||
def reset_temp_vars(self):
|
||||
del self.temp_scale_mask
|
||||
self.temp_scale_mask = None
|
||||
self.prev_hidden_states_batch = 0
|
||||
|
||||
def get_scale_mask(self, hidden_states: Tensor) -> Union[Tensor, None]:
|
||||
# if no raw mask, return None
|
||||
if self.raw_scale_mask is None:
|
||||
return None
|
||||
shape = hidden_states.shape
|
||||
batch, channel, height, width = shape
|
||||
# if temp mask already calculated, return it
|
||||
if self.temp_scale_mask != None:
|
||||
# check if hidden_states batch matches
|
||||
if batch == self.prev_hidden_states_batch:
|
||||
if self.sub_idxs is not None:
|
||||
return self.temp_scale_mask[:, self.sub_idxs, :]
|
||||
return self.temp_scale_mask
|
||||
# if does not match, reset cached temp_scale_mask and recalculate it
|
||||
del self.temp_scale_mask
|
||||
self.temp_scale_mask = None
|
||||
# otherwise, calculate temp mask
|
||||
self.prev_hidden_states_batch = batch
|
||||
mask = prepare_mask_batch(self.raw_scale_mask, shape=(self.full_length, 1, height, width))
|
||||
mask = repeat_to_batch_size(mask, self.full_length)
|
||||
# if mask not the same amount length as full length, make it match
|
||||
if self.full_length != mask.shape[0]:
|
||||
mask = broadcast_image_to(mask, self.full_length, 1)
|
||||
# reshape mask to attention K shape (h*w, latent_count, 1)
|
||||
batch, channel, height, width = mask.shape
|
||||
# first, perform same operations as on hidden_states,
|
||||
# turning (b, c, h, w) -> (b, h*w, c)
|
||||
mask = mask.permute(0, 2, 3, 1).reshape(batch, height*width, channel)
|
||||
# then, make it the same shape as attention's k, (h*w, b, c)
|
||||
mask = mask.permute(1, 0, 2)
|
||||
# make masks match the expected length of h*w
|
||||
batched_number = shape[0] // self.video_length
|
||||
if batched_number > 1:
|
||||
mask = torch.cat([mask] * batched_number, dim=0)
|
||||
# cache mask and set to proper device
|
||||
self.temp_scale_mask = mask
|
||||
# move temp_scale_mask to proper dtype + device
|
||||
self.temp_scale_mask = self.temp_scale_mask.to(dtype=hidden_states.dtype, device=hidden_states.device)
|
||||
# return subset of masks, if needed
|
||||
if self.sub_idxs is not None:
|
||||
return self.temp_scale_mask[:, self.sub_idxs, :]
|
||||
return self.temp_scale_mask
|
||||
|
||||
def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, view_options: ContextOptions=None):
|
||||
batch, channel, height, width = hidden_states.shape
|
||||
residual = hidden_states
|
||||
scale_mask = self.get_scale_mask(hidden_states)
|
||||
# add some casts for fp8 purposes - does not affect speed otherwise
|
||||
hidden_states = self.norm(hidden_states).to(hidden_states.dtype)
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
|
||||
batch, height * width, inner_dim
|
||||
)
|
||||
hidden_states = self.proj_in(hidden_states).to(hidden_states.dtype)
|
||||
|
||||
# Transformer Blocks
|
||||
for block in self.transformer_blocks:
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
video_length=self.video_length,
|
||||
scale_mask=scale_mask,
|
||||
view_options=view_options
|
||||
)
|
||||
|
||||
# output
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
hidden_states = (
|
||||
hidden_states.reshape(batch, height, width, inner_dim)
|
||||
.permute(0, 3, 1, 2)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
output = hidden_states + residual
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class TemporalTransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
attention_block_types=(
|
||||
"Temporal_Self",
|
||||
"Temporal_Self",
|
||||
),
|
||||
dropout=0.0,
|
||||
norm_num_groups=32,
|
||||
cross_attention_dim=768,
|
||||
activation_fn="geglu",
|
||||
attention_bias=False,
|
||||
upcast_attention=False,
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=24,
|
||||
ops=comfy.ops.disable_weight_init,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
attention_blocks = []
|
||||
norms = []
|
||||
|
||||
for block_name in attention_block_types:
|
||||
attention_blocks.append(
|
||||
VersatileAttention(
|
||||
attention_mode=block_name.split("_")[0],
|
||||
context_dim=cross_attention_dim # called context_dim for ComfyUI impl
|
||||
if block_name.endswith("_Cross")
|
||||
else None,
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
#bias=attention_bias, # remove for Comfy CrossAttention
|
||||
#upcast_attention=upcast_attention, # remove for Comfy CrossAttention
|
||||
cross_frame_attention_mode=cross_frame_attention_mode,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
ops=ops,
|
||||
)
|
||||
)
|
||||
norms.append(ops.LayerNorm(dim))
|
||||
|
||||
self.attention_blocks: Iterable[VersatileAttention] = nn.ModuleList(attention_blocks)
|
||||
self.norms = nn.ModuleList(norms)
|
||||
|
||||
self.ff = FeedForward(dim, dropout=dropout, glu=(activation_fn == "geglu"), operations=ops)
|
||||
self.ff_norm = ops.LayerNorm(dim)
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
for block in self.attention_blocks:
|
||||
block.set_scale_multiplier(multiplier)
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
for block in self.attention_blocks:
|
||||
block.set_sub_idxs(sub_idxs)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: Tensor,
|
||||
encoder_hidden_states: Tensor=None,
|
||||
attention_mask: Tensor=None,
|
||||
video_length: int=None,
|
||||
scale_mask: Tensor=None,
|
||||
view_options: ContextOptions=None,
|
||||
):
|
||||
# make view_options None if context_length > video_length, or if equal and equal not allowed
|
||||
if view_options:
|
||||
if view_options.context_length > video_length:
|
||||
view_options = None
|
||||
elif view_options.context_length == video_length and not view_options.use_on_equal_length:
|
||||
view_options = None
|
||||
if not view_options:
|
||||
for attention_block, norm in zip(self.attention_blocks, self.norms):
|
||||
norm_hidden_states = norm(hidden_states).to(hidden_states.dtype)
|
||||
hidden_states = (
|
||||
attention_block(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states
|
||||
if attention_block.is_cross_attention
|
||||
else None,
|
||||
attention_mask=attention_mask,
|
||||
video_length=video_length,
|
||||
scale_mask=scale_mask
|
||||
) + hidden_states
|
||||
)
|
||||
else:
|
||||
# views idea gotten from diffusers AnimateDiff FreeNoise implementation:
|
||||
# https://github.com/arthur-qiu/FreeNoise-AnimateDiff/blob/main/animatediff/models/motion_module.py
|
||||
# apply sliding context windows (views)
|
||||
views = get_context_windows(num_frames=video_length, opts=view_options)
|
||||
hidden_states = rearrange(hidden_states, "(b f) d c -> b f d c", f=video_length)
|
||||
value_final = torch.zeros_like(hidden_states)
|
||||
count_final = torch.zeros_like(hidden_states)
|
||||
# bias_final = [0.0] * video_length
|
||||
batched_conds = hidden_states.size(1) // video_length
|
||||
for sub_idxs in views:
|
||||
sub_hidden_states = rearrange(hidden_states[:, sub_idxs], "b f d c -> (b f) d c")
|
||||
for attention_block, norm in zip(self.attention_blocks, self.norms):
|
||||
norm_hidden_states = norm(sub_hidden_states).to(sub_hidden_states.dtype)
|
||||
sub_hidden_states = (
|
||||
attention_block(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states # do these need to be changed for sub_idxs too?
|
||||
if attention_block.is_cross_attention
|
||||
else None,
|
||||
attention_mask=attention_mask,
|
||||
video_length=len(sub_idxs),
|
||||
scale_mask=scale_mask[:, sub_idxs, :] if scale_mask is not None else scale_mask
|
||||
) + sub_hidden_states
|
||||
)
|
||||
sub_hidden_states = rearrange(sub_hidden_states, "(b f) d c -> b f d c", f=len(sub_idxs))
|
||||
|
||||
# if view_options.fuse_method == ContextFuseMethod.RELATIVE:
|
||||
# for pos, idx in enumerate(sub_idxs):
|
||||
# # bias is the influence of a specific index in relation to the whole context window
|
||||
# bias = 1 - abs(idx - (sub_idxs[0] + sub_idxs[-1]) / 2) / ((sub_idxs[-1] - sub_idxs[0] + 1e-2) / 2)
|
||||
# bias = max(1e-2, bias)
|
||||
# # take weighted averate relative to total bias of current idx
|
||||
# bias_total = bias_final[idx]
|
||||
# prev_weight = torch.tensor([bias_total / (bias_total + bias)],
|
||||
# dtype=value_final.dtype, device=value_final.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
|
||||
# #prev_weight = torch.cat([prev_weight]*value_final.shape[1], dim=1)
|
||||
# new_weight = torch.tensor([bias / (bias_total + bias)],
|
||||
# dtype=value_final.dtype, device=value_final.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
|
||||
# #new_weight = torch.cat([new_weight]*value_final.shape[1], dim=1)
|
||||
# test = value_final[:, idx:idx+1, :, :]
|
||||
# value_final[:, idx:idx+1, :, :] = value_final[:, idx:idx+1, :, :] * prev_weight + sub_hidden_states[:, pos:pos+1, : ,:] * new_weight
|
||||
# bias_final[idx] = bias_total + bias
|
||||
# else:
|
||||
weights = get_context_weights(len(sub_idxs), view_options.fuse_method) * batched_conds
|
||||
weights_tensor = torch.Tensor(weights).to(device=hidden_states.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
|
||||
value_final[:, sub_idxs] += sub_hidden_states * weights_tensor
|
||||
count_final[:, sub_idxs] += weights_tensor
|
||||
|
||||
# get weighted average of sub_hidden_states, if fuse method requires it
|
||||
# if view_options.fuse_method != ContextFuseMethod.RELATIVE:
|
||||
hidden_states = value_final / count_final
|
||||
hidden_states = rearrange(hidden_states, "b f d c -> (b f) d c")
|
||||
del value_final
|
||||
del count_final
|
||||
# del bias_final
|
||||
|
||||
hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states
|
||||
|
||||
output = hidden_states
|
||||
return output
|
||||
|
||||
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(self, d_model, dropout=0.0, max_len=24):
|
||||
super().__init__()
|
||||
self.dropout = nn.Dropout(p=dropout)
|
||||
position = torch.arange(max_len).unsqueeze(1)
|
||||
div_term = torch.exp(
|
||||
torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
|
||||
)
|
||||
pe = torch.zeros(1, max_len, d_model)
|
||||
pe[0, :, 0::2] = torch.sin(position * div_term)
|
||||
pe[0, :, 1::2] = torch.cos(position * div_term)
|
||||
self.register_buffer("pe", pe)
|
||||
self.sub_idxs = None
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
self.sub_idxs = sub_idxs
|
||||
|
||||
def forward(self, x):
|
||||
#if self.sub_idxs is not None:
|
||||
# x = x + self.pe[:, self.sub_idxs]
|
||||
#else:
|
||||
x = x + self.pe[:, : x.size(1)]
|
||||
return self.dropout(x)
|
||||
|
||||
|
||||
class VersatileAttention(CrossAttentionMM):
|
||||
def __init__(
|
||||
self,
|
||||
attention_mode=None,
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=24,
|
||||
ops=comfy.ops.disable_weight_init,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(operations=ops, *args, **kwargs)
|
||||
assert attention_mode == "Temporal"
|
||||
|
||||
self.attention_mode = attention_mode
|
||||
self.is_cross_attention = kwargs["context_dim"] is not None
|
||||
|
||||
self.pos_encoder = (
|
||||
PositionalEncoding(
|
||||
kwargs["query_dim"],
|
||||
dropout=0.0,
|
||||
max_len=temporal_position_encoding_max_len,
|
||||
)
|
||||
if (temporal_position_encoding and attention_mode == "Temporal")
|
||||
else None
|
||||
)
|
||||
|
||||
def extra_repr(self):
|
||||
return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}"
|
||||
|
||||
def set_scale_multiplier(self, multiplier: Union[float, None]):
|
||||
if multiplier is None or math.isclose(multiplier, 1.0):
|
||||
self.scale = 1.0
|
||||
else:
|
||||
self.scale = multiplier
|
||||
|
||||
def set_sub_idxs(self, sub_idxs: list[int]):
|
||||
if self.pos_encoder != None:
|
||||
self.pos_encoder.set_sub_idxs(sub_idxs)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: Tensor,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
video_length=None,
|
||||
scale_mask=None,
|
||||
):
|
||||
if self.attention_mode != "Temporal":
|
||||
raise NotImplementedError
|
||||
|
||||
d = hidden_states.shape[1]
|
||||
hidden_states = rearrange(
|
||||
hidden_states, "(b f) d c -> (b d) f c", f=video_length
|
||||
)
|
||||
|
||||
if self.pos_encoder is not None:
|
||||
hidden_states = self.pos_encoder(hidden_states).to(hidden_states.dtype)
|
||||
|
||||
encoder_hidden_states = (
|
||||
repeat(encoder_hidden_states, "b n c -> (b d) n c", d=d)
|
||||
if encoder_hidden_states is not None
|
||||
else encoder_hidden_states
|
||||
)
|
||||
|
||||
hidden_states = super().forward(
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
value=None,
|
||||
mask=attention_mask,
|
||||
scale_mask=scale_mask,
|
||||
)
|
||||
|
||||
hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
|
||||
|
||||
return hidden_states
|
||||
+149
@@ -0,0 +1,149 @@
|
||||
import comfy.sample as comfy_sample
|
||||
|
||||
from .sampling import motion_sample_factory
|
||||
|
||||
from .nodes_gen1 import (AnimateDiffLoaderGen1, LegacyAnimateDiffLoaderWithContext, AnimateDiffModelSettings,
|
||||
AnimateDiffModelSettingsSimple, AnimateDiffModelSettingsAdvanced, AnimateDiffModelSettingsAdvancedAttnStrengths)
|
||||
from .nodes_gen2 import UseEvolvedSamplingNode, ApplyAnimateDiffModelNode, ApplyAnimateDiffModelBasicNode, LoadAnimateDiffModelNode, ADKeyframeNode
|
||||
from .nodes_multival import MultivalDynamicNode, MultivalScaledMaskNode
|
||||
from .nodes_sample import (FreeInitOptionsNode, NoiseLayerAddWeightedNode, SampleSettingsNode, NoiseLayerAddNode, NoiseLayerReplaceNode, IterationOptionsNode,
|
||||
CustomCFGNode, CustomCFGKeyframeNode)
|
||||
from .nodes_sigma_schedule import (SigmaScheduleNode, RawSigmaScheduleNode, WeightedAverageSigmaScheduleNode, InterpolatedWeightedAverageSigmaScheduleNode, SplitAndCombineSigmaScheduleNode)
|
||||
from .nodes_context import (LegacyLoopedUniformContextOptionsNode, LoopedUniformContextOptionsNode, LoopedUniformViewOptionsNode, StandardUniformContextOptionsNode, StandardStaticContextOptionsNode, BatchedContextOptionsNode,
|
||||
StandardStaticViewOptionsNode, StandardUniformViewOptionsNode, ViewAsContextOptionsNode)
|
||||
from .nodes_ad_settings import AnimateDiffSettingsNode, ManualAdjustPENode, SweetspotStretchPENode, FullStretchPENode
|
||||
from .nodes_extras import AnimateDiffUnload, EmptyLatentImageLarge, CheckpointLoaderSimpleWithNoiseSelect
|
||||
from .nodes_deprecated import AnimateDiffLoader_Deprecated, AnimateDiffLoaderAdvanced_Deprecated, AnimateDiffCombine_Deprecated
|
||||
from .nodes_lora import AnimateDiffLoraLoader, MaskedLoraLoader
|
||||
|
||||
from .logger import logger
|
||||
|
||||
# override comfy_sample.sample with animatediff-support version
|
||||
comfy_sample.sample = motion_sample_factory(comfy_sample.sample)
|
||||
comfy_sample.sample_custom = motion_sample_factory(comfy_sample.sample_custom, is_custom=True)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
# Unencapsulated
|
||||
"ADE_AnimateDiffLoRALoader": AnimateDiffLoraLoader,
|
||||
"ADE_AnimateDiffSamplingSettings": SampleSettingsNode,
|
||||
"ADE_AnimateDiffKeyframe": ADKeyframeNode,
|
||||
# Multival Nodes
|
||||
"ADE_MultivalDynamic": MultivalDynamicNode,
|
||||
"ADE_MultivalScaledMask": MultivalScaledMaskNode,
|
||||
# Context Opts
|
||||
"ADE_StandardStaticContextOptions": StandardStaticContextOptionsNode,
|
||||
"ADE_StandardUniformContextOptions": StandardUniformContextOptionsNode,
|
||||
"ADE_LoopedUniformContextOptions": LoopedUniformContextOptionsNode,
|
||||
"ADE_ViewsOnlyContextOptions": ViewAsContextOptionsNode,
|
||||
"ADE_BatchedContextOptions": BatchedContextOptionsNode,
|
||||
"ADE_AnimateDiffUniformContextOptions": LegacyLoopedUniformContextOptionsNode, # Legacy
|
||||
# View Opts
|
||||
"ADE_StandardStaticViewOptions": StandardStaticViewOptionsNode,
|
||||
"ADE_StandardUniformViewOptions": StandardUniformViewOptionsNode,
|
||||
"ADE_LoopedUniformViewOptions": LoopedUniformViewOptionsNode,
|
||||
# Iteration Opts
|
||||
"ADE_IterationOptsDefault": IterationOptionsNode,
|
||||
"ADE_IterationOptsFreeInit": FreeInitOptionsNode,
|
||||
# Noise Layer Nodes
|
||||
"ADE_NoiseLayerAdd": NoiseLayerAddNode,
|
||||
"ADE_NoiseLayerAddWeighted": NoiseLayerAddWeightedNode,
|
||||
"ADE_NoiseLayerReplace": NoiseLayerReplaceNode,
|
||||
# AnimateDiff Settings
|
||||
"ADE_AnimateDiffSettings": AnimateDiffSettingsNode,
|
||||
"ADE_AdjustPESweetspotStretch": SweetspotStretchPENode,
|
||||
"ADE_AdjustPEFullStretch": FullStretchPENode,
|
||||
"ADE_AdjustPEManual": ManualAdjustPENode,
|
||||
# Sample Settings
|
||||
"ADE_CustomCFG": CustomCFGNode,
|
||||
"ADE_CustomCFGKeyframe": CustomCFGKeyframeNode,
|
||||
"ADE_SigmaSchedule": SigmaScheduleNode,
|
||||
"ADE_RawSigmaSchedule": RawSigmaScheduleNode,
|
||||
"ADE_SigmaScheduleWeightedAverage": WeightedAverageSigmaScheduleNode,
|
||||
"ADE_SigmaScheduleWeightedAverageInterp": InterpolatedWeightedAverageSigmaScheduleNode,
|
||||
"ADE_SigmaScheduleSplitAndCombine": SplitAndCombineSigmaScheduleNode,
|
||||
# Extras Nodes
|
||||
"ADE_AnimateDiffUnload": AnimateDiffUnload,
|
||||
"ADE_EmptyLatentImageLarge": EmptyLatentImageLarge,
|
||||
"CheckpointLoaderSimpleWithNoiseSelect": CheckpointLoaderSimpleWithNoiseSelect,
|
||||
# Gen1 Nodes
|
||||
"ADE_AnimateDiffLoaderGen1": AnimateDiffLoaderGen1,
|
||||
"ADE_AnimateDiffLoaderWithContext": LegacyAnimateDiffLoaderWithContext,
|
||||
"ADE_AnimateDiffModelSettings_Release": AnimateDiffModelSettings,
|
||||
"ADE_AnimateDiffModelSettingsSimple": AnimateDiffModelSettingsSimple,
|
||||
"ADE_AnimateDiffModelSettings": AnimateDiffModelSettingsAdvanced,
|
||||
"ADE_AnimateDiffModelSettingsAdvancedAttnStrengths": AnimateDiffModelSettingsAdvancedAttnStrengths,
|
||||
# Gen2 Nodes
|
||||
"ADE_UseEvolvedSampling": UseEvolvedSamplingNode,
|
||||
"ADE_ApplyAnimateDiffModelSimple": ApplyAnimateDiffModelBasicNode,
|
||||
"ADE_ApplyAnimateDiffModel": ApplyAnimateDiffModelNode,
|
||||
"ADE_LoadAnimateDiffModel": LoadAnimateDiffModelNode,
|
||||
# MaskedLoraLoader
|
||||
#"ADE_MaskedLoadLora": MaskedLoraLoader,
|
||||
# Deprecated Nodes
|
||||
"AnimateDiffLoaderV1": AnimateDiffLoader_Deprecated,
|
||||
"ADE_AnimateDiffLoaderV1Advanced": AnimateDiffLoaderAdvanced_Deprecated,
|
||||
"ADE_AnimateDiffCombine": AnimateDiffCombine_Deprecated,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
# Unencapsulated
|
||||
"ADE_AnimateDiffLoRALoader": "Load AnimateDiff LoRA 🎭🅐🅓",
|
||||
"ADE_AnimateDiffSamplingSettings": "Sample Settings 🎭🅐🅓",
|
||||
"ADE_AnimateDiffKeyframe": "AnimateDiff Keyframe 🎭🅐🅓",
|
||||
# Multival Nodes
|
||||
"ADE_MultivalDynamic": "Multival Dynamic 🎭🅐🅓",
|
||||
"ADE_MultivalScaledMask": "Multival Scaled Mask 🎭🅐🅓",
|
||||
# Context Opts
|
||||
"ADE_StandardStaticContextOptions": "Context Options◆Standard Static 🎭🅐🅓",
|
||||
"ADE_StandardUniformContextOptions": "Context Options◆Standard Uniform 🎭🅐🅓",
|
||||
"ADE_LoopedUniformContextOptions": "Context Options◆Looped Uniform 🎭🅐🅓",
|
||||
"ADE_ViewsOnlyContextOptions": "Context Options◆Views Only [VRAM⇈] 🎭🅐🅓",
|
||||
"ADE_BatchedContextOptions": "Context Options◆Batched [Non-AD] 🎭🅐🅓",
|
||||
"ADE_AnimateDiffUniformContextOptions": "Context Options◆Looped Uniform 🎭🅐🅓", # Legacy
|
||||
# View Opts
|
||||
"ADE_StandardStaticViewOptions": "View Options◆Standard Static 🎭🅐🅓",
|
||||
"ADE_StandardUniformViewOptions": "View Options◆Standard Uniform 🎭🅐🅓",
|
||||
"ADE_LoopedUniformViewOptions": "View Options◆Looped Uniform 🎭🅐🅓",
|
||||
# Iteration Opts
|
||||
"ADE_IterationOptsDefault": "Default Iteration Options 🎭🅐🅓",
|
||||
"ADE_IterationOptsFreeInit": "FreeInit Iteration Options 🎭🅐🅓",
|
||||
# Noise Layer Nodes
|
||||
"ADE_NoiseLayerAdd": "Noise Layer [Add] 🎭🅐🅓",
|
||||
"ADE_NoiseLayerAddWeighted": "Noise Layer [Add Weighted] 🎭🅐🅓",
|
||||
"ADE_NoiseLayerReplace": "Noise Layer [Replace] 🎭🅐🅓",
|
||||
# AnimateDiff Settings
|
||||
"ADE_AnimateDiffSettings": "AnimateDiff Settings 🎭🅐🅓",
|
||||
"ADE_AdjustPESweetspotStretch": "Adjust PE [Sweetspot Stretch] 🎭🅐🅓",
|
||||
"ADE_AdjustPEFullStretch": "Adjust PE [Full Stretch] 🎭🅐🅓",
|
||||
"ADE_AdjustPEManual": "Adjust PE [Manual] 🎭🅐🅓",
|
||||
# Sample Settings
|
||||
"ADE_CustomCFG": "Custom CFG 🎭🅐🅓",
|
||||
"ADE_CustomCFGKeyframe": "Custom CFG Keyframe 🎭🅐🅓",
|
||||
"ADE_SigmaSchedule": "Create Sigma Schedule 🎭🅐🅓",
|
||||
"ADE_RawSigmaSchedule": "Create Raw Sigma Schedule 🎭🅐🅓",
|
||||
"ADE_SigmaScheduleWeightedAverage": "Sigma Schedule Weighted Mean 🎭🅐🅓",
|
||||
"ADE_SigmaScheduleWeightedAverageInterp": "Sigma Schedule Interpolated Mean 🎭🅐🅓",
|
||||
"ADE_SigmaScheduleSplitAndCombine": "Sigma Schedule Split Combine 🎭🅐🅓",
|
||||
# Extras Nodes
|
||||
"ADE_AnimateDiffUnload": "AnimateDiff Unload 🎭🅐🅓",
|
||||
"ADE_EmptyLatentImageLarge": "Empty Latent Image (Big Batch) 🎭🅐🅓",
|
||||
"CheckpointLoaderSimpleWithNoiseSelect": "Load Checkpoint w/ Noise Select 🎭🅐🅓",
|
||||
# Gen1 Nodes
|
||||
"ADE_AnimateDiffLoaderGen1": "AnimateDiff Loader 🎭🅐🅓①",
|
||||
"ADE_AnimateDiffLoaderWithContext": "AnimateDiff Loader [Legacy] 🎭🅐🅓①",
|
||||
"ADE_AnimateDiffModelSettings_Release": "[DEPR] Motion Model Settings 🎭🅐🅓①",
|
||||
"ADE_AnimateDiffModelSettingsSimple": "[DEPR] Motion Model Settings (Simple) 🎭🅐🅓①",
|
||||
"ADE_AnimateDiffModelSettings": "[DEPR] Motion Model Settings (Advanced) 🎭🅐🅓①",
|
||||
"ADE_AnimateDiffModelSettingsAdvancedAttnStrengths": "[DEPR] Motion Model Settings (Adv. Attn) 🎭🅐🅓①",
|
||||
# Gen2 Nodes
|
||||
"ADE_UseEvolvedSampling": "Use Evolved Sampling 🎭🅐🅓②",
|
||||
"ADE_ApplyAnimateDiffModelSimple": "Apply AnimateDiff Model 🎭🅐🅓②",
|
||||
"ADE_ApplyAnimateDiffModel": "Apply AnimateDiff Model (Adv.) 🎭🅐🅓②",
|
||||
"ADE_LoadAnimateDiffModel": "Load AnimateDiff Model 🎭🅐🅓②",
|
||||
# MaskedLoraLoader
|
||||
#"ADE_MaskedLoadLora": "Load LoRA (Masked) 🎭🅐🅓",
|
||||
# Deprecated Nodes
|
||||
"AnimateDiffLoaderV1": "AnimateDiff Loader [DEPRECATED] 🎭🅐🅓",
|
||||
"ADE_AnimateDiffLoaderV1Advanced": "AnimateDiff Loader (Advanced) [DEPRECATED] 🎭🅐🅓",
|
||||
"ADE_AnimateDiffCombine": "AnimateDiff Combine [DEPRECATED, Use Video Combine (VHS) Instead!] 🎭🅐🅓",
|
||||
}
|
||||
+107
@@ -0,0 +1,107 @@
|
||||
from .ad_settings import AdjustPE, AdjustPEGroup, AnimateDiffSettings
|
||||
from .utils_model import BIGMAX
|
||||
|
||||
|
||||
class AnimateDiffSettingsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"optional": {
|
||||
"pe_adjust": ("PE_ADJUST",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("AD_SETTINGS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/ad settings"
|
||||
FUNCTION = "get_ad_settings"
|
||||
|
||||
def get_ad_settings(self, pe_adjust: AdjustPEGroup=None):
|
||||
return (AnimateDiffSettings(adjust_pe=pe_adjust),)
|
||||
|
||||
|
||||
class ManualAdjustPENode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"cap_initial_pe_length": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"interpolate_pe_to_length": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"initial_pe_idx_offset": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"final_pe_idx_offset": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"print_adjustment": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {
|
||||
"prev_pe_adjust": ("PE_ADJUST",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("PE_ADJUST",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/ad settings/pe adjust"
|
||||
FUNCTION = "get_pe_adjust"
|
||||
|
||||
def get_pe_adjust(self, cap_initial_pe_length: int, interpolate_pe_to_length: int,
|
||||
initial_pe_idx_offset: int, final_pe_idx_offset: int, print_adjustment: bool,
|
||||
prev_pe_adjust: AdjustPEGroup=None):
|
||||
if prev_pe_adjust is None:
|
||||
prev_pe_adjust = AdjustPEGroup()
|
||||
prev_pe_adjust = prev_pe_adjust.clone()
|
||||
adjust = AdjustPE(cap_initial_pe_length=cap_initial_pe_length, interpolate_pe_to_length=interpolate_pe_to_length,
|
||||
initial_pe_idx_offset=initial_pe_idx_offset, final_pe_idx_offset=final_pe_idx_offset,
|
||||
print_adjustment=print_adjustment)
|
||||
prev_pe_adjust.add(adjust)
|
||||
return (prev_pe_adjust,)
|
||||
|
||||
|
||||
class SweetspotStretchPENode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"sweetspot": ("INT", {"default": 16, "min": 0, "max": BIGMAX},),
|
||||
"new_sweetspot": ("INT", {"default": 16, "min": 0, "max": BIGMAX},),
|
||||
"print_adjustment": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {
|
||||
"prev_pe_adjust": ("PE_ADJUST",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("PE_ADJUST",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/ad settings/pe adjust"
|
||||
FUNCTION = "get_pe_adjust"
|
||||
|
||||
def get_pe_adjust(self, sweetspot: int, new_sweetspot: int, print_adjustment: bool, prev_pe_adjust: AdjustPEGroup=None):
|
||||
if prev_pe_adjust is None:
|
||||
prev_pe_adjust = AdjustPEGroup()
|
||||
prev_pe_adjust = prev_pe_adjust.clone()
|
||||
adjust = AdjustPE(cap_initial_pe_length=sweetspot, interpolate_pe_to_length=new_sweetspot,
|
||||
print_adjustment=print_adjustment)
|
||||
prev_pe_adjust.add(adjust)
|
||||
return (prev_pe_adjust,)
|
||||
|
||||
|
||||
class FullStretchPENode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"pe_stretch": ("INT", {"default": 0, "min": 0, "max": BIGMAX},),
|
||||
"print_adjustment": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {
|
||||
"prev_pe_adjust": ("PE_ADJUST",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("PE_ADJUST",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/ad settings/pe adjust"
|
||||
FUNCTION = "get_pe_adjust"
|
||||
|
||||
def get_pe_adjust(self, pe_stretch: int, print_adjustment: bool, prev_pe_adjust: AdjustPEGroup=None):
|
||||
if prev_pe_adjust is None:
|
||||
prev_pe_adjust = AdjustPEGroup()
|
||||
prev_pe_adjust = prev_pe_adjust.clone()
|
||||
adjust = AdjustPE(motion_pe_stretch=pe_stretch,
|
||||
print_adjustment=print_adjustment)
|
||||
prev_pe_adjust.add(adjust)
|
||||
return (prev_pe_adjust,)
|
||||
+347
@@ -0,0 +1,347 @@
|
||||
from .context import ContextFuseMethod, ContextOptions, ContextOptionsGroup, ContextSchedules
|
||||
from .utils_model import BIGMAX
|
||||
|
||||
|
||||
LENGTH_MAX = 128 # keep an eye on these max values;
|
||||
STRIDE_MAX = 32 # would need to be updated
|
||||
OVERLAP_MAX = 128 # if new motion modules come out
|
||||
|
||||
|
||||
class LoopedUniformContextOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"context_length": ("INT", {"default": 16, "min": 1, "max": LENGTH_MAX}),
|
||||
"context_stride": ("INT", {"default": 1, "min": 1, "max": STRIDE_MAX}),
|
||||
"context_overlap": ("INT", {"default": 4, "min": 0, "max": OVERLAP_MAX}),
|
||||
"closed_loop": ("BOOLEAN", {"default": False},),
|
||||
#"sync_context_to_pe": ("BOOLEAN", {"default": False},),
|
||||
},
|
||||
"optional": {
|
||||
"fuse_method": (ContextFuseMethod.LIST,),
|
||||
"use_on_equal_length": ("BOOLEAN", {"default": False},),
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"guarantee_steps": ("INT", {"default": 1, "min": 0, "max": BIGMAX}),
|
||||
"prev_context": ("CONTEXT_OPTIONS",),
|
||||
"view_opts": ("VIEW_OPTS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTEXT_OPTIONS",)
|
||||
RETURN_NAMES = ("CONTEXT_OPTS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/context opts"
|
||||
FUNCTION = "create_options"
|
||||
|
||||
def create_options(self, context_length: int, context_stride: int, context_overlap: int, closed_loop: bool,
|
||||
fuse_method: str=ContextFuseMethod.FLAT, use_on_equal_length=False, start_percent: float=0.0, guarantee_steps: int=1,
|
||||
view_opts: ContextOptions=None, prev_context: ContextOptionsGroup=None):
|
||||
if prev_context is None:
|
||||
prev_context = ContextOptionsGroup()
|
||||
prev_context = prev_context.clone()
|
||||
|
||||
context_options = ContextOptions(
|
||||
context_length=context_length,
|
||||
context_stride=context_stride,
|
||||
context_overlap=context_overlap,
|
||||
context_schedule=ContextSchedules.UNIFORM_LOOPED,
|
||||
closed_loop=closed_loop,
|
||||
fuse_method=fuse_method,
|
||||
use_on_equal_length=use_on_equal_length,
|
||||
start_percent=start_percent,
|
||||
guarantee_steps=guarantee_steps,
|
||||
view_options=view_opts,
|
||||
)
|
||||
#context_options.set_sync_context_to_pe(sync_context_to_pe)
|
||||
prev_context.add(context_options)
|
||||
return (prev_context,)
|
||||
|
||||
|
||||
# This Legacy version exists to maintain compatiblity with old workflows
|
||||
class LegacyLoopedUniformContextOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"context_length": ("INT", {"default": 16, "min": 1, "max": LENGTH_MAX}),
|
||||
"context_stride": ("INT", {"default": 1, "min": 1, "max": STRIDE_MAX}),
|
||||
"context_overlap": ("INT", {"default": 4, "min": 0, "max": OVERLAP_MAX}),
|
||||
"context_schedule": (ContextSchedules.LEGACY_UNIFORM_SCHEDULE_LIST,),
|
||||
"closed_loop": ("BOOLEAN", {"default": False},),
|
||||
#"sync_context_to_pe": ("BOOLEAN", {"default": False},),
|
||||
},
|
||||
"optional": {
|
||||
"fuse_method": (ContextFuseMethod.LIST, {"default": ContextFuseMethod.FLAT}),
|
||||
"use_on_equal_length": ("BOOLEAN", {"default": False},),
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"guarantee_steps": ("INT", {"default": 1, "min": 0, "max": BIGMAX}),
|
||||
"prev_context": ("CONTEXT_OPTIONS",),
|
||||
"view_opts": ("VIEW_OPTS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTEXT_OPTIONS",)
|
||||
RETURN_NAMES = ("CONTEXT_OPTS",)
|
||||
CATEGORY = "" # No Category, so will not appear in menu
|
||||
FUNCTION = "create_options"
|
||||
|
||||
def create_options(self, fuse_method: str=ContextFuseMethod.FLAT, context_schedule: str=None, **kwargs):
|
||||
return LoopedUniformContextOptionsNode.create_options(self, fuse_method=fuse_method, **kwargs)
|
||||
|
||||
|
||||
class StandardUniformContextOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"context_length": ("INT", {"default": 16, "min": 1, "max": LENGTH_MAX}),
|
||||
"context_stride": ("INT", {"default": 1, "min": 1, "max": STRIDE_MAX}),
|
||||
"context_overlap": ("INT", {"default": 4, "min": 0, "max": OVERLAP_MAX}),
|
||||
},
|
||||
"optional": {
|
||||
"fuse_method": (ContextFuseMethod.LIST,),
|
||||
"use_on_equal_length": ("BOOLEAN", {"default": False},),
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"guarantee_steps": ("INT", {"default": 1, "min": 0, "max": BIGMAX}),
|
||||
"prev_context": ("CONTEXT_OPTIONS",),
|
||||
"view_opts": ("VIEW_OPTS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTEXT_OPTIONS",)
|
||||
RETURN_NAMES = ("CONTEXT_OPTS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/context opts"
|
||||
FUNCTION = "create_options"
|
||||
|
||||
def create_options(self, context_length: int, context_stride: int, context_overlap: int,
|
||||
fuse_method: str=ContextFuseMethod.PYRAMID, use_on_equal_length=False, start_percent: float=0.0, guarantee_steps: int=1,
|
||||
view_opts: ContextOptions=None, prev_context: ContextOptionsGroup=None):
|
||||
if prev_context is None:
|
||||
prev_context = ContextOptionsGroup()
|
||||
prev_context = prev_context.clone()
|
||||
|
||||
context_options = ContextOptions(
|
||||
context_length=context_length,
|
||||
context_stride=context_stride,
|
||||
context_overlap=context_overlap,
|
||||
context_schedule=ContextSchedules.UNIFORM_STANDARD,
|
||||
closed_loop=False,
|
||||
fuse_method=fuse_method,
|
||||
use_on_equal_length=use_on_equal_length,
|
||||
start_percent=start_percent,
|
||||
guarantee_steps=guarantee_steps,
|
||||
view_options=view_opts,
|
||||
)
|
||||
prev_context.add(context_options)
|
||||
return (prev_context,)
|
||||
|
||||
|
||||
class StandardStaticContextOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"context_length": ("INT", {"default": 16, "min": 1, "max": LENGTH_MAX}),
|
||||
"context_overlap": ("INT", {"default": 4, "min": 0, "max": OVERLAP_MAX}),
|
||||
},
|
||||
"optional": {
|
||||
"fuse_method": (ContextFuseMethod.LIST_STATIC,),
|
||||
"use_on_equal_length": ("BOOLEAN", {"default": False},),
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"guarantee_steps": ("INT", {"default": 1, "min": 0, "max": BIGMAX}),
|
||||
"prev_context": ("CONTEXT_OPTIONS",),
|
||||
"view_opts": ("VIEW_OPTS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTEXT_OPTIONS",)
|
||||
RETURN_NAMES = ("CONTEXT_OPTS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/context opts"
|
||||
FUNCTION = "create_options"
|
||||
|
||||
def create_options(self, context_length: int, context_overlap: int,
|
||||
fuse_method: str=ContextFuseMethod.PYRAMID, use_on_equal_length=False, start_percent: float=0.0, guarantee_steps: int=1,
|
||||
view_opts: ContextOptions=None, prev_context: ContextOptionsGroup=None):
|
||||
if prev_context is None:
|
||||
prev_context = ContextOptionsGroup()
|
||||
prev_context = prev_context.clone()
|
||||
|
||||
context_options = ContextOptions(
|
||||
context_length=context_length,
|
||||
context_stride=None,
|
||||
context_overlap=context_overlap,
|
||||
context_schedule=ContextSchedules.STATIC_STANDARD,
|
||||
fuse_method=fuse_method,
|
||||
use_on_equal_length=use_on_equal_length,
|
||||
start_percent=start_percent,
|
||||
guarantee_steps=guarantee_steps,
|
||||
view_options=view_opts,
|
||||
)
|
||||
prev_context.add(context_options)
|
||||
return (prev_context,)
|
||||
|
||||
|
||||
class BatchedContextOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"context_length": ("INT", {"default": 16, "min": 1, "max": LENGTH_MAX}),
|
||||
},
|
||||
"optional": {
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"guarantee_steps": ("INT", {"default": 1, "min": 0, "max": BIGMAX}),
|
||||
"prev_context": ("CONTEXT_OPTIONS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTEXT_OPTIONS",)
|
||||
RETURN_NAMES = ("CONTEXT_OPTS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/context opts"
|
||||
FUNCTION = "create_options"
|
||||
|
||||
def create_options(self, context_length: int, start_percent: float=0.0, guarantee_steps: int=1,
|
||||
prev_context: ContextOptionsGroup=None):
|
||||
if prev_context is None:
|
||||
prev_context = ContextOptionsGroup()
|
||||
prev_context = prev_context.clone()
|
||||
|
||||
context_options = ContextOptions(
|
||||
context_length=context_length,
|
||||
context_overlap=0,
|
||||
context_schedule=ContextSchedules.BATCHED,
|
||||
start_percent=start_percent,
|
||||
guarantee_steps=guarantee_steps,
|
||||
)
|
||||
prev_context.add(context_options)
|
||||
return (prev_context,)
|
||||
|
||||
|
||||
class ViewAsContextOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"view_opts_req": ("VIEW_OPTS",),
|
||||
},
|
||||
"optional": {
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"guarantee_steps": ("INT", {"default": 1, "min": 0, "max": BIGMAX}),
|
||||
"prev_context": ("CONTEXT_OPTIONS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONTEXT_OPTIONS",)
|
||||
RETURN_NAMES = ("CONTEXT_OPTS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/context opts"
|
||||
FUNCTION = "create_options"
|
||||
|
||||
def create_options(self, view_opts_req: ContextOptions, start_percent: float=0.0, guarantee_steps: int=1,
|
||||
prev_context: ContextOptionsGroup=None):
|
||||
if prev_context is None:
|
||||
prev_context = ContextOptionsGroup()
|
||||
prev_context = prev_context.clone()
|
||||
context_options = ContextOptions(
|
||||
context_schedule=ContextSchedules.VIEW_AS_CONTEXT,
|
||||
start_percent=start_percent,
|
||||
guarantee_steps=guarantee_steps,
|
||||
view_options=view_opts_req,
|
||||
use_on_equal_length=True
|
||||
)
|
||||
prev_context.add(context_options)
|
||||
return (prev_context,)
|
||||
|
||||
|
||||
#########################
|
||||
# View Options
|
||||
class StandardStaticViewOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"view_length": ("INT", {"default": 16, "min": 1, "max": LENGTH_MAX}),
|
||||
"view_overlap": ("INT", {"default": 4, "min": 0, "max": OVERLAP_MAX}),
|
||||
},
|
||||
"optional": {
|
||||
"fuse_method": (ContextFuseMethod.LIST,),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIEW_OPTS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/context opts/view opts"
|
||||
FUNCTION = "create_options"
|
||||
|
||||
def create_options(self, view_length: int, view_overlap: int,
|
||||
fuse_method: str=ContextFuseMethod.FLAT,):
|
||||
view_options = ContextOptions(
|
||||
context_length=view_length,
|
||||
context_stride=None,
|
||||
context_overlap=view_overlap,
|
||||
context_schedule=ContextSchedules.STATIC_STANDARD,
|
||||
fuse_method=fuse_method,
|
||||
)
|
||||
return (view_options,)
|
||||
|
||||
|
||||
class StandardUniformViewOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"view_length": ("INT", {"default": 16, "min": 1, "max": LENGTH_MAX}),
|
||||
"view_stride": ("INT", {"default": 1, "min": 1, "max": STRIDE_MAX}),
|
||||
"view_overlap": ("INT", {"default": 4, "min": 0, "max": OVERLAP_MAX}),
|
||||
},
|
||||
"optional": {
|
||||
"fuse_method": (ContextFuseMethod.LIST,),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIEW_OPTS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/context opts/view opts"
|
||||
FUNCTION = "create_options"
|
||||
|
||||
def create_options(self, view_length: int, view_overlap: int, view_stride: int,
|
||||
fuse_method: str=ContextFuseMethod.PYRAMID,):
|
||||
view_options = ContextOptions(
|
||||
context_length=view_length,
|
||||
context_stride=view_stride,
|
||||
context_overlap=view_overlap,
|
||||
context_schedule=ContextSchedules.UNIFORM_STANDARD,
|
||||
fuse_method=fuse_method,
|
||||
)
|
||||
return (view_options,)
|
||||
|
||||
|
||||
class LoopedUniformViewOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"view_length": ("INT", {"default": 16, "min": 1, "max": LENGTH_MAX}),
|
||||
"view_stride": ("INT", {"default": 1, "min": 1, "max": STRIDE_MAX}),
|
||||
"view_overlap": ("INT", {"default": 4, "min": 0, "max": OVERLAP_MAX}),
|
||||
"closed_loop": ("BOOLEAN", {"default": False},),
|
||||
},
|
||||
"optional": {
|
||||
"fuse_method": (ContextFuseMethod.LIST,),
|
||||
"use_on_equal_length": ("BOOLEAN", {"default": False},),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIEW_OPTS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/context opts/view opts"
|
||||
FUNCTION = "create_options"
|
||||
|
||||
def create_options(self, view_length: int, view_overlap: int, view_stride: int, closed_loop: bool,
|
||||
fuse_method: str=ContextFuseMethod.PYRAMID, use_on_equal_length=False):
|
||||
view_options = ContextOptions(
|
||||
context_length=view_length,
|
||||
context_stride=view_stride,
|
||||
context_overlap=view_overlap,
|
||||
context_schedule=ContextSchedules.UNIFORM_LOOPED,
|
||||
closed_loop=closed_loop,
|
||||
fuse_method=fuse_method,
|
||||
use_on_equal_length=use_on_equal_length,
|
||||
)
|
||||
return (view_options,)
|
||||
+277
@@ -0,0 +1,277 @@
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from typing import Dict, List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
import folder_paths
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from .context import ContextOptionsGroup, ContextOptions, ContextSchedules
|
||||
from .logger import logger
|
||||
from .utils_model import Folders, BetaSchedules, get_available_motion_models
|
||||
from .model_injection import ModelPatcherAndInjector, InjectionParams, MotionModelGroup, load_motion_module_gen1
|
||||
|
||||
|
||||
class AnimateDiffLoader_Deprecated:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"latents": ("LATENT",),
|
||||
"model_name": (get_available_motion_models(),),
|
||||
"unlimited_area_hack": ("BOOLEAN", {"default": False},),
|
||||
"beta_schedule": (BetaSchedules.get_alias_list_with_first_element(BetaSchedules.SQRT_LINEAR),),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "LATENT")
|
||||
CATEGORY = ""
|
||||
FUNCTION = "load_mm_and_inject_params"
|
||||
|
||||
def load_mm_and_inject_params(
|
||||
self,
|
||||
model: ModelPatcher,
|
||||
latents: Dict[str, torch.Tensor],
|
||||
model_name: str, unlimited_area_hack: bool, beta_schedule: str,
|
||||
):
|
||||
# load motion module
|
||||
motion_model = load_motion_module_gen1(model_name, model)
|
||||
# get total frames
|
||||
init_frames_len = len(latents["samples"]) # deprecated - no longer used for anything lol
|
||||
# set injection params
|
||||
params = InjectionParams(
|
||||
unlimited_area_hack=unlimited_area_hack,
|
||||
apply_mm_groupnorm_hack=True,
|
||||
model_name=model_name,
|
||||
apply_v2_properly=False,
|
||||
)
|
||||
# inject for use in sampling code
|
||||
model = ModelPatcherAndInjector(model)
|
||||
model.motion_models = MotionModelGroup(motion_model)
|
||||
model.motion_injection_params = params
|
||||
|
||||
# save model sampling from BetaSchedule as object patch
|
||||
# if autoselect, get suggested beta_schedule from motion model
|
||||
if beta_schedule == BetaSchedules.AUTOSELECT and not model.motion_models.is_empty():
|
||||
beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
|
||||
new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
|
||||
if new_model_sampling is not None:
|
||||
model.add_object_patch("model_sampling", new_model_sampling)
|
||||
|
||||
del motion_model
|
||||
return (model, latents)
|
||||
|
||||
|
||||
class AnimateDiffLoaderAdvanced_Deprecated:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"latents": ("LATENT",),
|
||||
"model_name": (get_available_motion_models(),),
|
||||
"unlimited_area_hack": ("BOOLEAN", {"default": False},),
|
||||
"context_length": ("INT", {"default": 16, "min": 0, "max": 1000}),
|
||||
"context_stride": ("INT", {"default": 1, "min": 1, "max": 1000}),
|
||||
"context_overlap": ("INT", {"default": 4, "min": 0, "max": 1000}),
|
||||
"context_schedule": (ContextSchedules.LEGACY_UNIFORM_SCHEDULE_LIST,),
|
||||
"closed_loop": ("BOOLEAN", {"default": False},),
|
||||
"beta_schedule": (BetaSchedules.get_alias_list_with_first_element(BetaSchedules.SQRT_LINEAR),),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "LATENT")
|
||||
CATEGORY = ""
|
||||
FUNCTION = "load_mm_and_inject_params"
|
||||
|
||||
def load_mm_and_inject_params(self,
|
||||
model: ModelPatcher,
|
||||
latents: Dict[str, torch.Tensor],
|
||||
model_name: str, unlimited_area_hack: bool,
|
||||
context_length: int, context_stride: int, context_overlap: int, context_schedule: str, closed_loop: bool,
|
||||
beta_schedule: str,
|
||||
):
|
||||
# load motion module
|
||||
motion_model = load_motion_module_gen1(model_name, model)
|
||||
# get total frames
|
||||
init_frames_len = len(latents["samples"]) # deprecated - no longer used for anything lol
|
||||
# set injection params
|
||||
params = InjectionParams(
|
||||
unlimited_area_hack=unlimited_area_hack,
|
||||
apply_mm_groupnorm_hack=True,
|
||||
model_name=model_name,
|
||||
apply_v2_properly=False,
|
||||
)
|
||||
context_group = ContextOptionsGroup()
|
||||
context_group.add(
|
||||
ContextOptions(
|
||||
context_length=context_length,
|
||||
context_stride=context_stride,
|
||||
context_overlap=context_overlap,
|
||||
context_schedule=context_schedule,
|
||||
closed_loop=closed_loop,
|
||||
)
|
||||
)
|
||||
# set context settings
|
||||
params.set_context(context_options=context_group)
|
||||
# inject for use in sampling code
|
||||
model = ModelPatcherAndInjector(model)
|
||||
model.motion_models = MotionModelGroup(motion_model)
|
||||
model.motion_injection_params = params
|
||||
|
||||
# save model sampling from BetaSchedule as object patch
|
||||
# if autoselect, get suggested beta_schedule from motion model
|
||||
if beta_schedule == BetaSchedules.AUTOSELECT and not model.motion_models.is_empty():
|
||||
beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
|
||||
new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
|
||||
if new_model_sampling is not None:
|
||||
model.add_object_patch("model_sampling", new_model_sampling)
|
||||
|
||||
del motion_model
|
||||
return (model, latents)
|
||||
|
||||
|
||||
class AnimateDiffCombine_Deprecated:
|
||||
ffmpeg_warning_already_shown = False
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
ffmpeg_path = shutil.which("ffmpeg")
|
||||
#Hide ffmpeg formats if ffmpeg isn't available
|
||||
if ffmpeg_path is not None:
|
||||
ffmpeg_formats = ["video/"+x[:-5] for x in folder_paths.get_filename_list(Folders.VIDEO_FORMATS)]
|
||||
else:
|
||||
ffmpeg_formats = []
|
||||
if not s.ffmpeg_warning_already_shown:
|
||||
# Deprecated node are now hidden, so no need to show warning unless node is used.
|
||||
# logger.warning("This warning can be ignored, you should not be using the deprecated AnimateDiff Combine node anyway. If you are, use Video Combine from ComfyUI-VideoHelperSuite instead. ffmpeg could not be found. Outputs that require it have been disabled")
|
||||
s.ffmpeg_warning_already_shown = True
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"frame_rate": (
|
||||
"INT",
|
||||
{"default": 8, "min": 1, "max": 24, "step": 1},
|
||||
),
|
||||
"loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
|
||||
"filename_prefix": ("STRING", {"default": "AnimateDiff"}),
|
||||
"format": (["image/gif", "image/webp"] + ffmpeg_formats,),
|
||||
"pingpong": ("BOOLEAN", {"default": False}),
|
||||
"save_image": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("GIF",)
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = ""
|
||||
FUNCTION = "generate_gif"
|
||||
|
||||
def generate_gif(
|
||||
self,
|
||||
images,
|
||||
frame_rate: int,
|
||||
loop_count: int,
|
||||
filename_prefix="AnimateDiff",
|
||||
format="image/gif",
|
||||
pingpong=False,
|
||||
save_image=True,
|
||||
prompt=None,
|
||||
extra_pnginfo=None,
|
||||
):
|
||||
logger.warning("Do not use AnimateDiff Combine node, it is deprecated. Use Video Combine node from ComfyUI-VideoHelperSuite instead. Video nodes from VideoHelperSuite are actively maintained, more feature-rich, and also automatically attempts to get ffmpeg.")
|
||||
# convert images to numpy
|
||||
frames: List[Image.Image] = []
|
||||
for image in images:
|
||||
img = 255.0 * image.cpu().numpy()
|
||||
img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
|
||||
frames.append(img)
|
||||
|
||||
# get output information
|
||||
output_dir = (
|
||||
folder_paths.get_output_directory()
|
||||
if save_image
|
||||
else folder_paths.get_temp_directory()
|
||||
)
|
||||
(
|
||||
full_output_folder,
|
||||
filename,
|
||||
counter,
|
||||
subfolder,
|
||||
_,
|
||||
) = folder_paths.get_save_image_path(filename_prefix, output_dir)
|
||||
|
||||
metadata = PngInfo()
|
||||
if prompt is not None:
|
||||
metadata.add_text("prompt", json.dumps(prompt))
|
||||
if extra_pnginfo is not None:
|
||||
for x in extra_pnginfo:
|
||||
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
||||
|
||||
# save first frame as png to keep metadata
|
||||
file = f"{filename}_{counter:05}_.png"
|
||||
file_path = os.path.join(full_output_folder, file)
|
||||
frames[0].save(
|
||||
file_path,
|
||||
pnginfo=metadata,
|
||||
compress_level=4,
|
||||
)
|
||||
if pingpong:
|
||||
frames = frames + frames[-2:0:-1]
|
||||
|
||||
format_type, format_ext = format.split("/")
|
||||
file = f"{filename}_{counter:05}_.{format_ext}"
|
||||
file_path = os.path.join(full_output_folder, file)
|
||||
if format_type == "image":
|
||||
# Use pillow directly to save an animated image
|
||||
frames[0].save(
|
||||
file_path,
|
||||
format=format_ext.upper(),
|
||||
save_all=True,
|
||||
append_images=frames[1:],
|
||||
duration=round(1000 / frame_rate),
|
||||
loop=loop_count,
|
||||
compress_level=4,
|
||||
)
|
||||
else:
|
||||
# Use ffmpeg to save a video
|
||||
ffmpeg_path = shutil.which("ffmpeg")
|
||||
if ffmpeg_path is None:
|
||||
#Should never be reachable
|
||||
raise ProcessLookupError("Could not find ffmpeg")
|
||||
|
||||
video_format_path = folder_paths.get_full_path("video_formats", format_ext + ".json")
|
||||
with open(video_format_path, 'r') as stream:
|
||||
video_format = json.load(stream)
|
||||
file = f"{filename}_{counter:05}_.{video_format['extension']}"
|
||||
file_path = os.path.join(full_output_folder, file)
|
||||
dimensions = f"{frames[0].width}x{frames[0].height}"
|
||||
args = [ffmpeg_path, "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24",
|
||||
"-s", dimensions, "-r", str(frame_rate), "-i", "-"] \
|
||||
+ video_format['main_pass'] + [file_path]
|
||||
|
||||
env=os.environ.copy()
|
||||
if "environment" in video_format:
|
||||
env.update(video_format["environment"])
|
||||
with subprocess.Popen(args, stdin=subprocess.PIPE, env=env) as proc:
|
||||
for frame in frames:
|
||||
proc.stdin.write(frame.tobytes())
|
||||
|
||||
previews = [
|
||||
{
|
||||
"filename": file,
|
||||
"subfolder": subfolder,
|
||||
"type": "output" if save_image else "temp",
|
||||
"format": format,
|
||||
}
|
||||
]
|
||||
return {"ui": {"gifs": previews}}
|
||||
+78
@@ -0,0 +1,78 @@
|
||||
import torch
|
||||
|
||||
import folder_paths
|
||||
import nodes as comfy_nodes
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
from comfy.sd import load_checkpoint_guess_config
|
||||
|
||||
from .logger import logger
|
||||
from .utils_model import BetaSchedules
|
||||
from .model_injection import get_vanilla_model_patcher
|
||||
|
||||
|
||||
class AnimateDiffUnload:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"model": ("MODEL",)}}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/extras"
|
||||
FUNCTION = "unload_motion_modules"
|
||||
|
||||
def unload_motion_modules(self, model: ModelPatcher):
|
||||
# return model clone with ejected params
|
||||
#model = eject_params_from_model(model)
|
||||
model = get_vanilla_model_patcher(model)
|
||||
return (model.clone(),)
|
||||
|
||||
|
||||
class CheckpointLoaderSimpleWithNoiseSelect:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
|
||||
"beta_schedule": (BetaSchedules.ALIAS_LIST, {"default": BetaSchedules.USE_EXISTING}, )
|
||||
},
|
||||
"optional": {
|
||||
"use_custom_scale_factor": ("BOOLEAN", {"default": False}),
|
||||
"scale_factor": ("FLOAT", {"default": 0.18215, "min": 0.0, "max": 1.0, "step": 0.00001})
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
FUNCTION = "load_checkpoint"
|
||||
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/extras"
|
||||
|
||||
def load_checkpoint(self, ckpt_name, beta_schedule, output_vae=True, output_clip=True, use_custom_scale_factor=False, scale_factor=0.18215):
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
out = load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
||||
# register chosen beta schedule on model - convert to beta_schedule name recognized by ComfyUI
|
||||
new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, out[0])
|
||||
if new_model_sampling is not None:
|
||||
out[0].model.model_sampling = new_model_sampling
|
||||
if use_custom_scale_factor:
|
||||
out[0].model.latent_format.scale_factor = scale_factor
|
||||
return out
|
||||
|
||||
|
||||
class EmptyLatentImageLarge:
|
||||
def __init__(self, device="cpu"):
|
||||
self.device = device
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "width": ("INT", {"default": 512, "min": 64, "max": comfy_nodes.MAX_RESOLUTION, "step": 8}),
|
||||
"height": ("INT", {"default": 512, "min": 64, "max": comfy_nodes.MAX_RESOLUTION, "step": 8}),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 262144})}}
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "generate"
|
||||
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/extras"
|
||||
|
||||
def generate(self, width, height, batch_size=1):
|
||||
latent = torch.zeros([batch_size, 4, height // 8, width // 8])
|
||||
return ({"samples":latent}, )
|
||||
+340
@@ -0,0 +1,340 @@
|
||||
from pathlib import Path
|
||||
import torch
|
||||
|
||||
import comfy.sample as comfy_sample
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from .ad_settings import AdjustPEGroup, AnimateDiffSettings, AdjustPE
|
||||
from .context import ContextOptions, ContextOptionsGroup, ContextSchedules
|
||||
from .logger import logger
|
||||
from .utils_model import BetaSchedules, get_available_motion_loras, get_available_motion_models, get_motion_lora_path
|
||||
from .utils_motion import ADKeyframeGroup, get_combined_multival
|
||||
from .motion_lora import MotionLoraInfo, MotionLoraList
|
||||
from .model_injection import InjectionParams, ModelPatcherAndInjector, MotionModelGroup, load_motion_lora_as_patches, load_motion_module_gen1, load_motion_module_gen2, validate_model_compatibility_gen2
|
||||
from .sample_settings import SampleSettings, SeedNoiseGeneration
|
||||
from .sampling import motion_sample_factory
|
||||
|
||||
|
||||
class AnimateDiffLoaderGen1:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"model_name": (get_available_motion_models(),),
|
||||
"beta_schedule": (BetaSchedules.ALIAS_LIST, {"default": BetaSchedules.AUTOSELECT}),
|
||||
#"apply_mm_groupnorm_hack": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
"optional": {
|
||||
"context_options": ("CONTEXT_OPTIONS",),
|
||||
"motion_lora": ("MOTION_LORA",),
|
||||
"ad_settings": ("AD_SETTINGS",),
|
||||
"ad_keyframes": ("AD_KEYFRAMES",),
|
||||
"sample_settings": ("SAMPLE_SETTINGS",),
|
||||
"scale_multival": ("MULTIVAL",),
|
||||
"effect_multival": ("MULTIVAL",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/① Gen1 nodes ①"
|
||||
FUNCTION = "load_mm_and_inject_params"
|
||||
|
||||
def load_mm_and_inject_params(self,
|
||||
model: ModelPatcher,
|
||||
model_name: str, beta_schedule: str,# apply_mm_groupnorm_hack: bool,
|
||||
context_options: ContextOptionsGroup=None, motion_lora: MotionLoraList=None, ad_settings: AnimateDiffSettings=None,
|
||||
sample_settings: SampleSettings=None, scale_multival=None, effect_multival=None, ad_keyframes: ADKeyframeGroup=None,
|
||||
):
|
||||
# load motion module and motion settings, if included
|
||||
motion_model = load_motion_module_gen2(model_name=model_name, motion_model_settings=ad_settings)
|
||||
# confirm that it is compatible with SD model
|
||||
validate_model_compatibility_gen2(model=model, motion_model=motion_model)
|
||||
# apply motion model to loaded_mm
|
||||
if motion_lora is not None:
|
||||
for lora in motion_lora.loras:
|
||||
load_motion_lora_as_patches(motion_model, lora)
|
||||
motion_model.scale_multival = scale_multival
|
||||
motion_model.effect_multival = effect_multival
|
||||
motion_model.keyframes = ad_keyframes.clone() if ad_keyframes else ADKeyframeGroup()
|
||||
|
||||
# create injection params
|
||||
params = InjectionParams(unlimited_area_hack=False, model_name=motion_model.model.mm_info.mm_name)
|
||||
# apply context options
|
||||
if context_options:
|
||||
params.set_context(context_options)
|
||||
|
||||
# set motion_scale and motion_model_settings
|
||||
if not ad_settings:
|
||||
ad_settings = AnimateDiffSettings()
|
||||
ad_settings.attn_scale = 1.0
|
||||
params.set_motion_model_settings(ad_settings)
|
||||
|
||||
# backwards compatibility to support old way of masking scale
|
||||
if params.motion_model_settings.mask_attn_scale is not None:
|
||||
motion_model.scale_multival = get_combined_multival(scale_multival, (params.motion_model_settings.mask_attn_scale * params.motion_model_settings.attn_scale))
|
||||
|
||||
# need to use a ModelPatcher that supports injection of motion modules into unet
|
||||
# need to use a ModelPatcher that supports injection of motion modules into unet
|
||||
model = ModelPatcherAndInjector(model)
|
||||
model.motion_models = MotionModelGroup(motion_model)
|
||||
model.sample_settings = sample_settings if sample_settings is not None else SampleSettings()
|
||||
model.motion_injection_params = params
|
||||
|
||||
if model.sample_settings.custom_cfg is not None:
|
||||
logger.info("[Sample Settings] custom_cfg is set; will override any KSampler cfg values or patches.")
|
||||
|
||||
if model.sample_settings.sigma_schedule is not None:
|
||||
logger.info("[Sample Settings] sigma_schedule is set; will override beta_schedule.")
|
||||
model.add_object_patch("model_sampling", model.sample_settings.sigma_schedule.clone().model_sampling)
|
||||
else:
|
||||
# save model sampling from BetaSchedule as object patch
|
||||
# if autoselect, get suggested beta_schedule from motion model
|
||||
if beta_schedule == BetaSchedules.AUTOSELECT and not model.motion_models.is_empty():
|
||||
beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
|
||||
new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
|
||||
if new_model_sampling is not None:
|
||||
model.add_object_patch("model_sampling", new_model_sampling)
|
||||
|
||||
del motion_model
|
||||
return (model,)
|
||||
|
||||
|
||||
class LegacyAnimateDiffLoaderWithContext:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"model_name": (get_available_motion_models(),),
|
||||
"beta_schedule": (BetaSchedules.ALIAS_LIST, {"default": BetaSchedules.AUTOSELECT}),
|
||||
#"apply_mm_groupnorm_hack": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
"optional": {
|
||||
"context_options": ("CONTEXT_OPTIONS",),
|
||||
"motion_lora": ("MOTION_LORA",),
|
||||
"ad_settings": ("AD_SETTINGS",),
|
||||
"sample_settings": ("SAMPLE_SETTINGS",),
|
||||
"motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
"apply_v2_models_properly": ("BOOLEAN", {"default": True}),
|
||||
"ad_keyframes": ("AD_KEYFRAMES",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/① Gen1 nodes ①"
|
||||
FUNCTION = "load_mm_and_inject_params"
|
||||
|
||||
|
||||
def load_mm_and_inject_params(self,
|
||||
model: ModelPatcher,
|
||||
model_name: str, beta_schedule: str,# apply_mm_groupnorm_hack: bool,
|
||||
context_options: ContextOptionsGroup=None, motion_lora: MotionLoraList=None, ad_settings: AnimateDiffSettings=None, motion_model_settings: AnimateDiffSettings=None,
|
||||
sample_settings: SampleSettings=None, motion_scale: float=1.0, apply_v2_models_properly: bool=False, ad_keyframes: ADKeyframeGroup=None,
|
||||
):
|
||||
if ad_settings is not None:
|
||||
motion_model_settings = ad_settings
|
||||
# load motion module
|
||||
motion_model = load_motion_module_gen1(model_name, model, motion_lora=motion_lora, motion_model_settings=motion_model_settings)
|
||||
# set injection params
|
||||
params = InjectionParams(
|
||||
unlimited_area_hack=False,
|
||||
model_name=model_name,
|
||||
apply_v2_properly=apply_v2_models_properly,
|
||||
)
|
||||
if context_options:
|
||||
params.set_context(context_options)
|
||||
# set motion_scale and motion_model_settings
|
||||
if not motion_model_settings:
|
||||
motion_model_settings = AnimateDiffSettings()
|
||||
motion_model_settings.attn_scale = motion_scale
|
||||
params.set_motion_model_settings(motion_model_settings)
|
||||
|
||||
if params.motion_model_settings.mask_attn_scale is not None:
|
||||
motion_model.scale_multival = params.motion_model_settings.mask_attn_scale * params.motion_model_settings.attn_scale
|
||||
else:
|
||||
motion_model.scale_multival = params.motion_model_settings.attn_scale
|
||||
|
||||
motion_model.keyframes = ad_keyframes.clone() if ad_keyframes else ADKeyframeGroup()
|
||||
|
||||
model = ModelPatcherAndInjector(model)
|
||||
model.motion_models = MotionModelGroup(motion_model)
|
||||
model.sample_settings = sample_settings if sample_settings is not None else SampleSettings()
|
||||
model.motion_injection_params = params
|
||||
|
||||
# save model sampling from BetaSchedule as object patch
|
||||
# if autoselect, get suggested beta_schedule from motion model
|
||||
if beta_schedule == BetaSchedules.AUTOSELECT and not model.motion_models.is_empty():
|
||||
beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
|
||||
new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
|
||||
if new_model_sampling is not None:
|
||||
model.add_object_patch("model_sampling", new_model_sampling)
|
||||
|
||||
del motion_model
|
||||
return (model,)
|
||||
|
||||
|
||||
class AnimateDiffModelSettings:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"min_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
"max_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
},
|
||||
"optional": {
|
||||
"mask_motion_scale": ("MASK",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("AD_SETTINGS",)
|
||||
CATEGORY = "" #"Animate Diff 🎭🅐🅓/① Gen1 nodes ①/motion settings"
|
||||
FUNCTION = "get_motion_model_settings"
|
||||
|
||||
def get_motion_model_settings(self, mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
|
||||
motion_model_settings = AnimateDiffSettings(
|
||||
mask_attn_scale=mask_motion_scale,
|
||||
mask_attn_scale_min=min_motion_scale,
|
||||
mask_attn_scale_max=max_motion_scale,
|
||||
)
|
||||
|
||||
return (motion_model_settings,)
|
||||
|
||||
|
||||
class AnimateDiffModelSettingsSimple:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"motion_pe_stretch": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"mask_motion_scale": ("MASK",),
|
||||
"min_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
"max_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("AD_SETTINGS",)
|
||||
CATEGORY = "" #"Animate Diff 🎭🅐🅓/① Gen1 nodes ①/motion settings/experimental"
|
||||
FUNCTION = "get_motion_model_settings"
|
||||
|
||||
def get_motion_model_settings(self, motion_pe_stretch: int,
|
||||
mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
|
||||
adjust_pe = AdjustPEGroup(AdjustPE(motion_pe_stretch=motion_pe_stretch))
|
||||
motion_model_settings = AnimateDiffSettings(
|
||||
adjust_pe=adjust_pe,
|
||||
mask_attn_scale=mask_motion_scale,
|
||||
mask_attn_scale_min=min_motion_scale,
|
||||
mask_attn_scale_max=max_motion_scale,
|
||||
)
|
||||
|
||||
return (motion_model_settings,)
|
||||
|
||||
|
||||
class AnimateDiffModelSettingsAdvanced:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"pe_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"attn_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"other_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"motion_pe_stretch": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"cap_initial_pe_length": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"interpolate_pe_to_length": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"initial_pe_idx_offset": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"final_pe_idx_offset": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"mask_motion_scale": ("MASK",),
|
||||
"min_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
"max_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("AD_SETTINGS",)
|
||||
CATEGORY = "" #"Animate Diff 🎭🅐🅓/① Gen1 nodes ①/motion settings/experimental"
|
||||
FUNCTION = "get_motion_model_settings"
|
||||
|
||||
def get_motion_model_settings(self, pe_strength: float, attn_strength: float, other_strength: float,
|
||||
motion_pe_stretch: int,
|
||||
cap_initial_pe_length: int, interpolate_pe_to_length: int,
|
||||
initial_pe_idx_offset: int, final_pe_idx_offset: int,
|
||||
mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
|
||||
adjust_pe = AdjustPEGroup(AdjustPE(motion_pe_stretch=motion_pe_stretch,
|
||||
cap_initial_pe_length=cap_initial_pe_length, interpolate_pe_to_length=interpolate_pe_to_length,
|
||||
initial_pe_idx_offset=initial_pe_idx_offset, final_pe_idx_offset=final_pe_idx_offset))
|
||||
motion_model_settings = AnimateDiffSettings(
|
||||
adjust_pe=adjust_pe,
|
||||
pe_strength=pe_strength,
|
||||
attn_strength=attn_strength,
|
||||
other_strength=other_strength,
|
||||
mask_attn_scale=mask_motion_scale,
|
||||
mask_attn_scale_min=min_motion_scale,
|
||||
mask_attn_scale_max=max_motion_scale,
|
||||
)
|
||||
|
||||
return (motion_model_settings,)
|
||||
|
||||
|
||||
class AnimateDiffModelSettingsAdvancedAttnStrengths:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"pe_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"attn_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"attn_q_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"attn_k_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"attn_v_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"attn_out_weight_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"attn_out_bias_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"other_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.0001}),
|
||||
"motion_pe_stretch": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"cap_initial_pe_length": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"interpolate_pe_to_length": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"initial_pe_idx_offset": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
"final_pe_idx_offset": ("INT", {"default": 0, "min": 0, "step": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"mask_motion_scale": ("MASK",),
|
||||
"min_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
"max_motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("AD_SETTINGS",)
|
||||
CATEGORY = "" #"Animate Diff 🎭🅐🅓/① Gen1 nodes ①/motion settings/experimental"
|
||||
FUNCTION = "get_motion_model_settings"
|
||||
|
||||
def get_motion_model_settings(self, pe_strength: float, attn_strength: float,
|
||||
attn_q_strength: float,
|
||||
attn_k_strength: float,
|
||||
attn_v_strength: float,
|
||||
attn_out_weight_strength: float,
|
||||
attn_out_bias_strength: float,
|
||||
other_strength: float,
|
||||
motion_pe_stretch: int,
|
||||
cap_initial_pe_length: int, interpolate_pe_to_length: int,
|
||||
initial_pe_idx_offset: int, final_pe_idx_offset: int,
|
||||
mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
|
||||
adjust_pe = AdjustPEGroup(AdjustPE(motion_pe_stretch=motion_pe_stretch,
|
||||
cap_initial_pe_length=cap_initial_pe_length, interpolate_pe_to_length=interpolate_pe_to_length,
|
||||
initial_pe_idx_offset=initial_pe_idx_offset, final_pe_idx_offset=final_pe_idx_offset))
|
||||
motion_model_settings = AnimateDiffSettings(
|
||||
adjust_pe=adjust_pe,
|
||||
pe_strength=pe_strength,
|
||||
attn_strength=attn_strength,
|
||||
attn_q_strength=attn_q_strength,
|
||||
attn_k_strength=attn_k_strength,
|
||||
attn_v_strength=attn_v_strength,
|
||||
attn_out_weight_strength=attn_out_weight_strength,
|
||||
attn_out_bias_strength=attn_out_bias_strength,
|
||||
other_strength=other_strength,
|
||||
mask_attn_scale=mask_motion_scale,
|
||||
mask_attn_scale_min=min_motion_scale,
|
||||
mask_attn_scale_max=max_motion_scale,
|
||||
)
|
||||
|
||||
return (motion_model_settings,)
|
||||
+212
@@ -0,0 +1,212 @@
|
||||
from pathlib import Path
|
||||
import torch
|
||||
|
||||
import comfy.sample as comfy_sample
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from .ad_settings import AnimateDiffSettings
|
||||
from .context import ContextOptions, ContextOptionsGroup, ContextSchedules
|
||||
from .logger import logger
|
||||
from .utils_model import BIGMAX, BetaSchedules, get_available_motion_loras, get_available_motion_models, get_motion_lora_path
|
||||
from .utils_motion import ADKeyframeGroup, ADKeyframe
|
||||
from .motion_lora import MotionLoraInfo, MotionLoraList
|
||||
from .model_injection import (InjectionParams, ModelPatcherAndInjector, MotionModelGroup, MotionModelPatcher, create_fresh_motion_module,
|
||||
load_motion_module_gen1, load_motion_module_gen2, load_motion_lora_as_patches, validate_model_compatibility_gen2)
|
||||
from .sample_settings import SampleSettings, SeedNoiseGeneration
|
||||
from .sampling import motion_sample_factory
|
||||
|
||||
|
||||
class UseEvolvedSamplingNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"beta_schedule": (BetaSchedules.ALIAS_LIST, {"default": BetaSchedules.AUTOSELECT}),
|
||||
},
|
||||
"optional": {
|
||||
"m_models": ("M_MODELS",),
|
||||
"context_options": ("CONTEXT_OPTIONS",),
|
||||
"sample_settings": ("SAMPLE_SETTINGS",),
|
||||
#"beta_schedule_override": ("BETA_SCHEDULE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/② Gen2 nodes ②"
|
||||
FUNCTION = "use_evolved_sampling"
|
||||
|
||||
def use_evolved_sampling(self, model: ModelPatcher, beta_schedule: str, m_models: MotionModelGroup=None, context_options: ContextOptionsGroup=None,
|
||||
sample_settings: SampleSettings=None, beta_schedule_override=None):
|
||||
if m_models is not None:
|
||||
m_models = m_models.clone()
|
||||
# for each motion model, confirm that it is compatible with SD model
|
||||
for motion_model in m_models.models:
|
||||
validate_model_compatibility_gen2(model=model, motion_model=motion_model)
|
||||
# create injection params
|
||||
model_name_list = [motion_model.model.mm_info.mm_name for motion_model in m_models.models]
|
||||
model_names = ",".join(model_name_list)
|
||||
# TODO: check if any apply_v2_properly is set to False
|
||||
params = InjectionParams(unlimited_area_hack=False, model_name=model_names)
|
||||
else:
|
||||
params = InjectionParams()
|
||||
# apply context options
|
||||
if context_options:
|
||||
params.set_context(context_options)
|
||||
# need to use a ModelPatcher that supports injection of motion modules into unet
|
||||
model = ModelPatcherAndInjector(model)
|
||||
model.motion_models = m_models
|
||||
model.sample_settings = sample_settings if sample_settings is not None else SampleSettings()
|
||||
model.motion_injection_params = params
|
||||
|
||||
if model.sample_settings.custom_cfg is not None:
|
||||
logger.info("[Sample Settings] custom_cfg is set; will override any KSampler cfg values or patches.")
|
||||
|
||||
if model.sample_settings.sigma_schedule is not None:
|
||||
logger.info("[Sample Settings] sigma_schedule is set; will override beta_schedule.")
|
||||
model.add_object_patch("model_sampling", model.sample_settings.sigma_schedule.clone().model_sampling)
|
||||
else:
|
||||
# save model_sampling from BetaSchedule as object patch
|
||||
# if autoselect, get suggested beta_schedule from motion model
|
||||
if beta_schedule == BetaSchedules.AUTOSELECT and not model.motion_models.is_empty():
|
||||
beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
|
||||
new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
|
||||
if new_model_sampling is not None:
|
||||
model.add_object_patch("model_sampling", new_model_sampling)
|
||||
|
||||
del m_models
|
||||
return (model,)
|
||||
|
||||
|
||||
class ApplyAnimateDiffModelNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"motion_model": ("MOTION_MODEL_ADE",),
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
},
|
||||
"optional": {
|
||||
"motion_lora": ("MOTION_LORA",),
|
||||
"scale_multival": ("MULTIVAL",),
|
||||
"effect_multival": ("MULTIVAL",),
|
||||
"ad_keyframes": ("AD_KEYFRAMES",),
|
||||
"prev_m_models": ("M_MODELS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("M_MODELS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/② Gen2 nodes ②"
|
||||
FUNCTION = "apply_motion_model"
|
||||
|
||||
def apply_motion_model(self, motion_model: MotionModelPatcher, start_percent: float=0.0, end_percent: float=1.0,
|
||||
motion_lora: MotionLoraList=None, ad_keyframes: ADKeyframeGroup=None,
|
||||
scale_multival=None, effect_multival=None,
|
||||
prev_m_models: MotionModelGroup=None,):
|
||||
# set up motion models list
|
||||
if prev_m_models is None:
|
||||
prev_m_models = MotionModelGroup()
|
||||
prev_m_models = prev_m_models.clone()
|
||||
motion_model = motion_model.clone()
|
||||
# check if internal motion model already present in previous model - create new if so
|
||||
for prev_model in prev_m_models.models:
|
||||
if motion_model.model is prev_model.model:
|
||||
# need to create new internal model based on same state_dict
|
||||
motion_model = create_fresh_motion_module(motion_model)
|
||||
# apply motion model to loaded_mm
|
||||
if motion_lora is not None:
|
||||
for lora in motion_lora.loras:
|
||||
load_motion_lora_as_patches(motion_model, lora)
|
||||
motion_model.scale_multival = scale_multival
|
||||
motion_model.effect_multival = effect_multival
|
||||
motion_model.keyframes = ad_keyframes.clone() if ad_keyframes else ADKeyframeGroup()
|
||||
motion_model.timestep_percent_range = (start_percent, end_percent)
|
||||
# add to beginning, so that after injection, it will be the earliest of prev_m_models to be run
|
||||
prev_m_models.add_to_start(mm=motion_model)
|
||||
return (prev_m_models,)
|
||||
|
||||
|
||||
class ApplyAnimateDiffModelBasicNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"motion_model": ("MOTION_MODEL_ADE",),
|
||||
},
|
||||
"optional": {
|
||||
"motion_lora": ("MOTION_LORA",),
|
||||
"scale_multival": ("MULTIVAL",),
|
||||
"effect_multival": ("MULTIVAL",),
|
||||
"ad_keyframes": ("AD_KEYFRAMES",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("M_MODELS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/② Gen2 nodes ②"
|
||||
FUNCTION = "apply_motion_model"
|
||||
|
||||
def apply_motion_model(self,
|
||||
motion_model: MotionModelPatcher, motion_lora: MotionLoraList=None,
|
||||
scale_multival=None, effect_multival=None, ad_keyframes=None):
|
||||
# just a subset of normal ApplyAnimateDiffModelNode inputs
|
||||
return ApplyAnimateDiffModelNode.apply_motion_model(self, motion_model, motion_lora=motion_lora,
|
||||
scale_multival=scale_multival, effect_multival=effect_multival,
|
||||
ad_keyframes=ad_keyframes)
|
||||
|
||||
|
||||
class LoadAnimateDiffModelNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model_name": (get_available_motion_models(),),
|
||||
},
|
||||
"optional": {
|
||||
"ad_settings": ("AD_SETTINGS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MOTION_MODEL_ADE",)
|
||||
RETURN_NAMES = ("MOTION_MODEL",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/② Gen2 nodes ②"
|
||||
FUNCTION = "load_motion_model"
|
||||
|
||||
def load_motion_model(self, model_name: str, ad_settings: AnimateDiffSettings=None):
|
||||
# load motion module and motion settings, if included
|
||||
motion_model = load_motion_module_gen2(model_name=model_name, motion_model_settings=ad_settings)
|
||||
return (motion_model,)
|
||||
|
||||
|
||||
class ADKeyframeNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}, ),
|
||||
},
|
||||
"optional": {
|
||||
"prev_ad_keyframes": ("AD_KEYFRAMES", ),
|
||||
"scale_multival": ("MULTIVAL",),
|
||||
"effect_multival": ("MULTIVAL",),
|
||||
"inherit_missing": ("BOOLEAN", {"default": True}, ),
|
||||
"guarantee_steps": ("INT", {"default": 1, "min": 0, "max": BIGMAX}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("AD_KEYFRAMES", )
|
||||
FUNCTION = "load_keyframe"
|
||||
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓"
|
||||
|
||||
def load_keyframe(self,
|
||||
start_percent: float, prev_ad_keyframes=None,
|
||||
scale_multival: [float, torch.Tensor]=None, effect_multival: [float, torch.Tensor]=None,
|
||||
inherit_missing: bool=True, guarantee_steps: int=1):
|
||||
if not prev_ad_keyframes:
|
||||
prev_ad_keyframes = ADKeyframeGroup()
|
||||
prev_ad_keyframes = prev_ad_keyframes.clone()
|
||||
keyframe = ADKeyframe(start_percent=start_percent, scale_multival=scale_multival, effect_multival=effect_multival,
|
||||
inherit_missing=inherit_missing, guarantee_steps=guarantee_steps)
|
||||
prev_ad_keyframes.add(keyframe)
|
||||
return (prev_ad_keyframes,)
|
||||
+90
@@ -0,0 +1,90 @@
|
||||
from pathlib import Path
|
||||
|
||||
import folder_paths
|
||||
import comfy.utils
|
||||
import comfy.sd
|
||||
|
||||
from .logger import logger
|
||||
from .utils_model import get_available_motion_loras, get_motion_lora_path
|
||||
from .motion_lora import MotionLoraInfo, MotionLoraList
|
||||
|
||||
|
||||
class AnimateDiffLoraLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"lora_name": (get_available_motion_loras(),),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}),
|
||||
},
|
||||
"optional": {
|
||||
"prev_motion_lora": ("MOTION_LORA",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MOTION_LORA",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓"
|
||||
FUNCTION = "load_motion_lora"
|
||||
|
||||
def load_motion_lora(self, lora_name: str, strength: float, prev_motion_lora: MotionLoraList=None):
|
||||
if prev_motion_lora is None:
|
||||
prev_motion_lora = MotionLoraList()
|
||||
else:
|
||||
prev_motion_lora = prev_motion_lora.clone()
|
||||
# check if motion lora with name exists
|
||||
lora_path = get_motion_lora_path(lora_name)
|
||||
if not Path(lora_path).is_file():
|
||||
raise FileNotFoundError(f"Motion lora with name '{lora_name}' not found.")
|
||||
# create motion lora info to be loaded in AnimateDiff Loader
|
||||
lora_info = MotionLoraInfo(name=lora_name, strength=strength)
|
||||
prev_motion_lora.add_lora(lora_info)
|
||||
|
||||
return (prev_motion_lora,)
|
||||
|
||||
|
||||
class MaskedLoraLoader:
|
||||
def __init__(self):
|
||||
self.loaded_lora = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "model": ("MODEL",),
|
||||
"clip": ("CLIP", ),
|
||||
"lora_name": (folder_paths.get_filename_list("loras"), ),
|
||||
"strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
|
||||
"strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
|
||||
}}
|
||||
#RETURN_TYPES = ()
|
||||
RETURN_TYPES = ("MODEL", "CLIP")
|
||||
FUNCTION = "load_lora"
|
||||
|
||||
CATEGORY = "loaders"
|
||||
|
||||
def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
|
||||
if strength_model == 0 and strength_clip == 0:
|
||||
return (model, clip)
|
||||
|
||||
lora_path = folder_paths.get_full_path("loras", lora_name)
|
||||
lora = None
|
||||
if self.loaded_lora is not None:
|
||||
if self.loaded_lora[0] == lora_path:
|
||||
lora = self.loaded_lora[1]
|
||||
else:
|
||||
temp = self.loaded_lora
|
||||
self.loaded_lora = None
|
||||
del temp
|
||||
|
||||
if lora is None:
|
||||
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||||
self.loaded_lora = (lora_path, lora)
|
||||
|
||||
from pathlib import Path
|
||||
with open(Path(__file__).parent.parent.parent / "sd_lora_keys.txt", "w") as lfile:
|
||||
for key in lora:
|
||||
lfile.write(f"{key}:\t{lora[key].size()}\n")
|
||||
|
||||
#model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip)
|
||||
#return (model_lora, clip_lora)
|
||||
return (model, clip)
|
||||
|
||||
|
||||
+136
@@ -0,0 +1,136 @@
|
||||
from collections.abc import Iterable
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from .utils_motion import linear_conversion, normalize_min_max, extend_to_batch_size
|
||||
|
||||
|
||||
class ScaleType:
|
||||
ABSOLUTE = "absolute"
|
||||
RELATIVE = "relative"
|
||||
LIST = [ABSOLUTE, RELATIVE]
|
||||
|
||||
|
||||
class MultivalDynamicNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"float_val": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001},),
|
||||
},
|
||||
"optional": {
|
||||
"mask_optional": ("MASK",)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MULTIVAL",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/multival"
|
||||
FUNCTION = "create_multival"
|
||||
|
||||
def create_multival(self, float_val: Union[float, list[float]]=1.0, mask_optional: Tensor=None):
|
||||
# first, normalize inputs
|
||||
# if float_val is iterable, treat as a list and assume inputs are floats
|
||||
float_is_iterable = False
|
||||
if isinstance(float_val, Iterable):
|
||||
float_is_iterable = True
|
||||
float_val = list(float_val)
|
||||
# if mask present, make sure float_val list can be applied to list - match lengths
|
||||
if mask_optional is not None:
|
||||
if len(float_val) < mask_optional.shape[0]:
|
||||
# copies last entry enough times to match mask shape
|
||||
float_val = float_val + float_val[-1]*(mask_optional.shape[0]-len(float_val))
|
||||
if mask_optional.shape[0] < len(float_val):
|
||||
mask_optional = extend_to_batch_size(mask_optional, len(float_val))
|
||||
float_val = float_val[:mask_optional.shape[0]]
|
||||
float_val: Tensor = torch.tensor(float_val).unsqueeze(-1).unsqueeze(-1)
|
||||
# now that inputs are normalized, figure out what value to actually return
|
||||
if mask_optional is not None:
|
||||
mask_optional = mask_optional.clone()
|
||||
if float_is_iterable:
|
||||
mask_optional = mask_optional[:] * float_val.to(mask_optional.dtype).to(mask_optional.device)
|
||||
else:
|
||||
mask_optional = mask_optional * float_val
|
||||
return (mask_optional,)
|
||||
else:
|
||||
if not float_is_iterable:
|
||||
return (float_val,)
|
||||
# create a dummy mask of b,h,w=float_len,1,1 (sigle pixel)
|
||||
# purpose is for float input to work with mask code, without special cases
|
||||
float_len = float_val.shape[0] if float_is_iterable else 1
|
||||
shape = (float_len,1,1)
|
||||
mask_optional = torch.ones(shape)
|
||||
mask_optional = mask_optional[:] * float_val.to(mask_optional.dtype).to(mask_optional.device)
|
||||
return (mask_optional,)
|
||||
|
||||
|
||||
class MultivalScaledMaskNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"min_float_val": ("FLOAT", {"default": 0.0, "min": 0.0, "step": 0.001}),
|
||||
"max_float_val": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
"mask": ("MASK",),
|
||||
},
|
||||
"optional": {
|
||||
"scaling": (ScaleType.LIST,),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MULTIVAL",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/multival"
|
||||
FUNCTION = "create_multival"
|
||||
|
||||
def create_multival(self, min_float_val: float, max_float_val: float, mask: Tensor, scaling: str=ScaleType.ABSOLUTE):
|
||||
# TODO: allow min_float_val and max_float_val to be list[float]
|
||||
if isinstance(min_float_val, Iterable):
|
||||
raise ValueError(f"min_float_val must be type float (no lists allowed here), not {type(min_float_val).__name__}.")
|
||||
if isinstance(max_float_val, Iterable):
|
||||
raise ValueError(f"max_float_val must be type float (no lists allowed here), not {type(max_float_val).__name__}.")
|
||||
|
||||
if scaling == ScaleType.ABSOLUTE:
|
||||
mask = linear_conversion(mask.clone(), new_min=min_float_val, new_max=max_float_val)
|
||||
elif scaling == ScaleType.RELATIVE:
|
||||
mask = normalize_min_max(mask.clone(), new_min=min_float_val, new_max=max_float_val)
|
||||
else:
|
||||
raise ValueError(f"scaling '{scaling}' not recognized.")
|
||||
return MultivalDynamicNode.create_multival(self, mask_optional=mask)
|
||||
|
||||
|
||||
class MultivalDynamicFloatInputNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"float_val": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001, "forceInput": True},),
|
||||
},
|
||||
"optional": {
|
||||
"mask_optional": ("MASK",)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MULTIVAL",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/multival"
|
||||
FUNCTION = "create_multival"
|
||||
|
||||
def create_multival(self, float_val: Union[float, list[float]]=None, mask_optional: Tensor=None):
|
||||
return MultivalDynamicNode.create_multival(self, float_val=float_val, mask_optional=mask_optional)
|
||||
|
||||
|
||||
class MultivalFloatNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"float_val": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001},),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MULTIVAL",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/multival"
|
||||
FUNCTION = "create_multival"
|
||||
|
||||
def create_multival(self, float_val: Union[float, list[float]]=None):
|
||||
return MultivalDynamicNode.create_multival(self, float_val=float_val)
|
||||
+255
@@ -0,0 +1,255 @@
|
||||
from typing import Union
|
||||
from torch import Tensor
|
||||
|
||||
from .freeinit import FreeInitFilter
|
||||
from .sample_settings import (FreeInitOptions, IterationOptions,
|
||||
NoiseLayerAdd, NoiseLayerAddWeighted, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType,
|
||||
SeedNoiseGeneration, SampleSettings, CustomCFGKeyframeGroup, CustomCFGKeyframe)
|
||||
from .utils_model import BIGMIN, BIGMAX, SigmaSchedule
|
||||
|
||||
|
||||
class SampleSettingsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"batch_offset": ("INT", {"default": 0, "min": 0, "max": BIGMAX}),
|
||||
"noise_type": (NoiseLayerType.LIST,),
|
||||
"seed_gen": (SeedNoiseGeneration.LIST,),
|
||||
"seed_offset": ("INT", {"default": 0, "min": BIGMIN, "max": BIGMAX}),
|
||||
},
|
||||
"optional": {
|
||||
"noise_layers": ("NOISE_LAYERS",),
|
||||
"iteration_opts": ("ITERATION_OPTS",),
|
||||
"seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
|
||||
"adapt_denoise_steps": ("BOOLEAN", {"default": False},),
|
||||
"custom_cfg": ("CUSTOM_CFG",),
|
||||
"sigma_schedule": ("SIGMA_SCHEDULE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLE_SETTINGS",)
|
||||
RETURN_NAMES = ("settings",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓"
|
||||
FUNCTION = "create_settings"
|
||||
|
||||
def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, seed_offset: int, noise_layers: NoiseLayerGroup=None,
|
||||
iteration_opts: IterationOptions=None, seed_override: int=None, adapt_denoise_steps=False,
|
||||
custom_cfg: CustomCFGKeyframeGroup=None, sigma_schedule: SigmaSchedule=None):
|
||||
sampling_settings = SampleSettings(batch_offset=batch_offset, noise_type=noise_type, seed_gen=seed_gen, seed_offset=seed_offset, noise_layers=noise_layers,
|
||||
iteration_opts=iteration_opts, seed_override=seed_override, adapt_denoise_steps=adapt_denoise_steps,
|
||||
custom_cfg=custom_cfg, sigma_schedule=sigma_schedule)
|
||||
return (sampling_settings,)
|
||||
|
||||
|
||||
class NoiseLayerReplaceNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"batch_offset": ("INT", {"default": 0, "min": 0, "max": BIGMAX}),
|
||||
"noise_type": (NoiseLayerType.LIST,),
|
||||
"seed_gen_override": (SeedNoiseGeneration.LIST_WITH_OVERRIDE,),
|
||||
"seed_offset": ("INT", {"default": 0, "min": BIGMIN, "max": BIGMAX}),
|
||||
},
|
||||
"optional": {
|
||||
"prev_noise_layers": ("NOISE_LAYERS",),
|
||||
"mask_optional": ("MASK",),
|
||||
"seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("NOISE_LAYERS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/noise layers"
|
||||
FUNCTION = "create_layers"
|
||||
|
||||
def create_layers(self, batch_offset: int, noise_type: str, seed_gen_override: str, seed_offset: int,
|
||||
prev_noise_layers: NoiseLayerGroup=None, mask_optional: Tensor=None, seed_override: int=None,):
|
||||
# prepare prev_noise_layers
|
||||
if prev_noise_layers is None:
|
||||
prev_noise_layers = NoiseLayerGroup()
|
||||
prev_noise_layers = prev_noise_layers.clone()
|
||||
# create layer
|
||||
layer = NoiseLayerReplace(noise_type=noise_type, batch_offset=batch_offset, seed_gen_override=seed_gen_override, seed_offset=seed_offset,
|
||||
seed_override=seed_override, mask=mask_optional)
|
||||
prev_noise_layers.add_to_start(layer)
|
||||
return (prev_noise_layers,)
|
||||
|
||||
|
||||
class NoiseLayerAddNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"batch_offset": ("INT", {"default": 0, "min": 0, "max": BIGMAX}),
|
||||
"noise_type": (NoiseLayerType.LIST,),
|
||||
"seed_gen_override": (SeedNoiseGeneration.LIST_WITH_OVERRIDE,),
|
||||
"seed_offset": ("INT", {"default": 0, "min": BIGMIN, "max": BIGMAX}),
|
||||
"noise_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.001}),
|
||||
},
|
||||
"optional": {
|
||||
"prev_noise_layers": ("NOISE_LAYERS",),
|
||||
"mask_optional": ("MASK",),
|
||||
"seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("NOISE_LAYERS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/noise layers"
|
||||
FUNCTION = "create_layers"
|
||||
|
||||
def create_layers(self, batch_offset: int, noise_type: str, seed_gen_override: str, seed_offset: int,
|
||||
noise_weight: float,
|
||||
prev_noise_layers: NoiseLayerGroup=None, mask_optional: Tensor=None, seed_override: int=None,):
|
||||
# prepare prev_noise_layers
|
||||
if prev_noise_layers is None:
|
||||
prev_noise_layers = NoiseLayerGroup()
|
||||
prev_noise_layers = prev_noise_layers.clone()
|
||||
# create layer
|
||||
layer = NoiseLayerAdd(noise_type=noise_type, batch_offset=batch_offset, seed_gen_override=seed_gen_override, seed_offset=seed_offset,
|
||||
seed_override=seed_override, mask=mask_optional,
|
||||
noise_weight=noise_weight)
|
||||
prev_noise_layers.add_to_start(layer)
|
||||
return (prev_noise_layers,)
|
||||
|
||||
|
||||
class NoiseLayerAddWeightedNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"batch_offset": ("INT", {"default": 0, "min": 0, "max": BIGMAX}),
|
||||
"noise_type": (NoiseLayerType.LIST,),
|
||||
"seed_gen_override": (SeedNoiseGeneration.LIST_WITH_OVERRIDE,),
|
||||
"seed_offset": ("INT", {"default": 0, "min": BIGMIN, "max": BIGMAX}),
|
||||
"noise_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.001}),
|
||||
"balance_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
|
||||
},
|
||||
"optional": {
|
||||
"prev_noise_layers": ("NOISE_LAYERS",),
|
||||
"mask_optional": ("MASK",),
|
||||
"seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("NOISE_LAYERS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/noise layers"
|
||||
FUNCTION = "create_layers"
|
||||
|
||||
def create_layers(self, batch_offset: int, noise_type: str, seed_gen_override: str, seed_offset: int,
|
||||
noise_weight: float, balance_multiplier: float,
|
||||
prev_noise_layers: NoiseLayerGroup=None, mask_optional: Tensor=None, seed_override: int=None,):
|
||||
# prepare prev_noise_layers
|
||||
if prev_noise_layers is None:
|
||||
prev_noise_layers = NoiseLayerGroup()
|
||||
prev_noise_layers = prev_noise_layers.clone()
|
||||
# create layer
|
||||
layer = NoiseLayerAddWeighted(noise_type=noise_type, batch_offset=batch_offset, seed_gen_override=seed_gen_override, seed_offset=seed_offset,
|
||||
seed_override=seed_override, mask=mask_optional,
|
||||
noise_weight=noise_weight, balance_multiplier=balance_multiplier)
|
||||
prev_noise_layers.add_to_start(layer)
|
||||
return (prev_noise_layers,)
|
||||
|
||||
|
||||
class IterationOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"iterations": ("INT", {"default": 1, "min": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"iter_batch_offset": ("INT", {"default": 0, "min": 0, "max": BIGMAX}),
|
||||
"iter_seed_offset": ("INT", {"default": 0, "min": BIGMIN, "max": BIGMAX}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("ITERATION_OPTS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/iteration opts"
|
||||
FUNCTION = "create_iter_opts"
|
||||
|
||||
def create_iter_opts(self, iterations: int, iter_batch_offset: int=0, iter_seed_offset: int=0):
|
||||
iter_opts = IterationOptions(iterations=iterations, iter_batch_offset=iter_batch_offset, iter_seed_offset=iter_seed_offset)
|
||||
return (iter_opts,)
|
||||
|
||||
|
||||
class FreeInitOptionsNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"iterations": ("INT", {"default": 2, "min": 1}),
|
||||
"filter": (FreeInitFilter.LIST,),
|
||||
"d_s": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"d_t": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"n_butterworth": ("INT", {"default": 4, "min": 1, "max": 100},),
|
||||
"sigma_step": ("INT", {"default": 999, "min": 1, "max": 999}),
|
||||
"apply_to_1st_iter": ("BOOLEAN", {"default": False}),
|
||||
"init_type": (FreeInitOptions.LIST,)
|
||||
},
|
||||
"optional": {
|
||||
"iter_batch_offset": ("INT", {"default": 0, "min": 0, "max": BIGMAX}),
|
||||
"iter_seed_offset": ("INT", {"default": 1, "min": BIGMIN, "max": BIGMAX}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("ITERATION_OPTS",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/iteration opts"
|
||||
FUNCTION = "create_iter_opts"
|
||||
|
||||
def create_iter_opts(self, iterations: int, filter: str, d_s: float, d_t: float, n_butterworth: int,
|
||||
sigma_step: int, apply_to_1st_iter: bool, init_type: str,
|
||||
iter_batch_offset: int=0, iter_seed_offset: int=1):
|
||||
# init_type does nothing for now, not until I add more methods of applying low+high freq noise
|
||||
iter_opts = FreeInitOptions(iterations=iterations, step=sigma_step, apply_to_1st_iter=apply_to_1st_iter,
|
||||
filter=filter, d_s=d_s, d_t=d_t, n=n_butterworth, init_type=init_type,
|
||||
iter_batch_offset=iter_batch_offset, iter_seed_offset=iter_seed_offset)
|
||||
return (iter_opts,)
|
||||
|
||||
|
||||
class CustomCFGNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"cfg_multival": ("MULTIVAL",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CUSTOM_CFG",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/sample settings"
|
||||
FUNCTION = "create_custom_cfg"
|
||||
|
||||
def create_custom_cfg(self, cfg_multival: Union[float, Tensor]):
|
||||
keyframe = CustomCFGKeyframe(cfg_multival=cfg_multival)
|
||||
cfg_custom = CustomCFGKeyframeGroup()
|
||||
cfg_custom.add(keyframe)
|
||||
return (cfg_custom,)
|
||||
|
||||
|
||||
class CustomCFGKeyframeNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"cfg_multival": ("MULTIVAL",),
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"guarantee_steps": ("INT", {"default": 1, "min": 0, "max": BIGMAX}),
|
||||
},
|
||||
"optional": {
|
||||
"prev_custom_cfg": ("CUSTOM_CFG",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CUSTOM_CFG",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/sample settings"
|
||||
FUNCTION = "create_custom_cfg"
|
||||
|
||||
def create_custom_cfg(self, cfg_multival: Union[float, Tensor], start_percent: float=0.0, guarantee_steps: int=1,
|
||||
prev_custom_cfg: CustomCFGKeyframeGroup=None):
|
||||
if not prev_custom_cfg:
|
||||
prev_custom_cfg = CustomCFGKeyframeGroup()
|
||||
prev_custom_cfg = prev_custom_cfg.clone()
|
||||
keyframe = CustomCFGKeyframe(cfg_multival=cfg_multival, start_percent=start_percent, guarantee_steps=guarantee_steps)
|
||||
prev_custom_cfg.add(keyframe)
|
||||
return (prev_custom_cfg,)
|
||||
+141
@@ -0,0 +1,141 @@
|
||||
import torch
|
||||
|
||||
from .utils_model import BetaSchedules, SigmaSchedule, ModelSamplingType, ModelSamplingConfig, InterpolationMethod
|
||||
|
||||
|
||||
def validate_sigma_schedule_compatibility(schedule_A: SigmaSchedule, schedule_B: SigmaSchedule,
|
||||
name_a: str="sigma_schedule_A", name_b: str="sigma_schedule_B"):
|
||||
if schedule_A.total_sigmas() != schedule_B.total_sigmas():
|
||||
raise Exception(f"Weighted Average cannot be taken of Sigma Schedules that do not have the same amount of sigmas; " +
|
||||
f"{name_a} has {schedule_A.total_sigmas()} sigmas (lcm={schedule_A.is_lcm()}), " +
|
||||
f"{name_b} has {schedule_B.total_sigmas()} sigmas (lcm={schedule_B.is_lcm()}).")
|
||||
|
||||
|
||||
class SigmaScheduleNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"beta_schedule": (BetaSchedules.ALIAS_ACTIVE_LIST,),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SIGMA_SCHEDULE",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/sample settings/sigma schedule"
|
||||
FUNCTION = "get_sigma_schedule"
|
||||
|
||||
def get_sigma_schedule(self, beta_schedule: str):
|
||||
model_type = ModelSamplingType.from_alias(ModelSamplingType.EPS)
|
||||
new_model_sampling = BetaSchedules._to_model_sampling(alias=beta_schedule,
|
||||
model_type=model_type)
|
||||
return (SigmaSchedule(model_sampling=new_model_sampling, model_type=model_type),)
|
||||
|
||||
|
||||
class RawSigmaScheduleNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"raw_beta_schedule": (BetaSchedules.RAW_BETA_SCHEDULE_LIST,),
|
||||
"linear_start": ("FLOAT", {"default": 0.00085, "min": 0.0, "max": 1.0, "step": 0.000001}),
|
||||
"linear_end": ("FLOAT", {"default": 0.012, "min": 0.0, "max": 1.0, "step": 0.000001}),
|
||||
#"cosine_s": ("FLOAT", {"default": 8e-3, "min": 0.0, "max": 1.0, "step": 0.000001}),
|
||||
"sampling": (ModelSamplingType._FULL_LIST,),
|
||||
"lcm_original_timesteps": ("INT", {"default": 50, "min": 1, "max": 1000}),
|
||||
"lcm_zsnr": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SIGMA_SCHEDULE",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/sample settings/sigma schedule"
|
||||
FUNCTION = "get_sigma_schedule"
|
||||
|
||||
def get_sigma_schedule(self, raw_beta_schedule: str, linear_start: float, linear_end: float,# cosine_s: float,
|
||||
sampling: str, lcm_original_timesteps: int, lcm_zsnr: bool):
|
||||
new_config = ModelSamplingConfig(beta_schedule=raw_beta_schedule, linear_start=linear_start, linear_end=linear_end)
|
||||
if sampling != ModelSamplingType.LCM:
|
||||
lcm_original_timesteps=None
|
||||
lcm_zsnr=False
|
||||
model_type = ModelSamplingType.from_alias(sampling)
|
||||
new_model_sampling = BetaSchedules._to_model_sampling(alias=BetaSchedules.AUTOSELECT, model_type=model_type, config_override=new_config, original_timesteps=lcm_original_timesteps)
|
||||
if lcm_zsnr:
|
||||
SigmaSchedule.apply_zsnr(new_model_sampling=new_model_sampling)
|
||||
return (SigmaSchedule(model_sampling=new_model_sampling, model_type=model_type),)
|
||||
|
||||
|
||||
class WeightedAverageSigmaScheduleNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"schedule_A": ("SIGMA_SCHEDULE",),
|
||||
"schedule_B": ("SIGMA_SCHEDULE",),
|
||||
"weight_A": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SIGMA_SCHEDULE",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/sample settings/sigma schedule"
|
||||
FUNCTION = "get_sigma_schedule"
|
||||
|
||||
def get_sigma_schedule(self, schedule_A: SigmaSchedule, schedule_B: SigmaSchedule, weight_A: float):
|
||||
validate_sigma_schedule_compatibility(schedule_A, schedule_B)
|
||||
new_sigmas = schedule_A.model_sampling.sigmas * weight_A + schedule_B.model_sampling.sigmas * (1-weight_A)
|
||||
combo_schedule = schedule_A.clone()
|
||||
combo_schedule.model_sampling.set_sigmas(new_sigmas)
|
||||
return (combo_schedule,)
|
||||
|
||||
|
||||
class InterpolatedWeightedAverageSigmaScheduleNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"schedule_A": ("SIGMA_SCHEDULE",),
|
||||
"schedule_B": ("SIGMA_SCHEDULE",),
|
||||
"weight_A_Start": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"weight_A_End": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"interpolation": (InterpolationMethod._LIST,),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SIGMA_SCHEDULE",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/sample settings/sigma schedule"
|
||||
FUNCTION = "get_sigma_schedule"
|
||||
|
||||
def get_sigma_schedule(self, schedule_A: SigmaSchedule, schedule_B: SigmaSchedule,
|
||||
weight_A_Start: float, weight_A_End: float, interpolation: str):
|
||||
validate_sigma_schedule_compatibility(schedule_A, schedule_B)
|
||||
# get reverse weights, since sigmas are currently reversed
|
||||
weights = InterpolationMethod.get_weights(num_from=weight_A_Start, num_to=weight_A_End,
|
||||
length=schedule_A.total_sigmas(), method=interpolation, reverse=True)
|
||||
weights = weights.to(schedule_A.model_sampling.sigmas.dtype).to(schedule_A.model_sampling.sigmas.device)
|
||||
new_sigmas = schedule_A.model_sampling.sigmas * weights + schedule_B.model_sampling.sigmas * (1.0-weights)
|
||||
combo_schedule = schedule_A.clone()
|
||||
combo_schedule.model_sampling.set_sigmas(new_sigmas)
|
||||
return (combo_schedule,)
|
||||
|
||||
|
||||
class SplitAndCombineSigmaScheduleNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"schedule_Start": ("SIGMA_SCHEDULE",),
|
||||
"schedule_End": ("SIGMA_SCHEDULE",),
|
||||
"idx_split_percent": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.001})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SIGMA_SCHEDULE",)
|
||||
CATEGORY = "Animate Diff 🎭🅐🅓/sample settings/sigma schedule"
|
||||
FUNCTION = "get_sigma_schedule"
|
||||
|
||||
def get_sigma_schedule(self, schedule_Start: SigmaSchedule, schedule_End: SigmaSchedule, idx_split_percent: float):
|
||||
validate_sigma_schedule_compatibility(schedule_Start, schedule_End)
|
||||
# first, calculate index to act as split; get diff from 1.0 since sigmas are flipped at this stage
|
||||
idx = int((1.0-idx_split_percent) * schedule_Start.total_sigmas())
|
||||
new_sigmas = torch.cat([schedule_End.model_sampling.sigmas[:idx], schedule_Start.model_sampling.sigmas[idx:]], dim=0)
|
||||
new_schedule = schedule_Start.clone()
|
||||
new_schedule.model_sampling.set_sigmas(new_sigmas)
|
||||
return (new_schedule,)
|
||||
+555
@@ -0,0 +1,555 @@
|
||||
from collections.abc import Iterable
|
||||
from typing import Union
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
import comfy.sample
|
||||
import comfy.samplers
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
from comfy.model_base import BaseModel
|
||||
|
||||
from . import freeinit
|
||||
from .context import ContextOptions, ContextOptionsGroup
|
||||
from .utils_model import SigmaSchedule
|
||||
from .utils_motion import extend_to_batch_size, get_sorted_list_via_attr, prepare_mask_batch
|
||||
from .logger import logger
|
||||
|
||||
|
||||
def prepare_mask_ad(noise_mask, shape, device):
|
||||
"""ensures noise mask is of proper dimensions"""
|
||||
noise_mask = torch.nn.functional.interpolate(noise_mask.reshape((-1, 1, noise_mask.shape[-2], noise_mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear")
|
||||
#noise_mask = noise_mask.round()
|
||||
noise_mask = torch.cat([noise_mask] * shape[1], dim=1)
|
||||
noise_mask = comfy.utils.repeat_to_batch_size(noise_mask, shape[0])
|
||||
noise_mask = noise_mask.to(device)
|
||||
return noise_mask
|
||||
|
||||
|
||||
class NoiseLayerType:
|
||||
DEFAULT = "default"
|
||||
CONSTANT = "constant"
|
||||
EMPTY = "empty"
|
||||
REPEATED_CONTEXT = "repeated_context"
|
||||
FREENOISE = "FreeNoise"
|
||||
|
||||
LIST = [DEFAULT, CONSTANT, EMPTY, REPEATED_CONTEXT, FREENOISE]
|
||||
|
||||
|
||||
class NoiseApplication:
|
||||
ADD = "add"
|
||||
ADD_WEIGHTED = "add_weighted"
|
||||
REPLACE = "replace"
|
||||
|
||||
LIST = [ADD, ADD_WEIGHTED, REPLACE]
|
||||
|
||||
|
||||
class NoiseNormalize:
|
||||
DISABLE = "disable"
|
||||
NORMAL = "normal"
|
||||
|
||||
LIST = [DISABLE, NORMAL]
|
||||
|
||||
|
||||
class SampleSettings:
|
||||
def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, seed_offset: int=0, noise_layers: 'NoiseLayerGroup'=None,
|
||||
iteration_opts=None, seed_override:int=None, negative_cond_flipflop=False, adapt_denoise_steps: bool=False,
|
||||
custom_cfg: 'CustomCFGKeyframeGroup'=None, sigma_schedule: SigmaSchedule=None):
|
||||
self.batch_offset = batch_offset
|
||||
self.noise_type = noise_type if noise_type is not None else NoiseLayerType.DEFAULT
|
||||
self.seed_gen = seed_gen if seed_gen is not None else SeedNoiseGeneration.COMFY
|
||||
self.noise_layers = noise_layers if noise_layers else NoiseLayerGroup()
|
||||
self.iteration_opts = iteration_opts if iteration_opts else IterationOptions()
|
||||
self.seed_offset = seed_offset
|
||||
self.seed_override = seed_override
|
||||
self.negative_cond_flipflop = negative_cond_flipflop
|
||||
self.adapt_denoise_steps = adapt_denoise_steps
|
||||
self.custom_cfg = custom_cfg.clone() if custom_cfg else custom_cfg
|
||||
self.sigma_schedule = sigma_schedule
|
||||
|
||||
def prepare_noise(self, seed: int, latents: Tensor, noise: Tensor, extra_seed_offset=0, extra_args:dict={}, force_create_noise=True):
|
||||
if self.seed_override is not None:
|
||||
seed = self.seed_override
|
||||
# if seed is iterable, attempt to do per-latent noises
|
||||
if isinstance(seed, Iterable):
|
||||
noise = SeedNoiseGeneration.create_noise_individual_seeds(seeds=seed, latents=latents, seed_offset=self.seed_offset+extra_seed_offset, extra_args=extra_args)
|
||||
seed = seed[0]+self.seed_offset
|
||||
else:
|
||||
seed += self.seed_offset
|
||||
# replace initial noise if not batch_offset 0 or Comfy seed_gen or not NoiseType default
|
||||
if self.batch_offset != 0 or self.seed_offset != 0 or self.noise_type != NoiseLayerType.DEFAULT or self.seed_gen != SeedNoiseGeneration.COMFY or force_create_noise:
|
||||
noise = SeedNoiseGeneration.create_noise(seed=seed+extra_seed_offset, latents=latents, existing_seed_gen=self.seed_gen, seed_gen=self.seed_gen,
|
||||
noise_type=self.noise_type, batch_offset=self.batch_offset, extra_args=extra_args)
|
||||
# apply noise layers
|
||||
for noise_layer in self.noise_layers.layers:
|
||||
# first, generate new noise matching seed gen override
|
||||
layer_noise = noise_layer.create_layer_noise(existing_seed_gen=self.seed_gen, seed=seed, latents=latents,
|
||||
extra_seed_offset=extra_seed_offset, extra_args=extra_args)
|
||||
# next, get noise after applying layer
|
||||
noise = noise_layer.apply_layer_noise(new_noise=layer_noise, old_noise=noise)
|
||||
# noise prepared now
|
||||
return noise
|
||||
|
||||
def pre_run(self, model: ModelPatcher):
|
||||
if self.custom_cfg is not None:
|
||||
self.custom_cfg.reset()
|
||||
|
||||
def cleanup(self):
|
||||
if self.custom_cfg is not None:
|
||||
self.custom_cfg.reset()
|
||||
|
||||
def clone(self):
|
||||
return SampleSettings(batch_offset=self.batch_offset, noise_type=self.noise_type, seed_gen=self.seed_gen, seed_offset=self.seed_offset,
|
||||
noise_layers=self.noise_layers.clone(), iteration_opts=self.iteration_opts, seed_override=self.seed_override,
|
||||
negative_cond_flipflop=self.negative_cond_flipflop, adapt_denoise_steps=self.adapt_denoise_steps, custom_cfg=self.custom_cfg, sigma_schedule=self.sigma_schedule)
|
||||
|
||||
|
||||
class NoiseLayer:
|
||||
def __init__(self, noise_type: str, batch_offset: int, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None):
|
||||
self.application: str = NoiseApplication.REPLACE
|
||||
self.noise_type = noise_type
|
||||
self.batch_offset = batch_offset
|
||||
self.seed_gen_override = seed_gen_override
|
||||
self.seed_offset = seed_offset
|
||||
self.seed_override = seed_override
|
||||
self.mask = mask
|
||||
|
||||
def create_layer_noise(self, existing_seed_gen: str, seed: int, latents: Tensor, extra_seed_offset=0, extra_args:dict={}) -> Tensor:
|
||||
if self.seed_override is not None:
|
||||
seed = self.seed_override
|
||||
# if seed is iterable, attempt to do per-latent noises
|
||||
if isinstance(seed, Iterable):
|
||||
return SeedNoiseGeneration.create_noise_individual_seeds(seeds=seed, latents=latents, seed_offset=self.seed_offset+extra_seed_offset, extra_args=extra_args)
|
||||
seed += self.seed_offset + extra_seed_offset
|
||||
return SeedNoiseGeneration.create_noise(seed=seed, latents=latents, existing_seed_gen=existing_seed_gen, seed_gen=self.seed_gen_override,
|
||||
noise_type=self.noise_type, batch_offset=self.batch_offset, extra_args=extra_args)
|
||||
|
||||
def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor:
|
||||
return old_noise
|
||||
|
||||
def get_noise_mask(self, noise: Tensor) -> Tensor:
|
||||
if self.mask is None:
|
||||
return 1
|
||||
noise_mask = self.mask.reshape((-1, 1, self.mask.shape[-2], self.mask.shape[-1]))
|
||||
return prepare_mask_ad(noise_mask, noise.shape, noise.device)
|
||||
|
||||
|
||||
class NoiseLayerReplace(NoiseLayer):
|
||||
def __init__(self, noise_type: str, batch_offset: int, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None):
|
||||
super().__init__(noise_type, batch_offset, seed_gen_override, seed_offset, seed_override, mask)
|
||||
self.application = NoiseApplication.REPLACE
|
||||
|
||||
def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor:
|
||||
noise_mask = self.get_noise_mask(old_noise)
|
||||
return (1-noise_mask)*old_noise + noise_mask*new_noise
|
||||
|
||||
|
||||
class NoiseLayerAdd(NoiseLayer):
|
||||
def __init__(self, noise_type: str, batch_offset: int, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None,
|
||||
noise_weight=1.0):
|
||||
super().__init__(noise_type, batch_offset, seed_gen_override, seed_offset, seed_override, mask)
|
||||
self.noise_weight = noise_weight
|
||||
self.application = NoiseApplication.ADD
|
||||
|
||||
def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor:
|
||||
noise_mask = self.get_noise_mask(old_noise)
|
||||
return (1-noise_mask)*old_noise + noise_mask*(old_noise + new_noise * self.noise_weight)
|
||||
|
||||
|
||||
class NoiseLayerAddWeighted(NoiseLayerAdd):
|
||||
def __init__(self, noise_type: str, batch_offset: int, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None,
|
||||
noise_weight=1.0, balance_multiplier=1.0):
|
||||
super().__init__(noise_type, batch_offset, seed_gen_override, seed_offset, seed_override, mask, noise_weight)
|
||||
self.balance_multiplier = balance_multiplier
|
||||
self.application = NoiseApplication.ADD_WEIGHTED
|
||||
|
||||
def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor:
|
||||
noise_mask = self.get_noise_mask(old_noise)
|
||||
return (1-noise_mask)*old_noise + noise_mask*(old_noise * (1.0-(self.noise_weight*self.balance_multiplier)) + new_noise * self.noise_weight)
|
||||
|
||||
|
||||
class NoiseLayerGroup:
|
||||
def __init__(self):
|
||||
self.layers: list[NoiseLayer] = []
|
||||
|
||||
def add(self, layer: NoiseLayer) -> None:
|
||||
# add to the end of list
|
||||
self.layers.append(layer)
|
||||
|
||||
def add_to_start(self, layer: NoiseLayer) -> None:
|
||||
# add to the beginning of list
|
||||
self.layers.insert(0, layer)
|
||||
|
||||
def __getitem__(self, index) -> NoiseLayer:
|
||||
return self.layers[index]
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return len(self.layers) == 0
|
||||
|
||||
def clone(self) -> 'NoiseLayerGroup':
|
||||
cloned = NoiseLayerGroup()
|
||||
for layer in self.layers:
|
||||
cloned.add(layer)
|
||||
return cloned
|
||||
|
||||
class SeedNoiseGeneration:
|
||||
COMFY = "comfy"
|
||||
AUTO1111 = "auto1111"
|
||||
AUTO1111GPU = "auto1111 [gpu]" # TODO: implement this
|
||||
USE_EXISTING = "use existing"
|
||||
|
||||
LIST = [COMFY, AUTO1111]
|
||||
LIST_WITH_OVERRIDE = [USE_EXISTING, COMFY, AUTO1111]
|
||||
|
||||
@classmethod
|
||||
def create_noise(cls, seed: int, latents: Tensor, existing_seed_gen: str=COMFY, seed_gen: str=USE_EXISTING, noise_type: str=NoiseLayerType.DEFAULT, batch_offset: int=0, extra_args: dict={}):
|
||||
# determine if should use existing type
|
||||
if seed_gen == cls.USE_EXISTING:
|
||||
seed_gen = existing_seed_gen
|
||||
if seed_gen == cls.COMFY:
|
||||
return cls.create_noise_comfy(seed, latents, noise_type, batch_offset, extra_args)
|
||||
elif seed_gen in [cls.AUTO1111, cls.AUTO1111GPU]:
|
||||
return cls.create_noise_auto1111(seed, latents, noise_type, batch_offset, extra_args)
|
||||
raise ValueError(f"Noise seed_gen {seed_gen} is not recognized.")
|
||||
|
||||
@staticmethod
|
||||
def create_noise_comfy(seed: int, latents: Tensor, noise_type: str=NoiseLayerType.DEFAULT, batch_offset: int=0, extra_args: dict={}):
|
||||
common_noise = SeedNoiseGeneration._create_common_noise(seed, latents, noise_type, batch_offset, extra_args)
|
||||
if common_noise is not None:
|
||||
return common_noise
|
||||
if noise_type == NoiseLayerType.CONSTANT:
|
||||
generator = torch.manual_seed(seed)
|
||||
length = latents.shape[0]
|
||||
single_shape = (1 + batch_offset, latents.shape[1], latents.shape[2], latents.shape[3])
|
||||
single_noise = torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu")
|
||||
return torch.cat([single_noise[batch_offset:]] * length, dim=0)
|
||||
# comfy creates noise with a single seed for the entire shape of the latents batched tensor
|
||||
generator = torch.manual_seed(seed)
|
||||
offset_shape = (latents.shape[0] + batch_offset, latents.shape[1], latents.shape[2], latents.shape[3])
|
||||
final_noise = torch.randn(offset_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu")
|
||||
final_noise = final_noise[batch_offset:]
|
||||
# convert to derivative noise type, if needed
|
||||
derivative_noise = SeedNoiseGeneration._create_derivative_noise(final_noise, noise_type=noise_type, seed=seed, extra_args=extra_args)
|
||||
if derivative_noise is not None:
|
||||
return derivative_noise
|
||||
return final_noise
|
||||
|
||||
@staticmethod
|
||||
def create_noise_auto1111(seed: int, latents: Tensor, noise_type: str=NoiseLayerType.DEFAULT, batch_offset: int=0, extra_args: dict={}):
|
||||
common_noise = SeedNoiseGeneration._create_common_noise(seed, latents, noise_type, batch_offset, extra_args)
|
||||
if common_noise is not None:
|
||||
return common_noise
|
||||
if noise_type == NoiseLayerType.CONSTANT:
|
||||
generator = torch.manual_seed(seed+batch_offset)
|
||||
length = latents.shape[0]
|
||||
single_shape = (1, latents.shape[1], latents.shape[2], latents.shape[3])
|
||||
single_noise = torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu")
|
||||
return torch.cat([single_noise] * length, dim=0)
|
||||
# auto1111 applies growing seeds for a batch
|
||||
length = latents.shape[0]
|
||||
single_shape = (1, latents.shape[1], latents.shape[2], latents.shape[3])
|
||||
all_noises = []
|
||||
# i starts at 0
|
||||
for i in range(length):
|
||||
generator = torch.manual_seed(seed+i+batch_offset)
|
||||
all_noises.append(torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu"))
|
||||
final_noise = torch.cat(all_noises, dim=0)
|
||||
# convert to derivative noise type, if needed
|
||||
derivative_noise = SeedNoiseGeneration._create_derivative_noise(final_noise, noise_type=noise_type, seed=seed, extra_args=extra_args)
|
||||
if derivative_noise is not None:
|
||||
return derivative_noise
|
||||
return final_noise
|
||||
|
||||
@staticmethod
|
||||
def create_noise_individual_seeds(seeds: list[int], latents: Tensor, seed_offset: int=0, extra_args: dict={}):
|
||||
length = latents.shape[0]
|
||||
if len(seeds) < length:
|
||||
raise ValueError(f"{len(seeds)} seeds in seed_override were provided, but at least {length} are required to work with the current latents.")
|
||||
seeds = seeds[:length]
|
||||
single_shape = (1, latents.shape[1], latents.shape[2], latents.shape[3])
|
||||
all_noises = []
|
||||
for seed in seeds:
|
||||
generator = torch.manual_seed(seed+seed_offset)
|
||||
all_noises.append(torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu"))
|
||||
return torch.cat(all_noises, dim=0)
|
||||
|
||||
@staticmethod
|
||||
def _create_common_noise(seed: int, latents: Tensor, noise_type: str=NoiseLayerType.DEFAULT, batch_offset: int=0, extra_args: dict={}):
|
||||
if noise_type == NoiseLayerType.EMPTY:
|
||||
return torch.zeros_like(latents)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _create_derivative_noise(noise: Tensor, noise_type: str, seed: int, extra_args: dict):
|
||||
derivative_func = DERIVATIVE_NOISE_FUNC_MAP.get(noise_type, None)
|
||||
if derivative_func is None:
|
||||
return None
|
||||
return derivative_func(noise=noise, seed=seed, extra_args=extra_args)
|
||||
|
||||
@staticmethod
|
||||
def _convert_to_repeated_context(noise: Tensor, extra_args: dict, **kwargs):
|
||||
# if no context_length, return unmodified noise
|
||||
opts: ContextOptionsGroup = extra_args["context_options"]
|
||||
context_length: int = opts.context_length if not opts.view_options else opts.view_options.context_length
|
||||
if context_length is None:
|
||||
return noise
|
||||
length = noise.shape[0]
|
||||
noise = noise[:context_length]
|
||||
cat_count = (length // context_length) + 1
|
||||
return torch.cat([noise] * cat_count, dim=0)[:length]
|
||||
|
||||
@staticmethod
|
||||
def _convert_to_freenoise(noise: Tensor, seed: int, extra_args: dict, **kwargs):
|
||||
# if no context_length, return unmodified noise
|
||||
opts: ContextOptionsGroup = extra_args["context_options"]
|
||||
context_length: int = opts.context_length if not opts.view_options else opts.view_options.context_length
|
||||
context_overlap: int = opts.context_overlap if not opts.view_options else opts.view_options.context_overlap
|
||||
video_length: int = noise.shape[0]
|
||||
if context_length is None:
|
||||
return noise
|
||||
delta = context_length - context_overlap
|
||||
generator = torch.manual_seed(seed)
|
||||
|
||||
for start_idx in range(0, video_length-context_length, delta):
|
||||
# start_idx corresponds to the beginning of a context window
|
||||
# goal: place shuffled in the delta region right after the end of the context window
|
||||
# if space after context window is not enough to place the noise, adjust and finish
|
||||
place_idx = start_idx + context_length
|
||||
# if place_idx is outside the valid indexes, we are already finished
|
||||
if place_idx >= video_length:
|
||||
break
|
||||
end_idx = place_idx - 1
|
||||
# if there is not enough room to copy delta amount of indexes, copy limited amount and finish
|
||||
if end_idx + delta >= video_length:
|
||||
final_delta = video_length - place_idx
|
||||
# generate list of indexes in final delta region
|
||||
list_idx = torch.Tensor(list(range(start_idx,start_idx+final_delta))).to(torch.long)
|
||||
# shuffle list
|
||||
list_idx = list_idx[torch.randperm(final_delta, generator=generator)]
|
||||
# apply shuffled indexes
|
||||
noise[place_idx:place_idx+final_delta] = noise[list_idx]
|
||||
break
|
||||
# otherwise, do normal behavior
|
||||
# generate list of indexes in delta region
|
||||
list_idx = torch.Tensor(list(range(start_idx,start_idx+delta))).to(torch.long)
|
||||
# shuffle list
|
||||
list_idx = list_idx[torch.randperm(delta, generator=generator)]
|
||||
# apply shuffled indexes
|
||||
noise[place_idx:place_idx+delta] = noise[list_idx]
|
||||
return noise
|
||||
|
||||
|
||||
DERIVATIVE_NOISE_FUNC_MAP = {
|
||||
NoiseLayerType.REPEATED_CONTEXT: SeedNoiseGeneration._convert_to_repeated_context,
|
||||
NoiseLayerType.FREENOISE: SeedNoiseGeneration._convert_to_freenoise,
|
||||
}
|
||||
|
||||
|
||||
class IterationOptions:
|
||||
SAMPLER = "sampler"
|
||||
|
||||
def __init__(self, iterations: int=1, cache_init_noise=False, cache_init_latents=False,
|
||||
iter_batch_offset: int=0, iter_seed_offset: int=0):
|
||||
self.iterations = iterations
|
||||
self.cache_init_noise = cache_init_noise
|
||||
self.cache_init_latents = cache_init_latents
|
||||
self.iter_batch_offset = iter_batch_offset
|
||||
self.iter_seed_offset = iter_seed_offset
|
||||
self.need_sampler = False
|
||||
|
||||
def get_sigma(self, model: ModelPatcher, step: int):
|
||||
model_sampling = model.model.model_sampling
|
||||
if "model_sampling" in model.object_patches:
|
||||
model_sampling = model.object_patches["model_sampling"]
|
||||
return model_sampling.sigmas[step]
|
||||
|
||||
def initialize(self, latents: Tensor):
|
||||
pass
|
||||
|
||||
def preprocess_latents(self, curr_i: int, model: ModelPatcher, latents: Tensor, noise: Tensor,
|
||||
seed: int, sample_settings: SampleSettings, noise_extra_args: dict, **kwargs):
|
||||
if curr_i == 0 or (self.iter_batch_offset == 0 and self.iter_seed_offset == 0):
|
||||
return latents, noise
|
||||
temp_sample_settings = sample_settings.clone()
|
||||
temp_sample_settings.batch_offset += self.iter_batch_offset * curr_i
|
||||
temp_sample_settings.seed_offset += self.iter_seed_offset * curr_i
|
||||
return latents, temp_sample_settings.prepare_noise(seed=seed, latents=latents, noise=None,
|
||||
extra_args=noise_extra_args, force_create_noise=True)
|
||||
|
||||
|
||||
class FreeInitOptions(IterationOptions):
|
||||
FREEINIT_SAMPLER = "FreeInit [sampler sigma]"
|
||||
FREEINIT_MODEL = "FreeInit [model sigma]"
|
||||
DINKINIT_V1 = "DinkInit_v1"
|
||||
|
||||
LIST = [FREEINIT_SAMPLER, FREEINIT_MODEL, DINKINIT_V1]
|
||||
|
||||
def __init__(self, iterations: int, step: int=999, apply_to_1st_iter: bool=False,
|
||||
filter=freeinit.FreeInitFilter.GAUSSIAN, d_s=0.25, d_t=0.25, n=4, init_type=FREEINIT_SAMPLER,
|
||||
iter_batch_offset: int=0, iter_seed_offset: int=1):
|
||||
super().__init__(iterations=iterations, cache_init_noise=True, cache_init_latents=True,
|
||||
iter_batch_offset=iter_batch_offset, iter_seed_offset=iter_seed_offset)
|
||||
self.apply_to_1st_iter = apply_to_1st_iter
|
||||
self.step = step
|
||||
self.filter = filter
|
||||
self.d_s = d_s
|
||||
self.d_t = d_t
|
||||
self.n = n
|
||||
self.freq_filter = None
|
||||
self.freq_filter2 = None
|
||||
self.need_sampler = True if init_type in [self.FREEINIT_SAMPLER] else False
|
||||
self.init_type = init_type
|
||||
|
||||
def initialize(self, latents: Tensor):
|
||||
self.freq_filter = freeinit.get_freq_filter(latents.shape, device=latents.device, filter_type=self.filter,
|
||||
n=self.n, d_s=self.d_s, d_t=self.d_t)
|
||||
|
||||
def preprocess_latents(self, curr_i: int, model: ModelPatcher, latents: Tensor, noise: Tensor, cached_latents: Tensor, cached_noise: Tensor,
|
||||
seed:int, sample_settings: SampleSettings, noise_extra_args: dict, sampler: comfy.samplers.KSampler=None, **kwargs):
|
||||
# if first iter and should not apply, do nothing
|
||||
if curr_i == 0 and not self.apply_to_1st_iter:
|
||||
return latents, noise
|
||||
# otherwise, do FreeInit stuff
|
||||
if self.init_type in [self.FREEINIT_SAMPLER, self.FREEINIT_MODEL]:
|
||||
# NOTE: This should be very close (if not exactly) to how FreeInit is intended to initialize noise the latents.
|
||||
# The trick is that FreeInit is dependent on the behavior of diffuser's DDIMScheduler.add_noise function.
|
||||
# The typical noising method of latents + noise * sigma will NOT work.
|
||||
# 1. apply initial noise with appropriate step sigma, normalized against scale_factor
|
||||
if sampler is not None:
|
||||
sigma = sampler.sigmas[999-self.step].to(latents.device) / (model.model.latent_format.scale_factor)
|
||||
else:
|
||||
sigma = self.get_sigma(model, self.step-1000).to(latents.device) / (model.model.latent_format.scale_factor)
|
||||
alpha_cumprod = 1 / ((sigma * sigma) + 1)
|
||||
sqrt_alpha_prod = alpha_cumprod ** 0.5
|
||||
sqrt_one_minus_alpha_prod = (1 - alpha_cumprod) ** 0.5
|
||||
noised_latents = latents * sqrt_alpha_prod + noise * sqrt_one_minus_alpha_prod
|
||||
# 2. create random noise z_rand for high frequency
|
||||
temp_sample_settings = sample_settings.clone()
|
||||
temp_sample_settings.batch_offset += self.iter_batch_offset * curr_i
|
||||
temp_sample_settings.seed_offset += self.iter_seed_offset * curr_i
|
||||
z_rand = temp_sample_settings.prepare_noise(seed=seed, latents=latents, noise=None,
|
||||
extra_args=noise_extra_args, force_create_noise=True)
|
||||
# 3. noise reinitialization - combines low freq. noise from noised_latents and high freq. noise from z_rand
|
||||
noised_latents = freeinit.freq_mix_3d(x=noised_latents, noise=z_rand.to(dtype=latents.dtype, device=latents.device), LPF=self.freq_filter)
|
||||
return cached_latents, noised_latents
|
||||
elif self.init_type == self.DINKINIT_V1:
|
||||
# NOTE: This was my first attempt at implementing FreeInit; it sorta works due to my alpha_cumprod shenanigans,
|
||||
# but completely by accident.
|
||||
# 1. apply initial noise with appropriate step sigma
|
||||
sigma = self.get_sigma(model, self.step-1000).to(latents.device)
|
||||
alpha_cumprod = 1 / ((sigma * sigma) + 1) #1 / ((sigma * sigma)) # 1 / ((sigma * sigma) + 1)
|
||||
noised_latents = (latents + (cached_noise * sigma)) * alpha_cumprod
|
||||
# 2. create random noise z_rand for high frequency
|
||||
temp_sample_settings = sample_settings.clone()
|
||||
temp_sample_settings.batch_offset += self.iter_batch_offset * curr_i
|
||||
temp_sample_settings.seed_offset += self.iter_seed_offset * curr_i
|
||||
z_rand = temp_sample_settings.prepare_noise(seed=seed, latents=latents, noise=None,
|
||||
extra_args=noise_extra_args, force_create_noise=True)
|
||||
####z_rand = torch.randn_like(latents, dtype=latents.dtype, device=latents.device)
|
||||
# 3. noise reinitialization - combines low freq. noise from noised_latents and high freq. noise from z_rand
|
||||
noised_latents = freeinit.freq_mix_3d(x=noised_latents, noise=z_rand.to(dtype=latents.dtype, device=latents.device), LPF=self.freq_filter)
|
||||
return cached_latents, noised_latents
|
||||
else:
|
||||
raise ValueError(f"FreeInit init_type '{self.init_type}' is not recognized.")
|
||||
|
||||
|
||||
class CustomCFGKeyframe:
|
||||
def __init__(self, cfg_multival: Union[float, Tensor], start_percent=0.0, guarantee_steps=1):
|
||||
self.cfg_multival = cfg_multival
|
||||
# scheduling
|
||||
self.start_percent = float(start_percent)
|
||||
self.start_t = 999999999.9
|
||||
self.guarantee_steps = guarantee_steps
|
||||
|
||||
def clone(self):
|
||||
c = CustomCFGKeyframe(cfg_multival=self.cfg_multival,
|
||||
start_percent=self.start_percent, guarantee_steps=self.guarantee_steps)
|
||||
c.start_t = self.start_t
|
||||
return c
|
||||
|
||||
|
||||
class CustomCFGKeyframeGroup:
|
||||
def __init__(self):
|
||||
self.keyframes: list[CustomCFGKeyframe] = []
|
||||
self._current_keyframe: CustomCFGKeyframe = None
|
||||
self._current_used_steps: int = 0
|
||||
self._current_index: int = 0
|
||||
|
||||
def reset(self):
|
||||
self._current_keyframe = None
|
||||
self._current_used_steps = 0
|
||||
self._current_index = 0
|
||||
self._set_first_as_current()
|
||||
|
||||
def add(self, keyframe: CustomCFGKeyframe):
|
||||
# add to end of list, then sort
|
||||
self.keyframes.append(keyframe)
|
||||
self.keyframes = get_sorted_list_via_attr(self.keyframes, "start_percent")
|
||||
self._set_first_as_current()
|
||||
|
||||
def _set_first_as_current(self):
|
||||
if len(self.keyframes) > 0:
|
||||
self._current_keyframe = self.keyframes[0]
|
||||
else:
|
||||
self._current_keyframe = None
|
||||
|
||||
def has_index(self, index: int) -> int:
|
||||
return index >=0 and index < len(self.keyframes)
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return len(self.keyframes) == 0
|
||||
|
||||
def clone(self):
|
||||
cloned = CustomCFGKeyframeGroup()
|
||||
for keyframe in self.keyframes:
|
||||
cloned.keyframes.append(keyframe)
|
||||
cloned._set_first_as_current()
|
||||
return cloned
|
||||
|
||||
def initialize_timesteps(self, model: BaseModel):
|
||||
for keyframe in self.keyframes:
|
||||
keyframe.start_t = model.model_sampling.percent_to_sigma(keyframe.start_percent)
|
||||
|
||||
def prepare_current_keyframe(self, t: Tensor):
|
||||
curr_t: float = t[0]
|
||||
prev_index = self._current_index
|
||||
# if met guaranteed steps, look for next keyframe in case need to switch
|
||||
if self._current_used_steps >= self._current_keyframe.guarantee_steps:
|
||||
# if has next index, loop through and see if need t oswitch
|
||||
if self.has_index(self._current_index+1):
|
||||
for i in range(self._current_index+1, len(self.keyframes)):
|
||||
eval_c = self.keyframes[i]
|
||||
# check if start_t is greater or equal to curr_t
|
||||
# NOTE: t is in terms of sigmas, not percent, so bigger number = earlier step in sampling
|
||||
if eval_c.start_t >= curr_t:
|
||||
self._current_index = i
|
||||
self._current_keyframe = eval_c
|
||||
self._current_used_steps = 0
|
||||
# if guarantee_steps greater than zero, stop searching for other keyframes
|
||||
if self._current_keyframe.guarantee_steps > 0:
|
||||
break
|
||||
# if eval_c is outside the percent range, stop looking further
|
||||
else: break
|
||||
# update steps current context is used
|
||||
self._current_used_steps += 1
|
||||
|
||||
def patch_model(self, model: ModelPatcher) -> ModelPatcher:
|
||||
def evolved_custom_cfg(args):
|
||||
cond: Tensor = args["cond"]
|
||||
uncond: Tensor = args["uncond"]
|
||||
# cond scale is based purely off of CustomCFG - cond_scale input in sampler is ignored!
|
||||
cond_scale = self.cfg_multival
|
||||
if isinstance(cond_scale, Tensor):
|
||||
cond_scale = prepare_mask_batch(cond_scale.to(cond.dtype).to(cond.device), cond.shape)
|
||||
cond_scale = extend_to_batch_size(cond_scale, cond.shape[0])
|
||||
return uncond + (cond - uncond) * cond_scale
|
||||
|
||||
model = model.clone()
|
||||
model.set_model_sampler_cfg_function(evolved_custom_cfg)
|
||||
return model
|
||||
|
||||
# properties shadow those of CustomCFGKeyframe
|
||||
@property
|
||||
def cfg_multival(self):
|
||||
if self._current_keyframe != None:
|
||||
return self._current_keyframe.cfg_multival
|
||||
return None
|
||||
+528
@@ -0,0 +1,528 @@
|
||||
from typing import Callable
|
||||
|
||||
import math
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn.functional import group_norm
|
||||
from einops import rearrange
|
||||
|
||||
import comfy.ldm.modules.attention as attention
|
||||
from comfy.ldm.modules.diffusionmodules import openaimodel
|
||||
import comfy.model_management as model_management
|
||||
import comfy.samplers
|
||||
import comfy.sample
|
||||
import comfy.utils
|
||||
from comfy.controlnet import ControlBase
|
||||
import comfy.ops
|
||||
|
||||
from .context import ContextFuseMethod, ContextSchedules, get_context_weights, get_context_windows
|
||||
from .sample_settings import IterationOptions, SampleSettings, SeedNoiseGeneration, prepare_mask_ad
|
||||
from .utils_model import ModelTypeSD, wrap_function_to_inject_xformers_bug_info
|
||||
from .model_injection import InjectionParams, ModelPatcherAndInjector, MotionModelGroup, MotionModelPatcher
|
||||
from .motion_module_ad import AnimateDiffFormat, AnimateDiffInfo, AnimateDiffVersion, VanillaTemporalModule
|
||||
from .logger import logger
|
||||
|
||||
|
||||
##################################################################################
|
||||
######################################################################
|
||||
# Global variable to use to more conveniently hack variable access into samplers
|
||||
class AnimateDiffHelper_GlobalState:
|
||||
def __init__(self):
|
||||
self.motion_models: MotionModelGroup = None
|
||||
self.params: InjectionParams = None
|
||||
self.sample_settings: SampleSettings = None
|
||||
self.reset()
|
||||
|
||||
def initialize(self, model):
|
||||
# this function is to be run in sampling func
|
||||
if not self.initialized:
|
||||
self.initialized = True
|
||||
if self.motion_models is not None:
|
||||
self.motion_models.initialize_timesteps(model)
|
||||
if self.params.context_options is not None:
|
||||
self.params.context_options.initialize_timesteps(model)
|
||||
if self.sample_settings.custom_cfg is not None:
|
||||
self.sample_settings.custom_cfg.initialize_timesteps(model)
|
||||
|
||||
def reset(self):
|
||||
self.initialized = False
|
||||
self.start_step: int = 0
|
||||
self.last_step: int = 0
|
||||
self.current_step: int = 0
|
||||
self.total_steps: int = 0
|
||||
if self.motion_models is not None:
|
||||
del self.motion_models
|
||||
self.motion_models = None
|
||||
if self.params is not None:
|
||||
del self.params
|
||||
self.params = None
|
||||
if self.sample_settings is not None:
|
||||
del self.sample_settings
|
||||
self.sample_settings = None
|
||||
|
||||
def update_with_inject_params(self, params: InjectionParams):
|
||||
self.params = params
|
||||
|
||||
def is_using_sliding_context(self):
|
||||
return self.params is not None and self.params.is_using_sliding_context()
|
||||
|
||||
def create_exposed_params(self):
|
||||
# This dict will be exposed to be used by other extensions
|
||||
# DO NOT change any of the key names
|
||||
# or I will find you 👁.👁
|
||||
return {
|
||||
"full_length": self.params.full_length,
|
||||
"context_length": self.params.context_options.context_length,
|
||||
"sub_idxs": self.params.sub_idxs,
|
||||
}
|
||||
|
||||
ADGS = AnimateDiffHelper_GlobalState()
|
||||
######################################################################
|
||||
##################################################################################
|
||||
|
||||
|
||||
##################################################################################
|
||||
#### Code Injection ##################################################
|
||||
|
||||
# refer to forward_timestep_embed in comfy/ldm/modules/diffusionmodules/openaimodel.py
|
||||
def forward_timestep_embed_factory() -> Callable:
|
||||
def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, output_shape=None, time_context=None, num_video_frames=None, image_only_indicator=None):
|
||||
for layer in ts:
|
||||
if isinstance(layer, openaimodel.VideoResBlock):
|
||||
x = layer(x, emb, num_video_frames, image_only_indicator)
|
||||
elif isinstance(layer, openaimodel.TimestepBlock):
|
||||
x = layer(x, emb)
|
||||
elif isinstance(layer, VanillaTemporalModule):
|
||||
x = layer(x, context)
|
||||
elif isinstance(layer, attention.SpatialVideoTransformer):
|
||||
x = layer(x, context, time_context, num_video_frames, image_only_indicator, transformer_options)
|
||||
if "transformer_index" in transformer_options:
|
||||
transformer_options["transformer_index"] += 1
|
||||
if "current_index" in transformer_options: # keep this for backward compat, for now
|
||||
transformer_options["current_index"] += 1
|
||||
elif isinstance(layer, attention.SpatialTransformer):
|
||||
x = layer(x, context, transformer_options)
|
||||
if "transformer_index" in transformer_options:
|
||||
transformer_options["transformer_index"] += 1
|
||||
if "current_index" in transformer_options: # keep this for backward compat, for now
|
||||
transformer_options["current_index"] += 1
|
||||
elif isinstance(layer, openaimodel.Upsample):
|
||||
x = layer(x, output_shape=output_shape)
|
||||
else:
|
||||
x = layer(x)
|
||||
return x
|
||||
return forward_timestep_embed
|
||||
|
||||
|
||||
def unlimited_memory_required(*args, **kwargs):
|
||||
return 0
|
||||
|
||||
|
||||
def groupnorm_mm_factory(params: InjectionParams, manual_cast=False):
|
||||
def groupnorm_mm_forward(self, input: Tensor) -> Tensor:
|
||||
# axes_factor normalizes batch based on total conds and unconds passed in batch;
|
||||
# the conds and unconds per batch can change based on VRAM optimizations that may kick in
|
||||
if not params.is_using_sliding_context():
|
||||
batched_conds = input.size(0)//params.full_length
|
||||
else:
|
||||
batched_conds = input.size(0)//params.context_options.context_length
|
||||
|
||||
input = rearrange(input, "(b f) c h w -> b c f h w", b=batched_conds)
|
||||
if manual_cast:
|
||||
weight, bias = comfy.ops.cast_bias_weight(self, input)
|
||||
else:
|
||||
weight, bias = self.weight, self.bias
|
||||
input = group_norm(input, self.num_groups, weight, bias, self.eps)
|
||||
input = rearrange(input, "b c f h w -> (b f) c h w", b=batched_conds)
|
||||
return input
|
||||
return groupnorm_mm_forward
|
||||
|
||||
|
||||
def get_additional_models_factory(orig_get_additional_models: Callable, motion_models: MotionModelGroup):
|
||||
def get_additional_models_with_motion(*args, **kwargs):
|
||||
models, inference_memory = orig_get_additional_models(*args, **kwargs)
|
||||
if motion_models is not None:
|
||||
for motion_model in motion_models.models:
|
||||
models.append(motion_model)
|
||||
# TODO: account for inference memory as well?
|
||||
return models, inference_memory
|
||||
return get_additional_models_with_motion
|
||||
######################################################################
|
||||
##################################################################################
|
||||
|
||||
|
||||
def apply_params_to_motion_models(motion_models: MotionModelGroup, params: InjectionParams):
|
||||
params = params.clone()
|
||||
for context in params.context_options.contexts:
|
||||
if context.context_schedule == ContextSchedules.VIEW_AS_CONTEXT:
|
||||
context.context_length = params.full_length
|
||||
# TODO: check (and message) should be different based on use_on_equal_length setting
|
||||
if params.context_options.context_length:
|
||||
pass
|
||||
|
||||
allow_equal = params.context_options.use_on_equal_length
|
||||
if params.context_options.context_length:
|
||||
enough_latents = params.full_length >= params.context_options.context_length if allow_equal else params.full_length > params.context_options.context_length
|
||||
else:
|
||||
enough_latents = False
|
||||
if params.context_options.context_length and enough_latents:
|
||||
logger.info(f"Sliding context window activated - latents passed in ({params.full_length}) greater than context_length {params.context_options.context_length}.")
|
||||
else:
|
||||
logger.info(f"Regular AnimateDiff activated - latents passed in ({params.full_length}) less or equal to context_length {params.context_options.context_length}.")
|
||||
params.reset_context()
|
||||
if motion_models is not None:
|
||||
# if no context_length, treat video length as intended AD frame window
|
||||
if not params.context_options.context_length:
|
||||
for motion_model in motion_models.models:
|
||||
if not motion_model.model.is_length_valid_for_encoding_max_len(params.full_length):
|
||||
raise ValueError(f"Without a context window, AnimateDiff model {motion_model.model.mm_info.mm_name} has upper limit of {motion_model.model.encoding_max_len} frames, but received {params.full_length} latents.")
|
||||
motion_models.set_video_length(params.full_length, params.full_length)
|
||||
# otherwise, treat context_length as intended AD frame window
|
||||
else:
|
||||
for motion_model in motion_models.models:
|
||||
view_options = params.context_options.view_options
|
||||
context_length = view_options.context_length if view_options else params.context_options.context_length
|
||||
if not motion_model.model.is_length_valid_for_encoding_max_len(context_length):
|
||||
raise ValueError(f"AnimateDiff model {motion_model.model.mm_info.mm_name} has upper limit of {motion_model.model.encoding_max_len} frames for a context window, but received context length of {params.context_options.context_length}.")
|
||||
motion_models.set_video_length(params.context_options.context_length, params.full_length)
|
||||
# inject model
|
||||
module_str = "modules" if len(motion_models.models) > 1 else "module"
|
||||
logger.info(f"Using motion {module_str} {motion_models.get_name_string(show_version=True)}.")
|
||||
return params
|
||||
|
||||
|
||||
class FunctionInjectionHolder:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def inject_functions(self, model: ModelPatcherAndInjector, params: InjectionParams):
|
||||
# Save Original Functions
|
||||
self.orig_forward_timestep_embed = openaimodel.forward_timestep_embed # needed to account for VanillaTemporalModule
|
||||
self.orig_memory_required = model.model.memory_required # allows for "unlimited area hack" to prevent halving of conds/unconds
|
||||
self.orig_groupnorm_forward = torch.nn.GroupNorm.forward # used to normalize latents to remove "flickering" of colors/brightness between frames
|
||||
self.orig_groupnorm_manual_cast_forward = comfy.ops.manual_cast.GroupNorm.forward_comfy_cast_weights
|
||||
self.orig_sampling_function = comfy.samplers.sampling_function # used to support sliding context windows in samplers
|
||||
self.orig_prepare_mask = comfy.sample.prepare_mask
|
||||
self.orig_get_additional_models = comfy.sample.get_additional_models
|
||||
# Inject Functions
|
||||
openaimodel.forward_timestep_embed = forward_timestep_embed_factory()
|
||||
if params.unlimited_area_hack:
|
||||
model.model.memory_required = unlimited_memory_required
|
||||
if model.motion_models is not None:
|
||||
# only apply groupnorm hack if not [v3 or ([not Hotshot] and SD1.5 and v2 and apply_v2_properly)]
|
||||
info: AnimateDiffInfo = model.motion_models[0].model.mm_info
|
||||
if not (info.mm_version == AnimateDiffVersion.V3 or
|
||||
(info.mm_format not in [AnimateDiffFormat.HOTSHOTXL] and info.sd_type == ModelTypeSD.SD1_5 and info.mm_version == AnimateDiffVersion.V2 and params.apply_v2_properly)):
|
||||
torch.nn.GroupNorm.forward = groupnorm_mm_factory(params)
|
||||
comfy.ops.manual_cast.GroupNorm.forward_comfy_cast_weights = groupnorm_mm_factory(params, manual_cast=True)
|
||||
# if mps device (Apple Silicon), disable batched conds to avoid black images with groupnorm hack
|
||||
try:
|
||||
if model.load_device.type == "mps":
|
||||
model.model.memory_required = unlimited_memory_required
|
||||
except Exception:
|
||||
pass
|
||||
del info
|
||||
comfy.samplers.sampling_function = evolved_sampling_function
|
||||
comfy.sample.prepare_mask = prepare_mask_ad
|
||||
comfy.sample.get_additional_models = get_additional_models_factory(self.orig_get_additional_models, model.motion_models)
|
||||
|
||||
def restore_functions(self, model: ModelPatcherAndInjector):
|
||||
# Restoration
|
||||
try:
|
||||
model.model.memory_required = self.orig_memory_required
|
||||
openaimodel.forward_timestep_embed = self.orig_forward_timestep_embed
|
||||
torch.nn.GroupNorm.forward = self.orig_groupnorm_forward
|
||||
comfy.ops.manual_cast.GroupNorm.forward_comfy_cast_weights = self.orig_groupnorm_manual_cast_forward
|
||||
comfy.samplers.sampling_function = self.orig_sampling_function
|
||||
comfy.sample.prepare_mask = self.orig_prepare_mask
|
||||
comfy.sample.get_additional_models = self.orig_get_additional_models
|
||||
except AttributeError:
|
||||
logger.error("Encountered AttributeError while attempting to restore functions - likely, an error occured while trying " + \
|
||||
"to save original functions before injection, and a more specific error was thrown by ComfyUI.")
|
||||
|
||||
|
||||
def motion_sample_factory(orig_comfy_sample: Callable, is_custom: bool=False) -> Callable:
|
||||
def motion_sample(model: ModelPatcherAndInjector, noise: Tensor, *args, **kwargs):
|
||||
# check if model is intended for injecting
|
||||
if type(model) != ModelPatcherAndInjector:
|
||||
return orig_comfy_sample(model, noise, *args, **kwargs)
|
||||
# otherwise, injection time
|
||||
latents = None
|
||||
cached_latents = None
|
||||
cached_noise = None
|
||||
function_injections = FunctionInjectionHolder()
|
||||
try:
|
||||
if model.sample_settings.custom_cfg is not None:
|
||||
model = model.sample_settings.custom_cfg.patch_model(model)
|
||||
# clone params from model
|
||||
params = model.motion_injection_params.clone()
|
||||
# get amount of latents passed in, and store in params
|
||||
latents: Tensor = args[-1]
|
||||
params.full_length = latents.size(0)
|
||||
# reset global state
|
||||
ADGS.reset()
|
||||
|
||||
# apply custom noise, if needed
|
||||
disable_noise = kwargs.get("disable_noise") or False
|
||||
seed = kwargs["seed"]
|
||||
|
||||
# apply params to motion model
|
||||
params = apply_params_to_motion_models(model.motion_models, params)
|
||||
|
||||
# store and inject functions
|
||||
function_injections.inject_functions(model, params)
|
||||
|
||||
# prepare noise_extra_args for noise generation purposes
|
||||
noise_extra_args = {"disable_noise": disable_noise}
|
||||
params.set_noise_extra_args(noise_extra_args)
|
||||
# if noise is not disabled, do noise stuff
|
||||
if not disable_noise:
|
||||
noise = model.sample_settings.prepare_noise(seed, latents, noise, extra_args=noise_extra_args, force_create_noise=False)
|
||||
|
||||
# callback setup
|
||||
original_callback = kwargs.get("callback", None)
|
||||
def ad_callback(step, x0, x, total_steps):
|
||||
if original_callback is not None:
|
||||
original_callback(step, x0, x, total_steps)
|
||||
# update GLOBALSTATE for next iteration
|
||||
ADGS.current_step = ADGS.start_step + step + 1
|
||||
kwargs["callback"] = ad_callback
|
||||
ADGS.motion_models = model.motion_models
|
||||
ADGS.sample_settings = model.sample_settings
|
||||
|
||||
# apply adapt_denoise_steps
|
||||
args = list(args)
|
||||
if model.sample_settings.adapt_denoise_steps and not is_custom:
|
||||
# only applicable when denoise and steps are provided (from simple KSampler nodes)
|
||||
denoise = kwargs.get("denoise", None)
|
||||
steps = args[0]
|
||||
if denoise is not None and type(steps) == int:
|
||||
args[0] = max(int(denoise * steps), 1)
|
||||
|
||||
|
||||
iter_opts = IterationOptions()
|
||||
if model.sample_settings is not None:
|
||||
iter_opts = model.sample_settings.iteration_opts
|
||||
iter_opts.initialize(latents)
|
||||
# cache initial noise and latents, if needed
|
||||
if iter_opts.cache_init_latents:
|
||||
cached_latents = latents.clone()
|
||||
if iter_opts.cache_init_noise:
|
||||
cached_noise = noise.clone()
|
||||
# prepare iter opts preprocess kwargs, if needed
|
||||
iter_kwargs = {}
|
||||
if iter_opts.need_sampler:
|
||||
# -5 for sampler_name (not custom) and sampler (custom)
|
||||
model_management.load_model_gpu(model)
|
||||
if is_custom:
|
||||
iter_kwargs[IterationOptions.SAMPLER] = None #args[-5]
|
||||
else:
|
||||
iter_kwargs[IterationOptions.SAMPLER] = comfy.samplers.KSampler(
|
||||
model.model, steps=999, #steps=args[-7],
|
||||
device=model.current_device, sampler=args[-5],
|
||||
scheduler=args[-4], denoise=kwargs.get("denoise", None),
|
||||
model_options=model.model_options)
|
||||
|
||||
for curr_i in range(iter_opts.iterations):
|
||||
# handle GLOBALSTATE vars and step tally
|
||||
ADGS.update_with_inject_params(params)
|
||||
ADGS.start_step = kwargs.get("start_step") or 0
|
||||
ADGS.current_step = ADGS.start_step
|
||||
ADGS.last_step = kwargs.get("last_step") or 0
|
||||
if iter_opts.iterations > 1:
|
||||
logger.info(f"Iteration {curr_i+1}/{iter_opts.iterations}")
|
||||
# perform any iter_opts preprocessing on latents
|
||||
latents, noise = iter_opts.preprocess_latents(curr_i=curr_i, model=model, latents=latents, noise=noise,
|
||||
cached_latents=cached_latents, cached_noise=cached_noise,
|
||||
seed=seed,
|
||||
sample_settings=model.sample_settings, noise_extra_args=noise_extra_args,
|
||||
**iter_kwargs)
|
||||
args[-1] = latents
|
||||
|
||||
if model.motion_models is not None:
|
||||
model.motion_models.pre_run(model)
|
||||
if model.sample_settings is not None:
|
||||
model.sample_settings.pre_run(model)
|
||||
latents = wrap_function_to_inject_xformers_bug_info(orig_comfy_sample)(model, noise, *args, **kwargs)
|
||||
return latents
|
||||
finally:
|
||||
del latents
|
||||
del noise
|
||||
del cached_latents
|
||||
del cached_noise
|
||||
# reset global state
|
||||
ADGS.reset()
|
||||
# restore injected functions
|
||||
function_injections.restore_functions(model)
|
||||
del function_injections
|
||||
return motion_sample
|
||||
|
||||
|
||||
def evolved_sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options: dict={}, seed=None):
|
||||
ADGS.initialize(model)
|
||||
if ADGS.motion_models is not None:
|
||||
ADGS.motion_models.prepare_current_keyframe(t=timestep)
|
||||
if ADGS.params.context_options is not None:
|
||||
ADGS.params.context_options.prepare_current_context(t=timestep)
|
||||
if ADGS.sample_settings.custom_cfg is not None:
|
||||
ADGS.sample_settings.custom_cfg.prepare_current_keyframe(t=timestep)
|
||||
|
||||
# never use cfg1 optimization if using custom_cfg (since can have timesteps and such)
|
||||
if ADGS.sample_settings.custom_cfg is None and math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
|
||||
uncond_ = None
|
||||
else:
|
||||
uncond_ = uncond
|
||||
|
||||
# add AD/evolved-sampling params to model_options (transformer_options)
|
||||
model_options = model_options.copy()
|
||||
if "tranformer_options" not in model_options:
|
||||
model_options["tranformer_options"] = {}
|
||||
model_options["transformer_options"]["ad_params"] = ADGS.create_exposed_params()
|
||||
|
||||
if not ADGS.is_using_sliding_context():
|
||||
cond_pred, uncond_pred = comfy.samplers.calc_cond_uncond_batch(model, cond, uncond_, x, timestep, model_options)
|
||||
else:
|
||||
cond_pred, uncond_pred = sliding_calc_cond_uncond_batch(model, cond, uncond_, x, timestep, model_options)
|
||||
|
||||
if "sampler_cfg_function" in model_options:
|
||||
args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
|
||||
"cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
|
||||
cfg_result = x - model_options["sampler_cfg_function"](args)
|
||||
else:
|
||||
cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
|
||||
|
||||
for fn in model_options.get("sampler_post_cfg_function", []):
|
||||
args = {"denoised": cfg_result, "cond": cond, "uncond": uncond, "model": model, "uncond_denoised": uncond_pred, "cond_denoised": cond_pred,
|
||||
"sigma": timestep, "model_options": model_options, "input": x}
|
||||
cfg_result = fn(args)
|
||||
|
||||
return cfg_result
|
||||
|
||||
|
||||
# sliding_calc_cond_uncond_batch inspired by ashen's initial hack for 16-frame sliding context:
|
||||
# https://github.com/comfyanonymous/ComfyUI/compare/master...ashen-sensored:ComfyUI:master
|
||||
def sliding_calc_cond_uncond_batch(model, cond, uncond, x_in: Tensor, timestep, model_options):
|
||||
def prepare_control_objects(control: ControlBase, full_idxs: list[int]):
|
||||
if control.previous_controlnet is not None:
|
||||
prepare_control_objects(control.previous_controlnet, full_idxs)
|
||||
control.sub_idxs = full_idxs
|
||||
control.full_latent_length = ADGS.params.full_length
|
||||
control.context_length = ADGS.params.context_options.context_length
|
||||
|
||||
def get_resized_cond(cond_in, full_idxs) -> list:
|
||||
# reuse or resize cond items to match context requirements
|
||||
resized_cond = []
|
||||
# cond object is a list containing a dict - outer list is irrelevant, so just loop through it
|
||||
for actual_cond in cond_in:
|
||||
resized_actual_cond = actual_cond.copy()
|
||||
# now we are in the inner dict - "pooled_output" is a tensor, "control" is a ControlBase object, "model_conds" is dictionary
|
||||
for key in actual_cond:
|
||||
try:
|
||||
cond_item = actual_cond[key]
|
||||
if isinstance(cond_item, Tensor):
|
||||
# check that tensor is the expected length - x.size(0)
|
||||
if cond_item.size(0) == x_in.size(0):
|
||||
# if so, it's subsetting time - tell controls the expected indeces so they can handle them
|
||||
actual_cond_item = cond_item[full_idxs]
|
||||
resized_actual_cond[key] = actual_cond_item
|
||||
else:
|
||||
resized_actual_cond[key] = cond_item
|
||||
# look for control
|
||||
elif key == "control":
|
||||
control_item = cond_item
|
||||
if hasattr(control_item, "sub_idxs"):
|
||||
prepare_control_objects(control_item, full_idxs)
|
||||
else:
|
||||
raise ValueError(f"Control type {type(control_item).__name__} may not support required features for sliding context window; \
|
||||
use Control objects from Kosinkadink/ComfyUI-Advanced-ControlNet nodes, or make sure Advanced-ControlNet is updated.")
|
||||
resized_actual_cond[key] = control_item
|
||||
del control_item
|
||||
elif isinstance(cond_item, dict):
|
||||
new_cond_item = cond_item.copy()
|
||||
# when in dictionary, look for tensors and CONDCrossAttn [comfy/conds.py] (has cond attr that is a tensor)
|
||||
for cond_key, cond_value in new_cond_item.items():
|
||||
if isinstance(cond_value, Tensor):
|
||||
if cond_value.size(0) == x_in.size(0):
|
||||
new_cond_item[cond_key] = cond_value[full_idxs]
|
||||
# if has cond that is a Tensor, check if needs to be subset
|
||||
elif hasattr(cond_value, "cond") and isinstance(cond_value.cond, Tensor):
|
||||
if cond_value.cond.size(0) == x_in.size(0):
|
||||
new_cond_item[cond_key] = cond_value._copy_with(cond_value.cond[full_idxs])
|
||||
resized_actual_cond[key] = new_cond_item
|
||||
else:
|
||||
resized_actual_cond[key] = cond_item
|
||||
finally:
|
||||
del cond_item # just in case to prevent VRAM issues
|
||||
resized_cond.append(resized_actual_cond)
|
||||
return resized_cond
|
||||
|
||||
# get context windows
|
||||
ADGS.params.context_options.step = ADGS.current_step
|
||||
context_windows = get_context_windows(ADGS.params.full_length, ADGS.params.context_options)
|
||||
# figure out how input is split
|
||||
batched_conds = x_in.size(0)//ADGS.params.full_length
|
||||
|
||||
if ADGS.motion_models is not None:
|
||||
ADGS.motion_models.set_view_options(ADGS.params.context_options.view_options)
|
||||
|
||||
# prepare final cond, uncond, and out_count
|
||||
cond_final = torch.zeros_like(x_in)
|
||||
uncond_final = torch.zeros_like(x_in)
|
||||
out_count_final = torch.zeros((x_in.shape[0], 1, 1, 1), device=x_in.device)
|
||||
bias_final = [0.0] * x_in.shape[0]
|
||||
|
||||
# perform calc_cond_uncond_batch per context window
|
||||
for ctx_idxs in context_windows:
|
||||
ADGS.params.sub_idxs = ctx_idxs
|
||||
if ADGS.motion_models is not None:
|
||||
ADGS.motion_models.set_sub_idxs(ctx_idxs)
|
||||
ADGS.motion_models.set_video_length(len(ctx_idxs), ADGS.params.full_length)
|
||||
# update exposed params
|
||||
model_options["transformer_options"]["ad_params"]["sub_idxs"] = ctx_idxs
|
||||
model_options["transformer_options"]["ad_params"]["context_length"] = len(ctx_idxs)
|
||||
# account for all portions of input frames
|
||||
full_idxs = []
|
||||
for n in range(batched_conds):
|
||||
for ind in ctx_idxs:
|
||||
full_idxs.append((ADGS.params.full_length*n)+ind)
|
||||
# get subsections of x, timestep, cond, uncond, cond_concat
|
||||
sub_x = x_in[full_idxs]
|
||||
sub_timestep = timestep[full_idxs]
|
||||
sub_cond = get_resized_cond(cond, full_idxs) if cond is not None else None
|
||||
sub_uncond = get_resized_cond(uncond, full_idxs) if uncond is not None else None
|
||||
|
||||
sub_cond_out, sub_uncond_out = comfy.samplers.calc_cond_uncond_batch(model, sub_cond, sub_uncond, sub_x, sub_timestep, model_options)
|
||||
|
||||
if ADGS.params.context_options.fuse_method == ContextFuseMethod.RELATIVE:
|
||||
full_length = ADGS.params.full_length
|
||||
for pos, idx in enumerate(ctx_idxs):
|
||||
# bias is the influence of a specific index in relation to the whole context window
|
||||
bias = 1 - abs(idx - (ctx_idxs[0] + ctx_idxs[-1]) / 2) / ((ctx_idxs[-1] - ctx_idxs[0] + 1e-2) / 2)
|
||||
bias = max(1e-2, bias)
|
||||
# take weighted average relative to total bias of current idx
|
||||
# and account for batched_conds
|
||||
for n in range(batched_conds):
|
||||
bias_total = bias_final[(full_length*n)+idx]
|
||||
prev_weight = (bias_total / (bias_total + bias))
|
||||
new_weight = (bias / (bias_total + bias))
|
||||
cond_final[(full_length*n)+idx] = cond_final[(full_length*n)+idx] * prev_weight + sub_cond_out[(full_length*n)+pos] * new_weight
|
||||
uncond_final[(full_length*n)+idx] = uncond_final[(full_length*n)+idx] * prev_weight + sub_uncond_out[(full_length*n)+pos] * new_weight
|
||||
bias_final[(full_length*n)+idx] = bias_total + bias
|
||||
else:
|
||||
# add conds and counts based on weights of fuse method
|
||||
weights = get_context_weights(len(ctx_idxs), ADGS.params.context_options.fuse_method) * batched_conds
|
||||
weights_tensor = torch.Tensor(weights).to(device=x_in.device).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
|
||||
cond_final[full_idxs] += sub_cond_out * weights_tensor
|
||||
uncond_final[full_idxs] += sub_uncond_out * weights_tensor
|
||||
out_count_final[full_idxs] += weights_tensor
|
||||
|
||||
if ADGS.params.context_options.fuse_method == ContextFuseMethod.RELATIVE:
|
||||
# already normalized, so return as is
|
||||
del out_count_final
|
||||
return cond_final, uncond_final
|
||||
else:
|
||||
# normalize cond and uncond via division by context usage counts
|
||||
cond_final /= out_count_final
|
||||
uncond_final /= out_count_final
|
||||
del out_count_final
|
||||
return cond_final, uncond_final
|
||||
+417
@@ -0,0 +1,417 @@
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
from typing import Callable, Union
|
||||
from collections.abc import Iterable
|
||||
from time import time
|
||||
import copy
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
import folder_paths
|
||||
from comfy.model_base import SD21UNCLIP, SDXL, BaseModel, SDXLRefiner, SVD_img2vid, model_sampling, ModelType
|
||||
from comfy.model_management import xformers_enabled
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
import comfy.model_sampling
|
||||
import comfy_extras.nodes_model_advanced
|
||||
|
||||
|
||||
BIGMIN = -(2**53-1)
|
||||
BIGMAX = (2**53-1)
|
||||
|
||||
|
||||
class ModelSamplingConfig:
|
||||
def __init__(self, beta_schedule: str, linear_start: float=None, linear_end: float=None):
|
||||
self.sampling_settings = {"beta_schedule": beta_schedule}
|
||||
if linear_start is not None:
|
||||
self.sampling_settings["linear_start"] = linear_start
|
||||
if linear_end is not None:
|
||||
self.sampling_settings["linear_end"] = linear_end
|
||||
self.beta_schedule = beta_schedule # keeping this for backwards compatibility
|
||||
|
||||
|
||||
class ModelSamplingType:
|
||||
EPS = "eps"
|
||||
V_PREDICTION = "v_prediction"
|
||||
LCM = "lcm"
|
||||
|
||||
_NON_LCM_LIST = [EPS, V_PREDICTION]
|
||||
_FULL_LIST = [EPS, V_PREDICTION, LCM]
|
||||
|
||||
MAP = {
|
||||
EPS: ModelType.EPS,
|
||||
V_PREDICTION: ModelType.V_PREDICTION,
|
||||
LCM: comfy_extras.nodes_model_advanced.LCM,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_alias(cls, alias: str):
|
||||
return cls.MAP[alias]
|
||||
|
||||
|
||||
def factory_model_sampling_discrete_distilled(original_timesteps=50):
|
||||
class ModelSamplingDiscreteDistilledEvolved(comfy_extras.nodes_model_advanced.ModelSamplingDiscreteDistilled):
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.original_timesteps = original_timesteps # normal LCM has 50
|
||||
super().__init__(*args, **kwargs)
|
||||
return ModelSamplingDiscreteDistilledEvolved
|
||||
|
||||
|
||||
# based on code in comfy_extras/nodes_model_advanced.py
|
||||
def evolved_model_sampling(model_config: ModelSamplingConfig, model_type: ModelType, alias: str, original_timesteps: int=None):
|
||||
# if LCM, need to handle manually
|
||||
if BetaSchedules.is_lcm(alias) or original_timesteps is not None:
|
||||
sampling_type = comfy_extras.nodes_model_advanced.LCM
|
||||
if original_timesteps is not None:
|
||||
sampling_base = factory_model_sampling_discrete_distilled(original_timesteps=original_timesteps)
|
||||
elif alias == BetaSchedules.LCM_100:
|
||||
sampling_base = factory_model_sampling_discrete_distilled(original_timesteps=100)
|
||||
elif alias == BetaSchedules.LCM_25:
|
||||
sampling_base = factory_model_sampling_discrete_distilled(original_timesteps=25)
|
||||
else:
|
||||
sampling_base = comfy_extras.nodes_model_advanced.ModelSamplingDiscreteDistilled
|
||||
class ModelSamplingAdvancedEvolved(sampling_base, sampling_type):
|
||||
pass
|
||||
# NOTE: if I want to support zsnr, this is where I would add that code
|
||||
return ModelSamplingAdvancedEvolved(model_config)
|
||||
# otherwise, use vanilla model_sampling function
|
||||
return model_sampling(model_config, model_type)
|
||||
|
||||
|
||||
class BetaSchedules:
|
||||
AUTOSELECT = "autoselect"
|
||||
SQRT_LINEAR = "sqrt_linear (AnimateDiff)"
|
||||
LINEAR_ADXL = "linear (AnimateDiff-SDXL)"
|
||||
LINEAR = "linear (HotshotXL/default)"
|
||||
AVG_LINEAR_SQRT_LINEAR = "avg(sqrt_linear,linear)"
|
||||
LCM_AVG_LINEAR_SQRT_LINEAR = "lcm avg(sqrt_linear,linear)"
|
||||
LCM = "lcm"
|
||||
LCM_100 = "lcm[100_ots]"
|
||||
LCM_25 = "lcm[25_ots]"
|
||||
LCM_SQRT_LINEAR = "lcm >> sqrt_linear"
|
||||
USE_EXISTING = "use existing"
|
||||
SQRT = "sqrt"
|
||||
COSINE = "cosine"
|
||||
SQUAREDCOS_CAP_V2 = "squaredcos_cap_v2"
|
||||
RAW_LINEAR = "linear"
|
||||
RAW_SQRT_LINEAR = "sqrt_linear"
|
||||
|
||||
RAW_BETA_SCHEDULE_LIST = [RAW_LINEAR, RAW_SQRT_LINEAR, SQRT, COSINE, SQUAREDCOS_CAP_V2]
|
||||
|
||||
ALIAS_LCM_LIST = [LCM, LCM_100, LCM_25, LCM_SQRT_LINEAR]
|
||||
|
||||
ALIAS_ACTIVE_LIST = [SQRT_LINEAR, LINEAR_ADXL, LINEAR, AVG_LINEAR_SQRT_LINEAR, LCM_AVG_LINEAR_SQRT_LINEAR, LCM, LCM_100, LCM_SQRT_LINEAR, # LCM_25 is purposely omitted
|
||||
SQRT, COSINE, SQUAREDCOS_CAP_V2]
|
||||
|
||||
ALIAS_LIST = [AUTOSELECT, USE_EXISTING] + ALIAS_ACTIVE_LIST
|
||||
|
||||
|
||||
|
||||
ALIAS_MAP = {
|
||||
SQRT_LINEAR: "sqrt_linear",
|
||||
LINEAR_ADXL: "linear", # also linear, but has different linear_end (0.020)
|
||||
LINEAR: "linear",
|
||||
LCM_100: "linear", # distilled, 100 original timesteps
|
||||
LCM_25: "linear", # distilled, 25 original timesteps
|
||||
LCM: "linear", # distilled
|
||||
LCM_SQRT_LINEAR: "sqrt_linear", # distilled, sqrt_linear
|
||||
SQRT: "sqrt",
|
||||
COSINE: "cosine",
|
||||
SQUAREDCOS_CAP_V2: "squaredcos_cap_v2",
|
||||
RAW_LINEAR: "linear",
|
||||
RAW_SQRT_LINEAR: "sqrt_linear"
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def is_lcm(cls, alias: str):
|
||||
return alias in cls.ALIAS_LCM_LIST
|
||||
|
||||
@classmethod
|
||||
def to_name(cls, alias: str):
|
||||
return cls.ALIAS_MAP[alias]
|
||||
|
||||
@classmethod
|
||||
def to_config(cls, alias: str) -> ModelSamplingConfig:
|
||||
linear_start = None
|
||||
linear_end = None
|
||||
if alias == cls.LINEAR_ADXL:
|
||||
# uses linear_end=0.020
|
||||
linear_end = 0.020
|
||||
return ModelSamplingConfig(cls.to_name(alias), linear_start=linear_start, linear_end=linear_end)
|
||||
|
||||
@classmethod
|
||||
def _to_model_sampling(cls, alias: str, model_type: ModelType, config_override: ModelSamplingConfig=None, original_timesteps: int=None):
|
||||
if alias == cls.USE_EXISTING:
|
||||
return None
|
||||
elif config_override != None:
|
||||
return evolved_model_sampling(config_override, model_type=model_type, alias=alias, original_timesteps=original_timesteps)
|
||||
elif alias == cls.AVG_LINEAR_SQRT_LINEAR:
|
||||
ms_linear = evolved_model_sampling(cls.to_config(cls.LINEAR), model_type=model_type, alias=cls.LINEAR)
|
||||
ms_sqrt_linear = evolved_model_sampling(cls.to_config(cls.SQRT_LINEAR), model_type=model_type, alias=cls.SQRT_LINEAR)
|
||||
avg_sigmas = (ms_linear.sigmas + ms_sqrt_linear.sigmas) / 2
|
||||
ms_linear.set_sigmas(avg_sigmas)
|
||||
return ms_linear
|
||||
elif alias == cls.LCM_AVG_LINEAR_SQRT_LINEAR:
|
||||
ms_linear = evolved_model_sampling(cls.to_config(cls.LCM), model_type=model_type, alias=cls.LCM)
|
||||
ms_sqrt_linear = evolved_model_sampling(cls.to_config(cls.LCM_SQRT_LINEAR), model_type=model_type, alias=cls.LCM_SQRT_LINEAR)
|
||||
avg_sigmas = (ms_linear.sigmas + ms_sqrt_linear.sigmas) / 2
|
||||
ms_linear.set_sigmas(avg_sigmas)
|
||||
return ms_linear
|
||||
# average out the sigmas
|
||||
ms_obj = evolved_model_sampling(cls.to_config(alias), model_type=model_type, alias=alias, original_timesteps=original_timesteps)
|
||||
return ms_obj
|
||||
|
||||
@classmethod
|
||||
def to_model_sampling(cls, alias: str, model: ModelPatcher):
|
||||
return cls._to_model_sampling(alias=alias, model_type=model.model.model_type)
|
||||
|
||||
@staticmethod
|
||||
def get_alias_list_with_first_element(first_element: str):
|
||||
new_list = BetaSchedules.ALIAS_LIST.copy()
|
||||
element_index = new_list.index(first_element)
|
||||
new_list[0], new_list[element_index] = new_list[element_index], new_list[0]
|
||||
return new_list
|
||||
|
||||
|
||||
class SigmaSchedule:
|
||||
def __init__(self, model_sampling: comfy.model_sampling.ModelSamplingDiscrete, model_type: ModelType):
|
||||
self.model_sampling = model_sampling
|
||||
#self.config = config
|
||||
self.model_type = model_type
|
||||
self.original_timesteps = getattr(self.model_sampling, "original_timesteps", None)
|
||||
|
||||
def is_lcm(self):
|
||||
return self.original_timesteps is not None
|
||||
|
||||
def total_sigmas(self):
|
||||
return len(self.model_sampling.sigmas)
|
||||
|
||||
def clone(self) -> 'SigmaSchedule':
|
||||
new_model_sampling = copy.deepcopy(self.model_sampling)
|
||||
#new_config = copy.deepcopy(self.config)
|
||||
return SigmaSchedule(model_sampling=new_model_sampling, model_type=self.model_type)
|
||||
|
||||
# def clone(self):
|
||||
# pass
|
||||
|
||||
@staticmethod
|
||||
def apply_zsnr(new_model_sampling: comfy.model_sampling.ModelSamplingDiscrete):
|
||||
new_model_sampling.set_sigmas(comfy_extras.nodes_model_advanced.rescale_zero_terminal_snr_sigmas(new_model_sampling.sigmas))
|
||||
|
||||
# def get_lcmified(self, original_timesteps=50, zsnr=False) -> 'SigmaSchedule':
|
||||
# new_model_sampling = evolved_model_sampling(model_config=self.config, model_type=self.model_type, alias=None, original_timesteps=original_timesteps)
|
||||
# if zsnr:
|
||||
# new_model_sampling.set_sigmas(comfy_extras.nodes_model_advanced.rescale_zero_terminal_snr_sigmas(new_model_sampling.sigmas))
|
||||
# return SigmaSchedule(model_sampling=new_model_sampling, config=self.config, model_type=self.model_type, is_lcm=True)
|
||||
|
||||
|
||||
class InterpolationMethod:
|
||||
LINEAR = "linear"
|
||||
EASE_IN = "ease_in"
|
||||
EASE_OUT = "ease_out"
|
||||
EASE_IN_OUT = "ease_in_out"
|
||||
|
||||
_LIST = [LINEAR, EASE_IN, EASE_OUT, EASE_IN_OUT]
|
||||
|
||||
@classmethod
|
||||
def get_weights(cls, num_from: float, num_to: float, length: int, method: str, reverse=False):
|
||||
diff = num_to - num_from
|
||||
if method == cls.LINEAR:
|
||||
weights = torch.linspace(num_from, num_to, length)
|
||||
elif method == cls.EASE_IN:
|
||||
index = torch.linspace(0, 1, length)
|
||||
weights = diff * np.power(index, 2) + num_from
|
||||
elif method == cls.EASE_OUT:
|
||||
index = torch.linspace(0, 1, length)
|
||||
weights = diff * (1 - np.power(1 - index, 2)) + num_from
|
||||
elif method == cls.EASE_IN_OUT:
|
||||
index = torch.linspace(0, 1, length)
|
||||
weights = diff * ((1 - np.cos(index * np.pi)) / 2) + num_from
|
||||
else:
|
||||
raise ValueError(f"Unrecognized interpolation method '{method}'.")
|
||||
if reverse:
|
||||
weights = weights.flip(dims=(0,))
|
||||
return weights
|
||||
|
||||
|
||||
class Folders:
|
||||
ANIMATEDIFF_MODELS = "animatediff_models"
|
||||
MOTION_LORA = "animatediff_motion_lora"
|
||||
VIDEO_FORMATS = "animatediff_video_formats"
|
||||
|
||||
|
||||
def add_extension_to_folder_path(folder_name: str, extensions: Union[str, list[str]]):
|
||||
if folder_name in folder_paths.folder_names_and_paths:
|
||||
if isinstance(extensions, str):
|
||||
folder_paths.folder_names_and_paths[folder_name][1].add(extensions)
|
||||
elif isinstance(extensions, Iterable):
|
||||
for ext in extensions:
|
||||
folder_paths.folder_names_and_paths[folder_name][1].add(ext)
|
||||
|
||||
|
||||
def try_mkdir(full_path: str):
|
||||
try:
|
||||
Path(full_path).mkdir()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# register motion models folder(s)
|
||||
folder_paths.add_model_folder_path(Folders.ANIMATEDIFF_MODELS, str(Path(__file__).parent.parent / "models"))
|
||||
folder_paths.add_model_folder_path(Folders.ANIMATEDIFF_MODELS, str(Path(folder_paths.models_dir) / Folders.ANIMATEDIFF_MODELS))
|
||||
add_extension_to_folder_path(Folders.ANIMATEDIFF_MODELS, folder_paths.supported_pt_extensions)
|
||||
try_mkdir(str(Path(folder_paths.models_dir) / Folders.ANIMATEDIFF_MODELS))
|
||||
|
||||
# register motion LoRA folder(s)
|
||||
folder_paths.add_model_folder_path(Folders.MOTION_LORA, str(Path(__file__).parent.parent / "motion_lora"))
|
||||
folder_paths.add_model_folder_path(Folders.MOTION_LORA, str(Path(folder_paths.models_dir) / Folders.MOTION_LORA))
|
||||
add_extension_to_folder_path(Folders.MOTION_LORA, folder_paths.supported_pt_extensions)
|
||||
try_mkdir(str(Path(folder_paths.models_dir) / Folders.MOTION_LORA))
|
||||
|
||||
# register video_formats folder
|
||||
folder_paths.add_model_folder_path(Folders.VIDEO_FORMATS, str(Path(__file__).parent.parent / "video_formats"))
|
||||
add_extension_to_folder_path(Folders.VIDEO_FORMATS, ".json")
|
||||
|
||||
|
||||
def get_available_motion_models():
|
||||
return folder_paths.get_filename_list(Folders.ANIMATEDIFF_MODELS)
|
||||
|
||||
|
||||
def get_motion_model_path(model_name: str):
|
||||
return folder_paths.get_full_path(Folders.ANIMATEDIFF_MODELS, model_name)
|
||||
|
||||
|
||||
def get_available_motion_loras():
|
||||
return folder_paths.get_filename_list(Folders.MOTION_LORA)
|
||||
|
||||
|
||||
def get_motion_lora_path(lora_name: str):
|
||||
return folder_paths.get_full_path(Folders.MOTION_LORA, lora_name)
|
||||
|
||||
|
||||
# modified from https://stackoverflow.com/questions/22058048/hashing-a-file-in-python
|
||||
def calculate_file_hash(filename: str, hash_every_n: int = 50):
|
||||
h = hashlib.sha256()
|
||||
b = bytearray(1024*1024)
|
||||
mv = memoryview(b)
|
||||
with open(filename, 'rb', buffering=0) as f:
|
||||
i = 0
|
||||
# don't hash entire file, only portions of it
|
||||
while n := f.readinto(mv):
|
||||
if i%hash_every_n == 0:
|
||||
h.update(mv[:n])
|
||||
i += 1
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def calculate_model_hash(model: ModelPatcher):
|
||||
unet = model.model.diff
|
||||
t = unet.input_blocks[1]
|
||||
m = hashlib.sha256()
|
||||
for buf in t.buffers():
|
||||
m.update(buf.cpu().numpy().view(np.uint8))
|
||||
return m.hexdigest()
|
||||
|
||||
|
||||
class ModelTypeSD:
|
||||
SD1_5 = "SD1.5"
|
||||
SD2_1 = "SD2.1"
|
||||
SDXL = "SDXL"
|
||||
SDXL_REFINER = "SDXL_Refiner"
|
||||
SVD = "SVD"
|
||||
|
||||
|
||||
def get_sd_model_type(model: ModelPatcher) -> str:
|
||||
if model is None:
|
||||
return None
|
||||
elif type(model.model) == BaseModel:
|
||||
return ModelTypeSD.SD1_5
|
||||
elif type(model.model) == SDXL:
|
||||
return ModelTypeSD.SDXL
|
||||
elif type(model.model) == SD21UNCLIP:
|
||||
return ModelTypeSD.SD2_1
|
||||
elif type(model.model) == SDXLRefiner:
|
||||
return ModelTypeSD.SDXL_REFINER
|
||||
elif type(model.model) == SVD_img2vid:
|
||||
return ModelTypeSD.SVD
|
||||
else:
|
||||
return str(type(model.model).__name__)
|
||||
|
||||
def is_checkpoint_sd1_5(model: ModelPatcher):
|
||||
return False if model is None else type(model.model) == BaseModel
|
||||
|
||||
def is_checkpoint_sdxl(model: ModelPatcher):
|
||||
return False if model is None else type(model.model) == SDXL
|
||||
|
||||
|
||||
def raise_if_not_checkpoint_sd1_5(model: ModelPatcher):
|
||||
if not is_checkpoint_sd1_5(model):
|
||||
raise ValueError(f"For AnimateDiff, SD Checkpoint (model) is expected to be SD1.5-based (BaseModel), but was: {type(model.model).__name__}")
|
||||
|
||||
|
||||
# TODO: remove this filth when xformers bug gets fixed in future xformers version
|
||||
def wrap_function_to_inject_xformers_bug_info(function_to_wrap: Callable) -> Callable:
|
||||
if not xformers_enabled:
|
||||
return function_to_wrap
|
||||
else:
|
||||
def wrapped_function(*args, **kwargs):
|
||||
try:
|
||||
return function_to_wrap(*args, **kwargs)
|
||||
except RuntimeError as e:
|
||||
if str(e).startswith("CUDA error: invalid configuration argument"):
|
||||
raise RuntimeError(f"An xformers bug was encountered in AnimateDiff - this is unexpected, \
|
||||
report this to Kosinkadink/ComfyUI-AnimateDiff-Evolved repo as an issue, \
|
||||
and a workaround for now is to run ComfyUI with the --disable-xformers argument.")
|
||||
raise
|
||||
return wrapped_function
|
||||
|
||||
|
||||
class Timer(object):
|
||||
__slots__ = ("start_time", "end_time")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.start_time = 0.0
|
||||
self.end_time = 0.0
|
||||
|
||||
def start(self) -> None:
|
||||
self.start_time = time()
|
||||
|
||||
def update(self) -> None:
|
||||
self.start()
|
||||
|
||||
def stop(self) -> float:
|
||||
self.end_time = time()
|
||||
return self.get_time_diff()
|
||||
|
||||
def get_time_diff(self) -> float:
|
||||
return self.end_time - self.start_time
|
||||
|
||||
def get_time_current(self) -> float:
|
||||
return time() - self.start_time
|
||||
|
||||
|
||||
# TODO: possibly add configuration file in future when needed?
|
||||
# # Load config settings
|
||||
# ADE_DIR = Path(__file__).parent.parent
|
||||
# ADE_CONFIG_FILE = ADE_DIR / "ade_config.json"
|
||||
|
||||
# class ADE_Settings:
|
||||
# USE_XFORMERS_IN_VERSATILE_ATTENTION = "use_xformers_in_VersatileAttention"
|
||||
|
||||
# # Create ADE config if not present
|
||||
# ABS_CONFIG = {
|
||||
# ADE_Settings.USE_XFORMERS_IN_VERSATILE_ATTENTION: True
|
||||
# }
|
||||
# if not ADE_CONFIG_FILE.exists():
|
||||
# with ADE_CONFIG_FILE.open("w") as f:
|
||||
# json.dumps(ABS_CONFIG, indent=4)
|
||||
# # otherwise, load it and use values
|
||||
# else:
|
||||
# loaded_values: dict = None
|
||||
# with ADE_CONFIG_FILE.open("r") as f:
|
||||
# loaded_values = json.load(f)
|
||||
# if loaded_values is not None:
|
||||
# for key, value in loaded_values.items():
|
||||
# if key in ABS_CONFIG:
|
||||
# ABS_CONFIG[key] = value
|
||||
+230
@@ -0,0 +1,230 @@
|
||||
from typing import Union
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor, nn
|
||||
|
||||
import comfy.model_management as model_management
|
||||
import comfy.ops
|
||||
import comfy.utils
|
||||
from comfy.cli_args import args
|
||||
from comfy.ldm.modules.attention import attention_basic, attention_pytorch, attention_split, attention_sub_quad, default
|
||||
|
||||
from .logger import logger
|
||||
|
||||
|
||||
# until xformers bug is fixed, do not use xformers for VersatileAttention! TODO: change this when fix is out
|
||||
# logic for choosing optimized_attention method taken from comfy/ldm/modules/attention.py
|
||||
optimized_attention_mm = attention_basic
|
||||
if model_management.xformers_enabled():
|
||||
pass
|
||||
#optimized_attention_mm = attention_xformers
|
||||
if model_management.pytorch_attention_enabled():
|
||||
optimized_attention_mm = attention_pytorch
|
||||
else:
|
||||
if args.use_split_cross_attention:
|
||||
optimized_attention_mm = attention_split
|
||||
else:
|
||||
optimized_attention_mm = attention_sub_quad
|
||||
|
||||
|
||||
class CrossAttentionMM(nn.Module):
|
||||
def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None,
|
||||
operations=comfy.ops.disable_weight_init):
|
||||
super().__init__()
|
||||
inner_dim = dim_head * heads
|
||||
context_dim = default(context_dim, query_dim)
|
||||
|
||||
self.heads = heads
|
||||
self.dim_head = dim_head
|
||||
self.scale = None
|
||||
self.default_scale = dim_head ** -0.5
|
||||
|
||||
self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device)
|
||||
self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
|
||||
self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
|
||||
|
||||
self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout))
|
||||
|
||||
def forward(self, x, context=None, value=None, mask=None, scale_mask=None):
|
||||
q = self.to_q(x)
|
||||
context = default(context, x)
|
||||
k: Tensor = self.to_k(context)
|
||||
if value is not None:
|
||||
v = self.to_v(value)
|
||||
del value
|
||||
else:
|
||||
v = self.to_v(context)
|
||||
|
||||
# apply custom scale by multiplying k by scale factor
|
||||
if self.scale is not None:
|
||||
k *= self.scale
|
||||
|
||||
# apply scale mask, if present
|
||||
if scale_mask is not None:
|
||||
k *= scale_mask
|
||||
|
||||
out = optimized_attention_mm(q, k, v, self.heads, mask)
|
||||
return self.to_out(out)
|
||||
|
||||
# TODO: set up comfy.ops style classes for groupnorm and other functions
|
||||
class GroupNormAD(torch.nn.GroupNorm):
|
||||
def __init__(self, num_groups: int, num_channels: int, eps: float = 1e-5, affine: bool = True,
|
||||
device=None, dtype=None) -> None:
|
||||
super().__init__(num_groups=num_groups, num_channels=num_channels, eps=eps, affine=affine, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return F.group_norm(
|
||||
input, self.num_groups, self.weight, self.bias, self.eps)
|
||||
|
||||
|
||||
# applies min-max normalization, from:
|
||||
# https://stackoverflow.com/questions/68791508/min-max-normalization-of-a-tensor-in-pytorch
|
||||
def normalize_min_max(x: Tensor, new_min = 0.0, new_max = 1.0):
|
||||
return linear_conversion(x, x_min=x.min(), x_max=x.max(), new_min=new_min, new_max=new_max)
|
||||
|
||||
|
||||
def linear_conversion(x, x_min=0.0, x_max=1.0, new_min=0.0, new_max=1.0):
|
||||
x_min = float(x_min)
|
||||
x_max = float(x_max)
|
||||
new_min = float(new_min)
|
||||
new_max = float(new_max)
|
||||
return (((x - x_min)/(x_max - x_min)) * (new_max - new_min)) + new_min
|
||||
|
||||
|
||||
# adapted from comfy/sample.py
|
||||
def prepare_mask_batch(mask: Tensor, shape: Tensor, multiplier: int=1, match_dim1=False):
|
||||
mask = mask.clone()
|
||||
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[2]*multiplier, shape[3]*multiplier), mode="bilinear")
|
||||
if match_dim1:
|
||||
mask = torch.cat([mask] * shape[1], dim=1)
|
||||
return mask
|
||||
|
||||
|
||||
def extend_to_batch_size(tensor: Tensor, batch_size: int):
|
||||
if tensor.shape[0] > batch_size:
|
||||
return tensor[:batch_size]
|
||||
elif tensor.shape[0] < batch_size:
|
||||
remainder = batch_size-tensor.shape[0]
|
||||
return torch.cat([tensor] + [tensor[-1:]]*remainder, dim=0)
|
||||
return tensor
|
||||
|
||||
|
||||
def get_sorted_list_via_attr(objects: list, attr: str) -> list:
|
||||
if not objects:
|
||||
return objects
|
||||
elif len(objects) <= 1:
|
||||
return [x for x in objects]
|
||||
# now that we know we have to sort, do it following these rules:
|
||||
# a) if objects have same value of attribute, maintain their relative order
|
||||
# b) perform sorting of the groups of objects with same attributes
|
||||
unique_attrs = {}
|
||||
for o in objects:
|
||||
val_attr = getattr(o, attr)
|
||||
attr_list = unique_attrs.get(val_attr, list())
|
||||
attr_list.append(o)
|
||||
if val_attr not in unique_attrs:
|
||||
unique_attrs[val_attr] = attr_list
|
||||
# now that we have the unique attr values grouped together in relative order, sort them by key
|
||||
sorted_attrs = dict(sorted(unique_attrs.items()))
|
||||
# now flatten out the dict into a list to return
|
||||
sorted_list = []
|
||||
for object_list in sorted_attrs.values():
|
||||
sorted_list.extend(object_list)
|
||||
return sorted_list
|
||||
|
||||
|
||||
class MotionCompatibilityError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
def get_combined_multival(multivalA: Union[float, Tensor], multivalB: Union[float, Tensor]) -> Union[float, Tensor]:
|
||||
# if one is None, use the other
|
||||
if multivalA == None:
|
||||
return multivalB
|
||||
elif multivalB == None:
|
||||
return multivalA
|
||||
# both have a value - combine them based on type
|
||||
# if both are Tensors, make dims match before multiplying
|
||||
if type(multivalA) == Tensor and type(multivalB) == Tensor:
|
||||
areaA = multivalA.shape[1]*multivalA.shape[2]
|
||||
areaB = multivalB.shape[1]*multivalB.shape[2]
|
||||
# match height/width to mask with larger area
|
||||
leader,follower = (multivalA,multivalB) if areaA >= areaB else (multivalB,multivalA)
|
||||
batch_size = multivalA.shape[0] if multivalA.shape[0] >= multivalB.shape[0] else multivalB.shape[0]
|
||||
# make follower same dimensions as leader
|
||||
follower = torch.unsqueeze(follower, 1)
|
||||
follower = comfy.utils.common_upscale(follower, leader.shape[2], leader.shape[1], "bilinear", "center")
|
||||
follower = torch.squeeze(follower, 1)
|
||||
# make sure batch size will match
|
||||
leader = extend_to_batch_size(leader, batch_size)
|
||||
follower = extend_to_batch_size(follower, batch_size)
|
||||
return leader * follower
|
||||
# otherwise, just multiply them together - one of them is a float
|
||||
return multivalA * multivalB
|
||||
|
||||
|
||||
class ADKeyframe:
|
||||
def __init__(self,
|
||||
start_percent: float = 0.0,
|
||||
scale_multival: Union[float, Tensor]=None,
|
||||
effect_multival: Union[float, Tensor]=None,
|
||||
inherit_missing: bool=True,
|
||||
guarantee_steps: int=1,
|
||||
default: bool=False,
|
||||
):
|
||||
self.start_percent = start_percent
|
||||
self.start_t = 999999999.9
|
||||
self.scale_multival = scale_multival
|
||||
self.effect_multival = effect_multival
|
||||
self.inherit_missing = inherit_missing
|
||||
self.guarantee_steps = guarantee_steps
|
||||
self.default = default
|
||||
|
||||
def has_scale(self):
|
||||
return self.scale_multival is not None
|
||||
|
||||
def has_effect(self):
|
||||
return self.effect_multival is not None
|
||||
|
||||
|
||||
class ADKeyframeGroup:
|
||||
def __init__(self):
|
||||
self.keyframes: list[ADKeyframe] = []
|
||||
self.keyframes.append(ADKeyframe(guarantee_steps=1, default=True))
|
||||
|
||||
def add(self, keyframe: ADKeyframe):
|
||||
# remove any default keyframes that match start_percent of new keyframe
|
||||
default_to_delete = []
|
||||
for i in range(len(self.keyframes)):
|
||||
if self.keyframes[i].default and self.keyframes[i].start_percent == keyframe.start_percent:
|
||||
default_to_delete.append(i)
|
||||
for i in reversed(default_to_delete):
|
||||
self.keyframes.pop(i)
|
||||
# add to end of list, then sort
|
||||
self.keyframes.append(keyframe)
|
||||
self.keyframes = get_sorted_list_via_attr(self.keyframes, "start_percent")
|
||||
|
||||
def get_index(self, index: int) -> Union[ADKeyframe, None]:
|
||||
try:
|
||||
return self.keyframes[index]
|
||||
except IndexError:
|
||||
return None
|
||||
|
||||
def has_index(self, index: int) -> int:
|
||||
return index >=0 and index < len(self.keyframes)
|
||||
|
||||
def __getitem__(self, index) -> ADKeyframe:
|
||||
return self.keyframes[index]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.keyframes)
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
return len(self.keyframes) == 0
|
||||
|
||||
def clone(self) -> 'ADKeyframeGroup':
|
||||
cloned = ADKeyframeGroup()
|
||||
for tk in self.keyframes:
|
||||
if not tk.default:
|
||||
cloned.add(tk)
|
||||
return cloned
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
{
|
||||
"main_pass":
|
||||
[
|
||||
"-n", "-c:v", "libsvtav1",
|
||||
"-pix_fmt", "yuv420p10le",
|
||||
"-crf", "23"
|
||||
],
|
||||
"extension": "webm",
|
||||
"environment": {"SVT_LOG": "1"}
|
||||
}
|
||||
+9
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"main_pass":
|
||||
[
|
||||
"-n", "-c:v", "libx264",
|
||||
"-pix_fmt", "yuv420p",
|
||||
"-crf", "19"
|
||||
],
|
||||
"extension": "mp4"
|
||||
}
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"main_pass":
|
||||
[
|
||||
"-n", "-c:v", "libx265",
|
||||
"-pix_fmt", "yuv420p10le",
|
||||
"-preset", "medium",
|
||||
"-crf", "22",
|
||||
"-x265-params", "log-level=quiet"
|
||||
],
|
||||
"extension": "mp4"
|
||||
}
|
||||
+9
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"main_pass":
|
||||
[
|
||||
"-n",
|
||||
"-pix_fmt", "yuv420p",
|
||||
"-crf", "23"
|
||||
],
|
||||
"extension": "webm"
|
||||
}
|
||||
+142
@@ -0,0 +1,142 @@
|
||||
import { app } from '../../../scripts/app.js'
|
||||
import { api } from '../../../scripts/api.js'
|
||||
|
||||
function offsetDOMWidget(
|
||||
widget,
|
||||
ctx,
|
||||
node,
|
||||
widgetWidth,
|
||||
widgetY,
|
||||
height
|
||||
) {
|
||||
const margin = 10
|
||||
const elRect = ctx.canvas.getBoundingClientRect()
|
||||
const transform = new DOMMatrix()
|
||||
.scaleSelf(
|
||||
elRect.width / ctx.canvas.width,
|
||||
elRect.height / ctx.canvas.height
|
||||
)
|
||||
.multiplySelf(ctx.getTransform())
|
||||
.translateSelf(0, widgetY + margin)
|
||||
|
||||
const scale = new DOMMatrix().scaleSelf(transform.a, transform.d)
|
||||
Object.assign(widget.inputEl.style, {
|
||||
transformOrigin: '0 0',
|
||||
transform: scale,
|
||||
left: `${transform.e}px`,
|
||||
top: `${transform.d + transform.f}px`,
|
||||
width: `${widgetWidth}px`,
|
||||
height: `${(height || widget.parent?.inputHeight || 32) - margin}px`,
|
||||
position: 'absolute',
|
||||
background: !node.color ? '' : node.color,
|
||||
color: !node.color ? '' : 'white',
|
||||
zIndex: 5, //app.graph._nodes.indexOf(node),
|
||||
})
|
||||
}
|
||||
|
||||
export const hasWidgets = (node) => {
|
||||
if (!node.widgets || !node.widgets?.[Symbol.iterator]) {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
export const cleanupNode = (node) => {
|
||||
if (!hasWidgets(node)) {
|
||||
return
|
||||
}
|
||||
|
||||
for (const w of node.widgets) {
|
||||
if (w.canvas) {
|
||||
w.canvas.remove()
|
||||
}
|
||||
if (w.inputEl) {
|
||||
w.inputEl.remove()
|
||||
}
|
||||
// calls the widget remove callback
|
||||
w.onRemoved?.()
|
||||
}
|
||||
}
|
||||
|
||||
const CreatePreviewElement = (name, val, format) => {
|
||||
const [type] = format.split('/')
|
||||
const w = {
|
||||
name,
|
||||
type,
|
||||
value: val,
|
||||
draw: function (ctx, node, widgetWidth, widgetY, height) {
|
||||
const [cw, ch] = this.computeSize(widgetWidth)
|
||||
offsetDOMWidget(this, ctx, node, widgetWidth, widgetY, ch)
|
||||
},
|
||||
computeSize: function (_) {
|
||||
const ratio = this.inputRatio || 1
|
||||
const width = Math.max(220, this.parent.size[0])
|
||||
return [width, (width / ratio + 10)]
|
||||
},
|
||||
onRemoved: function () {
|
||||
if (this.inputEl) {
|
||||
this.inputEl.remove()
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
w.inputEl = document.createElement(type === 'video' ? 'video' : 'img')
|
||||
w.inputEl.src = w.value
|
||||
if (type === 'video') {
|
||||
w.inputEl.setAttribute('type', 'video/webm');
|
||||
w.inputEl.autoplay = true
|
||||
w.inputEl.loop = true
|
||||
w.inputEl.controls = false;
|
||||
}
|
||||
w.inputEl.onload = function () {
|
||||
w.inputRatio = w.inputEl.naturalWidth / w.inputEl.naturalHeight
|
||||
}
|
||||
document.body.appendChild(w.inputEl)
|
||||
return w
|
||||
}
|
||||
|
||||
const gif_preview = {
|
||||
name: 'AnimateDiff.gif_preview',
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
||||
switch (nodeData.name) {
|
||||
case 'ADE_AnimateDiffCombine':{
|
||||
const onExecuted = nodeType.prototype.onExecuted
|
||||
nodeType.prototype.onExecuted = function (message) {
|
||||
const prefix = 'ad_gif_preview_'
|
||||
const r = onExecuted ? onExecuted.apply(this, message) : undefined
|
||||
|
||||
if (this.widgets) {
|
||||
const pos = this.widgets.findIndex((w) => w.name === `${prefix}_0`)
|
||||
if (pos !== -1) {
|
||||
for (let i = pos; i < this.widgets.length; i++) {
|
||||
this.widgets[i].onRemoved?.()
|
||||
}
|
||||
this.widgets.length = pos
|
||||
}
|
||||
if (message?.gifs) {
|
||||
message.gifs.forEach((params, i) => {
|
||||
const previewUrl = api.apiURL(
|
||||
'/view?' + new URLSearchParams(params).toString()
|
||||
)
|
||||
const w = this.addCustomWidget(
|
||||
CreatePreviewElement(`${prefix}_${i}`, previewUrl, params.format || 'image/gif')
|
||||
)
|
||||
w.parent = this
|
||||
})
|
||||
}
|
||||
const onRemoved = this.onRemoved
|
||||
this.onRemoved = () => {
|
||||
cleanupNode(this)
|
||||
return onRemoved?.()
|
||||
}
|
||||
}
|
||||
this.setSize([this.size[0], this.computeSize([this.size[0], this.size[1]])[1]])
|
||||
return r
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
app.registerExtension(gif_preview)
|
||||
@@ -0,0 +1,20 @@
|
||||
# EditorConfig helps developers define and maintain consistent
|
||||
# coding styles between different editors and IDEs
|
||||
# editorconfig.org
|
||||
|
||||
root = true
|
||||
|
||||
[*]
|
||||
# Change these settings to your own preference
|
||||
indent_style = space
|
||||
indent_size = 2
|
||||
max_line_length = 120
|
||||
|
||||
end_of_line = lf
|
||||
charset = utf-8
|
||||
trim_trailing_whitespace = true
|
||||
insert_final_newline = true
|
||||
|
||||
[*.md]
|
||||
max_line_length = off
|
||||
trim_trailing_whitespace = false
|
||||
@@ -0,0 +1,164 @@
|
||||
module.exports = {
|
||||
root: true,
|
||||
env: { browser: true, es2020: true },
|
||||
extends: [
|
||||
'eslint:recommended',
|
||||
'plugin:@typescript-eslint/recommended',
|
||||
'plugin:@typescript-eslint/recommended-type-checked',
|
||||
'plugin:@typescript-eslint/stylistic-type-checked'
|
||||
],
|
||||
parser: '@typescript-eslint/parser',
|
||||
plugins: ['@typescript-eslint', 'import'],
|
||||
parserOptions: {
|
||||
// tsconfigRootDir: __dirname,
|
||||
sourceType: 'module',
|
||||
ecmaVersion: 'latest',
|
||||
ecmaFeatures: {
|
||||
// experimentalObjectRestSpread: true,
|
||||
modules: true,
|
||||
legacyDecorators: true
|
||||
},
|
||||
project: ['./tsconfig.json'],
|
||||
},
|
||||
rules: {
|
||||
// 'import/extensions': [
|
||||
// 'error',
|
||||
// "ignorePackages",
|
||||
// {
|
||||
// "js": "always",
|
||||
// "ts": "never"
|
||||
// }
|
||||
// ],
|
||||
'import/no-unresolved': ['off'],
|
||||
'class-methods-use-this': ['off'],
|
||||
'radix': ['off'],
|
||||
'import/prefer-default-export': ['off'],
|
||||
'implicit-arrow-linebreak': ['off'],
|
||||
'no-trailing-spaces': [
|
||||
'error', {
|
||||
'skipBlankLines': true,
|
||||
'ignoreComments': true,
|
||||
},
|
||||
],
|
||||
'max-len': [
|
||||
'error', {
|
||||
'code': 120,
|
||||
},
|
||||
],
|
||||
'prefer-rest-params': ['off'],
|
||||
|
||||
'import/no-extraneous-dependencies': [
|
||||
'error', {
|
||||
'devDependencies': ['**/__tests__/*.ts', '**/__mocks__/*.ts'],
|
||||
},
|
||||
],
|
||||
|
||||
'one-var': ['off'],
|
||||
|
||||
// false positive https://github.com/typescript-eslint/typescript-eslint/issues/2483
|
||||
'no-shadow': 'off',
|
||||
'object-curly-newline': ["error", {
|
||||
"ObjectExpression": { "multiline": true, "minProperties": 1 },
|
||||
"ObjectPattern": { "multiline": true, "minProperties": 1 },
|
||||
"ImportDeclaration": {"multiline": true, "minProperties": 8 },
|
||||
"ExportDeclaration": { "multiline": true, "minProperties": 1 }
|
||||
}],
|
||||
"@typescript-eslint/explicit-function-return-type": 'error',
|
||||
'@typescript-eslint/triple-slash-reference': 'off',
|
||||
'@typescript-eslint/no-shadow': 'error',
|
||||
'@typescript-eslint/no-unsafe-argument': 'off',
|
||||
'@typescript-eslint/no-unsafe-assignment': 'off',
|
||||
'@typescript-eslint/no-unsafe-member-access': 'off',
|
||||
'@typescript-eslint/no-unsafe-return': 'off',
|
||||
'@typescript-eslint/no-unsafe-call': 'off',
|
||||
'@typescript-eslint/prefer-nullish-coalescing': 'off',
|
||||
'@typescript-eslint/consistent-type-definitions': 'off',
|
||||
'@typescript-eslint/no-unused-expressions': [
|
||||
'error',
|
||||
{
|
||||
'allowShortCircuit': true,
|
||||
'allowTernary': true,
|
||||
},
|
||||
],
|
||||
'@typescript-eslint/no-use-before-define': ['off'],
|
||||
'@typescript-eslint/no-unused-vars': [
|
||||
'error',
|
||||
{'argsIgnorePattern': '^_'},
|
||||
],
|
||||
'@typescript-eslint/explicit-module-boundary-types': [
|
||||
'off', {
|
||||
allowArgumentsExplicitlyTypedAsAny: true,
|
||||
},
|
||||
],
|
||||
'@typescript-eslint/no-explicit-any': 'off',
|
||||
'@typescript-eslint/ban-ts-comment': [
|
||||
'off',
|
||||
{'ts-expect-error': 'allow-with-description'},
|
||||
],
|
||||
'@typescript-eslint/no-empty-function': [
|
||||
'error',
|
||||
{
|
||||
'allow': [
|
||||
'methods',
|
||||
'asyncMethods',
|
||||
],
|
||||
},
|
||||
],
|
||||
|
||||
'lines-between-class-members': 'off',
|
||||
'@typescript-eslint/lines-between-class-members': 'off',
|
||||
|
||||
'curly': ['error', 'multi-line', 'consistent'],
|
||||
'no-unused-vars': 'off',
|
||||
'jest/no-test-prefixes': 'off',
|
||||
'jest/no-focused-tests': 'off',
|
||||
'comma-dangle': ['error', 'only-multiline'],
|
||||
'linebreak-style': ['error', 'unix'],
|
||||
'quotes': ['error', 'single'],
|
||||
'semi': ['error', 'always'],
|
||||
'eqeqeq': ['error', 'always'],
|
||||
'complexity': [
|
||||
'error', {
|
||||
max: 8,
|
||||
},
|
||||
],
|
||||
'block-scoped-var': 'error',
|
||||
'no-else-return': [
|
||||
'error', {
|
||||
allowElseIf: true,
|
||||
},
|
||||
],
|
||||
'no-debugger': 'off',
|
||||
'no-eval': 'error',
|
||||
'no-lone-blocks': 'error',
|
||||
'no-multi-spaces': 'error',
|
||||
'no-useless-return': 'error',
|
||||
'no-var': 'error',
|
||||
'no-console': [
|
||||
// 'error', {
|
||||
// allow: ['warn', 'error'],
|
||||
// },
|
||||
'off',
|
||||
],
|
||||
'no-throw-literal': 'error',
|
||||
'newline-per-chained-call': [
|
||||
'error', {
|
||||
ignoreChainWithDepth: 4,
|
||||
},
|
||||
],
|
||||
'no-extra-boolean-cast': [
|
||||
'error', {
|
||||
enforceForLogicalOperands: true,
|
||||
},
|
||||
],
|
||||
'no-fallthrough': 'error',
|
||||
'no-use-before-define': 'off',
|
||||
'no-case-declarations': 'off',
|
||||
},
|
||||
ignorePatterns: [
|
||||
'node_modules',
|
||||
'web/**/*.d.ts',
|
||||
'web/**/*.js'
|
||||
],
|
||||
globals: {},
|
||||
}
|
||||
Vendored
+47
@@ -0,0 +1,47 @@
|
||||
---
|
||||
name: Bug report
|
||||
about: Create a report to help us improve
|
||||
title: ''
|
||||
labels: ''
|
||||
assignees: ''
|
||||
|
||||
---
|
||||
|
||||
**Describe the bug**
|
||||
A clear and concise description of what the bug is.
|
||||
|
||||
**To Reproduce**
|
||||
__Please attach a workflow file to make it easier for others to reproduce the error!__
|
||||
|
||||
**Expected behavior**
|
||||
A clear and concise description of what you expected to happen.
|
||||
|
||||
**Screenshots**
|
||||
If applicable, add screenshots to help explain your problem.
|
||||
|
||||
**Error in console:**
|
||||
If applicable, add the error message from the console.
|
||||
|
||||
**Versions:**
|
||||
Copy the output of the console (4 parts), like this:
|
||||
```
|
||||
** Python version: 3.10.11 (tags/v3.10.11:7d4cc5a, Apr 5 2023, 00:38:17) [MSC v.1929 64 bit (AMD64)]
|
||||
```
|
||||
```
|
||||
Total VRAM 12287 MB, total RAM 130996 MB
|
||||
Set vram state to: NORMAL_VRAM
|
||||
```
|
||||
```
|
||||
[Crystools INFO] Crystools version: 1.9.2
|
||||
[Crystools INFO] CPU: AMD Ryzen 9 5950X 16-Core Processor - Arch: AMD64 - OS: Windows 10
|
||||
[Crystools INFO] GPU/s:
|
||||
[Crystools INFO] 0) NVIDIA GeForce RTX 3080 Ti
|
||||
[Crystools INFO] NVIDIA Driver: 546.33
|
||||
```
|
||||
```
|
||||
### Loading: ComfyUI-Manager (V1.25.3)
|
||||
### ComfyUI Revision: 1886 [6a10640f] | Released on '2024-01-08'
|
||||
```
|
||||
|
||||
**Additional context**
|
||||
Add any other context about the problem here.
|
||||
Vendored
+20
@@ -0,0 +1,20 @@
|
||||
---
|
||||
name: Feature request
|
||||
about: Suggest an idea for this project
|
||||
title: ''
|
||||
labels: ''
|
||||
assignees: ''
|
||||
|
||||
---
|
||||
|
||||
**Is your feature request related to a problem? Please describe.**
|
||||
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
|
||||
|
||||
**Describe the solution you'd like**
|
||||
A clear and concise description of what you want to happen.
|
||||
|
||||
**Describe alternatives you've considered**
|
||||
A clear and concise description of any alternative solutions or features you've considered.
|
||||
|
||||
**Additional context**
|
||||
Add any other context or screenshots about the feature request here.
|
||||
@@ -0,0 +1,155 @@
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintainted in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
*.pyc
|
||||
|
||||
.idea
|
||||
/node_modules
|
||||
@@ -0,0 +1,26 @@
|
||||
<component name="ProjectRunConfigurationManager">
|
||||
<configuration default="false" name="ComfyUI" type="PythonConfigurationType" factoryName="Python">
|
||||
<module name="ComfyUI" />
|
||||
<option name="ENV_FILES" value="" />
|
||||
<option name="INTERPRETER_OPTIONS" value="" />
|
||||
<option name="PARENT_ENVS" value="true" />
|
||||
<envs>
|
||||
<env name="PYTHONUNBUFFERED" value="1" />
|
||||
<env name="CRYSTOOLS_LOGLEVEL" value="10" />
|
||||
</envs>
|
||||
<option name="SDK_HOME" value="" />
|
||||
<option name="SDK_NAME" value="Python 3.11 (.env)" />
|
||||
<option name="WORKING_DIRECTORY" value="E:\AI\SD\comfyUI\ComfyUI" />
|
||||
<option name="IS_MODULE_SDK" value="false" />
|
||||
<option name="ADD_CONTENT_ROOTS" value="true" />
|
||||
<option name="ADD_SOURCE_ROOTS" value="true" />
|
||||
<option name="SCRIPT_NAME" value="$PROJECT_DIR$/main.py" />
|
||||
<option name="PARAMETERS" value="--windows-standalone-build --extra-model-paths-config ..\extra_model_paths.yaml --output-directory ..\..\comfyUI.output --input-directory ..\..\comfyUI.input --disable-auto-launch" />
|
||||
<option name="SHOW_COMMAND_LINE" value="false" />
|
||||
<option name="EMULATE_TERMINAL" value="false" />
|
||||
<option name="MODULE_MODE" value="false" />
|
||||
<option name="REDIRECT_INPUT" value="false" />
|
||||
<option name="INPUT_FILE" value="" />
|
||||
<method v="2" />
|
||||
</configuration>
|
||||
</component>
|
||||
@@ -0,0 +1,12 @@
|
||||
<component name="ProjectRunConfigurationManager">
|
||||
<configuration default="false" name="lint" type="js.build_tools.npm">
|
||||
<package-json value="$PROJECT_DIR$/custom_nodes/ComfyUI-Crystools/package.json" />
|
||||
<command value="run" />
|
||||
<scripts>
|
||||
<script value="lint" />
|
||||
</scripts>
|
||||
<node-interpreter value="project" />
|
||||
<envs />
|
||||
<method v="2" />
|
||||
</configuration>
|
||||
</component>
|
||||
@@ -0,0 +1,32 @@
|
||||
#!/bin/bash
|
||||
|
||||
cd ../comfyUI/ComfyUI
|
||||
rm ./custom_nodes/ComfyUI-Crystools/web/*.js -f
|
||||
|
||||
LOCAL_COMMAND="npm --prefix ./custom_nodes/ComfyUI-Crystools run validate"
|
||||
OUTPUT=$($LOCAL_COMMAND 2>&1)
|
||||
EXIT_CODE=$?
|
||||
|
||||
if [ $EXIT_CODE -eq 0 ]; then
|
||||
echo "Validate ok"
|
||||
cd ../../ComfyUI-Crystools
|
||||
else
|
||||
echo "Validate KO!"
|
||||
echo $OUTPUT
|
||||
exit 1
|
||||
fi
|
||||
|
||||
LOCAL_COMMAND="python .others/version.py"
|
||||
|
||||
OUTPUT=$($LOCAL_COMMAND 2>&1)
|
||||
EXIT_CODE=$?
|
||||
|
||||
if [ $EXIT_CODE -eq 0 ]; then
|
||||
echo "Version changed"
|
||||
else
|
||||
echo "Can't change the version! check the .other/version.py"
|
||||
echo $OUTPUT
|
||||
exit 1
|
||||
fi
|
||||
|
||||
git add __init__.py core/version.py web/*.js
|
||||
@@ -0,0 +1,12 @@
|
||||
<component name="ProjectRunConfigurationManager">
|
||||
<configuration default="false" name="typescript" type="js.build_tools.npm">
|
||||
<package-json value="$PROJECT_DIR$/custom_nodes/ComfyUI-Crystools/package.json" />
|
||||
<command value="run" />
|
||||
<scripts>
|
||||
<script value="start" />
|
||||
</scripts>
|
||||
<node-interpreter value="project" />
|
||||
<envs />
|
||||
<method v="2" />
|
||||
</configuration>
|
||||
</component>
|
||||
@@ -0,0 +1,40 @@
|
||||
import os
|
||||
|
||||
print("-- Setting up the version --")
|
||||
|
||||
versionLabel = "@version:"
|
||||
version_file = os.path.join(os.path.dirname(__file__), "../version")
|
||||
|
||||
with open(version_file, "r") as f:
|
||||
version = f.read().strip()
|
||||
|
||||
if version is None or version == "":
|
||||
raise RuntimeError(f"File without version: {version_file}")
|
||||
|
||||
print(f"Version to set: {version}")
|
||||
|
||||
# modify the __init__.py file with @version
|
||||
file1 = os.path.join(os.path.dirname(__file__), "../__init__.py")
|
||||
print(f"> File1: {file1}")
|
||||
|
||||
file1VersionFound = False
|
||||
# lines = None
|
||||
with open(file1, "r") as f:
|
||||
lines = f.readlines()
|
||||
for line in lines:
|
||||
if versionLabel in line:
|
||||
file1VersionFound = True
|
||||
lines[lines.index(line)] = f"{versionLabel} {version}\n"
|
||||
|
||||
if not file1VersionFound:
|
||||
raise RuntimeError(f"Version label not found in file: {file1}")
|
||||
|
||||
with open(file1, "w") as f2:
|
||||
f2.writelines(lines)
|
||||
|
||||
file2 = os.path.join(os.path.dirname(__file__), "../core/version.py")
|
||||
print(f"> File2 {file2}")
|
||||
|
||||
with open(file2, "w") as f2:
|
||||
f2.writelines([f"version = \"{version}\"\n"])
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2023 Crystian
|
||||
|
||||
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.
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user