Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,4 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+# Auto detect text files and perform LF normalization
+* text=auto
+HomeImage.png filter=lfs diff=lfs merge=lfs -text
+stable_fast-1.0.5+torch222cu121-cp310-cp310-manylinux2014_x86_64.whl filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,15 @@
+
+*.pyc
+*.pth
+*.pt
+*.safetensors
+*.gguf
+*.png
+/.idea
+/htmlcov
+.coverage
+.toml
+__pycache__
+.venv
+!HomeImage.png
+*.txt

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

.vscode/settings.json

@@ -0,0 +1,13 @@
+{
+    "python.testing.unittestArgs": [
+        "-v",
+        "-s",
+        ".",
+        "-p",
+        "*test.py"
+    ],
+    "python.testing.pytestEnabled": false,
+    "python.testing.unittestEnabled": true,
+    "python.analysis.autoImportCompletions": true,
+    "python.analysis.typeCheckingMode": "off"
+}

Compiler.py

@@ -0,0 +1,106 @@
+import os
+import re
+
+files_ordered = [
+    "./modules/Utilities/util.py",
+    "./modules/sample/sampling_util.py",
+    "./modules/Device/Device.py",
+    "./modules/cond/cond_util.py",
+    "./modules/cond/cond.py",
+    "./modules/sample/ksampler_util.py",
+    "./modules/cond/cast.py",
+    "./modules/Attention/AttentionMethods.py",
+    "./modules/AutoEncoders/taesd.py",
+    "./modules/cond/cond.py",
+    "./modules/cond/Activation.py",
+    "./modules/Attention/Attention.py",
+    "./modules/sample/samplers.py",
+    "./modules/sample/CFG.py",
+    "./modules/NeuralNetwork/transformer.py",
+    "./modules/sample/sampling.py",
+    "./modules/clip/CLIPTextModel.py",
+    "./modules/AutoEncoders/ResBlock.py",
+    "./modules/AutoDetailer/mask_util.py",
+    "./modules/NeuralNetwork/unet.py",
+    "./modules/SD15/SDClip.py",
+    "./modules/SD15/SDToken.py",
+    "./modules/UltimateSDUpscale/USDU_util.py",
+    "./modules/StableFast/SF_util.py",
+    "./modules/Utilities/Latent.py",
+    "./modules/AutoDetailer/SEGS.py",
+    "./modules/AutoDetailer/tensor_util.py",
+    "./modules/AutoDetailer/AD_util.py",
+    "./modules/clip/FluxClip.py",
+    "./modules/Model/ModelPatcher.py",
+    "./modules/Model/ModelBase.py",
+    "./modules/UltimateSDUpscale/image_util.py",
+    "./modules/UltimateSDUpscale/RDRB.py",
+    "./modules/StableFast/ModuleFactory.py",
+    "./modules/AutoDetailer/bbox.py",
+    "./modules/AutoEncoders/VariationalAE.py",
+    "./modules/clip/Clip.py",
+    "./modules/Model/LoRas.py",
+    "./modules/BlackForest/Flux.py",
+    "./modules/UltimateSDUpscale/USDU_upscaler.py",
+    "./modules/StableFast/ModuleTracing.py",
+    "./modules/hidiffusion/utils.py",
+    "./modules/FileManaging/Downloader.py",
+    "./modules/AutoDetailer/SAM.py",
+    "./modules/AutoDetailer/ADetailer.py",
+    "./modules/Quantize/Quantizer.py",
+    "./modules/FileManaging/Loader.py",
+    "./modules/SD15/SD15.py",
+    "./modules/UltimateSDUpscale/UltimateSDUpscale.py",
+    "./modules/StableFast/StableFast.py",
+    "./modules/hidiffusion/msw_msa_attention.py",
+    "./modules/FileManaging/ImageSaver.py",
+    "./modules/Utilities/Enhancer.py",
+    "./modules/Utilities/upscale.py",
+    "./modules/user/pipeline.py",
+]
+
+def get_file_patterns():
+    patterns = []
+    seen = set()
+    for path in files_ordered:
+        filename = os.path.basename(path)
+        name = os.path.splitext(filename)[0]
+        if name not in seen:
+            # Pattern 1: matches module name when not in brackets or after a dot
+            pattern1 = rf'(?<![a-zA-Z0-9_\.])({name}\.)(?![)\]])'
+            # Pattern 2: matches module name inside brackets while preserving them  
+            pattern2 = rf'(\[|\()({name}\.)([^\]\)]+?)(\]|\))'
+            pattern3 = 'cond_util\.'
+            patterns.extend([
+                (pattern1, ''),  # Remove module name and dot outside brackets
+                (pattern2, r'\1\3\4'),  # Keep brackets, remove only module name
+                (pattern3, '')
+            ])
+            seen.add(name)
+    return patterns
+
+def remove_file_names(line):
+    patterns = get_file_patterns()
+    result = line
+    for pattern, replacement in patterns:
+        result = re.sub(pattern, replacement, result)
+    return result
+
+try:
+    with open("./compiled.py", "w") as output_file:
+        for file_path in files_ordered:
+            try:
+                with open(file_path, "r") as input_file:
+                    for line in input_file:
+                        if not line.lstrip().startswith("from modules."):
+                            # Apply the file name removal before writing
+                            modified_line = remove_file_names(line)
+                            output_file.write(modified_line)
+                    output_file.write("\n\n")
+                print(f"Processed: {file_path}")
+            except FileNotFoundError:
+                print(f"Error: Could not find file {file_path}")
+            except Exception as e:
+                print(f"Error processing {file_path}: {str(e)}")
+except Exception as e:
+    print(f"Error creating compiled.py: {str(e)}")

LICENSE

@@ -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>.

README.md

@@ -1,12 +1,144 @@
 ---
-title: LightDiffusion Next
-emoji: 😻
-colorFrom: gray
-colorTo: blue
+title: LightDiffusion-Next
+app_file: app.py
 sdk: gradio
 sdk_version: 5.14.0
-app_file: app.py
-pinned: false
+---
+<div align="center">
+
+# Say hi to LightDiffusion-Next 👋
+
+
+**LightDiffusion-Next**  is the fastest AI-powered image generation GUI/CLI, combining speed, precision, and flexibility in one cohesive tool.
+</br>
+</br>
+  <a href="https://github.com/LightDiffusion/LightDiffusion-Next">
+    <img src="./HomeImage.png" alt="Logo">
+
+  </a>
+</br>
+</div>
+
+
+As a refactored and improved version of the original [LightDiffusion repository](https://github.com/Aatrick/LightDiffusion), this project enhances usability, maintainability, and functionality while introducing a host of new features to streamline your creative workflows.
+
+## Motivation:
+
+**LightDiffusion** was originally meant to be made in Rust, but due to the lack of support for the Rust language in the AI community, it was made in Python with the goal of being the simplest and fastest AI image generation tool.
+
+That's when the first version of LightDiffusion was born which only counted [3000 lines of code](https://github.com/LightDiffusion/LightDiffusion-original), only using Pytorch. With time, the [project](https://github.com/Aatrick/LightDiffusion) grew and became more complex, and the need for a refactor was evident. This is where **LightDiffusion-Next** comes in, with a more modular and maintainable codebase, and a plethora of new features and optimizations.
+
+📚 Learn more in the [official documentation](https://aatrick.github.io/LightDiffusion/).
+
+---
+
+## 🌟 Highlights
+
+![image](https://github.com/user-attachments/assets/b994fe0d-3a2e-44ff-93a4-46919cf865e3)
+
+**LightDiffusion-Next** offers a powerful suite of tools to cater to creators at every level. At its core, it supports **Text-to-Image** (Txt2Img) and **Image-to-Image** (Img2Img) generation, offering a variety of upscale methods and samplers, to make it easier to create stunning images with minimal effort.
+
+Advanced users can take advantage of features like **attention syntax**, **Hires-Fix** or **ADetailer**. These tools provide better quality and flexibility for generating complex and high-resolution outputs.
+
+**LightDiffusion-Next** is fine-tuned for **performance**. Features such as **Xformers** acceleration, **BFloat16** precision support, **WaveSpeed** dynamic caching, and **Stable-Fast** model compilation (which offers up to a 70% speed boost) ensure smooth and efficient operation, even on demanding workloads.
+
+---
+
+## ✨ Feature Showcase
+
+Here’s what makes LightDiffusion-Next stand out:
+
+- **Speed and Efficiency**:
+  Enjoy industry-leading performance with built-in Xformers, Pytorch, Wavespeed and Stable-Fast optimizations, achieving up to 30% faster speeds compared to the rest of the AI image generation backends in SD1.5 and up to 2x for Flux.
+
+- **Automatic Detailing**:
+  Effortlessly enhance faces and body details with AI-driven tools based on the [Impact Pack](https://github.com/ltdrdata/ComfyUI-Impact-Pack).
+
+- **State Preservation**:
+  Save and resume your progress with saved states, ensuring seamless transitions between sessions.
+
+- **Advanced GUI and CLI**:
+  Work through a user-friendly graphical interface or leverage the streamlined pipeline for CLI-based workflows.
+
+- **Integration-Ready**:
+  Collaborate and create directly in Discord with [Boubou](https://github.com/Aatrick/Boubou), or preview images dynamically with the optional **TAESD preview mode**.
+
+- **Image Previewing**:
+  Get a real-time preview of your generated images with TAESD, allowing for user-friendly and interactive workflows.
+
+- **Image Upscaling**:
+  Enhance your images with advanced upscaling options like UltimateSDUpscaling, ensuring high-quality results every time.
+
+- **Prompt Refinement**:
+    Use the Ollama-powered automatic prompt enhancer to refine your prompts and generate more accurate and detailed outputs.
+
+- **LoRa and Textual Inversion Embeddings**:
+    Leverage LoRa and textual inversion embeddings for highly customized and nuanced results, adding a new dimension to your creative process.
+
+- **Low-End Device Support**:
+    Run LightDiffusion-Next on low-end devices with as little as 2GB of VRAM or even no GPU, ensuring accessibility for all users.
+
+---
+
+## ⚡ Performance Benchmarks
+
+**LightDiffusion-Next** dominates in performance:
+
+| **Tool**                           | **Speed (it/s)** |
+|------------------------------------|------------------|
+| **LightDiffusion with Stable-Fast** | 2.8              |
+| **LightDiffusion**                 | 1.8              |
+| **ComfyUI**                        | 1.4              |
+| **SDForge**                        | 1.3              |
+| **SDWebUI**                        | 0.9              |
+
+(All benchmarks are based on a 1024x1024 resolution with a batch size of 1 using BFloat16 precision without tweaking installations. Made with a 3060 mobile GPU using SD1.5.)
+
+With its unmatched speed and efficiency, LightDiffusion-Next sets the benchmark for AI image generation tools.
+
+---
+
+## 🛠 Installation
+
+### Quick Start
+
+1. Download a release or clone this repository.
+2. Run `run.bat` in a terminal.
+3. Start creating!
+
+### Command-Line Pipeline
+
+For a GUI-free experience, use the pipeline:
+```bash
+pipeline.bat <prompt> <width> <height> <num_images> <batch_size>
+```
+Use `pipeline.bat -h` for more options.
+
+---
+
+### Advanced Setup
+
+- **Install from Source**:
+  Install dependencies via:
+  ```bash
+  pip install -r requirements.txt
+  ```
+  Add your SD1/1.5 safetensors model to the `checkpoints` directory, then launch the application.
+
+- **⚡Stable-Fast Optimization**:
+  Follow [this guide](https://github.com/chengzeyi/stable-fast?tab=readme-ov-file#installation) to enable Stable-Fast mode for optimal performance.
+
+- **🦙 Prompt Enhancer**:
+  Refine your prompts with Ollama:
+  ```bash
+  pip install ollama
+  ollama run deepseek-r1
+  ```
+  See the [Ollama guide](https://github.com/ollama/ollama?tab=readme-ov-file) for details.
+
+- **🤖 Discord Integration**:
+  Set up the Discord bot by following the [Boubou installation guide](https://github.com/Aatrick/Boubou).
+
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+🎨 Enjoy exploring the powerful features of LightDiffusion-Next!

_internal/clip/sd1_clip_config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "openai/clip-vit-large-patch14",
+  "architectures": [
+    "CLIPTextModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 768,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "projection_dim": 768,
+  "torch_dtype": "float32",
+  "transformers_version": "4.24.0",
+  "vocab_size": 49408
+}

_internal/sd1_tokenizer/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

_internal/sd1_tokenizer/tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+  "add_prefix_space": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "do_lower_case": true,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "errors": "replace",
+  "model_max_length": 77,
+  "name_or_path": "openai/clip-vit-large-patch14",
+  "pad_token": "<|endoftext|>",
+  "special_tokens_map_file": "./special_tokens_map.json",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

app.py

@@ -0,0 +1,171 @@
+import glob
+import gradio as gr
+import sys
+import os
+from PIL import Image
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+
+from modules.user.pipeline import pipeline
+import torch
+
+def load_generated_images():
+    """Load generated images with given prefix from disk"""
+    image_files = glob.glob("./_internal/output/*")
+
+    # If there are no image files, return
+    if not image_files:
+        return []
+
+    # Sort files by modification time in descending order
+    image_files.sort(key=os.path.getmtime, reverse=True)
+
+    # Get most recent timestamp
+    latest_time = os.path.getmtime(image_files[0])
+
+    # Get all images from same batch (within 1 second of most recent)
+    batch_images = []
+    for file in image_files:
+        if abs(os.path.getmtime(file) - latest_time) < 1.0:
+            try:
+                img = Image.open(file)
+                batch_images.append(img)
+            except:
+                continue
+
+    if not batch_images:
+        return []
+    return batch_images
+
+def generate_images(
+    prompt: str,
+    width: int = 512,
+    height: int = 512,
+    num_images: int = 1,
+    batch_size: int = 1,
+    hires_fix: bool = False,
+    adetailer: bool = False,
+    enhance_prompt: bool = False,
+    img2img_enabled: bool = False,
+    img2img_image: str = None,
+    stable_fast: bool = False,
+    reuse_seed: bool = False,
+    flux_enabled: bool = False,
+    prio_speed: bool = False,
+    progress=gr.Progress()
+):
+    """Generate images using the LightDiffusion pipeline"""
+    try:
+        if img2img_enabled and img2img_image is not None:
+            # Save uploaded image temporarily and pass path to pipeline
+            img2img_image.save("temp_img2img.png")
+            prompt = "temp_img2img.png"
+        
+        # Run pipeline and capture saved images
+        with torch.inference_mode():
+            images = pipeline(
+                prompt=prompt,
+                w=width,
+                h=height, 
+                number=num_images,
+                batch=batch_size,
+                hires_fix=hires_fix,
+                adetailer=adetailer,
+                enhance_prompt=enhance_prompt,
+                img2img=img2img_enabled,
+                stable_fast=stable_fast,
+                reuse_seed=reuse_seed,
+                flux_enabled=flux_enabled,
+                prio_speed=prio_speed
+            )
+            
+        return load_generated_images()
+
+    except Exception as e:
+        import traceback
+        print(traceback.format_exc())
+        return [Image.new('RGB', (512, 512), color='black')]
+
+# Create Gradio interface
+with gr.Blocks(title="LightDiffusion Web UI") as demo:
+    gr.Markdown("# LightDiffusion Web UI")
+    gr.Markdown("Generate AI images using LightDiffusion")
+    
+    with gr.Row():
+        with gr.Column():
+            # Input components
+            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
+            
+            with gr.Row():
+                width = gr.Slider(minimum=64, maximum=2048, value=512, step=64, label="Width")
+                height = gr.Slider(minimum=64, maximum=2048, value=512, step=64, label="Height")
+            
+            with gr.Row():
+                num_images = gr.Slider(minimum=1, maximum=10, value=1, step=1, label="Number of Images")
+                batch_size = gr.Slider(minimum=1, maximum=4, value=1, step=1, label="Batch Size")
+            
+            with gr.Row():
+                hires_fix = gr.Checkbox(label="HiRes Fix")
+                adetailer = gr.Checkbox(label="Auto Face/Body Enhancement")
+                enhance_prompt = gr.Checkbox(label="Enhance Prompt")
+                stable_fast = gr.Checkbox(label="Stable Fast Mode")
+            
+            with gr.Row():
+                reuse_seed = gr.Checkbox(label="Reuse Seed")
+                flux_enabled = gr.Checkbox(label="Flux Mode") 
+                prio_speed = gr.Checkbox(label="Prioritize Speed")
+            
+            with gr.Row():
+                img2img_enabled = gr.Checkbox(label="Image to Image Mode")
+                img2img_image = gr.Image(label="Input Image for img2img", visible=False)
+            
+            # Make input image visible only when img2img is enabled
+            img2img_enabled.change(
+                fn=lambda x: gr.update(visible=x),
+                inputs=[img2img_enabled],
+                outputs=[img2img_image]
+            )
+            
+            generate_btn = gr.Button("Generate")
+
+        # Output gallery
+        gallery = gr.Gallery(
+            label="Generated Images",
+            show_label=True,
+            elem_id="gallery",
+            columns=[2], 
+            rows=[2],
+            object_fit="contain",
+            height="auto"
+        )
+    
+    # Connect generate button to pipeline
+    generate_btn.click(
+        fn=generate_images,
+        inputs=[
+            prompt,
+            width, 
+            height,
+            num_images,
+            batch_size,
+            hires_fix,
+            adetailer, 
+            enhance_prompt,
+            img2img_enabled,
+            img2img_image,
+            stable_fast,
+            reuse_seed,
+            flux_enabled,
+            prio_speed
+        ],
+        outputs=gallery
+    )
+
+# For local testing
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=8000,
+        auth=None,
+        share=True,
+        debug=True
+    )

modules/Attention/Attention.py

@@ -0,0 +1,191 @@
+import torch
+import torch.nn as nn
+import logging
+
+from modules.Utilities import util
+from modules.Attention import AttentionMethods
+from modules.Device import Device
+from modules.cond import cast
+
+
+def Normalize(
+    in_channels: int, dtype: torch.dtype = None, device: torch.device = None
+) -> torch.nn.GroupNorm:
+    """#### Normalize the input channels.
+
+    #### Args:
+        - `in_channels` (int): The input channels.
+        - `dtype` (torch.dtype, optional): The data type. Defaults to `None`.
+        - `device` (torch.device, optional): The device. Defaults to `None`.
+
+    #### Returns:
+        - `torch.nn.GroupNorm`: The normalized input channels
+    """
+    return torch.nn.GroupNorm(
+        num_groups=32,
+        num_channels=in_channels,
+        eps=1e-6,
+        affine=True,
+        dtype=dtype,
+        device=device,
+    )
+
+
+if Device.xformers_enabled():
+    logging.info("Using xformers cross attention")
+    optimized_attention = AttentionMethods.attention_xformers
+else:
+    logging.info("Using pytorch cross attention")
+    optimized_attention = AttentionMethods.attention_pytorch
+
+optimized_attention_masked = optimized_attention
+
+
+def optimized_attention_for_device() -> AttentionMethods.attention_pytorch:
+    """#### Get the optimized attention for a device.
+
+    #### Returns:
+        - `function`: The optimized attention function.
+    """
+    return AttentionMethods.attention_pytorch
+
+
+class CrossAttention(nn.Module):
+    """#### Cross attention module, which applies attention across the query and context.
+
+    #### Args:
+        - `query_dim` (int): The query dimension.
+        - `context_dim` (int, optional): The context dimension. Defaults to `None`.
+        - `heads` (int, optional): The number of heads. Defaults to `8`.
+        - `dim_head` (int, optional): The head dimension. Defaults to `64`.
+        - `dropout` (float, optional): The dropout rate. Defaults to `0.0`.
+        - `dtype` (torch.dtype, optional): The data type. Defaults to `None`.
+        - `device` (torch.device, optional): The device. Defaults to `None`.
+        - `operations` (cast.disable_weight_init, optional): The operations. Defaults to `cast.disable_weight_init`.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        context_dim: int = None,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        dtype: torch.dtype = None,
+        device: torch.device = None,
+        operations: cast.disable_weight_init = cast.disable_weight_init,
+    ):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = util.default(context_dim, query_dim)
+
+        self.heads = heads
+        self.dim_head = dim_head
+
+        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: torch.Tensor,
+        context: torch.Tensor = None,
+        value: torch.Tensor = None,
+        mask: torch.Tensor = None,
+    ) -> torch.Tensor:
+        """#### Forward pass of the cross attention module.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `context` (torch.Tensor, optional): The context tensor. Defaults to `None`.
+            - `value` (torch.Tensor, optional): The value tensor. Defaults to `None`.
+            - `mask` (torch.Tensor, optional): The mask tensor. Defaults to `None`.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        q = self.to_q(x)
+        context = util.default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        out = optimized_attention(q, k, v, self.heads)
+        return self.to_out(out)
+
+
+class AttnBlock(nn.Module):
+    """#### Attention block, which applies attention to the input tensor.
+
+    #### Args:
+        - `in_channels` (int): The input channels.
+    """
+
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = cast.disable_weight_init.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.k = cast.disable_weight_init.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.v = cast.disable_weight_init.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.proj_out = cast.disable_weight_init.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+
+        if Device.xformers_enabled_vae():
+            logging.info("Using xformers attention in VAE")
+            self.optimized_attention = AttentionMethods.xformers_attention
+        else:
+            logging.info("Using pytorch attention in VAE")
+            self.optimized_attention = AttentionMethods.pytorch_attention
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass of the attention block.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        h_ = self.optimized_attention(q, k, v)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+def make_attn(in_channels: int, attn_type: str = "vanilla") -> AttnBlock:
+    """#### Make an attention block.
+
+    #### Args:
+        - `in_channels` (int): The input channels.
+        - `attn_type` (str, optional): The attention type. Defaults to "vanilla".
+
+    #### Returns:
+        - `AttnBlock`: A class instance of the attention block.
+    """
+    return AttnBlock(in_channels)

modules/Attention/AttentionMethods.py

@@ -0,0 +1,120 @@
+try :
+    import xformers
+except ImportError:
+    pass
+import torch
+
+
+def attention_xformers(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask=None
+) -> torch.Tensor:
+    """#### Make an attention call using xformers. Fastest attention implementation.
+
+    #### Args:
+        - `q` (torch.Tensor): The query tensor.
+        - `k` (torch.Tensor): The key tensor, must have the same shape as `q`.
+        - `v` (torch.Tensor): The value tensor, must have the same shape as `q`.
+        - `heads` (int): The number of heads, must be a divisor of the hidden dimension.
+        - `mask` (torch.Tensor, optional): The mask tensor. Defaults to `None`.
+
+    #### Returns:
+        - `torch.Tensor`: The output tensor.
+    """
+    b, _, dim_head = q.shape
+    dim_head //= heads
+
+    q, k, v = map(
+        lambda t: t.unsqueeze(3)
+        .reshape(b, -1, heads, dim_head)
+        .permute(0, 2, 1, 3)
+        .reshape(b * heads, -1, dim_head)
+        .contiguous(),
+        (q, k, v),
+    )
+
+    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
+
+    out = (
+        out.unsqueeze(0)
+        .reshape(b, heads, -1, dim_head)
+        .permute(0, 2, 1, 3)
+        .reshape(b, -1, heads * dim_head)
+    )
+    return out
+
+
+def attention_pytorch(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask=None
+) -> torch.Tensor:
+    """#### Make an attention call using PyTorch.
+
+    #### Args:
+        - `q` (torch.Tensor): The query tensor.
+        - `k` (torch.Tensor): The key tensor, must have the same shape as `q.
+        - `v` (torch.Tensor): The value tensor, must have the same shape as `q.
+        - `heads` (int): The number of heads, must be a divisor of the hidden dimension.
+        - `mask` (torch.Tensor, optional): The mask tensor. Defaults to `None`.
+
+    #### Returns:
+        - `torch.Tensor`: The output tensor.
+    """
+    b, _, dim_head = q.shape
+    dim_head //= heads
+    q, k, v = map(
+        lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
+        (q, k, v),
+    )
+
+    out = torch.nn.functional.scaled_dot_product_attention(
+        q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
+    )
+    out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+    return out
+
+
+def xformers_attention(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
+) -> torch.Tensor:
+    """#### Compute attention using xformers.
+
+    #### Args:
+        - `q` (torch.Tensor): The query tensor.
+        - `k` (torch.Tensor): The key tensor, must have the same shape as `q`.
+        - `v` (torch.Tensor): The value tensor, must have the same shape as `q`.
+
+    Returns:
+        - `torch.Tensor`: The output tensor.
+    """
+    B, C, H, W = q.shape
+    q, k, v = map(
+        lambda t: t.view(B, C, -1).transpose(1, 2).contiguous(),
+        (q, k, v),
+    )
+    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
+    out = out.transpose(1, 2).reshape(B, C, H, W)
+    return out
+
+
+def pytorch_attention(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
+) -> torch.Tensor:
+    """#### Compute attention using PyTorch.
+
+    #### Args:
+        - `q` (torch.Tensor): The query tensor.
+        - `k` (torch.Tensor): The key tensor, must have the same shape as `q.
+        - `v` (torch.Tensor): The value tensor, must have the same shape as `q.
+
+    #### Returns:
+        - `torch.Tensor`: The output tensor.
+    """
+    B, C, H, W = q.shape
+    q, k, v = map(
+        lambda t: t.view(B, 1, C, -1).transpose(2, 3).contiguous(),
+        (q, k, v),
+    )
+    out = torch.nn.functional.scaled_dot_product_attention(
+        q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False
+    )
+    out = out.transpose(2, 3).reshape(B, C, H, W)
+    return out

modules/AutoDetailer/AD_util.py

@@ -0,0 +1,245 @@
+from typing import List
+import cv2
+import numpy as np
+import torch
+from ultralytics import YOLO
+from PIL import Image
+
+orig_torch_load = torch.load
+
+# importing YOLO breaking original torch.load capabilities
+torch.load = orig_torch_load
+
+
+def load_yolo(model_path: str) -> YOLO:
+    """#### Load YOLO model.
+
+    #### Args:
+        - `model_path` (str): The path to the YOLO model.
+
+    #### Returns:
+        - `YOLO`: The YOLO model initialized with the specified model path.
+    """
+    try:
+        return YOLO(model_path)
+    except ModuleNotFoundError:
+        print("please download yolo model")
+
+
+def inference_bbox(
+    model: YOLO,
+    image: Image.Image,
+    confidence: float = 0.3,
+    device: str = "",
+) -> List:
+    """#### Perform inference on an image and return bounding boxes.
+
+    #### Args:
+        - `model` (YOLO): The YOLO model.
+        - `image` (Image.Image): The image to perform inference on.
+        - `confidence` (float): The confidence threshold for the bounding boxes.
+        - `device` (str): The device to run the model on.
+
+    #### Returns:
+        - `List[List[str, List[int], np.ndarray, float]]`: The list of bounding boxes.
+    """
+    pred = model(image, conf=confidence, device=device)
+
+    bboxes = pred[0].boxes.xyxy.cpu().numpy()
+    cv2_image = np.array(image)
+    cv2_image = cv2_image[:, :, ::-1].copy()  # Convert RGB to BGR for cv2 processing
+    cv2_gray = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2GRAY)
+
+    segms = []
+    for x0, y0, x1, y1 in bboxes:
+        cv2_mask = np.zeros(cv2_gray.shape, np.uint8)
+        cv2.rectangle(cv2_mask, (int(x0), int(y0)), (int(x1), int(y1)), 255, -1)
+        cv2_mask_bool = cv2_mask.astype(bool)
+        segms.append(cv2_mask_bool)
+
+    results = [[], [], [], []]
+    for i in range(len(bboxes)):
+        results[0].append(pred[0].names[int(pred[0].boxes[i].cls.item())])
+        results[1].append(bboxes[i])
+        results[2].append(segms[i])
+        results[3].append(pred[0].boxes[i].conf.cpu().numpy())
+
+    return results
+
+
+def create_segmasks(results: List) -> List:
+    """#### Create segmentation masks from the results of the inference.
+
+    #### Args:
+        - `results` (List[List[str, List[int], np.ndarray, float]]): The results of the inference.
+
+    #### Returns:
+        - `List[List[int], np.ndarray, float]`: The list of segmentation masks.
+    """
+    bboxs = results[1]
+    segms = results[2]
+    confidence = results[3]
+
+    results = []
+    for i in range(len(segms)):
+        item = (bboxs[i], segms[i].astype(np.float32), confidence[i])
+        results.append(item)
+    return results
+
+
+def dilate_masks(segmasks: List, dilation_factor: int, iter: int = 1) -> List:
+    """#### Dilate the segmentation masks.
+
+    #### Args:
+        - `segmasks` (List[List[int], np.ndarray, float]): The segmentation masks.
+        - `dilation_factor` (int): The dilation factor.
+        - `iter` (int): The number of iterations.
+
+    #### Returns:
+        - `List[List[int], np.ndarray, float]`: The dilated segmentation masks.
+    """
+    dilated_masks = []
+    kernel = np.ones((abs(dilation_factor), abs(dilation_factor)), np.uint8)
+
+    for i in range(len(segmasks)):
+        cv2_mask = segmasks[i][1]
+
+        dilated_mask = cv2.dilate(cv2_mask, kernel, iter)
+
+        item = (segmasks[i][0], dilated_mask, segmasks[i][2])
+        dilated_masks.append(item)
+
+    return dilated_masks
+
+
+def normalize_region(limit: int, startp: int, size: int) -> List:
+    """#### Normalize the region.
+
+    #### Args:
+        - `limit` (int): The limit.
+        - `startp` (int): The start point.
+        - `size` (int): The size.
+
+    #### Returns:
+        - `List[int]`: The normalized start and end points.
+    """
+    if startp < 0:
+        new_endp = min(limit, size)
+        new_startp = 0
+    elif startp + size > limit:
+        new_startp = max(0, limit - size)
+        new_endp = limit
+    else:
+        new_startp = startp
+        new_endp = min(limit, startp + size)
+
+    return int(new_startp), int(new_endp)
+
+
+def make_crop_region(w: int, h: int, bbox: List, crop_factor: float) -> List:
+    """#### Make the crop region.
+
+    #### Args:
+        - `w` (int): The width.
+        - `h` (int): The height.
+        - `bbox` (List[int]): The bounding box.
+        - `crop_factor` (float): The crop factor.
+
+    #### Returns:
+        - `List[x1: int, y1: int, x2: int, y2: int]`: The crop region.
+    """
+    x1 = bbox[0]
+    y1 = bbox[1]
+    x2 = bbox[2]
+    y2 = bbox[3]
+
+    bbox_w = x2 - x1
+    bbox_h = y2 - y1
+
+    crop_w = bbox_w * crop_factor
+    crop_h = bbox_h * crop_factor
+
+    kernel_x = x1 + bbox_w / 2
+    kernel_y = y1 + bbox_h / 2
+
+    new_x1 = int(kernel_x - crop_w / 2)
+    new_y1 = int(kernel_y - crop_h / 2)
+
+    # make sure position in (w,h)
+    new_x1, new_x2 = normalize_region(w, new_x1, crop_w)
+    new_y1, new_y2 = normalize_region(h, new_y1, crop_h)
+
+    return [new_x1, new_y1, new_x2, new_y2]
+
+
+def crop_ndarray2(npimg: np.ndarray, crop_region: List) -> np.ndarray:
+    """#### Crop the ndarray in 2 dimensions.
+
+    #### Args:
+        - `npimg` (np.ndarray): The ndarray to crop.
+        - `crop_region` (List[int]): The crop region.
+
+    #### Returns:
+        - `np.ndarray`: The cropped ndarray.
+    """
+    x1 = crop_region[0]
+    y1 = crop_region[1]
+    x2 = crop_region[2]
+    y2 = crop_region[3]
+
+    cropped = npimg[y1:y2, x1:x2]
+
+    return cropped
+
+
+def crop_ndarray4(npimg: np.ndarray, crop_region: List) -> np.ndarray:
+    """#### Crop the ndarray in 4 dimensions.
+
+    #### Args:
+        - `npimg` (np.ndarray): The ndarray to crop.
+        - `crop_region` (List[int]): The crop region.
+
+    #### Returns:
+        - `np.ndarray`: The cropped ndarray.
+    """
+    x1 = crop_region[0]
+    y1 = crop_region[1]
+    x2 = crop_region[2]
+    y2 = crop_region[3]
+
+    cropped = npimg[:, y1:y2, x1:x2, :]
+
+    return cropped
+
+
+def crop_image(image: Image.Image, crop_region: List) -> Image.Image:
+    """#### Crop the image.
+
+    #### Args:
+        - `image` (Image.Image): The image to crop.
+        - `crop_region` (List[int]): The crop region.
+
+    #### Returns:
+        - `Image.Image`: The cropped image.
+    """
+    return crop_ndarray4(image, crop_region)
+
+
+def segs_scale_match(segs: List[np.ndarray], target_shape: List) -> List:
+    """#### Match the scale of the segmentation masks.
+
+    #### Args:
+        - `segs` (List[np.ndarray]): The segmentation masks.
+        - `target_shape` (List[int]): The target shape.
+
+    #### Returns:
+        - `List[np.ndarray]`: The matched segmentation masks.
+    """
+    h = segs[0][0]
+    w = segs[0][1]
+
+    th = target_shape[1]
+    tw = target_shape[2]
+
+    if (h == th and w == tw) or h == 0 or w == 0:
+        return segs

modules/AutoDetailer/ADetailer.py

@@ -0,0 +1,952 @@
+import math
+import torch
+from typing import Any, Dict, Optional, Tuple
+
+from modules.AutoDetailer import AD_util, bbox, tensor_util
+from modules.AutoDetailer import SEGS
+from modules.Utilities import util
+from modules.AutoEncoders import VariationalAE
+from modules.Device import Device
+from modules.sample import ksampler_util, samplers, sampling, sampling_util
+
+# FIXME: Improve slow inference times
+
+
+class DifferentialDiffusion:
+    """#### Class for applying differential diffusion to a model."""
+
+    def apply(self, model: torch.nn.Module) -> Tuple[torch.nn.Module]:
+        """#### Apply differential diffusion to a model.
+
+        #### Args:
+            - `model` (torch.nn.Module): The input model.
+
+        #### Returns:
+            - `Tuple[torch.nn.Module]`: The modified model.
+        """
+        model = model.clone()
+        model.set_model_denoise_mask_function(self.forward)
+        return (model,)
+
+    def forward(
+        self,
+        sigma: torch.Tensor,
+        denoise_mask: torch.Tensor,
+        extra_options: Dict[str, Any],
+    ) -> torch.Tensor:
+        """#### Forward function for differential diffusion.
+
+        #### Args:
+            - `sigma` (torch.Tensor): The sigma tensor.
+            - `denoise_mask` (torch.Tensor): The denoise mask tensor.
+            - `extra_options` (Dict[str, Any]): Additional options.
+
+        #### Returns:
+            - `torch.Tensor`: The processed denoise mask tensor.
+        """
+        model = extra_options["model"]
+        step_sigmas = extra_options["sigmas"]
+        sigma_to = model.inner_model.model_sampling.sigma_min
+        sigma_from = step_sigmas[0]
+
+        ts_from = model.inner_model.model_sampling.timestep(sigma_from)
+        ts_to = model.inner_model.model_sampling.timestep(sigma_to)
+        current_ts = model.inner_model.model_sampling.timestep(sigma[0])
+
+        threshold = (current_ts - ts_to) / (ts_from - ts_to)
+
+        return (denoise_mask >= threshold).to(denoise_mask.dtype)
+
+
+def to_latent_image(pixels: torch.Tensor, vae: VariationalAE.VAE) -> torch.Tensor:
+    """#### Convert pixels to a latent image using a VAE.
+
+    #### Args:
+        - `pixels` (torch.Tensor): The input pixel tensor.
+        - `vae` (VariationalAE.VAE): The VAE model.
+
+    #### Returns:
+        - `torch.Tensor`: The latent image tensor.
+    """
+    pixels.shape[1]
+    pixels.shape[2]
+    return VariationalAE.VAEEncode().encode(vae, pixels)[0]
+
+
+def calculate_sigmas2(
+    model: torch.nn.Module, sampler: str, scheduler: str, steps: int
+) -> torch.Tensor:
+    """#### Calculate sigmas for a model.
+
+    #### Args:
+        - `model` (torch.nn.Module): The input model.
+        - `sampler` (str): The sampler name.
+        - `scheduler` (str): The scheduler name.
+        - `steps` (int): The number of steps.
+
+    #### Returns:
+        - `torch.Tensor`: The calculated sigmas.
+    """
+    return ksampler_util.calculate_sigmas(
+        model.get_model_object("model_sampling"), scheduler, steps
+    )
+
+
+def get_noise_sampler(
+    x: torch.Tensor, cpu: bool, total_sigmas: torch.Tensor, **kwargs
+) -> Optional[sampling_util.BrownianTreeNoiseSampler]:
+    """#### Get a noise sampler.
+
+    #### Args:
+        - `x` (torch.Tensor): The input tensor.
+        - `cpu` (bool): Whether to use CPU.
+        - `total_sigmas` (torch.Tensor): The total sigmas tensor.
+        - `kwargs` (dict): Additional arguments.
+
+    #### Returns:
+        - `Optional[sampling_util.BrownianTreeNoiseSampler]`: The noise sampler.
+    """
+    if "extra_args" in kwargs and "seed" in kwargs["extra_args"]:
+        sigma_min, sigma_max = total_sigmas[total_sigmas > 0].min(), total_sigmas.max()
+        seed = kwargs["extra_args"].get("seed", None)
+        return sampling_util.BrownianTreeNoiseSampler(
+            x, sigma_min, sigma_max, seed=seed, cpu=cpu
+        )
+    return None
+
+
+def ksampler2(
+    sampler_name: str,
+    total_sigmas: torch.Tensor,
+    extra_options: Dict[str, Any] = {},
+    inpaint_options: Dict[str, Any] = {},
+    pipeline: bool = False,
+) -> sampling.KSAMPLER:
+    """#### Get a ksampler.
+
+    #### Args:
+        - `sampler_name` (str): The sampler name.
+        - `total_sigmas` (torch.Tensor): The total sigmas tensor.
+        - `extra_options` (Dict[str, Any], optional): Additional options. Defaults to {}.
+        - `inpaint_options` (Dict[str, Any], optional): Inpaint options. Defaults to {}.
+        - `pipeline` (bool, optional): Whether to use pipeline. Defaults to False.
+
+    #### Returns:
+        - `sampling.KSAMPLER`: The ksampler.
+    """
+    if sampler_name == "dpmpp_2m_sde":
+
+        def sample_dpmpp_sde(model, x, sigmas, pipeline, **kwargs):
+            noise_sampler = get_noise_sampler(x, True, total_sigmas, **kwargs)
+            if noise_sampler is not None:
+                kwargs["noise_sampler"] = noise_sampler
+
+            return samplers.sample_dpmpp_2m_sde(
+                model, x, sigmas, pipeline=pipeline, **kwargs
+            )
+
+        sampler_function = sample_dpmpp_sde
+
+    else:
+        return sampling.sampler_object(sampler_name, pipeline=pipeline)
+
+    return sampling.KSAMPLER(sampler_function, extra_options, inpaint_options)
+
+
+class Noise_RandomNoise:
+    """#### Class for generating random noise."""
+
+    def __init__(self, seed: int):
+        """#### Initialize the Noise_RandomNoise class.
+
+        #### Args:
+            - `seed` (int): The seed for random noise.
+        """
+        self.seed = seed
+
+    def generate_noise(self, input_latent: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """#### Generate random noise.
+
+        #### Args:
+            - `input_latent` (Dict[str, torch.Tensor]): The input latent tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The generated noise tensor.
+        """
+        latent_image = input_latent["samples"]
+        batch_inds = (
+            input_latent["batch_index"] if "batch_index" in input_latent else None
+        )
+        return ksampler_util.prepare_noise(latent_image, self.seed, batch_inds)
+
+
+def sample_with_custom_noise(
+    model: torch.nn.Module,
+    add_noise: bool,
+    noise_seed: int,
+    cfg: int,
+    positive: Any,
+    negative: Any,
+    sampler: Any,
+    sigmas: torch.Tensor,
+    latent_image: Dict[str, torch.Tensor],
+    noise: Optional[torch.Tensor] = None,
+    callback: Optional[callable] = None,
+    pipeline: bool = False,
+) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+    """#### Sample with custom noise.
+
+    #### Args:
+        - `model` (torch.nn.Module): The input model.
+        - `add_noise` (bool): Whether to add noise.
+        - `noise_seed` (int): The noise seed.
+        - `cfg` (int): Classifier-Free Guidance Scale
+        - `positive` (Any): The positive prompt.
+        - `negative` (Any): The negative prompt.
+        - `sampler` (Any): The sampler.
+        - `sigmas` (torch.Tensor): The sigmas tensor.
+        - `latent_image` (Dict[str, torch.Tensor]): The latent image tensor.
+        - `noise` (Optional[torch.Tensor], optional): The noise tensor. Defaults to None.
+        - `callback` (Optional[callable], optional): The callback function. Defaults to None.
+        - `pipeline` (bool, optional): Whether to use pipeline. Defaults to False.
+
+    #### Returns:
+        - `Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]`: The sampled and denoised tensors.
+    """
+    latent = latent_image
+    latent_image = latent["samples"]
+
+    out = latent.copy()
+    out["samples"] = latent_image
+
+    if noise is None:
+        noise = Noise_RandomNoise(noise_seed).generate_noise(out)
+
+    noise_mask = None
+    if "noise_mask" in latent:
+        noise_mask = latent["noise_mask"]
+
+    disable_pbar = not util.PROGRESS_BAR_ENABLED
+
+    device = Device.get_torch_device()
+
+    noise = noise.to(device)
+    latent_image = latent_image.to(device)
+    if noise_mask is not None:
+        noise_mask = noise_mask.to(device)
+
+    samples = sampling.sample_custom(
+        model,
+        noise,
+        cfg,
+        sampler,
+        sigmas,
+        positive,
+        negative,
+        latent_image,
+        noise_mask=noise_mask,
+        disable_pbar=disable_pbar,
+        seed=noise_seed,
+        pipeline=pipeline,
+    )
+
+    samples = samples.to(Device.intermediate_device())
+
+    out["samples"] = samples
+    out_denoised = out
+    return out, out_denoised
+
+
+def separated_sample(
+    model: torch.nn.Module,
+    add_noise: bool,
+    seed: int,
+    steps: int,
+    cfg: int,
+    sampler_name: str,
+    scheduler: str,
+    positive: Any,
+    negative: Any,
+    latent_image: Dict[str, torch.Tensor],
+    start_at_step: Optional[int],
+    end_at_step: Optional[int],
+    return_with_leftover_noise: bool,
+    sigma_ratio: float = 1.0,
+    sampler_opt: Optional[Dict[str, Any]] = None,
+    noise: Optional[torch.Tensor] = None,
+    callback: Optional[callable] = None,
+    scheduler_func: Optional[callable] = None,
+    pipeline: bool = False,
+) -> Dict[str, torch.Tensor]:
+    """#### Perform separated sampling.
+
+    #### Args:
+        - `model` (torch.nn.Module): The input model.
+        - `add_noise` (bool): Whether to add noise.
+        - `seed` (int): The seed for random noise.
+        - `steps` (int): The number of steps.
+        - `cfg` (int): Classifier-Free Guidance Scale
+        - `sampler_name` (str): The sampler name.
+        - `scheduler` (str): The scheduler name.
+        - `positive` (Any): The positive prompt.
+        - `negative` (Any): The negative prompt.
+        - `latent_image` (Dict[str, torch.Tensor]): The latent image tensor.
+        - `start_at_step` (Optional[int]): The step to start at.
+        - `end_at_step` (Optional[int]): The step to end at.
+        - `return_with_leftover_noise` (bool): Whether to return with leftover noise.
+        - `sigma_ratio` (float, optional): The sigma ratio. Defaults to 1.0.
+        - `sampler_opt` (Optional[Dict[str, Any]], optional): The sampler options. Defaults to None.
+        - `noise` (Optional[torch.Tensor], optional): The noise tensor. Defaults to None.
+        - `callback` (Optional[callable], optional): The callback function. Defaults to None.
+        - `scheduler_func` (Optional[callable], optional): The scheduler function. Defaults to None.
+        - `pipeline` (bool, optional): Whether to use pipeline. Defaults to False.
+
+    #### Returns:
+        - `Dict[str, torch.Tensor]`: The sampled tensor.
+    """
+    total_sigmas = calculate_sigmas2(model, sampler_name, scheduler, steps)
+
+    sigmas = total_sigmas
+
+    if start_at_step is not None:
+        sigmas = sigmas[start_at_step:] * sigma_ratio
+
+    impact_sampler = ksampler2(sampler_name, total_sigmas, pipeline=pipeline)
+
+    res = sample_with_custom_noise(
+        model,
+        add_noise,
+        seed,
+        cfg,
+        positive,
+        negative,
+        impact_sampler,
+        sigmas,
+        latent_image,
+        noise=noise,
+        callback=callback,
+        pipeline=pipeline,
+    )
+
+    return res[1]
+
+
+def ksampler_wrapper(
+    model: torch.nn.Module,
+    seed: int,
+    steps: int,
+    cfg: int,
+    sampler_name: str,
+    scheduler: str,
+    positive: Any,
+    negative: Any,
+    latent_image: Dict[str, torch.Tensor],
+    denoise: float,
+    refiner_ratio: Optional[float] = None,
+    refiner_model: Optional[torch.nn.Module] = None,
+    refiner_clip: Optional[Any] = None,
+    refiner_positive: Optional[Any] = None,
+    refiner_negative: Optional[Any] = None,
+    sigma_factor: float = 1.0,
+    noise: Optional[torch.Tensor] = None,
+    scheduler_func: Optional[callable] = None,
+    pipeline: bool = False,
+) -> Dict[str, torch.Tensor]:
+    """#### Wrapper for ksampler.
+
+    #### Args:
+        - `model` (torch.nn.Module): The input model.
+        - `seed` (int): The seed for random noise.
+        - `steps` (int): The number of steps.
+        - `cfg` (int): Classifier-Free Guidance Scale
+        - `sampler_name` (str): The sampler name.
+        - `scheduler` (str): The scheduler name.
+        - `positive` (Any): The positive prompt.
+        - `negative` (Any): The negative prompt.
+        - `latent_image` (Dict[str, torch.Tensor]): The latent image tensor.
+        - `denoise` (float): The denoise factor.
+        - `refiner_ratio` (Optional[float], optional): The refiner ratio. Defaults to None.
+        - `refiner_model` (Optional[torch.nn.Module], optional): The refiner model. Defaults to None.
+        - `refiner_clip` (Optional[Any], optional): The refiner clip. Defaults to None.
+        - `refiner_positive` (Optional[Any], optional): The refiner positive prompt. Defaults to None.
+        - `refiner_negative` (Optional[Any], optional): The refiner negative prompt. Defaults to None.
+        - `sigma_factor` (float, optional): The sigma factor. Defaults to 1.0.
+        - `noise` (Optional[torch.Tensor], optional): The noise tensor. Defaults to None.
+        - `scheduler_func` (Optional[callable], optional): The scheduler function. Defaults to None.
+        - `pipeline` (bool, optional): Whether to use pipeline. Defaults to False.
+
+    #### Returns:
+        - `Dict[str, torch.Tensor]`: The refined latent tensor.
+    """
+    advanced_steps = math.floor(steps / denoise)
+    start_at_step = advanced_steps - steps
+    end_at_step = start_at_step + steps
+    refined_latent = separated_sample(
+        model,
+        True,
+        seed,
+        advanced_steps,
+        cfg,
+        sampler_name,
+        scheduler,
+        positive,
+        negative,
+        latent_image,
+        start_at_step,
+        end_at_step,
+        False,
+        sigma_ratio=sigma_factor,
+        noise=noise,
+        scheduler_func=scheduler_func,
+        pipeline=pipeline,
+    )
+
+    return refined_latent
+
+
+def enhance_detail(
+    image: torch.Tensor,
+    model: torch.nn.Module,
+    clip: Any,
+    vae: VariationalAE.VAE,
+    guide_size: int,
+    guide_size_for_bbox: bool,
+    max_size: int,
+    bbox: Tuple[int, int, int, int],
+    seed: int,
+    steps: int,
+    cfg: int,
+    sampler_name: str,
+    scheduler: str,
+    positive: Any,
+    negative: Any,
+    denoise: float,
+    noise_mask: Optional[torch.Tensor],
+    force_inpaint: bool,
+    wildcard_opt: Optional[Any] = None,
+    wildcard_opt_concat_mode: Optional[Any] = None,
+    detailer_hook: Optional[callable] = None,
+    refiner_ratio: Optional[float] = None,
+    refiner_model: Optional[torch.nn.Module] = None,
+    refiner_clip: Optional[Any] = None,
+    refiner_positive: Optional[Any] = None,
+    refiner_negative: Optional[Any] = None,
+    control_net_wrapper: Optional[Any] = None,
+    cycle: int = 1,
+    inpaint_model: bool = False,
+    noise_mask_feather: int = 0,
+    scheduler_func: Optional[callable] = None,
+    pipeline: bool = False,
+) -> Tuple[torch.Tensor, Optional[Any]]:
+    """#### Enhance detail of an image.
+
+    #### Args:
+        - `image` (torch.Tensor): The input image tensor.
+        - `model` (torch.nn.Module): The model.
+        - `clip` (Any): The clip model.
+        - `vae` (VariationalAE.VAE): The VAE model.
+        - `guide_size` (int): The guide size.
+        - `guide_size_for_bbox` (bool): Whether to use guide size for bbox.
+        - `max_size` (int): The maximum size.
+        - `bbox` (Tuple[int, int, int, int]): The bounding box.
+        - `seed` (int): The seed for random noise.
+        - `steps` (int): The number of steps.
+        - `cfg` (int): Classifier-Free Guidance Scale
+        - `sampler_name` (str): The sampler name.
+        - `scheduler` (str): The scheduler name.
+        - `positive` (Any): The positive prompt.
+        - `negative` (Any): The negative prompt.
+        - `denoise` (float): The denoise factor.
+        - `noise_mask` (Optional[torch.Tensor]): The noise mask tensor.
+        - `force_inpaint` (bool): Whether to force inpaint.
+        - `wildcard_opt` (Optional[Any], optional): The wildcard options. Defaults to None.
+        - `wildcard_opt_concat_mode` (Optional[Any], optional): The wildcard concat mode. Defaults to None.
+        - `detailer_hook` (Optional[callable], optional): The detailer hook. Defaults to None.
+        - `refiner_ratio` (Optional[float], optional): The refiner ratio. Defaults to None.
+        - `refiner_model` (Optional[torch.nn.Module], optional): The refiner model. Defaults to None.
+        - `refiner_clip` (Optional[Any], optional): The refiner clip. Defaults to None.
+        - `refiner_positive` (Optional[Any], optional): The refiner positive prompt. Defaults to None.
+        - `refiner_negative` (Optional[Any], optional): The refiner negative prompt. Defaults to None.
+        - `control_net_wrapper` (Optional[Any], optional): The control net wrapper. Defaults to None.
+        - `cycle` (int, optional): The number of cycles. Defaults to 1.
+        - `inpaint_model` (bool, optional): Whether to use inpaint model. Defaults to False.
+        - `noise_mask_feather` (int, optional): The noise mask feather. Defaults to 0.
+        - `scheduler_func` (Optional[callable], optional): The scheduler function. Defaults to None.
+        - `pipeline` (bool, optional): Whether to use pipeline. Defaults to False.
+
+    #### Returns:
+        - `Tuple[torch.Tensor, Optional[Any]]`: The refined image tensor and optional cnet_pils.
+    """
+    if noise_mask is not None:
+        noise_mask = tensor_util.tensor_gaussian_blur_mask(
+            noise_mask, noise_mask_feather
+        )
+        noise_mask = noise_mask.squeeze(3)
+
+    h = image.shape[1]
+    w = image.shape[2]
+
+    bbox_h = bbox[3] - bbox[1]
+    bbox_w = bbox[2] - bbox[0]
+
+    # for cropped_size
+    upscale = guide_size / min(w, h)
+
+    new_w = int(w * upscale)
+    new_h = int(h * upscale)
+
+    if new_w > max_size or new_h > max_size:
+        upscale *= max_size / max(new_w, new_h)
+        new_w = int(w * upscale)
+        new_h = int(h * upscale)
+
+    if upscale <= 1.0 or new_w == 0 or new_h == 0:
+        print("Detailer: force inpaint")
+        upscale = 1.0
+        new_w = w
+        new_h = h
+
+    print(
+        f"Detailer: segment upscale for ({bbox_w, bbox_h}) | crop region {w, h} x {upscale} -> {new_w, new_h}"
+    )
+
+    # upscale
+    upscaled_image = tensor_util.tensor_resize(image, new_w, new_h)
+
+    cnet_pils = None
+
+    # prepare mask
+    latent_image = to_latent_image(upscaled_image, vae)
+    if noise_mask is not None:
+        latent_image["noise_mask"] = noise_mask
+
+    refined_latent = latent_image
+
+    # ksampler
+    for i in range(0, cycle):
+        (
+            model2,
+            seed2,
+            steps2,
+            cfg2,
+            sampler_name2,
+            scheduler2,
+            positive2,
+            negative2,
+            _upscaled_latent2,
+            denoise2,
+        ) = (
+            model,
+            seed + i,
+            steps,
+            cfg,
+            sampler_name,
+            scheduler,
+            positive,
+            negative,
+            latent_image,
+            denoise,
+        )
+        noise = None
+
+        refined_latent = ksampler_wrapper(
+            model2,
+            seed2,
+            steps2,
+            cfg2,
+            sampler_name2,
+            scheduler2,
+            positive2,
+            negative2,
+            refined_latent,
+            denoise2,
+            refiner_ratio,
+            refiner_model,
+            refiner_clip,
+            refiner_positive,
+            refiner_negative,
+            noise=noise,
+            scheduler_func=scheduler_func,
+            pipeline=pipeline,
+        )
+
+    # non-latent downscale - latent downscale cause bad quality
+    try:
+        # try to decode image normally
+        refined_image = vae.decode(refined_latent["samples"])
+    except Exception:
+        # usually an out-of-memory exception from the decode, so try a tiled approach
+        refined_image = vae.decode_tiled(
+            refined_latent["samples"],
+            tile_x=64,
+            tile_y=64,
+        )
+
+    # downscale
+    refined_image = tensor_util.tensor_resize(refined_image, w, h)
+
+    # prevent mixing of device
+    refined_image = refined_image.cpu()
+
+    # don't convert to latent - latent break image
+    # preserving pil is much better
+    return refined_image, cnet_pils
+
+
+class DetailerForEach:
+    """#### Class for detailing each segment of an image."""
+
+    @staticmethod
+    def do_detail(
+        image: torch.Tensor,
+        segs: Tuple[torch.Tensor, Any],
+        model: torch.nn.Module,
+        clip: Any,
+        vae: VariationalAE.VAE,
+        guide_size: int,
+        guide_size_for_bbox: bool,
+        max_size: int,
+        seed: int,
+        steps: int,
+        cfg: int,
+        sampler_name: str,
+        scheduler: str,
+        positive: Any,
+        negative: Any,
+        denoise: float,
+        feather: int,
+        noise_mask: Optional[torch.Tensor],
+        force_inpaint: bool,
+        wildcard_opt: Optional[Any] = None,
+        detailer_hook: Optional[callable] = None,
+        refiner_ratio: Optional[float] = None,
+        refiner_model: Optional[torch.nn.Module] = None,
+        refiner_clip: Optional[Any] = None,
+        refiner_positive: Optional[Any] = None,
+        refiner_negative: Optional[Any] = None,
+        cycle: int = 1,
+        inpaint_model: bool = False,
+        noise_mask_feather: int = 0,
+        scheduler_func_opt: Optional[callable] = None,
+        pipeline: bool = False,
+    ) -> Tuple[torch.Tensor, list, list, list, list, Tuple[torch.Tensor, list]]:
+        """#### Perform detailing on each segment of an image.
+
+        #### Args:
+            - `image` (torch.Tensor): The input image tensor.
+            - `segs` (Tuple[torch.Tensor, Any]): The segments.
+            - `model` (torch.nn.Module): The model.
+            - `clip` (Any): The clip model.
+            - `vae` (VariationalAE.VAE): The VAE model.
+            - `guide_size` (int): The guide size.
+            - `guide_size_for_bbox` (bool): Whether to use guide size for bbox.
+            - `max_size` (int): The maximum size.
+            - `seed` (int): The seed for random noise.
+            - `steps` (int): The number of steps.
+            - `cfg` (int): Classifier-Free Guidance Scale.
+            - `sampler_name` (str): The sampler name.
+            - `scheduler` (str): The scheduler name.
+            - `positive` (Any): The positive prompt.
+            - `negative` (Any): The negative prompt.
+            - `denoise` (float): The denoise factor.
+            - `feather` (int): The feather value.
+            - `noise_mask` (Optional[torch.Tensor]): The noise mask tensor.
+            - `force_inpaint` (bool): Whether to force inpaint.
+            - `wildcard_opt` (Optional[Any], optional): The wildcard options. Defaults to None.
+            - `detailer_hook` (Optional[callable], optional): The detailer hook. Defaults to None.
+            - `refiner_ratio` (Optional[float], optional): The refiner ratio. Defaults to None.
+            - `refiner_model` (Optional[torch.nn.Module], optional): The refiner model. Defaults to None.
+            - `refiner_clip` (Optional[Any], optional): The refiner clip. Defaults to None.
+            - `refiner_positive` (Optional[Any], optional): The refiner positive prompt. Defaults to None.
+            - `refiner_negative` (Optional[Any], optional): The refiner negative prompt. Defaults to None.
+            - `cycle` (int, optional): The number of cycles. Defaults to 1.
+            - `inpaint_model` (bool, optional): Whether to use inpaint model. Defaults to False.
+            - `noise_mask_feather` (int, optional): The noise mask feather. Defaults to 0.
+            - `scheduler_func_opt` (Optional[callable], optional): The scheduler function. Defaults to None.
+            - `pipeline` (bool, optional): Whether to use pipeline. Defaults to False.
+
+        #### Returns:
+            - `Tuple[torch.Tensor, list, list, list, list, Tuple[torch.Tensor, list]]`: The detailed image tensor, cropped list, enhanced list, enhanced alpha list, cnet PIL list, and new segments.
+        """
+        image = image.clone()
+        enhanced_alpha_list = []
+        enhanced_list = []
+        cropped_list = []
+        cnet_pil_list = []
+
+        segs = AD_util.segs_scale_match(segs, image.shape)
+        new_segs = []
+
+        wildcard_concat_mode = None
+        wmode, wildcard_chooser = bbox.process_wildcard_for_segs(wildcard_opt)
+
+        ordered_segs = segs[1]
+
+        if (
+            noise_mask_feather > 0
+            and "denoise_mask_function" not in model.model_options
+        ):
+            model = DifferentialDiffusion().apply(model)[0]
+
+        for i, seg in enumerate(ordered_segs):
+            cropped_image = AD_util.crop_ndarray4(
+                image.cpu().numpy(), seg.crop_region
+            )  # Never use seg.cropped_image to handle overlapping area
+            cropped_image = tensor_util.to_tensor(cropped_image)
+            mask = tensor_util.to_tensor(seg.cropped_mask)
+            mask = tensor_util.tensor_gaussian_blur_mask(mask, feather)
+
+            is_mask_all_zeros = (seg.cropped_mask == 0).all().item()
+            if is_mask_all_zeros:
+                print("Detailer: segment skip [empty mask]")
+                continue
+
+            cropped_mask = seg.cropped_mask
+
+            seg_seed, wildcard_item = wildcard_chooser.get(seg)
+
+            seg_seed = seed + i if seg_seed is None else seg_seed
+
+            cropped_positive = [
+                [
+                    condition,
+                    {
+                        k: (
+                            crop_condition_mask(v, image, seg.crop_region)
+                            if k == "mask"
+                            else v
+                        )
+                        for k, v in details.items()
+                    },
+                ]
+                for condition, details in positive
+            ]
+
+            cropped_negative = [
+                [
+                    condition,
+                    {
+                        k: (
+                            crop_condition_mask(v, image, seg.crop_region)
+                            if k == "mask"
+                            else v
+                        )
+                        for k, v in details.items()
+                    },
+                ]
+                for condition, details in negative
+            ]
+
+            orig_cropped_image = cropped_image.clone()
+            enhanced_image, cnet_pils = enhance_detail(
+                cropped_image,
+                model,
+                clip,
+                vae,
+                guide_size,
+                guide_size_for_bbox,
+                max_size,
+                seg.bbox,
+                seg_seed,
+                steps,
+                cfg,
+                sampler_name,
+                scheduler,
+                cropped_positive,
+                cropped_negative,
+                denoise,
+                cropped_mask,
+                force_inpaint,
+                wildcard_opt=wildcard_item,
+                wildcard_opt_concat_mode=wildcard_concat_mode,
+                detailer_hook=detailer_hook,
+                refiner_ratio=refiner_ratio,
+                refiner_model=refiner_model,
+                refiner_clip=refiner_clip,
+                refiner_positive=refiner_positive,
+                refiner_negative=refiner_negative,
+                control_net_wrapper=seg.control_net_wrapper,
+                cycle=cycle,
+                inpaint_model=inpaint_model,
+                noise_mask_feather=noise_mask_feather,
+                scheduler_func=scheduler_func_opt,
+                pipeline=pipeline,
+            )
+
+            if enhanced_image is not None:
+                # don't latent composite-> converting to latent caused poor quality
+                # use image paste
+                image = image.cpu()
+                enhanced_image = enhanced_image.cpu()
+                tensor_util.tensor_paste(
+                    image,
+                    enhanced_image,
+                    (seg.crop_region[0], seg.crop_region[1]),
+                    mask,
+                )  # this code affecting to `cropped_image`.
+                enhanced_list.append(enhanced_image)
+
+            # Convert enhanced_pil_alpha to RGBA mode
+            enhanced_image_alpha = tensor_util.tensor_convert_rgba(enhanced_image)
+            new_seg_image = (
+                enhanced_image.numpy()
+            )  # alpha should not be applied to seg_image
+            # Apply the mask
+            mask = tensor_util.tensor_resize(
+                mask, *tensor_util.tensor_get_size(enhanced_image)
+            )
+            tensor_util.tensor_putalpha(enhanced_image_alpha, mask)
+            enhanced_alpha_list.append(enhanced_image_alpha)
+
+            cropped_list.append(orig_cropped_image)  # NOTE: Don't use `cropped_image`
+
+            new_seg = SEGS.SEG(
+                new_seg_image,
+                seg.cropped_mask,
+                seg.confidence,
+                seg.crop_region,
+                seg.bbox,
+                seg.label,
+                seg.control_net_wrapper,
+            )
+            new_segs.append(new_seg)
+
+        image_tensor = tensor_util.tensor_convert_rgb(image)
+
+        cropped_list.sort(key=lambda x: x.shape, reverse=True)
+        enhanced_list.sort(key=lambda x: x.shape, reverse=True)
+        enhanced_alpha_list.sort(key=lambda x: x.shape, reverse=True)
+
+        return (
+            image_tensor,
+            cropped_list,
+            enhanced_list,
+            enhanced_alpha_list,
+            cnet_pil_list,
+            (segs[0], new_segs),
+        )
+
+
+def empty_pil_tensor(w: int = 64, h: int = 64) -> torch.Tensor:
+    """#### Create an empty PIL tensor.
+
+    #### Args:
+        - `w` (int, optional): The width of the tensor. Defaults to 64.
+        - `h` (int, optional): The height of the tensor. Defaults to 64.
+
+    #### Returns:
+        - `torch.Tensor`: The empty tensor.
+    """
+    return torch.zeros((1, h, w, 3), dtype=torch.float32)
+
+
+class DetailerForEachTest(DetailerForEach):
+    """#### Test class for DetailerForEach."""
+
+    def doit(
+        self,
+        image: torch.Tensor,
+        segs: Any,
+        model: torch.nn.Module,
+        clip: Any,
+        vae: VariationalAE.VAE,
+        guide_size: int,
+        guide_size_for: bool,
+        max_size: int,
+        seed: int,
+        steps: int,
+        cfg: Any,
+        sampler_name: str,
+        scheduler: str,
+        positive: Any,
+        negative: Any,
+        denoise: float,
+        feather: int,
+        noise_mask: Optional[torch.Tensor],
+        force_inpaint: bool,
+        wildcard: Optional[Any],
+        detailer_hook: Optional[callable] = None,
+        cycle: int = 1,
+        inpaint_model: bool = False,
+        noise_mask_feather: int = 0,
+        scheduler_func_opt: Optional[callable] = None,
+        pipeline: bool = False,
+    ) -> Tuple[torch.Tensor, list, list, list, list]:
+        """#### Perform detail enhancement for testing.
+
+        #### Args:
+            - `image` (torch.Tensor): The input image tensor.
+            - `segs` (Any): The segments.
+            - `model` (torch.nn.Module): The model.
+            - `clip` (Any): The clip model.
+            - `vae` (VariationalAE.VAE): The VAE model.
+            - `guide_size` (int): The guide size.
+            - `guide_size_for` (bool): Whether to use guide size for.
+            - `max_size` (int): The maximum size.
+            - `seed` (int): The seed for random noise.
+            - `steps` (int): The number of steps.
+            - `cfg` (Any): The configuration.
+            - `sampler_name` (str): The sampler name.
+            - `scheduler` (str): The scheduler name.
+            - `positive` (Any): The positive prompt.
+            - `negative` (Any): The negative prompt.
+            - `denoise` (float): The denoise factor.
+            - `feather` (int): The feather value.
+            - `noise_mask` (Optional[torch.Tensor]): The noise mask tensor.
+            - `force_inpaint` (bool): Whether to force inpaint.
+            - `wildcard` (Optional[Any]): The wildcard options.
+            - `detailer_hook` (Optional[callable], optional): The detailer hook. Defaults to None.
+            - `cycle` (int, optional): The number of cycles. Defaults to 1.
+            - `inpaint_model` (bool, optional): Whether to use inpaint model. Defaults to False.
+            - `noise_mask_feather` (int, optional): The noise mask feather. Defaults to 0.
+            - `scheduler_func_opt` (Optional[callable], optional): The scheduler function. Defaults to None.
+            - `pipeline` (bool, optional): Whether to use pipeline. Defaults to False.
+
+        #### Returns:
+            - `Tuple[torch.Tensor, list, list, list, list]`: The enhanced image tensor, cropped list, cropped enhanced list, cropped enhanced alpha list, and cnet PIL list.
+        """
+        (
+            enhanced_img,
+            cropped,
+            cropped_enhanced,
+            cropped_enhanced_alpha,
+            cnet_pil_list,
+            new_segs,
+        ) = DetailerForEach.do_detail(
+            image,
+            segs,
+            model,
+            clip,
+            vae,
+            guide_size,
+            guide_size_for,
+            max_size,
+            seed,
+            steps,
+            cfg,
+            sampler_name,
+            scheduler,
+            positive,
+            negative,
+            denoise,
+            feather,
+            noise_mask,
+            force_inpaint,
+            wildcard,
+            detailer_hook,
+            cycle=cycle,
+            inpaint_model=inpaint_model,
+            noise_mask_feather=noise_mask_feather,
+            scheduler_func_opt=scheduler_func_opt,
+            pipeline=pipeline,
+        )
+
+        cnet_pil_list = [empty_pil_tensor()]
+
+        return (
+            enhanced_img,
+            cropped,
+            cropped_enhanced,
+            cropped_enhanced_alpha,
+            cnet_pil_list,
+        )

modules/AutoDetailer/SAM.py

@@ -0,0 +1,300 @@
+import os
+import numpy as np
+from segment_anything import SamPredictor, sam_model_registry
+import torch
+
+from modules.AutoDetailer import mask_util
+from modules.Device import Device
+
+
+def sam_predict(
+    predictor: SamPredictor, points: list, plabs: list, bbox: list, threshold: float
+) -> list:
+    """#### Predict masks using SAM.
+
+    #### Args:
+        - `predictor` (SamPredictor): The SAM predictor.
+        - `points` (list): List of points.
+        - `plabs` (list): List of point labels.
+        - `bbox` (list): Bounding box.
+        - `threshold` (float): Threshold for mask selection.
+
+    #### Returns:
+        - `list`: List of predicted masks.
+    """
+    point_coords = None if not points else np.array(points)
+    point_labels = None if not plabs else np.array(plabs)
+
+    box = np.array([bbox]) if bbox is not None else None
+
+    cur_masks, scores, _ = predictor.predict(
+        point_coords=point_coords, point_labels=point_labels, box=box
+    )
+
+    total_masks = []
+
+    selected = False
+    max_score = 0
+    max_mask = None
+    for idx in range(len(scores)):
+        if scores[idx] > max_score:
+            max_score = scores[idx]
+            max_mask = cur_masks[idx]
+
+        if scores[idx] >= threshold:
+            selected = True
+            total_masks.append(cur_masks[idx])
+        else:
+            pass
+
+    if not selected and max_mask is not None:
+        total_masks.append(max_mask)
+
+    return total_masks
+
+
+def is_same_device(a: torch.device, b: torch.device) -> bool:
+    """#### Check if two devices are the same.
+
+    #### Args:
+        - `a` (torch.device): The first device.
+        - `b` (torch.device): The second device.
+
+    #### Returns:
+        - `bool`: Whether the devices are the same.
+    """
+    a_device = torch.device(a) if isinstance(a, str) else a
+    b_device = torch.device(b) if isinstance(b, str) else b
+    return a_device.type == b_device.type and a_device.index == b_device.index
+
+
+class SafeToGPU:
+    """#### Class to safely move objects to GPU."""
+
+    def __init__(self, size: int):
+        self.size = size
+
+    def to_device(self, obj: torch.nn.Module, device: torch.device) -> None:
+        """#### Move an object to a device.
+
+        #### Args:
+            - `obj` (torch.nn.Module): The object to move.
+            - `device` (torch.device): The target device.
+        """
+        if is_same_device(device, "cpu"):
+            obj.to(device)
+        else:
+            if is_same_device(obj.device, "cpu"):  # cpu to gpu
+                Device.free_memory(self.size * 1.3, device)
+                if Device.get_free_memory(device) > self.size * 1.3:
+                    try:
+                        obj.to(device)
+                    except:
+                        print(
+                            f"WARN: The model is not moved to the '{device}' due to insufficient memory. [1]"
+                        )
+                else:
+                    print(
+                        f"WARN: The model is not moved to the '{device}' due to insufficient memory. [2]"
+                    )
+
+
+class SAMWrapper:
+    """#### Wrapper class for SAM model."""
+
+    def __init__(
+        self, model: torch.nn.Module, is_auto_mode: bool, safe_to_gpu: SafeToGPU = None
+    ):
+        self.model = model
+        self.safe_to_gpu = safe_to_gpu if safe_to_gpu is not None else SafeToGPU()
+        self.is_auto_mode = is_auto_mode
+
+    def prepare_device(self) -> None:
+        """#### Prepare the device for the model."""
+        if self.is_auto_mode:
+            device = Device.get_torch_device()
+            self.safe_to_gpu.to_device(self.model, device=device)
+
+    def release_device(self) -> None:
+        """#### Release the device from the model."""
+        if self.is_auto_mode:
+            self.model.to(device="cpu")
+
+    def predict(
+        self, image: np.ndarray, points: list, plabs: list, bbox: list, threshold: float
+    ) -> list:
+        """#### Predict masks using the SAM model.
+
+        #### Args:
+            - `image` (np.ndarray): The input image.
+            - `points` (list): List of points.
+            - `plabs` (list): List of point labels.
+            - `bbox` (list): Bounding box.
+            - `threshold` (float): Threshold for mask selection.
+
+        #### Returns:
+            - `list`: List of predicted masks.
+        """
+        predictor = SamPredictor(self.model)
+        predictor.set_image(image, "RGB")
+
+        return sam_predict(predictor, points, plabs, bbox, threshold)
+
+
+class SAMLoader:
+    """#### Class to load SAM models."""
+
+    def load_model(self, model_name: str, device_mode: str = "auto") -> tuple:
+        """#### Load a SAM model.
+
+        #### Args:
+            - `model_name` (str): The name of the model.
+            - `device_mode` (str, optional): The device mode. Defaults to "auto".
+
+        #### Returns:
+            - `tuple`: The loaded SAM model.
+        """
+        modelname = "./_internal/yolos/" + model_name
+
+        if "vit_h" in model_name:
+            model_kind = "vit_h"
+        elif "vit_l" in model_name:
+            model_kind = "vit_l"
+        else:
+            model_kind = "vit_b"
+
+        sam = sam_model_registry[model_kind](checkpoint=modelname)
+        size = os.path.getsize(modelname)
+        safe_to = SafeToGPU(size)
+
+        # Unless user explicitly wants to use CPU, we use GPU
+        device = Device.get_torch_device() if device_mode == "Prefer GPU" else "CPU"
+
+        if device_mode == "Prefer GPU":
+            safe_to.to_device(sam, device)
+
+        is_auto_mode = device_mode == "AUTO"
+
+        sam_obj = SAMWrapper(sam, is_auto_mode=is_auto_mode, safe_to_gpu=safe_to)
+        sam.sam_wrapper = sam_obj
+
+        print(f"Loads SAM model: {modelname} (device:{device_mode})")
+        return (sam,)
+
+
+def make_sam_mask(
+    sam: SAMWrapper,
+    segs: tuple,
+    image: torch.Tensor,
+    detection_hint: bool,
+    dilation: int,
+    threshold: float,
+    bbox_expansion: int,
+    mask_hint_threshold: float,
+    mask_hint_use_negative: bool,
+) -> torch.Tensor:
+    """#### Create a SAM mask.
+
+    #### Args:
+        - `sam` (SAMWrapper): The SAM wrapper.
+        - `segs` (tuple): Segmentation information.
+        - `image` (torch.Tensor): The input image.
+        - `detection_hint` (bool): Whether to use detection hint.
+        - `dilation` (int): Dilation value.
+        - `threshold` (float): Threshold for mask selection.
+        - `bbox_expansion` (int): Bounding box expansion value.
+        - `mask_hint_threshold` (float): Mask hint threshold.
+        - `mask_hint_use_negative` (bool): Whether to use negative mask hint.
+
+    #### Returns:
+        - `torch.Tensor`: The created SAM mask.
+    """
+    sam_obj = sam.sam_wrapper
+    sam_obj.prepare_device()
+
+    try:
+        image = np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
+
+        total_masks = []
+        # seg_shape = segs[0]
+        segs = segs[1]
+        for i in range(len(segs)):
+            bbox = segs[i].bbox
+            center = mask_util.center_of_bbox(bbox)
+            x1 = max(bbox[0] - bbox_expansion, 0)
+            y1 = max(bbox[1] - bbox_expansion, 0)
+            x2 = min(bbox[2] + bbox_expansion, image.shape[1])
+            y2 = min(bbox[3] + bbox_expansion, image.shape[0])
+            dilated_bbox = [x1, y1, x2, y2]
+            points = []
+            plabs = []
+            points.append(center)
+            plabs = [1]  # 1 = foreground point, 0 = background point
+            detected_masks = sam_obj.predict(
+                image, points, plabs, dilated_bbox, threshold
+            )
+            total_masks += detected_masks
+
+        # merge every collected masks
+        mask = mask_util.combine_masks2(total_masks)
+
+    finally:
+        sam_obj.release_device()
+
+    if mask is not None:
+        mask = mask.float()
+        mask = mask_util.dilate_mask(mask.cpu().numpy(), dilation)
+        mask = torch.from_numpy(mask)
+
+        mask = mask_util.make_3d_mask(mask)
+        return mask
+    else:
+        return None
+
+
+class SAMDetectorCombined:
+    """#### Class to combine SAM detection."""
+
+    def doit(
+        self,
+        sam_model: SAMWrapper,
+        segs: tuple,
+        image: torch.Tensor,
+        detection_hint: bool,
+        dilation: int,
+        threshold: float,
+        bbox_expansion: int,
+        mask_hint_threshold: float,
+        mask_hint_use_negative: bool,
+    ) -> tuple:
+        """#### Combine SAM detection.
+
+        #### Args:
+            - `sam_model` (SAMWrapper): The SAM wrapper.
+            - `segs` (tuple): Segmentation information.
+            - `image` (torch.Tensor): The input image.
+            - `detection_hint` (bool): Whether to use detection hint.
+            - `dilation` (int): Dilation value.
+            - `threshold` (float): Threshold for mask selection.
+            - `bbox_expansion` (int): Bounding box expansion value.
+            - `mask_hint_threshold` (float): Mask hint threshold.
+            - `mask_hint_use_negative` (bool): Whether to use negative mask hint.
+
+        #### Returns:
+            - `tuple`: The combined SAM detection result.
+        """
+        sam = make_sam_mask(
+            sam_model,
+            segs,
+            image,
+            detection_hint,
+            dilation,
+            threshold,
+            bbox_expansion,
+            mask_hint_threshold,
+            mask_hint_use_negative,
+        )
+        if sam is not None:
+            return (sam,)
+        else:
+            return None

modules/AutoDetailer/SEGS.py

@@ -0,0 +1,95 @@
+from collections import namedtuple
+import numpy as np
+import torch
+from modules.AutoDetailer import mask_util
+
+SEG = namedtuple(
+    "SEG",
+    [
+        "cropped_image",
+        "cropped_mask",
+        "confidence",
+        "crop_region",
+        "bbox",
+        "label",
+        "control_net_wrapper",
+    ],
+    defaults=[None],
+)
+
+
+def segs_bitwise_and_mask(segs: tuple, mask: torch.Tensor) -> tuple:
+    """#### Apply bitwise AND operation between segmentation masks and a given mask.
+
+    #### Args:
+        - `segs` (tuple): A tuple containing segmentation information.
+        - `mask` (torch.Tensor): The mask tensor.
+
+    #### Returns:
+        - `tuple`: A tuple containing the original segmentation and the updated items.
+    """
+    mask = mask_util.make_2d_mask(mask)
+    items = []
+
+    mask = (mask.cpu().numpy() * 255).astype(np.uint8)
+
+    for seg in segs[1]:
+        cropped_mask = (seg.cropped_mask * 255).astype(np.uint8)
+        crop_region = seg.crop_region
+
+        cropped_mask2 = mask[
+            crop_region[1] : crop_region[3], crop_region[0] : crop_region[2]
+        ]
+
+        new_mask = np.bitwise_and(cropped_mask.astype(np.uint8), cropped_mask2)
+        new_mask = new_mask.astype(np.float32) / 255.0
+
+        item = SEG(
+            seg.cropped_image,
+            new_mask,
+            seg.confidence,
+            seg.crop_region,
+            seg.bbox,
+            seg.label,
+            None,
+        )
+        items.append(item)
+
+    return segs[0], items
+
+
+class SegsBitwiseAndMask:
+    """#### Class to apply bitwise AND operation between segmentation masks and a given mask."""
+
+    def doit(self, segs: tuple, mask: torch.Tensor) -> tuple:
+        """#### Apply bitwise AND operation between segmentation masks and a given mask.
+
+        #### Args:
+            - `segs` (tuple): A tuple containing segmentation information.
+            - `mask` (torch.Tensor): The mask tensor.
+
+        #### Returns:
+            - `tuple`: A tuple containing the original segmentation and the updated items.
+        """
+        return (segs_bitwise_and_mask(segs, mask),)
+
+
+class SEGSLabelFilter:
+    """#### Class to filter segmentation labels."""
+
+    @staticmethod
+    def filter(segs: tuple, labels: list) -> tuple:
+        """#### Filter segmentation labels.
+
+        #### Args:
+            - `segs` (tuple): A tuple containing segmentation information.
+            - `labels` (list): A list of labels to filter.
+
+        #### Returns:
+            - `tuple`: A tuple containing the original segmentation and an empty list.
+        """
+        labels = set([label.strip() for label in labels])
+        return (
+            segs,
+            (segs[0], []),
+        )

modules/AutoDetailer/bbox.py

@@ -0,0 +1,203 @@
+import torch
+from ultralytics import YOLO
+from modules.AutoDetailer import SEGS, AD_util, tensor_util
+from typing import List, Tuple, Optional
+
+
+class UltraBBoxDetector:
+    """#### Class to detect bounding boxes using a YOLO model."""
+
+    bbox_model: Optional[YOLO] = None
+
+    def __init__(self, bbox_model: YOLO):
+        """#### Initialize the UltraBBoxDetector with a YOLO model.
+
+        #### Args:
+            - `bbox_model` (YOLO): The YOLO model to use for detection.
+        """
+        self.bbox_model = bbox_model
+
+    def detect(
+        self,
+        image: torch.Tensor,
+        threshold: float,
+        dilation: int,
+        crop_factor: float,
+        drop_size: int = 1,
+        detailer_hook: Optional[callable] = None,
+    ) -> Tuple[Tuple[int, int], List[SEGS.SEG]]:
+        """#### Detect bounding boxes in an image.
+
+        #### Args:
+            - `image` (torch.Tensor): The input image tensor.
+            - `threshold` (float): The detection threshold.
+            - `dilation` (int): The dilation factor for masks.
+            - `crop_factor` (float): The crop factor for bounding boxes.
+            - `drop_size` (int, optional): The minimum size of bounding boxes to keep. Defaults to 1.
+            - `detailer_hook` (callable, optional): A hook function for additional processing. Defaults to None.
+
+        #### Returns:
+            - `Tuple[Tuple[int, int], List[SEGS.SEG]]`: The shape of the image and a list of detected segments.
+        """
+        drop_size = max(drop_size, 1)
+        detected_results = AD_util.inference_bbox(
+            self.bbox_model, tensor_util.tensor2pil(image), threshold
+        )
+        segmasks = AD_util.create_segmasks(detected_results)
+
+        if dilation > 0:
+            segmasks = AD_util.dilate_masks(segmasks, dilation)
+
+        items = []
+        h = image.shape[1]
+        w = image.shape[2]
+
+        for x, label in zip(segmasks, detected_results[0]):
+            item_bbox = x[0]
+            item_mask = x[1]
+
+            y1, x1, y2, x2 = item_bbox
+
+            if (
+                x2 - x1 > drop_size and y2 - y1 > drop_size
+            ):  # minimum dimension must be (2,2) to avoid squeeze issue
+                crop_region = AD_util.make_crop_region(w, h, item_bbox, crop_factor)
+
+                cropped_image = AD_util.crop_image(image, crop_region)
+                cropped_mask = AD_util.crop_ndarray2(item_mask, crop_region)
+                confidence = x[2]
+
+                item = SEGS.SEG(
+                    cropped_image,
+                    cropped_mask,
+                    confidence,
+                    crop_region,
+                    item_bbox,
+                    label,
+                    None,
+                )
+
+                items.append(item)
+
+        shape = image.shape[1], image.shape[2]
+        segs = shape, items
+
+        return segs
+
+
+class UltraSegmDetector:
+    """#### Class to detect segments using a YOLO model."""
+
+    bbox_model: Optional[YOLO] = None
+
+    def __init__(self, bbox_model: YOLO):
+        """#### Initialize the UltraSegmDetector with a YOLO model.
+
+        #### Args:
+            - `bbox_model` (YOLO): The YOLO model to use for detection.
+        """
+        self.bbox_model = bbox_model
+
+
+class NO_SEGM_DETECTOR:
+    """#### Placeholder class for no segment detector."""
+
+    pass
+
+
+class UltralyticsDetectorProvider:
+    """#### Class to provide YOLO models for detection."""
+
+    def doit(self, model_name: str) -> Tuple[UltraBBoxDetector, UltraSegmDetector]:
+        """#### Load a YOLO model and return detectors.
+
+        #### Args:
+            - `model_name` (str): The name of the YOLO model to load.
+
+        #### Returns:
+            - `Tuple[UltraBBoxDetector, UltraSegmDetector]`: The bounding box and segment detectors.
+        """
+        model = AD_util.load_yolo("./_internal/yolos/" + model_name)
+        return UltraBBoxDetector(model), UltraSegmDetector(model)
+
+
+class BboxDetectorForEach:
+    """#### Class to detect bounding boxes for each segment."""
+
+    def doit(
+        self,
+        bbox_detector: UltraBBoxDetector,
+        image: torch.Tensor,
+        threshold: float,
+        dilation: int,
+        crop_factor: float,
+        drop_size: int,
+        labels: Optional[str] = None,
+        detailer_hook: Optional[callable] = None,
+    ) -> Tuple[Tuple[int, int], List[SEGS.SEG]]:
+        """#### Detect bounding boxes for each segment in an image.
+
+        #### Args:
+            - `bbox_detector` (UltraBBoxDetector): The bounding box detector.
+            - `image` (torch.Tensor): The input image tensor.
+            - `threshold` (float): The detection threshold.
+            - `dilation` (int): The dilation factor for masks.
+            - `crop_factor` (float): The crop factor for bounding boxes.
+            - `drop_size` (int): The minimum size of bounding boxes to keep.
+            - `labels` (str, optional): The labels to filter. Defaults to None.
+            - `detailer_hook` (callable, optional): A hook function for additional processing. Defaults to None.
+
+        #### Returns:
+            - `Tuple[Tuple[int, int], List[SEGS.SEG]]`: The shape of the image and a list of detected segments.
+        """
+        segs = bbox_detector.detect(
+            image, threshold, dilation, crop_factor, drop_size, detailer_hook
+        )
+
+        if labels is not None and labels != "":
+            labels = labels.split(",")
+            if len(labels) > 0:
+                segs, _ = SEGS.SEGSLabelFilter.filter(segs, labels)
+
+        return segs
+
+
+class WildcardChooser:
+    """#### Class to choose wildcards for segments."""
+
+    def __init__(self, items: List[Tuple[None, str]], randomize_when_exhaust: bool):
+        """#### Initialize the WildcardChooser.
+
+        #### Args:
+            - `items` (List[Tuple[None, str]]): The list of items to choose from.
+            - `randomize_when_exhaust` (bool): Whether to randomize when the list is exhausted.
+        """
+        self.i = 0
+        self.items = items
+        self.randomize_when_exhaust = randomize_when_exhaust
+
+    def get(self, seg: SEGS.SEG) -> Tuple[None, str]:
+        """#### Get the next item from the list.
+
+        #### Args:
+            - `seg` (SEGS.SEG): The segment.
+
+        #### Returns:
+            - `Tuple[None, str]`: The next item from the list.
+        """
+        item = self.items[self.i]
+        self.i += 1
+
+        return item
+
+
+def process_wildcard_for_segs(wildcard: str) -> Tuple[None, WildcardChooser]:
+    """#### Process a wildcard for segments.
+
+    #### Args:
+        - `wildcard` (str): The wildcard.
+
+    #### Returns:
+        - `Tuple[None, WildcardChooser]`: The processed wildcard and a WildcardChooser.
+    """
+    return None, WildcardChooser([(None, wildcard)], False)

modules/AutoDetailer/mask_util.py

@@ -0,0 +1,80 @@
+import numpy as np
+import torch
+
+
+def center_of_bbox(bbox: list) -> tuple[float, float]:
+    """#### Calculate the center of a bounding box.
+
+    #### Args:
+        - `bbox` (list): The bounding box coordinates [x1, y1, x2, y2].
+
+    #### Returns:
+        - `tuple[float, float]`: The center coordinates (x, y).
+    """
+    w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
+    return bbox[0] + w / 2, bbox[1] + h / 2
+
+
+def make_2d_mask(mask: torch.Tensor) -> torch.Tensor:
+    """#### Convert a mask to 2D.
+
+    #### Args:
+        - `mask` (torch.Tensor): The input mask tensor.
+
+    #### Returns:
+        - `torch.Tensor`: The 2D mask tensor.
+    """
+    if len(mask.shape) == 4:
+        return mask.squeeze(0).squeeze(0)
+    elif len(mask.shape) == 3:
+        return mask.squeeze(0)
+    return mask
+
+
+def combine_masks2(masks: list) -> torch.Tensor | None:
+    """#### Combine multiple masks into one.
+
+    #### Args:
+        - `masks` (list): A list of mask tensors.
+
+    #### Returns:
+        - `torch.Tensor | None`: The combined mask tensor or None if no masks are provided.
+    """
+    try:
+        mask = torch.from_numpy(np.array(masks[0]).astype(np.uint8))
+    except:
+        print("No Human Detected")
+        return None
+    return mask
+
+
+def dilate_mask(
+    mask: torch.Tensor, dilation_factor: int, iter: int = 1
+) -> torch.Tensor:
+    """#### Dilate a mask.
+
+    #### Args:
+        - `mask` (torch.Tensor): The input mask tensor.
+        - `dilation_factor` (int): The dilation factor.
+        - `iter` (int, optional): The number of iterations. Defaults to 1.
+
+    #### Returns:
+        - `torch.Tensor`: The dilated mask tensor.
+    """
+    return make_2d_mask(mask)
+
+
+def make_3d_mask(mask: torch.Tensor) -> torch.Tensor:
+    """#### Convert a mask to 3D.
+
+    #### Args:
+        - `mask` (torch.Tensor): The input mask tensor.
+
+    #### Returns:
+        - `torch.Tensor`: The 3D mask tensor.
+    """
+    if len(mask.shape) == 4:
+        return mask.squeeze(0)
+    elif len(mask.shape) == 2:
+        return mask.unsqueeze(0)
+    return mask

modules/AutoDetailer/tensor_util.py

@@ -0,0 +1,253 @@
+import numpy as np
+import torch
+from PIL import Image
+import torchvision
+
+from modules.Device import Device
+
+
+def _tensor_check_image(image: torch.Tensor) -> None:
+    """#### Check if the input is a valid tensor image.
+
+    #### Args:
+        - `image` (torch.Tensor): The input tensor image.
+    """
+    return
+
+
+def tensor2pil(image: torch.Tensor) -> Image.Image:
+    """#### Convert a tensor to a PIL image.
+
+    #### Args:
+        - `image` (torch.Tensor): The input tensor.
+
+    #### Returns:
+        - `Image.Image`: The converted PIL image.
+    """
+    _tensor_check_image(image)
+    return Image.fromarray(
+        np.clip(255.0 * image.cpu().numpy().squeeze(0), 0, 255).astype(np.uint8)
+    )
+
+
+def general_tensor_resize(image: torch.Tensor, w: int, h: int) -> torch.Tensor:
+    """#### Resize a tensor image using bilinear interpolation.
+
+    #### Args:
+        - `image` (torch.Tensor): The input tensor image.
+        - `w` (int): The target width.
+        - `h` (int): The target height.
+
+    #### Returns:
+        - `torch.Tensor`: The resized tensor image.
+    """
+    _tensor_check_image(image)
+    image = image.permute(0, 3, 1, 2)
+    image = torch.nn.functional.interpolate(image, size=(h, w), mode="bilinear")
+    image = image.permute(0, 2, 3, 1)
+    return image
+
+
+def pil2tensor(image: Image.Image) -> torch.Tensor:
+    """#### Convert a PIL image to a tensor.
+
+    #### Args:
+        - `image` (Image.Image): The input PIL image.
+
+    #### Returns:
+        - `torch.Tensor`: The converted tensor.
+    """
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+
+class TensorBatchBuilder:
+    """#### Class for building a batch of tensors."""
+
+    def __init__(self):
+        self.tensor: torch.Tensor | None = None
+
+    def concat(self, new_tensor: torch.Tensor) -> None:
+        """#### Concatenate a new tensor to the batch.
+
+        #### Args:
+            - `new_tensor` (torch.Tensor): The new tensor to concatenate.
+        """
+        self.tensor = new_tensor
+
+
+LANCZOS = Image.Resampling.LANCZOS if hasattr(Image, "Resampling") else Image.LANCZOS
+
+
+def tensor_resize(image: torch.Tensor, w: int, h: int) -> torch.Tensor:
+    """#### Resize a tensor image.
+
+    #### Args:
+        - `image` (torch.Tensor): The input tensor image.
+        - `w` (int): The target width.
+        - `h` (int): The target height.
+
+    #### Returns:
+        - `torch.Tensor`: The resized tensor image.
+    """
+    _tensor_check_image(image)
+    if image.shape[3] >= 3:
+        scaled_images = TensorBatchBuilder()
+        for single_image in image:
+            single_image = single_image.unsqueeze(0)
+            single_pil = tensor2pil(single_image)
+            scaled_pil = single_pil.resize((w, h), resample=LANCZOS)
+
+            single_image = pil2tensor(scaled_pil)
+            scaled_images.concat(single_image)
+
+        return scaled_images.tensor
+    else:
+        return general_tensor_resize(image, w, h)
+
+
+def tensor_paste(
+    image1: torch.Tensor,
+    image2: torch.Tensor,
+    left_top: tuple[int, int],
+    mask: torch.Tensor,
+) -> None:
+    """#### Paste one tensor image onto another using a mask.
+
+    #### Args:
+        - `image1` (torch.Tensor): The base tensor image.
+        - `image2` (torch.Tensor): The tensor image to paste.
+        - `left_top` (tuple[int, int]): The top-left corner where the image2 will be pasted.
+        - `mask` (torch.Tensor): The mask tensor.
+    """
+    _tensor_check_image(image1)
+    _tensor_check_image(image2)
+    _tensor_check_mask(mask)
+
+    x, y = left_top
+    _, h1, w1, _ = image1.shape
+    _, h2, w2, _ = image2.shape
+
+    # calculate image patch size
+    w = min(w1, x + w2) - x
+    h = min(h1, y + h2) - y
+
+    mask = mask[:, :h, :w, :]
+    image1[:, y : y + h, x : x + w, :] = (1 - mask) * image1[
+        :, y : y + h, x : x + w, :
+    ] + mask * image2[:, :h, :w, :]
+    return
+
+
+def tensor_convert_rgba(image: torch.Tensor, prefer_copy: bool = True) -> torch.Tensor:
+    """#### Convert a tensor image to RGBA format.
+
+    #### Args:
+        - `image` (torch.Tensor): The input tensor image.
+        - `prefer_copy` (bool, optional): Whether to prefer copying the tensor. Defaults to True.
+
+    #### Returns:
+        - `torch.Tensor`: The converted RGBA tensor image.
+    """
+    _tensor_check_image(image)
+    alpha = torch.ones((*image.shape[:-1], 1))
+    return torch.cat((image, alpha), axis=-1)
+
+
+def tensor_convert_rgb(image: torch.Tensor, prefer_copy: bool = True) -> torch.Tensor:
+    """#### Convert a tensor image to RGB format.
+
+    #### Args:
+        - `image` (torch.Tensor): The input tensor image.
+        - `prefer_copy` (bool, optional): Whether to prefer copying the tensor. Defaults to True.
+
+    #### Returns:
+        - `torch.Tensor`: The converted RGB tensor image.
+    """
+    _tensor_check_image(image)
+    return image
+
+
+def tensor_get_size(image: torch.Tensor) -> tuple[int, int]:
+    """#### Get the size of a tensor image.
+
+    #### Args:
+        - `image` (torch.Tensor): The input tensor image.
+
+    #### Returns:
+        - `tuple[int, int]`: The width and height of the tensor image.
+    """
+    _tensor_check_image(image)
+    _, h, w, _ = image.shape
+    return (w, h)
+
+
+def tensor_putalpha(image: torch.Tensor, mask: torch.Tensor) -> None:
+    """#### Add an alpha channel to a tensor image using a mask.
+
+    #### Args:
+        - `image` (torch.Tensor): The input tensor image.
+        - `mask` (torch.Tensor): The mask tensor.
+    """
+    _tensor_check_image(image)
+    _tensor_check_mask(mask)
+    image[..., -1] = mask[..., 0]
+
+
+def _tensor_check_mask(mask: torch.Tensor) -> None:
+    """#### Check if the input is a valid tensor mask.
+
+    #### Args:
+        - `mask` (torch.Tensor): The input tensor mask.
+    """
+    return
+
+
+def tensor_gaussian_blur_mask(
+    mask: torch.Tensor | np.ndarray, kernel_size: int, sigma: float = 10.0
+) -> torch.Tensor:
+    """#### Apply Gaussian blur to a tensor mask.
+
+    #### Args:
+        - `mask` (torch.Tensor | np.ndarray): The input tensor mask.
+        - `kernel_size` (int): The size of the Gaussian kernel.
+        - `sigma` (float, optional): The standard deviation of the Gaussian kernel. Defaults to 10.0.
+
+    #### Returns:
+        - `torch.Tensor`: The blurred tensor mask.
+    """
+    if isinstance(mask, np.ndarray):
+        mask = torch.from_numpy(mask)
+
+    if mask.ndim == 2:
+        mask = mask[None, ..., None]
+
+    _tensor_check_mask(mask)
+
+    kernel_size = kernel_size * 2 + 1
+
+    prev_device = mask.device
+    device = Device.get_torch_device()
+    mask.to(device)
+
+    # apply gaussian blur
+    mask = mask[:, None, ..., 0]
+    blurred_mask = torchvision.transforms.GaussianBlur(
+        kernel_size=kernel_size, sigma=sigma
+    )(mask)
+    blurred_mask = blurred_mask[:, 0, ..., None]
+
+    blurred_mask.to(prev_device)
+
+    return blurred_mask
+
+
+def to_tensor(image: np.ndarray) -> torch.Tensor:
+    """#### Convert a numpy array to a tensor.
+
+    #### Args:
+        - `image` (np.ndarray): The input numpy array.
+
+    #### Returns:
+        - `torch.Tensor`: The converted tensor.
+    """
+    return torch.from_numpy(image)

modules/AutoEncoders/ResBlock.py

@@ -0,0 +1,406 @@
+from abc import abstractmethod
+from typing import Optional, Any, Dict
+
+import torch
+from modules.NeuralNetwork import transformer
+import torch.nn as nn
+import torch.nn.functional as F
+
+from modules.Attention import Attention
+from modules.cond import cast
+from modules.sample import sampling_util
+
+
+oai_ops = cast.disable_weight_init
+
+
+class TimestepBlock1(nn.Module):
+    """#### Abstract class representing a timestep block."""
+
+    @abstractmethod
+    def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the timestep block.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `emb` (torch.Tensor): The embedding tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        pass
+
+
+def forward_timestep_embed1(
+    ts: nn.ModuleList,
+    x: torch.Tensor,
+    emb: torch.Tensor,
+    context: Optional[torch.Tensor] = None,
+    transformer_options: Optional[Dict[str, Any]] = {},
+    output_shape: Optional[torch.Size] = None,
+    time_context: Optional[torch.Tensor] = None,
+    num_video_frames: Optional[int] = None,
+    image_only_indicator: Optional[bool] = None,
+) -> torch.Tensor:
+    """#### Forward pass for timestep embedding.
+
+    #### Args:
+        - `ts` (nn.ModuleList): The list of timestep blocks.
+        - `x` (torch.Tensor): The input tensor.
+        - `emb` (torch.Tensor): The embedding tensor.
+        - `context` (torch.Tensor, optional): The context tensor. Defaults to None.
+        - `transformer_options` (dict, optional): The transformer options. Defaults to {}.
+        - `output_shape` (torch.Size, optional): The output shape. Defaults to None.
+        - `time_context` (torch.Tensor, optional): The time context tensor. Defaults to None.
+        - `num_video_frames` (int, optional): The number of video frames. Defaults to None.
+        - `image_only_indicator` (bool, optional): The image only indicator. Defaults to None.
+
+    #### Returns:
+        - `torch.Tensor`: The output tensor.
+    """
+    for layer in ts:
+        if isinstance(layer, TimestepBlock1):
+            x = layer(x, emb)
+        elif isinstance(layer, transformer.SpatialTransformer):
+            x = layer(x, context, transformer_options)
+            if "transformer_index" in transformer_options:
+                transformer_options["transformer_index"] += 1
+        elif isinstance(layer, Upsample1):
+            x = layer(x, output_shape=output_shape)
+        else:
+            x = layer(x)
+    return x
+
+
+class Upsample1(nn.Module):
+    """#### Class representing an upsample layer."""
+
+    def __init__(
+        self,
+        channels: int,
+        use_conv: bool,
+        dims: int = 2,
+        out_channels: Optional[int] = None,
+        padding: int = 1,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        operations: Any = oai_ops,
+    ):
+        """#### Initialize the upsample layer.
+
+        #### Args:
+            - `channels` (int): The number of input channels.
+            - `use_conv` (bool): Whether to use convolution.
+            - `dims` (int, optional): The number of dimensions. Defaults to 2.
+            - `out_channels` (int, optional): The number of output channels. Defaults to None.
+            - `padding` (int, optional): The padding size. Defaults to 1.
+            - `dtype` (torch.dtype, optional): The data type. Defaults to None.
+            - `device` (torch.device, optional): The device. Defaults to None.
+            - `operations` (any, optional): The operations. Defaults to oai_ops.
+        """
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        if use_conv:
+            self.conv = operations.conv_nd(
+                dims,
+                self.channels,
+                self.out_channels,
+                3,
+                padding=padding,
+                dtype=dtype,
+                device=device,
+            )
+
+    def forward(
+        self, x: torch.Tensor, output_shape: Optional[torch.Size] = None
+    ) -> torch.Tensor:
+        """#### Forward pass for the upsample layer.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `output_shape` (torch.Size, optional): The output shape. Defaults to None.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        assert x.shape[1] == self.channels
+        shape = [x.shape[2] * 2, x.shape[3] * 2]
+        if output_shape is not None:
+            shape[0] = output_shape[2]
+            shape[1] = output_shape[3]
+
+        x = F.interpolate(x, size=shape, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample1(nn.Module):
+    """#### Class representing a downsample layer."""
+
+    def __init__(
+        self,
+        channels: int,
+        use_conv: bool,
+        dims: int = 2,
+        out_channels: Optional[int] = None,
+        padding: int = 1,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        operations: Any = oai_ops,
+    ):
+        """#### Initialize the downsample layer.
+
+        #### Args:
+            - `channels` (int): The number of input channels.
+            - `use_conv` (bool): Whether to use convolution.
+            - `dims` (int, optional): The number of dimensions. Defaults to 2.
+            - `out_channels` (int, optional): The number of output channels. Defaults to None.
+            - `padding` (int, optional): The padding size. Defaults to 1.
+            - `dtype` (torch.dtype, optional): The data type. Defaults to None.
+            - `device` (torch.device, optional): The device. Defaults to None.
+            - `operations` (any, optional): The operations. Defaults to oai_ops.
+        """
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        stride = 2 if dims != 3 else (1, 2, 2)
+        self.op = operations.conv_nd(
+            dims,
+            self.channels,
+            self.out_channels,
+            3,
+            stride=stride,
+            padding=padding,
+            dtype=dtype,
+            device=device,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the downsample layer.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class ResBlock1(TimestepBlock1):
+    """#### Class representing a residual block layer."""
+
+    def __init__(
+        self,
+        channels: int,
+        emb_channels: int,
+        dropout: float,
+        out_channels: Optional[int] = None,
+        use_conv: bool = False,
+        use_scale_shift_norm: bool = False,
+        dims: int = 2,
+        use_checkpoint: bool = False,
+        up: bool = False,
+        down: bool = False,
+        kernel_size: int = 3,
+        exchange_temb_dims: bool = False,
+        skip_t_emb: bool = False,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        operations: Any = oai_ops,
+    ):
+        """#### Initialize the residual block layer.
+
+        #### Args:
+            - `channels` (int): The number of input channels.
+            - `emb_channels` (int): The number of embedding channels.
+            - `dropout` (float): The dropout rate.
+            - `out_channels` (int, optional): The number of output channels. Defaults to None.
+            - `use_conv` (bool, optional): Whether to use convolution. Defaults to False.
+            - `use_scale_shift_norm` (bool, optional): Whether to use scale shift normalization. Defaults to False.
+            - `dims` (int, optional): The number of dimensions. Defaults to 2.
+            - `use_checkpoint` (bool, optional): Whether to use checkpointing. Defaults to False.
+            - `up` (bool, optional): Whether to use upsampling. Defaults to False.
+            - `down` (bool, optional): Whether to use downsampling. Defaults to False.
+            - `kernel_size` (int, optional): The kernel size. Defaults to 3.
+            - `exchange_temb_dims` (bool, optional): Whether to exchange embedding dimensions. Defaults to False.
+            - `skip_t_emb` (bool, optional): Whether to skip embedding. Defaults to False.
+            - `dtype` (torch.dtype, optional): The data type. Defaults to None.
+            - `device` (torch.device, optional): The device. Defaults to None.
+            - `operations` (any, optional): The operations. Defaults to oai_ops.
+        """
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_checkpoint = use_checkpoint
+        self.use_scale_shift_norm = use_scale_shift_norm
+        self.exchange_temb_dims = exchange_temb_dims
+
+        padding = kernel_size // 2
+
+        self.in_layers = nn.Sequential(
+            operations.GroupNorm(32, channels, dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.conv_nd(
+                dims,
+                channels,
+                self.out_channels,
+                kernel_size,
+                padding=padding,
+                dtype=dtype,
+                device=device,
+            ),
+        )
+
+        self.updown = up or down
+
+        self.h_upd = self.x_upd = nn.Identity()
+
+        self.skip_t_emb = skip_t_emb
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            operations.Linear(
+                emb_channels,
+                (2 * self.out_channels if use_scale_shift_norm else self.out_channels),
+                dtype=dtype,
+                device=device,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            operations.GroupNorm(32, self.out_channels, dtype=dtype, device=device),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            operations.conv_nd(
+                dims,
+                self.out_channels,
+                self.out_channels,
+                kernel_size,
+                padding=padding,
+                dtype=dtype,
+                device=device,
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        else:
+            self.skip_connection = operations.conv_nd(
+                dims, channels, self.out_channels, 1, dtype=dtype, device=device
+            )
+
+    def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the residual block layer.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `emb` (torch.Tensor): The embedding tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        return sampling_util.checkpoint(
+            self._forward, (x, emb), self.parameters(), self.use_checkpoint
+        )
+
+    def _forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
+        """#### Internal forward pass for the residual block layer.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `emb` (torch.Tensor): The embedding tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        h = self.in_layers(x)
+
+        emb_out = None
+        if not self.skip_t_emb:
+            emb_out = self.emb_layers(emb).type(h.dtype)
+            while len(emb_out.shape) < len(h.shape):
+                emb_out = emb_out[..., None]
+        if emb_out is not None:
+            h = h + emb_out
+        h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+ops = cast.disable_weight_init
+
+
+class ResnetBlock(nn.Module):
+    """#### Class representing a ResNet block layer."""
+
+    def __init__(
+        self,
+        *,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float,
+        temb_channels: int = 512,
+    ):
+        """#### Initialize the ResNet block layer.
+
+        #### Args:
+            - `in_channels` (int): The number of input channels.
+            - `out_channels` (int, optional): The number of output channels. Defaults to None.
+            - `conv_shortcut` (bool, optional): Whether to use convolution shortcut. Defaults to False.
+            - `dropout` (float): The dropout rate.
+            - `temb_channels` (int, optional): The number of embedding channels. Defaults to 512.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.swish = torch.nn.SiLU(inplace=True)
+        self.norm1 = Attention.Normalize(in_channels)
+        self.conv1 = ops.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        self.norm2 = Attention.Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout, inplace=True)
+        self.conv2 = ops.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = ops.Conv2d(
+                in_channels, out_channels, kernel_size=1, stride=1, padding=0
+            )
+
+    def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the ResNet block layer.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `temb` (torch.Tensor): The embedding tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        h = x
+        h = self.norm1(h)
+        h = self.swish(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = self.swish(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x + h

modules/AutoEncoders/VariationalAE.py

@@ -0,0 +1,824 @@
+import logging
+from typing import Dict, Optional, Tuple, Union
+import numpy as np
+import torch
+from modules.Model import ModelPatcher
+import torch.nn as nn
+
+from modules.Attention import Attention
+from modules.AutoEncoders import ResBlock
+from modules.Device import Device
+from modules.Utilities import util
+from modules.cond import cast
+
+
+class DiagonalGaussianDistribution(object):
+    """#### Represents a diagonal Gaussian distribution parameterized by mean and log-variance.
+
+    #### Attributes:
+        - `parameters` (torch.Tensor): The concatenated mean and log-variance of the distribution.
+        - `mean` (torch.Tensor): The mean of the distribution.
+        - `logvar` (torch.Tensor): The log-variance of the distribution, clamped between -30.0 and 20.0.
+        - `std` (torch.Tensor): The standard deviation of the distribution, computed as exp(0.5 * logvar).
+        - `var` (torch.Tensor): The variance of the distribution, computed as exp(logvar).
+        - `deterministic` (bool): If True, the distribution is deterministic.
+
+    #### Methods:
+        - `sample() -> torch.Tensor`:
+            Samples from the distribution using the reparameterization trick.
+        - `kl(other: DiagonalGaussianDistribution = None) -> torch.Tensor`:
+            Computes the Kullback-Leibler divergence between this distribution and a standard normal distribution.
+            If `other` is provided, computes the KL divergence between this distribution and `other`.
+    """
+
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+
+    def sample(self) -> torch.Tensor:
+        """#### Samples from the distribution using the reparameterization trick.
+
+        #### Returns:
+            - `torch.Tensor`: A sample from the distribution.
+        """
+        x = self.mean + self.std * torch.randn(self.mean.shape).to(
+            device=self.parameters.device
+        )
+        return x
+
+    def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor:
+        """#### Computes the Kullback-Leibler divergence between this distribution and a standard normal distribution.
+
+        If `other` is provided, computes the KL divergence between this distribution and `other`.
+
+        #### Args:
+            - `other` (DiagonalGaussianDistribution, optional): Another distribution to compute the KL divergence with.
+
+        #### Returns:
+            - `torch.Tensor`: The KL divergence.
+        """
+        return 0.5 * torch.sum(
+            torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+            dim=[1, 2, 3],
+        )
+
+
+class DiagonalGaussianRegularizer(torch.nn.Module):
+    """#### Regularizer for diagonal Gaussian distributions."""
+
+    def __init__(self, sample: bool = True):
+        """#### Initialize the regularizer.
+
+        #### Args:
+            - `sample` (bool, optional): Whether to sample from the distribution. Defaults to True.
+        """
+        super().__init__()
+        self.sample = sample
+
+    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
+        """#### Forward pass for the regularizer.
+
+        #### Args:
+            - `z` (torch.Tensor): The input tensor.
+
+        #### Returns:
+            - `Tuple[torch.Tensor, dict]`: The regularized tensor and a log dictionary.
+        """
+        log = dict()
+        posterior = DiagonalGaussianDistribution(z)
+        if self.sample:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+        kl_loss = posterior.kl()
+        kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
+        log["kl_loss"] = kl_loss
+        return z, log
+
+
+class AutoencodingEngine(nn.Module):
+    """#### Class representing an autoencoding engine."""
+
+    def __init__(self, encoder: nn.Module, decoder: nn.Module, regularizer: nn.Module, flux: bool = False):
+        """#### Initialize the autoencoding engine.
+
+        #### Args:
+            - `encoder` (nn.Module): The encoder module.
+            - `decoder` (nn.Module): The decoder module.
+            - `regularizer` (nn.Module): The regularizer module.
+        """
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.regularization = regularizer
+        if not flux:
+            self.post_quant_conv = cast.disable_weight_init.Conv2d(4, 4, 1)
+            self.quant_conv = cast.disable_weight_init.Conv2d(8, 8, 1)
+        
+    def get_last_layer(self):
+        """#### Get the last layer of the decoder.
+
+        Returns:
+            - `nn.Module`: The last layer of the decoder.
+        """
+        return self.decoder.get_last_layer()
+    
+    def decode(self, z: torch.Tensor, flux:bool = False, **kwargs) -> torch.Tensor:
+        """#### Decode the latent tensor.
+
+        #### Args:
+            - `z` (torch.Tensor): The latent tensor.
+            - `decoder_kwargs` (dict): Additional arguments for the decoder.
+
+        #### Returns:
+            - `torch.Tensor`: The decoded tensor.
+        """
+        if flux:
+            x = self.decoder(z, **kwargs)
+            return x
+        dec = self.post_quant_conv(z)
+        dec = self.decoder(dec, **kwargs)
+        return dec
+
+
+    def encode(
+        self,
+        x: torch.Tensor,
+        return_reg_log: bool = False,
+        unregularized: bool = False,
+        flux: bool = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, dict]]:
+        """#### Encode the input tensor.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `return_reg_log` (bool, optional): Whether to return the regularization log. Defaults to False.
+
+        #### Returns:
+            - `Union[torch.Tensor, Tuple[torch.Tensor, dict]]`: The encoded tensor and optionally the regularization log.
+        """
+        z = self.encoder(x)
+        if not flux:
+            z = self.quant_conv(z)
+        if unregularized:
+            return z, dict()
+        z, reg_log = self.regularization(z)
+        if return_reg_log:
+            return z, reg_log
+        return z
+
+ops = cast.disable_weight_init
+
+if Device.xformers_enabled_vae():
+    pass
+
+
+def nonlinearity(x: torch.Tensor) -> torch.Tensor:
+    """#### Apply the swish nonlinearity.
+
+    #### Args:
+        - `x` (torch.Tensor): The input tensor.
+
+    #### Returns:
+        - `torch.Tensor`: The output tensor.
+    """
+    return x * torch.sigmoid(x)
+
+
+class Upsample(nn.Module):
+    """#### Class representing an upsample layer."""
+
+    def __init__(self, in_channels: int, with_conv: bool):
+        """#### Initialize the upsample layer.
+
+        #### Args:
+            - `in_channels` (int): The number of input channels.
+            - `with_conv` (bool): Whether to use convolution.
+        """
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = ops.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=1, padding=1
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the upsample layer.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """#### Class representing a downsample layer."""
+
+    def __init__(self, in_channels: int, with_conv: bool):
+        """#### Initialize the downsample layer.
+
+        #### Args:
+            - `in_channels` (int): The number of input channels.
+            - `with_conv` (bool): Whether to use convolution.
+        """
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = ops.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=2, padding=0
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the downsample layer.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        pad = (0, 1, 0, 1)
+        x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    """#### Class representing an encoder."""
+
+    def __init__(
+        self,
+        *,
+        ch: int,
+        out_ch: int,
+        ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
+        num_res_blocks: int,
+        attn_resolutions: Tuple[int, ...],
+        dropout: float = 0.0,
+        resamp_with_conv: bool = True,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+        double_z: bool = True,
+        use_linear_attn: bool = False,
+        attn_type: str = "vanilla",
+        **ignore_kwargs,
+    ):
+        """#### Initialize the encoder.
+
+        #### Args:
+            - `ch` (int): The base number of channels.
+            - `out_ch` (int): The number of output channels.
+            - `ch_mult` (Tuple[int, ...], optional): Channel multiplier at each resolution. Defaults to (1, 2, 4, 8).
+            - `num_res_blocks` (int): The number of residual blocks.
+            - `attn_resolutions` (Tuple[int, ...]): The resolutions at which to apply attention.
+            - `dropout` (float, optional): The dropout rate. Defaults to 0.0.
+            - `resamp_with_conv` (bool, optional): Whether to use convolution for resampling. Defaults to True.
+            - `in_channels` (int): The number of input channels.
+            - `resolution` (int): The resolution of the input.
+            - `z_channels` (int): The number of latent channels.
+            - `double_z` (bool, optional): Whether to double the latent channels. Defaults to True.
+            - `use_linear_attn` (bool, optional): Whether to use linear attention. Defaults to False.
+            - `attn_type` (str, optional): The type of attention. Defaults to "vanilla".
+        """
+        super().__init__()
+        if use_linear_attn:
+            attn_type = "linear"
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        # downsampling
+        self.conv_in = ops.Conv2d(
+            in_channels, self.ch, kernel_size=3, stride=1, padding=1
+        )
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(
+                    ResBlock.ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResBlock.ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = Attention.make_attn(block_in, attn_type=attn_type)
+        self.mid.block_2 = ResBlock.ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+        # end
+        self.norm_out = Attention.Normalize(block_in)
+        self.conv_out = ops.Conv2d(
+            block_in,
+            2 * z_channels if double_z else z_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+        self._device = torch.device("cpu")
+        self._dtype = torch.float32
+
+    def to(self, device=None, dtype=None):
+        """#### Move the encoder to a device and data type.
+        
+        #### Args:
+            - `device` (torch.device, optional): The device to move to. Defaults to None.
+            - `dtype` (torch.dtype, optional): The data type to move to. Defaults to None.
+        
+        #### Returns:
+            - `nn.Module`: The encoder.
+        """
+        if device is not None:
+            self._device = device
+        if dtype is not None:
+            self._dtype = dtype
+        return super().to(device=device, dtype=dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the encoder.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The encoded tensor.
+        """
+        if x.device != self._device or x.dtype != self._dtype:
+            self.to(device=x.device, dtype=x.dtype)
+        # timestep embedding
+        temb = None
+        # downsampling
+        h = self.conv_in(x)
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](h, temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+            if i_level != self.num_resolutions - 1:
+                h = self.down[i_level].downsample(h)
+
+        # middle
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    """#### Class representing a decoder."""
+
+    def __init__(
+        self,
+        *,
+        ch: int,
+        out_ch: int,
+        ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
+        num_res_blocks: int,
+        attn_resolutions: Tuple[int, ...],
+        dropout: float = 0.0,
+        resamp_with_conv: bool = True,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+        give_pre_end: bool = False,
+        tanh_out: bool = False,
+        use_linear_attn: bool = False,
+        conv_out_op: nn.Module = ops.Conv2d,
+        resnet_op: nn.Module = ResBlock.ResnetBlock,
+        attn_op: nn.Module = Attention.AttnBlock,
+        **ignorekwargs,
+    ):
+        """#### Initialize the decoder.
+
+        #### Args:
+            - `ch` (int): The base number of channels.
+            - `out_ch` (int): The number of output channels.
+            - `ch_mult` (Tuple[int, ...], optional): Channel multiplier at each resolution. Defaults to (1, 2, 4, 8).
+            - `num_res_blocks` (int): The number of residual blocks.
+            - `attn_resolutions` (Tuple[int, ...]): The resolutions at which to apply attention.
+            - `dropout` (float, optional): The dropout rate. Defaults to 0.0.
+            - `resamp_with_conv` (bool, optional): Whether to use convolution for resampling. Defaults to True.
+            - `in_channels` (int): The number of input channels.
+            - `resolution` (int): The resolution of the input.
+            - `z_channels` (int): The number of latent channels.
+            - `give_pre_end` (bool, optional): Whether to give pre-end. Defaults to False.
+            - `tanh_out` (bool, optional): Whether to use tanh activation at the output. Defaults to False.
+            - `use_linear_attn` (bool, optional): Whether to use linear attention. Defaults to False.
+            - `conv_out_op` (nn.Module, optional): The convolution output operation. Defaults to ops.Conv2d.
+            - `resnet_op` (nn.Module, optional): The residual block operation. Defaults to ResBlock.ResnetBlock.
+            - `attn_op` (nn.Module, optional): The attention block operation. Defaults to Attention.AttnBlock.
+        """
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
+        self.tanh_out = tanh_out
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        (1,) + tuple(ch_mult)
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+        logging.debug(
+            "Working with z of shape {} = {} dimensions.".format(
+                self.z_shape, np.prod(self.z_shape)
+            )
+        )
+
+        # z to block_in
+        self.conv_in = ops.Conv2d(
+            z_channels, block_in, kernel_size=3, stride=1, padding=1
+        )
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = resnet_op(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = attn_op(block_in)
+        self.mid.block_2 = resnet_op(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(
+                    resnet_op(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = Attention.Normalize(block_in)
+        self.conv_out = conv_out_op(
+            block_in, out_ch, kernel_size=3, stride=1, padding=1
+        )
+
+    def forward(self, z: torch.Tensor, **kwargs) -> torch.Tensor:
+        """#### Forward pass for the decoder.
+
+        #### Args:
+            - `z` (torch.Tensor): The input tensor.
+            - `**kwargs`: Additional arguments.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+
+        """
+        # assert z.shape[1:] == self.z_shape[1:]
+        self.last_z_shape = z.shape
+
+        # timestep embedding
+        temb = None
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h, temb, **kwargs)
+        h = self.mid.attn_1(h, **kwargs)
+        h = self.mid.block_2(h, temb, **kwargs)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h, temb, **kwargs)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h, **kwargs)
+        return h
+
+
+class VAE:
+    """#### Class representing a Variational Autoencoder (VAE)."""
+
+    def __init__(
+        self,
+        sd: Optional[dict] = None,
+        device: Optional[torch.device] = None,
+        config: Optional[dict] = None,
+        dtype: Optional[torch.dtype] = None,
+        flux: Optional[bool] = False,
+    ):
+        """#### Initialize the VAE.
+
+        #### Args:
+            - `sd` (dict, optional): The state dictionary. Defaults to None.
+            - `device` (torch.device, optional): The device to use. Defaults to None.
+            - `config` (dict, optional): The configuration dictionary. Defaults to None.
+            - `dtype` (torch.dtype, optional): The data type. Defaults to None.
+        """
+        self.memory_used_encode = lambda shape, dtype: (
+            1767 * shape[2] * shape[3]
+        ) * Device.dtype_size(
+            dtype
+        )  # These are for AutoencoderKL and need tweaking (should be lower)
+        self.memory_used_decode = lambda shape, dtype: (
+            2178 * shape[2] * shape[3] * 64
+        ) * Device.dtype_size(dtype)
+        self.downscale_ratio = 8
+        self.upscale_ratio = 8
+        self.latent_channels = 4
+        self.output_channels = 3
+        self.process_input = lambda image: image * 2.0 - 1.0
+        self.process_output = lambda image: torch.clamp(
+            (image + 1.0) / 2.0, min=0.0, max=1.0
+        )
+        self.working_dtypes = [torch.bfloat16, torch.float32]
+
+        if config is None:
+            if "decoder.conv_in.weight" in sd:
+                # default SD1.x/SD2.x VAE parameters
+                ddconfig = {
+                    "double_z": True,
+                    "z_channels": 4,
+                    "resolution": 256,
+                    "in_channels": 3,
+                    "out_ch": 3,
+                    "ch": 128,
+                    "ch_mult": [1, 2, 4, 4],
+                    "num_res_blocks": 2,
+                    "attn_resolutions": [],
+                    "dropout": 0.0,
+                }
+
+                if (
+                    "encoder.down.2.downsample.conv.weight" not in sd
+                    and "decoder.up.3.upsample.conv.weight" not in sd
+                ):  # Stable diffusion x4 upscaler VAE
+                    ddconfig["ch_mult"] = [1, 2, 4]
+                    self.downscale_ratio = 4
+                    self.upscale_ratio = 4
+
+                self.latent_channels = ddconfig["z_channels"] = sd[
+                    "decoder.conv_in.weight"
+                ].shape[1]
+                # Initialize model
+                self.first_stage_model = AutoencodingEngine(
+                    Encoder(**ddconfig),
+                    Decoder(**ddconfig), 
+                    DiagonalGaussianRegularizer(),
+                    flux=flux
+                )
+            else:
+                logging.warning("WARNING: No VAE weights detected, VAE not initalized.")
+                self.first_stage_model = None
+                return
+        
+        self.first_stage_model = self.first_stage_model.eval()
+
+        m, u = self.first_stage_model.load_state_dict(sd, strict=False)
+        if len(m) > 0:
+            logging.warning("Missing VAE keys {}".format(m))
+
+        if len(u) > 0:
+            logging.debug("Leftover VAE keys {}".format(u))
+
+        if device is None:
+            device = Device.vae_device()
+        self.device = device
+        offload_device = Device.vae_offload_device()
+        if dtype is None:
+            dtype = Device.vae_dtype()
+        self.vae_dtype = dtype
+        self.first_stage_model.to(self.vae_dtype)
+        self.output_device = Device.intermediate_device()
+
+        self.patcher = ModelPatcher.ModelPatcher(
+            self.first_stage_model,
+            load_device=self.device,
+            offload_device=offload_device,
+        )
+        logging.debug(
+            "VAE load device: {}, offload device: {}, dtype: {}".format(
+                self.device, offload_device, self.vae_dtype
+            )
+        )
+
+
+    def vae_encode_crop_pixels(self, pixels: torch.Tensor) -> torch.Tensor:
+        """#### Crop the input pixels to be compatible with the VAE.
+
+        #### Args:
+            - `pixels` (torch.Tensor): The input pixel tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The cropped pixel tensor.
+        """
+        (pixels.shape[1] // self.downscale_ratio) * self.downscale_ratio
+        (pixels.shape[2] // self.downscale_ratio) * self.downscale_ratio
+        return pixels
+
+    def decode(self, samples_in: torch.Tensor, flux:bool = False) -> torch.Tensor:
+        """#### Decode the latent samples to pixel samples.
+
+        #### Args:
+            - `samples_in` (torch.Tensor): The input latent samples.
+
+        #### Returns:
+            - `torch.Tensor`: The decoded pixel samples.
+        """
+        memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
+        Device.load_models_gpu([self.patcher], memory_required=memory_used)
+        free_memory = Device.get_free_memory(self.device)
+        batch_number = int(free_memory / memory_used)
+        batch_number = max(1, batch_number)
+
+        pixel_samples = torch.empty(
+            (
+                samples_in.shape[0],
+                3,
+                round(samples_in.shape[2] * self.upscale_ratio),
+                round(samples_in.shape[3] * self.upscale_ratio),
+            ),
+            device=self.output_device,
+        )
+        for x in range(0, samples_in.shape[0], batch_number):
+            samples = (
+                samples_in[x : x + batch_number].to(self.vae_dtype).to(self.device)
+            )
+            pixel_samples[x : x + batch_number] = self.process_output(
+                self.first_stage_model.decode(samples, flux=flux).to(self.output_device).float()
+            )
+        pixel_samples = pixel_samples.to(self.output_device).movedim(1, -1)
+        return pixel_samples
+
+
+    def encode(self, pixel_samples: torch.Tensor, flux:bool = False) -> torch.Tensor:
+        """#### Encode the pixel samples to latent samples.
+
+        #### Args:
+            - `pixel_samples` (torch.Tensor): The input pixel samples.
+
+        #### Returns:
+            - `torch.Tensor`: The encoded latent samples.
+        """
+        pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
+        pixel_samples = pixel_samples.movedim(-1, 1)
+        memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)
+        Device.load_models_gpu([self.patcher], memory_required=memory_used)
+        free_memory = Device.get_free_memory(self.device)
+        batch_number = int(free_memory / memory_used)
+        batch_number = max(1, batch_number)
+        samples = torch.empty(
+            (
+                pixel_samples.shape[0],
+                self.latent_channels,
+                round(pixel_samples.shape[2] // self.downscale_ratio),
+                round(pixel_samples.shape[3] // self.downscale_ratio),
+            ),
+            device=self.output_device,
+        )
+        for x in range(0, pixel_samples.shape[0], batch_number):
+            pixels_in = (
+                self.process_input(pixel_samples[x : x + batch_number])
+                .to(self.vae_dtype)
+                .to(self.device)
+            )
+            samples[x : x + batch_number] = (
+                self.first_stage_model.encode(pixels_in, flux=flux).to(self.output_device).float()
+            )
+
+        return samples
+
+    def get_sd(self):
+        """#### Get the state dictionary.
+        
+        #### Returns:
+            - `dict`: The state dictionary.
+        """
+        return self.first_stage_model.state_dict()
+
+
+class VAEDecode:
+    """#### Class for decoding VAE samples."""
+
+    def decode(self, vae: VAE, samples: dict, flux:bool = False) -> Tuple[torch.Tensor]:
+        """#### Decode the VAE samples.
+
+        #### Args:
+            - `vae` (VAE): The VAE instance.
+            - `samples` (dict): The samples dictionary.
+
+        #### Returns:
+            - `Tuple[torch.Tensor]`: The decoded samples.
+        """
+        return (vae.decode(samples["samples"], flux=flux),)
+
+
+class VAEEncode:
+    """#### Class for encoding VAE samples."""
+
+    def encode(self, vae: VAE, pixels: torch.Tensor, flux:bool = False) -> Tuple[dict]:
+        """#### Encode the VAE samples.
+
+        #### Args:
+            - `vae` (VAE): The VAE instance.
+            - `pixels` (torch.Tensor): The input pixel tensor.
+
+        #### Returns:
+            - `Tuple[dict]`: The encoded samples dictionary.
+        """
+        t = vae.encode(pixels[:, :, :, :3], flux=flux)
+        return ({"samples": t},)
+
+
+class VAELoader:
+    """#### Class for loading VAEs."""
+    # TODO: scale factor?
+    def load_vae(self, vae_name):
+        """#### Load the VAE.
+        
+        #### Args:
+            - `vae_name`: The name of the VAE.
+        
+        #### Returns:
+            - `Tuple[VAE]`: The VAE instance.
+        """
+        if vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]:
+            sd = self.load_taesd(vae_name)
+        else:
+            vae_path = "./_internal/vae/" + vae_name
+            sd = util.load_torch_file(vae_path)
+        vae = VAE(sd=sd)
+        return (vae,)
+
+

modules/AutoEncoders/taesd.py

@@ -0,0 +1,310 @@
+"""
+Tiny AutoEncoder for Stable Diffusion
+(DNN for encoding / decoding SD's latent space)
+"""
+
+# TODO: Check if multiprocessing is possible for this module
+from PIL import Image
+import numpy as np
+from sympy import im
+import torch
+from modules.Utilities import util
+import torch.nn as nn
+
+from modules.cond import cast
+from modules.user import app_instance
+
+
+def conv(n_in: int, n_out: int, **kwargs) -> cast.disable_weight_init.Conv2d:
+    """#### Create a convolutional layer.
+
+    #### Args:
+        - `n_in` (int): The number of input channels.
+        - `n_out` (int): The number of output channels.
+
+    #### Returns:
+        - `torch.nn.Module`: The convolutional layer.
+    """
+    return cast.disable_weight_init.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
+
+
+class Clamp(nn.Module):
+    """#### Class representing a clamping layer."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass of the clamping layer.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The clamped tensor.
+        """
+        return torch.tanh(x / 3) * 3
+
+
+class Block(nn.Module):
+    """#### Class representing a block layer."""
+
+    def __init__(self, n_in: int, n_out: int):
+        """#### Initialize the block layer.
+
+        #### Args:
+            - `n_in` (int): The number of input channels.
+            - `n_out` (int): The number of output channels.
+
+        #### Returns:
+            - `Block`: The block layer.
+        """
+        super().__init__()
+        self.conv = nn.Sequential(
+            conv(n_in, n_out),
+            nn.ReLU(),
+            conv(n_out, n_out),
+            nn.ReLU(),
+            conv(n_out, n_out),
+        )
+        self.skip = (
+            cast.disable_weight_init.Conv2d(n_in, n_out, 1, bias=False)
+            if n_in != n_out
+            else nn.Identity()
+        )
+        self.fuse = nn.ReLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fuse(self.conv(x) + self.skip(x))
+
+
+def Encoder2(latent_channels: int = 4) -> nn.Sequential:
+    """#### Create an encoder.
+
+    #### Args:
+        - `latent_channels` (int, optional): The number of latent channels. Defaults to 4.
+
+    #### Returns:
+        - `torch.nn.Module`: The encoder.
+    """
+    return nn.Sequential(
+        conv(3, 64),
+        Block(64, 64),
+        conv(64, 64, stride=2, bias=False),
+        Block(64, 64),
+        Block(64, 64),
+        Block(64, 64),
+        conv(64, 64, stride=2, bias=False),
+        Block(64, 64),
+        Block(64, 64),
+        Block(64, 64),
+        conv(64, 64, stride=2, bias=False),
+        Block(64, 64),
+        Block(64, 64),
+        Block(64, 64),
+        conv(64, latent_channels),
+    )
+
+
+def Decoder2(latent_channels: int = 4) -> nn.Sequential:
+    """#### Create a decoder.
+
+    #### Args:
+        - `latent_channels` (int, optional): The number of latent channels. Defaults to 4.
+
+    #### Returns:
+        - `torch.nn.Module`: The decoder.
+    """
+    return nn.Sequential(
+        Clamp(),
+        conv(latent_channels, 64),
+        nn.ReLU(),
+        Block(64, 64),
+        Block(64, 64),
+        Block(64, 64),
+        nn.Upsample(scale_factor=2),
+        conv(64, 64, bias=False),
+        Block(64, 64),
+        Block(64, 64),
+        Block(64, 64),
+        nn.Upsample(scale_factor=2),
+        conv(64, 64, bias=False),
+        Block(64, 64),
+        Block(64, 64),
+        Block(64, 64),
+        nn.Upsample(scale_factor=2),
+        conv(64, 64, bias=False),
+        Block(64, 64),
+        conv(64, 3),
+    )
+
+
+class TAESD(nn.Module):
+    """#### Class representing a Tiny AutoEncoder for Stable Diffusion.
+
+    #### Attributes:
+        - `latent_magnitude` (float): Magnitude of the latent space.
+        - `latent_shift` (float): Shift value for the latent space.
+        - `vae_shift` (torch.nn.Parameter): Shift parameter for the VAE.
+        - `vae_scale` (torch.nn.Parameter): Scale parameter for the VAE.
+        - `taesd_encoder` (Encoder2): Encoder network for the TAESD.
+        - `taesd_decoder` (Decoder2): Decoder network for the TAESD.
+
+    #### Args:
+        - `encoder_path` (str, optional): Path to the encoder model file. Defaults to None.
+        - `decoder_path` (str, optional): Path to the decoder model file. Defaults to "./_internal/vae_approx/taesd_decoder.safetensors".
+        - `latent_channels` (int, optional): Number of channels in the latent space. Defaults to 4.
+
+    #### Methods:
+        - `scale_latents(x)`:
+            Scales raw latents to the range [0, 1].
+        - `unscale_latents(x)`:
+            Unscales latents from the range [0, 1] to raw latents.
+        - `decode(x)`:
+            Decodes the given latent representation to the original space.
+        - `encode(x)`:
+            Encodes the given input to the latent space.
+    """
+
+    latent_magnitude = 3
+    latent_shift = 0.5
+
+    def __init__(
+        self,
+        encoder_path: str = None,
+        decoder_path: str = None,
+        latent_channels: int = 4,
+    ):
+        """#### Initialize the TAESD model.
+
+        #### Args:
+            - `encoder_path` (str, optional): Path to the encoder model file. Defaults to None.
+            - `decoder_path` (str, optional): Path to the decoder model file. Defaults to "./_internal/vae_approx/taesd_decoder.safetensors".
+            - `latent_channels` (int, optional): Number of channels in the latent space. Defaults to 4.
+        """
+        super().__init__()
+        self.vae_shift = torch.nn.Parameter(torch.tensor(0.0))
+        self.vae_scale = torch.nn.Parameter(torch.tensor(1.0))
+        self.taesd_encoder = Encoder2(latent_channels)
+        self.taesd_decoder = Decoder2(latent_channels)
+        decoder_path = (
+            "./_internal/vae_approx/taesd_decoder.safetensors"
+            if decoder_path is None
+            else decoder_path
+        )
+        if encoder_path is not None:
+            self.taesd_encoder.load_state_dict(
+                util.load_torch_file(encoder_path, safe_load=True)
+            )
+        if decoder_path is not None:
+            self.taesd_decoder.load_state_dict(
+                util.load_torch_file(decoder_path, safe_load=True)
+            )
+
+    @staticmethod
+    def scale_latents(x: torch.Tensor) -> torch.Tensor:
+        """#### Scales raw latents to the range [0, 1].
+
+        #### Args:
+            - `x` (torch.Tensor): The raw latents.
+
+        #### Returns:
+            - `torch.Tensor`: The scaled latents.
+        """
+        return x.div(2 * TAESD.latent_magnitude).add(TAESD.latent_shift).clamp(0, 1)
+
+    @staticmethod
+    def unscale_latents(x: torch.Tensor) -> torch.Tensor:
+        """#### Unscales latents from the range [0, 1] to raw latents.
+
+        #### Args:
+            - `x` (torch.Tensor): The scaled latents.
+
+        #### Returns:
+            - `torch.Tensor`: The raw latents.
+        """
+        return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude)
+
+    def decode(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Decodes the given latent representation to the original space.
+
+        #### Args:
+            - `x` (torch.Tensor): The latent representation.
+
+        #### Returns:
+            - `torch.Tensor`: The decoded representation.
+        """
+        device = next(self.taesd_decoder.parameters()).device
+        x = x.to(device)
+        x_sample = self.taesd_decoder((x - self.vae_shift) * self.vae_scale)
+        x_sample = x_sample.sub(0.5).mul(2)
+        return x_sample
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Encodes the given input to the latent space.
+
+        #### Args:
+            - `x` (torch.Tensor): The input.
+
+        #### Returns:
+            - `torch.Tensor`: The latent representation.
+        """
+        device = next(self.taesd_encoder.parameters()).device
+        x = x.to(device)
+        return (self.taesd_encoder(x * 0.5 + 0.5) / self.vae_scale) + self.vae_shift
+
+
+def taesd_preview(x: torch.Tensor, flux: bool = False):
+    """#### Preview the batched latent tensors as images.
+    
+    #### Args:
+        - `x` (torch.Tensor): Input latent tensor with shape [B,C,H,W] 
+        - `flux` (bool, optional): Whether using flux model (for channel ordering). Defaults to False.
+    """
+    if app_instance.app.previewer_var.get() is True:
+        taesd_instance = TAESD()
+        
+        # Handle channel dimension
+        if x.shape[1] != 4:
+            desired_channels = 4
+            current_channels = x.shape[1]
+            
+            if current_channels > desired_channels:
+                x = x[:, :desired_channels, :, :]
+            else:
+                padding = torch.zeros(x.shape[0], desired_channels - current_channels, 
+                                   x.shape[2], x.shape[3], device=x.device)
+                x = torch.cat([x, padding], dim=1)
+
+        # Process entire batch at once
+        decoded_batch = taesd_instance.decode(x)
+        
+        images = []
+        
+        # Convert each image in batch 
+        for decoded in decoded_batch:
+            # Handle channel dimension
+            if decoded.shape[0] == 1:
+                decoded = decoded.repeat(3, 1, 1)
+                
+            # Apply different normalization for flux vs standard mode
+            if flux:
+                # For flux: Assume BGR ordering and different normalization
+                decoded = decoded[[2,1,0], :, :] # BGR -> RGB
+                # Adjust normalization for flux model range
+                decoded = decoded.clamp(-1, 1)
+                decoded = (decoded + 1.0) * 0.5 # Scale from [-1,1] to [0,1]
+            else:
+                # Standard normalization
+                decoded = (decoded + 1.0) / 2.0
+            
+            # Convert to numpy and uint8
+            image_np = (decoded.cpu().detach().numpy() * 255.0)
+            image_np = np.transpose(image_np, (1, 2, 0))
+            image_np = np.clip(image_np, 0, 255).astype(np.uint8)
+            
+            # Create PIL Image
+            img = Image.fromarray(image_np, mode='RGB')
+            images.append(img)
+            
+        # Update display with all images
+        app_instance.app.update_image(images)
+    else:
+        pass

modules/BlackForest/Flux.py

@@ -0,0 +1,853 @@
+# Original code can be found on: https://github.com/black-forest-labs/flux
+
+
+from dataclasses import dataclass
+from einops import rearrange, repeat
+import torch
+import torch.nn as nn
+
+from modules.Attention import Attention
+from modules.Device import Device
+from modules.Model import ModelBase
+from modules.Utilities import Latent
+from modules.cond import cast, cond
+from modules.sample import sampling, sampling_util
+
+
+# Define the attention mechanism
+def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, pe: torch.Tensor) -> torch.Tensor:
+    """#### Compute the attention mechanism.
+
+    #### Args:
+        - `q` (Tensor): The query tensor.
+        - `k` (Tensor): The key tensor.
+        - `v` (Tensor): The value tensor.
+        - `pe` (Tensor): The positional encoding tensor.
+
+    #### Returns:
+        - `Tensor`: The attention tensor.
+    """
+    q, k = apply_rope(q, k, pe)
+    heads = q.shape[1]
+    x = Attention.optimized_attention(q, k, v, heads, skip_reshape=True)
+    return x
+
+# Define the rotary positional encoding (RoPE)
+def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
+    """#### Compute the rotary positional encoding.
+
+    #### Args:
+        - `pos` (Tensor): The position tensor.
+        - `dim` (int): The dimension of the tensor.
+        - `theta` (int): The theta value for scaling.
+
+    #### Returns:
+        - `Tensor`: The rotary positional encoding tensor.
+    """
+    assert dim % 2 == 0
+    if Device.is_device_mps(pos.device) or Device.is_intel_xpu():
+        device = torch.device("cpu")
+    else:
+        device = pos.device
+
+    scale = torch.linspace(
+        0, (dim - 2) / dim, steps=dim // 2, dtype=torch.float64, device=device
+    )
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum(
+        "...n,d->...nd", pos.to(dtype=torch.float32, device=device), omega
+    )
+    out = torch.stack(
+        [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1
+    )
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.to(dtype=torch.float32, device=pos.device)
+
+# Apply the rotary positional encoding to the query and key tensors
+def apply_rope(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor) -> tuple:
+    """#### Apply the rotary positional encoding to the query and key tensors.
+
+    #### Args:
+        - `xq` (Tensor): The query tensor.
+        - `xk` (Tensor): The key tensor.
+        - `freqs_cis` (Tensor): The frequency tensor.
+
+    #### Returns:
+        - `tuple`: The modified query and key tensors.
+    """
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+# Define the embedding class
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: list):
+        """#### Initialize the EmbedND class.
+
+        #### Args:
+            - `dim` (int): The dimension of the tensor.
+            - `theta` (int): The theta value for scaling.
+            - `axes_dim` (list): The list of axis dimensions.
+        """
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the EmbedND class.
+
+        #### Args:
+            - `ids` (Tensor): The input tensor.
+
+        #### Returns:
+            - `Tensor`: The embedded tensor.
+        """
+        n_axes = ids.shape[-1]
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+            dim=-3,
+        )
+        return emb.unsqueeze(1)
+
+# Define the MLP embedder class
+class MLPEmbedder(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, dtype=None, device=None, operations=None):
+        """#### Initialize the MLPEmbedder class.
+
+        #### Args:
+            - `in_dim` (int): The input dimension.
+            - `hidden_dim` (int): The hidden dimension.
+            - `dtype` (optional): The data type.
+            - `device` (optional): The device.
+            - `operations` (optional): The operations module.
+        """
+        super().__init__()
+        self.in_layer = operations.Linear(
+            in_dim, hidden_dim, bias=True, dtype=dtype, device=device
+        )
+        self.silu = nn.SiLU()
+        self.out_layer = operations.Linear(
+            hidden_dim, hidden_dim, bias=True, dtype=dtype, device=device
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the MLPEmbedder class.
+
+        #### Args:
+            - `x` (Tensor): The input tensor.
+
+        #### Returns:
+            - `Tensor`: The output tensor.
+        """
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+# Define the RMS normalization class
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, dtype=None, device=None, operations=None):
+        """#### Initialize the RMSNorm class.
+
+        #### Args:
+            - `dim` (int): The dimension of the tensor.
+            - `dtype` (optional): The data type.
+            - `device` (optional): The device.
+            - `operations` (optional): The operations module.
+        """
+        super().__init__()
+        self.scale = nn.Parameter(torch.empty((dim), dtype=dtype, device=device))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the RMSNorm class.
+
+        #### Args:
+            - `x` (Tensor): The input tensor.
+
+        #### Returns:
+            - `Tensor`: The normalized tensor.
+        """
+        return rms_norm(x, self.scale, 1e-6)
+
+# Define the query-key normalization class
+class QKNorm(nn.Module):
+    def __init__(self, dim: int, dtype=None, device=None, operations=None):
+        """#### Initialize the QKNorm class.
+
+        #### Args:
+            - `dim` (int): The dimension of the tensor.
+            - `dtype` (optional): The data type.
+            - `device` (optional): The device.
+            - `operations` (optional): The operations module.
+        """
+        super().__init__()
+        self.query_norm = RMSNorm(dim, dtype=dtype, device=device, operations=operations)
+        self.key_norm = RMSNorm(dim, dtype=dtype, device=device, operations=operations)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> tuple:
+        """#### Forward pass for the QKNorm class.
+
+        #### Args:
+            - `q` (Tensor): The query tensor.
+            - `k` (Tensor): The key tensor.
+            - `v` (Tensor): The value tensor.
+
+        #### Returns:
+            - `tuple`: The normalized query and key tensors.
+        """
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q.to(v), k.to(v)
+
+# Define the self-attention class
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False, dtype=None, device=None, operations=None):
+        """#### Initialize the SelfAttention class.
+
+        #### Args:
+            - `dim` (int): The dimension of the tensor.
+            - `num_heads` (int, optional): The number of attention heads. Defaults to 8.
+            - `qkv_bias` (bool, optional): Whether to use bias in QKV projection. Defaults to False.
+            - `dtype` (optional): The data type.
+            - `device` (optional): The device.
+            - `operations` (optional): The operations module.
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.qkv = operations.Linear(dim, dim * 3, bias=qkv_bias, dtype=dtype, device=device)
+        self.norm = QKNorm(head_dim, dtype=dtype, device=device, operations=operations)
+        self.proj = operations.Linear(dim, dim, dtype=dtype, device=device)
+
+# Define the modulation output dataclass
+@dataclass
+class ModulationOut:
+    shift: torch.Tensor
+    scale: torch.Tensor
+    gate: torch.Tensor
+
+# Define the modulation class
+class Modulation(nn.Module):
+    def __init__(self, dim: int, double: bool, dtype=None, device=None, operations=None):
+        """#### Initialize the Modulation class.
+
+        #### Args:
+            - `dim` (int): The dimension of the tensor.
+            - `double` (bool): Whether to use double modulation.
+            - `dtype` (optional): The data type.
+            - `device` (optional): The device.
+            - `operations` (optional): The operations module.
+        """
+        super().__init__()
+        self.is_double = double
+        self.multiplier = 6 if double else 3
+        self.lin = operations.Linear(dim, self.multiplier * dim, bias=True, dtype=dtype, device=device)
+
+    def forward(self, vec: torch.Tensor) -> tuple:
+        """#### Forward pass for the Modulation class.
+
+        #### Args:
+            - `vec` (Tensor): The input tensor.
+
+        #### Returns:
+            - `tuple`: The modulation output.
+        """
+        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
+        return (ModulationOut(*out[:3]), ModulationOut(*out[3:]) if self.is_double else None)
+
+# Define the double stream block class
+class DoubleStreamBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, dtype=None, device=None, operations=None):
+        """#### Initialize the DoubleStreamBlock class.
+
+        #### Args:
+            - `hidden_size` (int): The hidden size.
+            - `num_heads` (int): The number of attention heads.
+            - `mlp_ratio` (float): The MLP ratio.
+            - `qkv_bias` (bool, optional): Whether to use bias in QKV projection. Defaults to False.
+            - `dtype` (optional): The data type.
+            - `device` (optional): The device.
+            - `operations` (optional): The operations module.
+        """
+        super().__init__()
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.img_mod = Modulation(hidden_size, double=True, dtype=dtype, device=device, operations=operations)
+        self.img_norm1 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, dtype=dtype, device=device, operations=operations)
+        self.img_norm2 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.img_mlp = nn.Sequential(
+            operations.Linear(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, device=device),
+            nn.GELU(approximate="tanh"),
+            operations.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
+        )
+
+        self.txt_mod = Modulation(hidden_size, double=True, dtype=dtype, device=device, operations=operations)
+        self.txt_norm1 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, dtype=dtype, device=device, operations=operations)
+        self.txt_norm2 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.txt_mlp = nn.Sequential(
+            operations.Linear(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, device=device),
+            nn.GELU(approximate="tanh"),
+            operations.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
+        )
+
+    def forward(self, img: torch.Tensor, txt: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor) -> tuple:
+        """#### Forward pass for the DoubleStreamBlock class.
+
+        #### Args:
+            - `img` (Tensor): The image tensor.
+            - `txt` (Tensor): The text tensor.
+            - `vec` (Tensor): The vector tensor.
+            - `pe` (Tensor): The positional encoding tensor.
+
+        #### Returns:
+            - `tuple`: The modified image and text tensors.
+        """
+        img_mod1, img_mod2 = self.img_mod(vec)
+        txt_mod1, txt_mod2 = self.txt_mod(vec)
+
+        # prepare image for attention
+        img_modulated = self.img_norm1(img)
+        img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
+        img_qkv = self.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = img_qkv.view(img_qkv.shape[0], img_qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+
+        # prepare txt for attention
+        txt_modulated = self.txt_norm1(txt)
+        txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+        txt_qkv = self.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = txt_qkv.view(txt_qkv.shape[0], txt_qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        # run actual attention
+        attn = attention(
+            torch.cat((txt_q, img_q), dim=2),
+            torch.cat((txt_k, img_k), dim=2),
+            torch.cat((txt_v, img_v), dim=2),
+            pe=pe,
+        )
+
+        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+        # calculate the img bloks
+        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
+        img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
+
+        # calculate the txt bloks
+        txt += txt_mod1.gate * self.txt_attn.proj(txt_attn)
+        txt += txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+
+        if txt.dtype == torch.float16:
+            txt = torch.nan_to_num(txt, nan=0.0, posinf=65504, neginf=-65504)
+
+        return img, txt
+
+# Define the single stream block class
+class SingleStreamBlock(nn.Module):
+    """
+    A DiT block with parallel linear layers as described in
+    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0, qk_scale: float = None, dtype=None, device=None, operations=None):
+        """#### Initialize the SingleStreamBlock class.
+
+        #### Args:
+            - `hidden_size` (int): The hidden size.
+            - `num_heads` (int): The number of attention heads.
+            - `mlp_ratio` (float, optional): The MLP ratio. Defaults to 4.0.
+            - `qk_scale` (float, optional): The QK scale. Defaults to None.
+            - `dtype` (optional): The data type.
+            - `device` (optional): The device.
+            - `operations` (optional): The operations module.
+        """
+        super().__init__()
+        self.hidden_dim = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        # qkv and mlp_in
+        self.linear1 = operations.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim, dtype=dtype, device=device)
+        # proj and mlp_out
+        self.linear2 = operations.Linear(hidden_size + self.mlp_hidden_dim, hidden_size, dtype=dtype, device=device)
+
+        self.norm = QKNorm(head_dim, dtype=dtype, device=device, operations=operations)
+
+        self.hidden_size = hidden_size
+        self.pre_norm = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+
+        self.mlp_act = nn.GELU(approximate="tanh")
+        self.modulation = Modulation(hidden_size, double=False, dtype=dtype, device=device, operations=operations)
+
+    def forward(self, x: torch.Tensor, vec: torch.Tensor, pe: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the SingleStreamBlock class.
+
+        #### Args:
+            - `x` (Tensor): The input tensor.
+            - `vec` (Tensor): The vector tensor.
+            - `pe` (Tensor): The positional encoding tensor.
+
+        #### Returns:
+            - `Tensor`: The modified tensor.
+        """
+        mod, _ = self.modulation(vec)
+        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
+        qkv, mlp = torch.split(
+            self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1
+        )
+
+        q, k, v = qkv.view(qkv.shape[0], qkv.shape[1], 3, self.num_heads, -1).permute(
+            2, 0, 3, 1, 4
+        )
+        q, k = self.norm(q, k, v)
+
+        # compute attention
+        attn = attention(q, k, v, pe=pe)
+        # compute activation in mlp stream, cat again and run second linear layer
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+        x += mod.gate * output
+        if x.dtype == torch.float16:
+            x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
+        return x
+
+class LastLayer(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        patch_size: int,
+        out_channels: int,
+        dtype=None,
+        device=None,
+        operations=None,
+    ):
+        """#### Initialize the LastLayer class.
+
+        #### Args:
+            - `hidden_size` (int): The hidden size.
+            - `patch_size` (int): The patch size.
+            - `out_channels` (int): The number of output channels.
+            - `dtype` (optional): The data type.
+            - `device` (optional): The device.
+            - `operations` (optional): The operations module.
+        """
+        super().__init__()
+        self.norm_final = operations.LayerNorm(
+            hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device
+        )
+        self.linear = operations.Linear(
+            hidden_size,
+            patch_size * patch_size * out_channels,
+            bias=True,
+            dtype=dtype,
+            device=device,
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            operations.Linear(
+                hidden_size, 2 * hidden_size, bias=True, dtype=dtype, device=device
+            ),
+        )
+
+    def forward(self, x: torch.Tensor, vec: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass for the LastLayer class.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `vec` (torch.Tensor): The vector tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
+        x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+        x = self.linear(x)
+        return x
+
+
+def pad_to_patch_size(img: torch.Tensor, patch_size: tuple = (2, 2), padding_mode: str = "circular") -> torch.Tensor:
+    """#### Pad the image to the specified patch size.
+
+    #### Args:
+        - `img` (torch.Tensor): The input image tensor.
+        - `patch_size` (tuple, optional): The patch size. Defaults to (2, 2).
+        - `padding_mode` (str, optional): The padding mode. Defaults to "circular".
+
+    #### Returns:
+        - `torch.Tensor`: The padded image tensor.
+    """
+    if (
+        padding_mode == "circular"
+        and torch.jit.is_tracing()
+        or torch.jit.is_scripting()
+    ):
+        padding_mode = "reflect"
+    pad_h = (patch_size[0] - img.shape[-2] % patch_size[0]) % patch_size[0]
+    pad_w = (patch_size[1] - img.shape[-1] % patch_size[1]) % patch_size[1]
+    return torch.nn.functional.pad(img, (0, pad_w, 0, pad_h), mode=padding_mode)
+
+
+try:
+    rms_norm_torch = torch.nn.functional.rms_norm
+except Exception:
+    rms_norm_torch = None
+
+
+def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    """#### Apply RMS normalization to the input tensor.
+
+    #### Args:
+        - `x` (torch.Tensor): The input tensor.
+        - `weight` (torch.Tensor): The weight tensor.
+        - `eps` (float, optional): The epsilon value for numerical stability. Defaults to 1e-6.
+
+    #### Returns:
+        - `torch.Tensor`: The normalized tensor.
+    """
+    if rms_norm_torch is not None and not (
+        torch.jit.is_tracing() or torch.jit.is_scripting()
+    ):
+        return rms_norm_torch(
+            x,
+            weight.shape,
+            weight=cast.cast_to(weight, dtype=x.dtype, device=x.device),
+            eps=eps,
+        )
+    else:
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + eps)
+        return (x * rrms) * cast.cast_to(weight, dtype=x.dtype, device=x.device)
+
+
+@dataclass
+class FluxParams:
+    in_channels: int
+    vec_in_dim: int
+    context_in_dim: int
+    hidden_size: int
+    mlp_ratio: float
+    num_heads: int
+    depth: int
+    depth_single_blocks: int
+    axes_dim: list
+    theta: int
+    qkv_bias: bool
+    guidance_embed: bool
+
+
+class Flux3(nn.Module):
+    """
+    Transformer model for flow matching on sequences.
+    """
+
+    def __init__(
+        self,
+        image_model=None,
+        final_layer: bool = True,
+        dtype=None,
+        device=None,
+        operations=None,
+        **kwargs,
+    ):
+        """#### Initialize the Flux3 class.
+
+        #### Args:
+            - `image_model` (optional): The image model.
+            - `final_layer` (bool, optional): Whether to include the final layer. Defaults to True.
+            - `dtype` (optional): The data type.
+            - `device` (optional): The device.
+            - `operations` (optional): The operations module.
+            - `**kwargs`: Additional keyword arguments.
+        """
+        super().__init__()
+        self.dtype = dtype
+        params = FluxParams(**kwargs)
+        self.params = params
+        self.in_channels = params.in_channels * 2 * 2
+        self.out_channels = self.in_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(
+                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
+            )
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(
+                f"Got {params.axes_dim} but expected positional dim {pe_dim}"
+            )
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.pe_embedder = EmbedND(
+            dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim
+        )
+        self.img_in = operations.Linear(
+            self.in_channels, self.hidden_size, bias=True, dtype=dtype, device=device
+        )
+        self.time_in = MLPEmbedder(
+            in_dim=256,
+            hidden_dim=self.hidden_size,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )
+        self.vector_in = MLPEmbedder(
+            params.vec_in_dim,
+            self.hidden_size,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )
+        self.guidance_in = (
+            MLPEmbedder(
+                in_dim=256,
+                hidden_dim=self.hidden_size,
+                dtype=dtype,
+                device=device,
+                operations=operations,
+            )
+            if params.guidance_embed
+            else nn.Identity()
+        )
+        self.txt_in = operations.Linear(
+            params.context_in_dim, self.hidden_size, dtype=dtype, device=device
+        )
+
+        self.double_blocks = nn.ModuleList(
+            [
+                DoubleStreamBlock(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                    dtype=dtype,
+                    device=device,
+                    operations=operations,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.single_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    dtype=dtype,
+                    device=device,
+                    operations=operations,
+                )
+                for _ in range(params.depth_single_blocks)
+            ]
+        )
+
+        if final_layer:
+            self.final_layer = LastLayer(
+                self.hidden_size,
+                1,
+                self.out_channels,
+                dtype=dtype,
+                device=device,
+                operations=operations,
+            )
+
+    def forward_orig(
+        self,
+        img: torch.Tensor,
+        img_ids: torch.Tensor,
+        txt: torch.Tensor,
+        txt_ids: torch.Tensor,
+        timesteps: torch.Tensor,
+        y: torch.Tensor,
+        guidance: torch.Tensor = None,
+        control=None,
+    ) -> torch.Tensor:
+        """#### Original forward pass for the Flux3 class.
+
+        #### Args:
+            - `img` (torch.Tensor): The image tensor.
+            - `img_ids` (torch.Tensor): The image IDs tensor.
+            - `txt` (torch.Tensor): The text tensor.
+            - `txt_ids` (torch.Tensor): The text IDs tensor.
+            - `timesteps` (torch.Tensor): The timesteps tensor.
+            - `y` (torch.Tensor): The vector tensor.
+            - `guidance` (torch.Tensor, optional): The guidance tensor. Defaults to None.
+            - `control` (optional): The control tensor. Defaults to None.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # running on sequences img
+        img = self.img_in(img)
+        vec = self.time_in(sampling_util.timestep_embedding_flux(timesteps, 256).to(img.dtype))
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError(
+                    "Didn't get guidance strength for guidance distilled model."
+                )
+            vec = vec + self.guidance_in(
+                sampling_util.timestep_embedding_flux(guidance, 256).to(img.dtype)
+            )
+
+        vec = vec + self.vector_in(y)
+        txt = self.txt_in(txt)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        for i, block in enumerate(self.double_blocks):
+            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+
+            if control is not None:  # Controlnet
+                control_i = control.get("input")
+                if i < len(control_i):
+                    add = control_i[i]
+                    if add is not None:
+                        img += add
+
+        img = torch.cat((txt, img), 1)
+
+        for i, block in enumerate(self.single_blocks):
+            img = block(img, vec=vec, pe=pe)
+
+            if control is not None:  # Controlnet
+                control_o = control.get("output")
+                if i < len(control_o):
+                    add = control_o[i]
+                    if add is not None:
+                        img[:, txt.shape[1] :, ...] += add
+
+        img = img[:, txt.shape[1] :, ...]
+
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+        return img
+
+    def forward(self, x: torch.Tensor, timestep: torch.Tensor, context: torch.Tensor, y: torch.Tensor, guidance: torch.Tensor, control=None, **kwargs) -> torch.Tensor:
+        """#### Forward pass for the Flux3 class.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `timestep` (torch.Tensor): The timestep tensor.
+            - `context` (torch.Tensor): The context tensor.
+            - `y` (torch.Tensor): The vector tensor.
+            - `guidance` (torch.Tensor): The guidance tensor.
+            - `control` (optional): The control tensor. Defaults to None.
+            - `**kwargs`: Additional keyword arguments.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        bs, c, h, w = x.shape
+        patch_size = 2
+        x = pad_to_patch_size(x, (patch_size, patch_size))
+
+        img = rearrange(
+            x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size
+        )
+
+        h_len = (h + (patch_size // 2)) // patch_size
+        w_len = (w + (patch_size // 2)) // patch_size
+        img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
+        img_ids[..., 1] = (
+            img_ids[..., 1]
+            + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype)[
+                :, None
+            ]
+        )
+        img_ids[..., 2] = (
+            img_ids[..., 2]
+            + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype)[
+                None, :
+            ]
+        )
+        img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+        txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
+        out = self.forward_orig(
+            img, img_ids, context, txt_ids, timestep, y, guidance, control
+        )
+        return rearrange(
+            out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2
+        )[:, :, :h, :w]
+
+
+class Flux2(ModelBase.BaseModel):
+    def __init__(self, model_config: dict, model_type=sampling.ModelType.FLUX, device=None):
+        """#### Initialize the Flux2 class.
+
+        #### Args:
+            - `model_config` (dict): The model configuration.
+            - `model_type` (sampling.ModelType, optional): The model type. Defaults to sampling.ModelType.FLUX.
+            - `device` (optional): The device.
+        """
+        super().__init__(model_config, model_type, device=device, unet_model=Flux3, flux=True)
+
+    def encode_adm(self, **kwargs) -> torch.Tensor:
+        """#### Encode the ADM.
+
+        #### Args:
+            - `**kwargs`: Additional keyword arguments.
+
+        #### Returns:
+            - `torch.Tensor`: The encoded ADM tensor.
+        """
+        return kwargs["pooled_output"]
+
+    def extra_conds(self, **kwargs) -> dict:
+        """#### Get extra conditions.
+
+        #### Args:
+            - `**kwargs`: Additional keyword arguments.
+
+        #### Returns:
+            - `dict`: The extra conditions.
+        """
+        out = super().extra_conds(**kwargs)
+        cross_attn = kwargs.get("cross_attn", None)
+        if cross_attn is not None:
+            out["c_crossattn"] = cond.CONDRegular(cross_attn)
+        out["guidance"] = cond.CONDRegular(torch.FloatTensor([kwargs.get("guidance", 3.5)]))
+        return out
+
+
+class Flux(ModelBase.BASE):
+    unet_config = {
+        "image_model": "flux",
+        "guidance_embed": True,
+    }
+
+    sampling_settings = {}
+
+    unet_extra_config = {}
+    latent_format = Latent.Flux1
+
+    memory_usage_factor = 2.8
+
+    supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
+
+    vae_key_prefix = ["vae."]
+    text_encoder_key_prefix = ["text_encoders."]
+
+    def get_model(self, state_dict: dict, prefix: str = "", device=None) -> Flux2:
+        """#### Get the model.
+
+        #### Args:
+            - `state_dict` (dict): The state dictionary.
+            - `prefix` (str, optional): The prefix. Defaults to "".
+            - `device` (optional): The device.
+
+        #### Returns:
+            - `Flux2`: The Flux2 model.
+        """
+        out = Flux2(self, device=device)
+        return out
+
+
+models = [Flux]

modules/Device/Device.py

@@ -0,0 +1,1602 @@
+import logging
+import platform
+import sys
+from enum import Enum
+from typing import Tuple, Union
+import packaging.version
+
+import psutil
+import torch
+
+if packaging.version.parse(torch.__version__) >= packaging.version.parse("1.12.0"):
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+
+class VRAMState(Enum):
+    """#### Enum for VRAM states.
+    """
+    DISABLED = 0  # No vram present: no need to move _internal to vram
+    NO_VRAM = 1  # Very low vram: enable all the options to save vram
+    LOW_VRAM = 2
+    NORMAL_VRAM = 3
+    HIGH_VRAM = 4
+    SHARED = 5  # No dedicated vram: memory shared between CPU and GPU but _internal still need to be moved between both.
+
+
+class CPUState(Enum):
+    """#### Enum for CPU states.
+    """
+    GPU = 0
+    CPU = 1
+    MPS = 2
+
+
+# Determine VRAM State
+vram_state = VRAMState.NORMAL_VRAM
+set_vram_to = VRAMState.NORMAL_VRAM
+cpu_state = CPUState.GPU
+
+total_vram = 0
+
+lowvram_available = True
+xpu_available = False
+
+directml_enabled = False
+try:
+    if torch.xpu.is_available():
+        xpu_available = True
+except:
+    pass
+
+try:
+    if torch.backends.mps.is_available():
+        cpu_state = CPUState.MPS
+        import torch.mps
+except:
+    pass
+
+
+def is_intel_xpu() -> bool:
+    """#### Check if Intel XPU is available.
+
+    #### Returns:
+        - `bool`: Whether Intel XPU is available.
+    """
+    global cpu_state
+    global xpu_available
+    if cpu_state == CPUState.GPU:
+        if xpu_available:
+            return True
+    return False
+
+
+def get_torch_device() -> torch.device:
+    """#### Get the torch device.
+    
+    #### Returns:
+        - `torch.device`: The torch device.
+    """
+    global directml_enabled
+    global cpu_state
+    if directml_enabled:
+        global directml_device
+        return directml_device
+    if cpu_state == CPUState.MPS:
+        return torch.device("mps")
+    if cpu_state == CPUState.CPU:
+        return torch.device("cpu")
+    else:
+        if is_intel_xpu():
+            return torch.device("xpu", torch.xpu.current_device())
+        else:
+            return torch.device(torch.cuda.current_device())
+
+
+def get_total_memory(dev: torch.device = None, torch_total_too: bool = False) -> int:
+    """#### Get the total memory.
+
+    #### Args:
+        - `dev` (torch.device, optional): The device. Defaults to None.
+        - `torch_total_too` (bool, optional): Whether to get the total memory in PyTorch. Defaults to False.
+
+    #### Returns:
+        - `int`: The total memory.
+    """
+    global directml_enabled
+    if dev is None:
+        dev = get_torch_device()
+
+    if hasattr(dev, "type") and (dev.type == "cpu" or dev.type == "mps"):
+        mem_total = psutil.virtual_memory().total
+        mem_total_torch = mem_total
+    else:
+        if directml_enabled:
+            mem_total = 1024 * 1024 * 1024
+            mem_total_torch = mem_total
+        elif is_intel_xpu():
+            stats = torch.xpu.memory_stats(dev)
+            mem_reserved = stats["reserved_bytes.all.current"]
+            mem_total_torch = mem_reserved
+            mem_total = torch.xpu.get_device_properties(dev).total_memory
+        else:
+            stats = torch.cuda.memory_stats(dev)
+            mem_reserved = stats["reserved_bytes.all.current"]
+            _, mem_total_cuda = torch.cuda.mem_get_info(dev)
+            mem_total_torch = mem_reserved
+            mem_total = mem_total_cuda
+
+    if torch_total_too:
+        return (mem_total, mem_total_torch)
+    else:
+        return mem_total
+
+
+total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
+total_ram = psutil.virtual_memory().total / (1024 * 1024)
+logging.info(
+    "Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram)
+)
+try:
+    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
+except:
+    OOM_EXCEPTION = Exception
+
+XFORMERS_VERSION = ""
+XFORMERS_ENABLED_VAE = True
+try:
+    import xformers
+    import xformers.ops
+
+    XFORMERS_IS_AVAILABLE = True
+    try:
+        XFORMERS_IS_AVAILABLE = xformers._has_cpp_library
+    except:
+        pass
+    try:
+        XFORMERS_VERSION = xformers.version.__version__
+        logging.info("xformers version: {}".format(XFORMERS_VERSION))
+        if XFORMERS_VERSION.startswith("0.0.18"):
+            logging.warning(
+                "\nWARNING: This version of xformers has a major bug where you will get black images when generating high resolution images."
+            )
+            logging.warning(
+                "Please downgrade or upgrade xformers to a different version.\n"
+            )
+            XFORMERS_ENABLED_VAE = False
+    except:
+        pass
+except:
+    XFORMERS_IS_AVAILABLE = False
+
+
+def is_nvidia() -> bool:
+    """#### Checks if user has an Nvidia GPU
+
+    #### Returns
+        - `bool`: Whether the GPU is Nvidia
+    """
+    global cpu_state
+    if cpu_state == CPUState.GPU:
+        if torch.version.cuda:
+            return True
+    return False
+
+
+ENABLE_PYTORCH_ATTENTION = False
+
+VAE_DTYPE = torch.float32
+
+try:
+    if is_nvidia():
+        torch_version = torch.version.__version__
+        if int(torch_version[0]) >= 2:
+            if ENABLE_PYTORCH_ATTENTION is False:
+                ENABLE_PYTORCH_ATTENTION = True
+            if (
+                torch.cuda.is_bf16_supported()
+                and torch.cuda.get_device_properties(torch.cuda.current_device()).major
+                >= 8
+            ):
+                VAE_DTYPE = torch.bfloat16
+except:
+    pass
+
+if is_intel_xpu():
+    VAE_DTYPE = torch.bfloat16
+
+if ENABLE_PYTORCH_ATTENTION:
+    torch.backends.cuda.enable_math_sdp(True)
+    torch.backends.cuda.enable_flash_sdp(True)
+    torch.backends.cuda.enable_mem_efficient_sdp(True)
+
+
+FORCE_FP32 = False
+FORCE_FP16 = False
+
+if lowvram_available:
+    if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
+        vram_state = set_vram_to
+
+if cpu_state != CPUState.GPU:
+    vram_state = VRAMState.DISABLED
+
+if cpu_state == CPUState.MPS:
+    vram_state = VRAMState.SHARED
+
+logging.info(f"Set vram state to: {vram_state.name}")
+
+DISABLE_SMART_MEMORY = False
+
+if DISABLE_SMART_MEMORY:
+    logging.info("Disabling smart memory management")
+
+
+def get_torch_device_name(device: torch.device) -> str:
+    """#### Get the name of the torch compatible device
+
+    #### Args:
+        - `device` (torch.device): the device
+
+    #### Returns:
+        - `str`: the name of the device
+    """
+    if hasattr(device, "type"):
+        if device.type == "cuda":
+            try:
+                allocator_backend = torch.cuda.get_allocator_backend()
+            except:
+                allocator_backend = ""
+            return "{} {} : {}".format(
+                device, torch.cuda.get_device_name(device), allocator_backend
+            )
+        else:
+            return "{}".format(device.type)
+    elif is_intel_xpu():
+        return "{} {}".format(device, torch.xpu.get_device_name(device))
+    else:
+        return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))
+
+
+try:
+    logging.info("Device: {}".format(get_torch_device_name(get_torch_device())))
+except:
+    logging.warning("Could not pick default device.")
+
+logging.info("VAE dtype: {}".format(VAE_DTYPE))
+
+current_loaded_models = []
+
+
+def module_size(module: torch.nn.Module) -> int:
+    """#### Get the size of a module
+    
+    #### Args:
+        - `module` (torch.nn.Module): The module
+    
+    #### Returns:
+        - `int`: The size of the module
+    """
+    module_mem = 0
+    sd = module.state_dict()
+    for k in sd:
+        t = sd[k]
+        module_mem += t.nelement() * t.element_size()
+    return module_mem
+
+
+class LoadedModel:
+    """#### Class to load a model
+    """
+    def __init__(self, model: torch.nn.Module):
+        """#### Initialize the class
+        
+        #### Args:
+            - `model`: The model
+        """
+        self.model = model
+        self.device = model.load_device
+        self.weights_loaded = False
+        self.real_model = None
+
+    def model_memory(self):
+        """#### Get the model memory
+        
+        #### Returns:
+            - `int`: The model memory
+        """
+        return self.model.model_size()
+
+    
+    def model_offloaded_memory(self):
+        """#### Get the offloaded model memory
+        
+        #### Returns:
+            - `int`: The offloaded model memory
+        """
+        return self.model.model_size() - self.model.loaded_size()
+
+    def model_memory_required(self, device: torch.device) -> int:
+        """#### Get the required model memory
+        
+        #### Args:
+            - `device`: The device
+        
+        #### Returns:
+            - `int`: The required model memory
+        """
+        if hasattr(self.model, 'current_loaded_device') and device == self.model.current_loaded_device():
+            return self.model_offloaded_memory()
+        else:
+            return self.model_memory()
+
+    def model_load(self, lowvram_model_memory: int = 0, force_patch_weights: bool = False) -> torch.nn.Module:
+        """#### Load the model
+        
+        #### Args:
+            - `lowvram_model_memory` (int, optional): The low VRAM model memory. Defaults to 0.
+            - `force_patch_weights` (bool, optional): Whether to force patch the weights. Defaults to False.
+        
+        #### Returns:
+            - `torch.nn.Module`: The real model
+        """
+        patch_model_to = self.device
+
+        self.model.model_patches_to(self.device)
+        self.model.model_patches_to(self.model.model_dtype())
+
+        load_weights = not self.weights_loaded
+
+        try:
+            if hasattr(self.model, "patch_model_lowvram") and  lowvram_model_memory > 0 and load_weights:
+                self.real_model = self.model.patch_model_lowvram(
+                    device_to=patch_model_to,
+                    lowvram_model_memory=lowvram_model_memory,
+                    force_patch_weights=force_patch_weights,
+                )
+            else:
+                self.real_model = self.model.patch_model(
+                    device_to=patch_model_to, patch_weights=load_weights
+                )
+        except Exception as e:
+            self.model.unpatch_model(self.model.offload_device)
+            self.model_unload()
+            raise e
+        self.weights_loaded = True
+        return self.real_model
+
+    def model_load_flux(self, lowvram_model_memory: int = 0, force_patch_weights: bool = False) -> torch.nn.Module:
+        """#### Load the model
+        
+        #### Args:
+            - `lowvram_model_memory` (int, optional): The low VRAM model memory. Defaults to 0.
+            - `force_patch_weights` (bool, optional): Whether to force patch the weights. Defaults to False.
+        
+        #### Returns:
+            - `torch.nn.Module`: The real model
+        """
+        patch_model_to = self.device
+
+        self.model.model_patches_to(self.device)
+        self.model.model_patches_to(self.model.model_dtype())
+
+        load_weights = not self.weights_loaded
+
+        if self.model.loaded_size() > 0:
+            use_more_vram = lowvram_model_memory
+            if use_more_vram == 0:
+                use_more_vram = 1e32
+            self.model_use_more_vram(use_more_vram)
+        else:
+            try:
+                self.real_model = self.model.patch_model_flux(
+                    device_to=patch_model_to,
+                    lowvram_model_memory=lowvram_model_memory,
+                    load_weights=load_weights,
+                    force_patch_weights=force_patch_weights,
+                )
+            except Exception as e:
+                self.model.unpatch_model(self.model.offload_device)
+                self.model_unload()
+                raise e
+
+        if (
+            is_intel_xpu()
+            and "ipex" in globals()
+            and self.real_model is not None
+        ):
+            import ipex
+            with torch.no_grad():
+                self.real_model = ipex.optimize(
+                    self.real_model.eval(),
+                    inplace=True,
+                    graph_mode=True,
+                    concat_linear=True,
+                )
+
+        self.weights_loaded = True
+        return self.real_model
+    
+    def should_reload_model(self, force_patch_weights: bool = False) -> bool:
+        """#### Checks if the model should be reloaded
+
+        #### Args:
+            - `force_patch_weights` (bool, optional): If model reloading should be enforced. Defaults to False.
+
+        #### Returns:
+            - `bool`: Whether the model should be reloaded
+        """
+        if force_patch_weights and self.model.lowvram_patch_counter > 0:
+            return True
+        return False
+
+    def model_unload(self, unpatch_weights: bool = True) -> None:
+        """#### Unloads the patched model
+
+        #### Args:
+            - `unpatch_weights` (bool, optional): Whether the weights should be unpatched. Defaults to True.
+        """
+        self.model.unpatch_model(
+            self.model.offload_device, unpatch_weights=unpatch_weights
+        )
+        self.model.model_patches_to(self.model.offload_device)
+        self.weights_loaded = self.weights_loaded and not unpatch_weights
+        self.real_model = None
+
+    def model_use_more_vram(self, extra_memory: int) -> int:
+        """#### Use more VRAM
+        
+        #### Args:
+            - `extra_memory`: The extra memory
+        """
+        return self.model.partially_load(self.device, extra_memory)
+    
+    def __eq__(self, other: torch.nn.Module) -> bool:
+        """#### Verify if the model is equal to another
+
+        #### Args:
+            - `other` (torch.nn.Module): the other model
+
+        #### Returns:
+            - `bool`: Whether the two models are equal
+        """
+        return self.model is other.model
+
+
+def minimum_inference_memory() -> int:
+    """#### The minimum memory requirement for inference, equals to 1024^3
+
+    #### Returns:
+        - `int`: the memory requirement
+    """
+    return 1024 * 1024 * 1024
+
+
+def unload_model_clones(model: torch.nn.Module, unload_weights_only:bool = True, force_unload: bool = True) -> bool:
+    """#### Unloads the model clones
+
+    #### Args:
+        - `model` (torch.nn.Module): The model
+        - `unload_weights_only` (bool, optional): Whether to unload only the weights. Defaults to True.
+        - `force_unload` (bool, optional): Whether to force the unload. Defaults to True.
+
+    #### Returns:
+        - `bool`: Whether the model was unloaded
+    """
+    to_unload = []
+    for i in range(len(current_loaded_models)):
+        if model.is_clone(current_loaded_models[i].model):
+            to_unload = [i] + to_unload
+
+    if len(to_unload) == 0:
+        return True
+
+    same_weights = 0
+
+    if same_weights == len(to_unload):
+        unload_weight = False
+    else:
+        unload_weight = True
+
+    if not force_unload:
+        if unload_weights_only and unload_weight is False:
+            return None
+
+    for i in to_unload:
+        logging.debug("unload clone {} {}".format(i, unload_weight))
+        current_loaded_models.pop(i).model_unload(unpatch_weights=unload_weight)
+
+    return unload_weight
+
+
+def free_memory(memory_required: int, device: torch.device, keep_loaded: list = []) -> None:
+    """#### Free memory
+    
+    #### Args:
+        - `memory_required` (int): The required memory
+        - `device` (torch.device): The device
+        - `keep_loaded` (list, optional): The list of loaded models to keep. Defaults to [].
+    """
+    unloaded_model = []
+    can_unload = []
+
+    for i in range(len(current_loaded_models) - 1, -1, -1):
+        shift_model = current_loaded_models[i]
+        if shift_model.device == device:
+            if shift_model not in keep_loaded:
+                can_unload.append(
+                    (sys.getrefcount(shift_model.model), shift_model.model_memory(), i)
+                )
+
+    for x in sorted(can_unload):
+        i = x[-1]
+        if not DISABLE_SMART_MEMORY:
+            if get_free_memory(device) > memory_required:
+                break
+        current_loaded_models[i].model_unload()
+        unloaded_model.append(i)
+
+    for i in sorted(unloaded_model, reverse=True):
+        current_loaded_models.pop(i)
+
+    if len(unloaded_model) > 0:
+        soft_empty_cache()
+    else:
+        if vram_state != VRAMState.HIGH_VRAM:
+            mem_free_total, mem_free_torch = get_free_memory(
+                device, torch_free_too=True
+            )
+            if mem_free_torch > mem_free_total * 0.25:
+                soft_empty_cache()
+
+def use_more_memory(extra_memory: int, loaded_models: list, device: torch.device) -> None:
+    """#### Use more memory
+    
+    #### Args:
+        - `extra_memory` (int): The extra memory
+        - `loaded_models` (list): The loaded models
+        - `device` (torch.device): The device
+    """
+    for m in loaded_models:
+        if m.device == device:
+            extra_memory -= m.model_use_more_vram(extra_memory)
+            if extra_memory <= 0:
+                break
+
+WINDOWS = any(platform.win32_ver())
+
+EXTRA_RESERVED_VRAM = 400 * 1024 * 1024
+if WINDOWS:
+    EXTRA_RESERVED_VRAM = (
+        600 * 1024 * 1024
+    )  # Windows is higher because of the shared vram issue
+
+def extra_reserved_memory() -> int:
+    """#### Extra reserved memory
+    
+    #### Returns:
+        - `int`: The extra reserved memory
+    """
+    return EXTRA_RESERVED_VRAM
+
+def offloaded_memory(loaded_models: list, device: torch.device) -> int:
+    """#### Offloaded memory
+    
+    #### Args:
+        - `loaded_models` (list): The loaded models
+        - `device` (torch.device): The device
+        
+    #### Returns:
+        - `int`: The offloaded memory
+    """
+    offloaded_mem = 0
+    for m in loaded_models:
+        if m.device == device:
+            offloaded_mem += m.model_offloaded_memory()
+    return offloaded_mem
+
+def load_models_gpu(models: list, memory_required: int = 0, force_patch_weights: bool = False, minimum_memory_required=None, force_full_load=False, flux_enabled: bool = False) -> None:
+    """#### Load models on the GPU
+    
+    #### Args:
+        - `models`(list): The models
+        - `memory_required` (int, optional): The required memory. Defaults to 0.
+        - `force_patch_weights` (bool, optional): Whether to force patch the weights. Defaults to False.
+        - `minimum_memory_required` (int, optional): The minimum memory required. Defaults to None.
+        - `force_full_load` (bool, optional
+        - `flux_enabled` (bool, optional): Whether flux is enabled. Defaults to False.
+    """
+    global vram_state
+    if not flux_enabled:
+    
+        inference_memory = minimum_inference_memory()
+        extra_mem = max(inference_memory, memory_required)
+
+        models = set(models)
+
+        models_to_load = []
+        models_already_loaded = []
+        for x in models:
+            loaded_model = LoadedModel(x)
+            loaded = None
+
+            try:
+                loaded_model_index = current_loaded_models.index(loaded_model)
+            except:
+                loaded_model_index = None
+
+            if loaded_model_index is not None:
+                loaded = current_loaded_models[loaded_model_index]
+                if loaded.should_reload_model(force_patch_weights=force_patch_weights):
+                    current_loaded_models.pop(loaded_model_index).model_unload(
+                        unpatch_weights=True
+                    )
+                    loaded = None
+                else:
+                    models_already_loaded.append(loaded)
+
+            if loaded is None:
+                if hasattr(x, "model"):
+                    logging.info(f"Requested to load {x.model.__class__.__name__}")
+                models_to_load.append(loaded_model)
+
+        if len(models_to_load) == 0:
+            devs = set(map(lambda a: a.device, models_already_loaded))
+            for d in devs:
+                if d != torch.device("cpu"):
+                    free_memory(extra_mem, d, models_already_loaded)
+            return
+
+        logging.info(
+            f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}"
+        )
+
+        total_memory_required = {}
+        for loaded_model in models_to_load:
+            if (
+                unload_model_clones(
+                    loaded_model.model, unload_weights_only=True, force_unload=False
+                )
+                is True
+            ):  # unload clones where the weights are different
+                total_memory_required[loaded_model.device] = total_memory_required.get(
+                    loaded_model.device, 0
+                ) + loaded_model.model_memory_required(loaded_model.device)
+
+        for device in total_memory_required:
+            if device != torch.device("cpu"):
+                free_memory(
+                    total_memory_required[device] * 1.3 + extra_mem,
+                    device,
+                    models_already_loaded,
+                )
+
+        for loaded_model in models_to_load:
+            weights_unloaded = unload_model_clones(
+                loaded_model.model, unload_weights_only=False, force_unload=False
+            )  # unload the rest of the clones where the weights can stay loaded
+            if weights_unloaded is not None:
+                loaded_model.weights_loaded = not weights_unloaded
+
+        for loaded_model in models_to_load:
+            model = loaded_model.model
+            torch_dev = model.load_device
+            if is_device_cpu(torch_dev):
+                vram_set_state = VRAMState.DISABLED
+            else:
+                vram_set_state = vram_state
+            lowvram_model_memory = 0
+            if lowvram_available and (
+                vram_set_state == VRAMState.LOW_VRAM
+                or vram_set_state == VRAMState.NORMAL_VRAM
+            ):
+                model_size = loaded_model.model_memory_required(torch_dev)
+                current_free_mem = get_free_memory(torch_dev)
+                lowvram_model_memory = int(
+                    max(64 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3)
+                )
+                if model_size > (
+                    current_free_mem - inference_memory
+                ):  # only switch to lowvram if really necessary
+                    vram_set_state = VRAMState.LOW_VRAM
+                else:
+                    lowvram_model_memory = 0
+
+            if vram_set_state == VRAMState.NO_VRAM:
+                lowvram_model_memory = 64 * 1024 * 1024
+
+            loaded_model.model_load(
+                lowvram_model_memory, force_patch_weights=force_patch_weights
+            )
+            current_loaded_models.insert(0, loaded_model)
+        return
+    else:
+        inference_memory = minimum_inference_memory()
+        extra_mem = max(inference_memory, memory_required + extra_reserved_memory())
+        if minimum_memory_required is None:
+            minimum_memory_required = extra_mem
+        else:
+            minimum_memory_required = max(
+                inference_memory, minimum_memory_required + extra_reserved_memory()
+            )
+
+        models = set(models)
+
+        models_to_load = []
+        models_already_loaded = []
+        for x in models:
+            loaded_model = LoadedModel(x)
+            loaded = None
+
+            try:
+                loaded_model_index = current_loaded_models.index(loaded_model)
+            except:
+                loaded_model_index = None
+
+            if loaded_model_index is not None:
+                loaded = current_loaded_models[loaded_model_index]
+                if loaded.should_reload_model(
+                    force_patch_weights=force_patch_weights
+                ):  # TODO: cleanup this model reload logic
+                    current_loaded_models.pop(loaded_model_index).model_unload(
+                        unpatch_weights=True
+                    )
+                    loaded = None
+                else:
+                    loaded.currently_used = True
+                    models_already_loaded.append(loaded)
+
+            if loaded is None:
+                if hasattr(x, "model"):
+                    logging.info(f"Requested to load {x.model.__class__.__name__}")
+                models_to_load.append(loaded_model)
+
+        if len(models_to_load) == 0:
+            devs = set(map(lambda a: a.device, models_already_loaded))
+            for d in devs:
+                if d != torch.device("cpu"):
+                    free_memory(
+                        extra_mem + offloaded_memory(models_already_loaded, d),
+                        d,
+                        models_already_loaded,
+                    )
+                    free_mem = get_free_memory(d)
+                    if free_mem < minimum_memory_required:
+                        logging.info(
+                            "Unloading models for lowram load."
+                        )  # TODO: partial model unloading when this case happens, also handle the opposite case where models can be unlowvramed.
+                        models_to_load = free_memory(minimum_memory_required, d)
+                        logging.info("{} models unloaded.".format(len(models_to_load)))
+                    else:
+                        use_more_memory(
+                            free_mem - minimum_memory_required, models_already_loaded, d
+                        )
+            if len(models_to_load) == 0:
+                return
+
+        logging.info(
+            f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}"
+        )
+
+        total_memory_required = {}
+        for loaded_model in models_to_load:
+            unload_model_clones(
+                loaded_model.model, unload_weights_only=True, force_unload=False
+            )  # unload clones where the weights are different
+            total_memory_required[loaded_model.device] = total_memory_required.get(
+                loaded_model.device, 0
+            ) + loaded_model.model_memory_required(loaded_model.device)
+
+        for loaded_model in models_already_loaded:
+            total_memory_required[loaded_model.device] = total_memory_required.get(
+                loaded_model.device, 0
+            ) + loaded_model.model_memory_required(loaded_model.device)
+
+        for loaded_model in models_to_load:
+            weights_unloaded = unload_model_clones(
+                loaded_model.model, unload_weights_only=False, force_unload=False
+            )  # unload the rest of the clones where the weights can stay loaded
+            if weights_unloaded is not None:
+                loaded_model.weights_loaded = not weights_unloaded
+
+        for device in total_memory_required:
+            if device != torch.device("cpu"):
+                free_memory(
+                    total_memory_required[device] * 1.1 + extra_mem,
+                    device,
+                    models_already_loaded,
+                )
+
+        for loaded_model in models_to_load:
+            model = loaded_model.model
+            torch_dev = model.load_device
+            if is_device_cpu(torch_dev):
+                vram_set_state = VRAMState.DISABLED
+            else:
+                vram_set_state = vram_state
+            lowvram_model_memory = 0
+            if (
+                lowvram_available
+                and (
+                    vram_set_state == VRAMState.LOW_VRAM
+                    or vram_set_state == VRAMState.NORMAL_VRAM
+                )
+                and not force_full_load
+            ):
+                model_size = loaded_model.model_memory_required(torch_dev)
+                current_free_mem = get_free_memory(torch_dev)
+                lowvram_model_memory = max(
+                    64 * (1024 * 1024),
+                    (current_free_mem - minimum_memory_required),
+                    min(
+                        current_free_mem * 0.4,
+                        current_free_mem - minimum_inference_memory(),
+                    ),
+                )
+                if (
+                    model_size <= lowvram_model_memory
+                ):  # only switch to lowvram if really necessary
+                    lowvram_model_memory = 0
+
+            if vram_set_state == VRAMState.NO_VRAM:
+                lowvram_model_memory = 64 * 1024 * 1024
+
+            loaded_model.model_load_flux(
+                lowvram_model_memory, force_patch_weights=force_patch_weights
+            )
+            current_loaded_models.insert(0, loaded_model)
+
+        devs = set(map(lambda a: a.device, models_already_loaded))
+        for d in devs:
+            if d != torch.device("cpu"):
+                free_mem = get_free_memory(d)
+                if free_mem > minimum_memory_required:
+                    use_more_memory(
+                        free_mem - minimum_memory_required, models_already_loaded, d
+                    )
+        return
+
+def load_model_gpu(model: torch.nn.Module, flux_enabled:bool = False) -> None:
+    """#### Load a model on the GPU
+    
+    #### Args:
+        - `model` (torch.nn.Module): The model
+        - `flux_enable` (bool, optional): Whether flux is enabled. Defaults to False.
+    """
+    return load_models_gpu([model], flux_enabled=flux_enabled)
+
+
+def cleanup_models(keep_clone_weights_loaded:bool = False):
+    """#### Cleanup the models
+    
+    #### Args:
+        - `keep_clone_weights_loaded` (bool, optional): Whether to keep the clone weights loaded. Defaults to False.
+    """
+    to_delete = []
+    for i in range(len(current_loaded_models)):
+        if sys.getrefcount(current_loaded_models[i].model) <= 2:
+            if not keep_clone_weights_loaded:
+                to_delete = [i] + to_delete
+            elif (
+                sys.getrefcount(current_loaded_models[i].real_model) <= 3
+            ):  # references from .real_model + the .model
+                to_delete = [i] + to_delete
+
+    for i in to_delete:
+        x = current_loaded_models.pop(i)
+        x.model_unload()
+        del x
+
+
+def dtype_size(dtype: torch.dtype) -> int:
+    """#### Get the size of a dtype
+
+    #### Args:
+        - `dtype` (torch.dtype): The dtype
+
+    #### Returns:
+        - `int`: The size of the dtype
+    """
+    dtype_size = 4
+    if dtype == torch.float16 or dtype == torch.bfloat16:
+        dtype_size = 2
+    elif dtype == torch.float32:
+        dtype_size = 4
+    else:
+        try:
+            dtype_size = dtype.itemsize
+        except:  # Old pytorch doesn't have .itemsize
+            pass
+    return dtype_size
+
+
+def unet_offload_device() -> torch.device:
+    """#### Get the offload device for UNet
+    
+    #### Returns:
+        - `torch.device`: The offload device
+    """
+    if vram_state == VRAMState.HIGH_VRAM:
+        return get_torch_device()
+    else:
+        return torch.device("cpu")
+
+
+def unet_inital_load_device(parameters, dtype) -> torch.device:
+    """#### Get the initial load device for UNet
+    
+    #### Args:
+        - `parameters` (int): The parameters
+        - `dtype` (torch.dtype): The dtype
+    
+    #### Returns:
+        - `torch.device`: The initial load device
+    """
+    torch_dev = get_torch_device()
+    if vram_state == VRAMState.HIGH_VRAM:
+        return torch_dev
+
+    cpu_dev = torch.device("cpu")
+    if DISABLE_SMART_MEMORY:
+        return cpu_dev
+
+    model_size = dtype_size(dtype) * parameters
+
+    mem_dev = get_free_memory(torch_dev)
+    mem_cpu = get_free_memory(cpu_dev)
+    if mem_dev > mem_cpu and model_size < mem_dev:
+        return torch_dev
+    else:
+        return cpu_dev
+
+
+def unet_dtype(
+    device: torch.dtype = None,
+    model_params: int = 0,
+    supported_dtypes: list = [torch.float16, torch.bfloat16, torch.float32],
+) -> torch.dtype:
+    """#### Get the dtype for UNet
+
+    #### Args:
+        - `device` (torch.dtype, optional): The device. Defaults to None.
+        - `model_params` (int, optional): The model parameters. Defaults to 0.
+        - `supported_dtypes` (list, optional): The supported dtypes. Defaults to [torch.float16, torch.bfloat16, torch.float32].
+
+    #### Returns:
+        - `torch.dtype`: The dtype
+    """
+    if should_use_fp16(device=device, model_params=model_params, manual_cast=True):
+        if torch.float16 in supported_dtypes:
+            return torch.float16
+    if should_use_bf16(device, model_params=model_params, manual_cast=True):
+        if torch.bfloat16 in supported_dtypes:
+            return torch.bfloat16
+    return torch.float32
+
+
+# None means no manual cast
+def unet_manual_cast(
+    weight_dtype: torch.dtype,
+    inference_device: torch.device,
+    supported_dtypes: list = [torch.float16, torch.bfloat16, torch.float32],
+) -> torch.dtype:
+    """#### Manual cast for UNet
+
+    #### Args:
+        - `weight_dtype` (torch.dtype): The dtype of the weights
+        - `inference_device` (torch.device): The device used for inference
+        - `supported_dtypes` (list, optional): The supported dtypes. Defaults to [torch.float16, torch.bfloat16, torch.float32].
+
+    #### Returns:
+        - `torch.dtype`: The dtype
+    """
+    if weight_dtype == torch.float32:
+        return None
+
+    fp16_supported = should_use_fp16(inference_device, prioritize_performance=False)
+    if fp16_supported and weight_dtype == torch.float16:
+        return None
+
+    bf16_supported = should_use_bf16(inference_device)
+    if bf16_supported and weight_dtype == torch.bfloat16:
+        return None
+
+    if fp16_supported and torch.float16 in supported_dtypes:
+        return torch.float16
+
+    elif bf16_supported and torch.bfloat16 in supported_dtypes:
+        return torch.bfloat16
+    else:
+        return torch.float32
+
+
+def text_encoder_offload_device() -> torch.device:
+    """#### Get the offload device for the text encoder
+    
+    #### Returns:
+        - `torch.device`: The offload device
+    """
+    return torch.device("cpu")
+
+
+def text_encoder_device() -> torch.device:
+    """#### Get the device for the text encoder
+    
+    #### Returns:
+        - `torch.device`: The device
+    """
+    if vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
+        if should_use_fp16(prioritize_performance=False):
+            return get_torch_device()
+        else:
+            return torch.device("cpu")
+    else:
+        return torch.device("cpu")
+    
+def text_encoder_initial_device(load_device: torch.device, offload_device: torch.device, model_size: int = 0) -> torch.device:
+    """#### Get the initial device for the text encoder
+    
+    #### Args:
+        - `load_device` (torch.device): The load device
+        - `offload_device` (torch.device): The offload device
+        - `model_size` (int, optional): The model size. Defaults to 0.
+    
+    #### Returns:
+        - `torch.device`: The initial device
+    """
+    if load_device == offload_device or model_size <= 1024 * 1024 * 1024:
+        return offload_device
+
+    if is_device_mps(load_device):
+        return offload_device
+
+    mem_l = get_free_memory(load_device)
+    mem_o = get_free_memory(offload_device)
+    if mem_l > (mem_o * 0.5) and model_size * 1.2 < mem_l:
+        return load_device
+    else:
+        return offload_device
+
+
+def text_encoder_dtype(device: torch.device = None) -> torch.dtype:
+    """#### Get the dtype for the text encoder
+
+    #### Args:
+        - `device` (torch.device, optional): The device used by the text encoder. Defaults to None.
+
+    Returns:
+        torch.dtype: The dtype
+    """
+    if is_device_cpu(device):
+        return torch.float16
+
+    return torch.float16
+
+
+def intermediate_device() -> torch.device:
+    """#### Get the intermediate device
+    
+    #### Returns:
+        - `torch.device`: The intermediate device
+    """
+    return torch.device("cpu")
+
+
+def vae_device() -> torch.device:
+    """#### Get the VAE device
+    
+    #### Returns:
+        - `torch.device`: The VAE device
+    """
+    return get_torch_device()
+
+
+def vae_offload_device() -> torch.device:
+    """#### Get the offload device for VAE
+    
+    #### Returns:
+        - `torch.device`: The offload device
+    """
+    return torch.device("cpu")
+
+
+def vae_dtype():
+    """#### Get the dtype for VAE
+    
+    #### Returns:
+        - `torch.dtype`: The dtype
+    """
+    global VAE_DTYPE
+    return VAE_DTYPE
+
+
+def get_autocast_device(dev: torch.device) -> str:
+    """#### Get the autocast device
+    
+    #### Args:
+        - `dev` (torch.device): The device
+    
+    #### Returns:
+        - `str`: The autocast device type
+    """
+    if hasattr(dev, "type"):
+        return dev.type
+    return "cuda"
+
+
+def supports_dtype(device: torch.device, dtype: torch.dtype) -> bool:
+    """#### Check if the device supports the dtype
+    
+    #### Args:
+        - `device` (torch.device): The device to check
+        - `dtype`  (torch.dtype): The dtype to check support
+        
+    #### Returns:
+        - `bool`: Whether the dtype is supported by the device
+    """
+    if dtype == torch.float32:
+        return True
+    if is_device_cpu(device):
+        return False
+    if dtype == torch.float16:
+        return True
+    if dtype == torch.bfloat16:
+        return True
+    return False
+
+
+def device_supports_non_blocking(device: torch.device) -> bool:
+    """#### Check if the device supports non-blocking
+
+    #### Args:
+        - `device` (torch.device): The device to check
+
+    #### Returns:
+        - `bool`: Whether the device supports non-blocking
+    """
+    if is_device_mps(device):
+        return False  # pytorch bug? mps doesn't support non blocking
+    return True
+
+def supports_cast(device: torch.device, dtype: torch.dtype):  # TODO
+    """#### Check if the device supports casting
+    
+    #### Args:
+        - `device`: The device
+        - `dtype`: The dtype
+    
+    #### Returns:
+        - `bool`: Whether the device supports casting
+    """
+    if dtype == torch.float32:
+        return True
+    if dtype == torch.float16:
+        return True
+    if directml_enabled:
+        return False
+    if dtype == torch.bfloat16:
+        return True
+    if is_device_mps(device):
+        return False
+    if dtype == torch.float8_e4m3fn:
+        return True
+    if dtype == torch.float8_e5m2:
+        return True
+    return False
+
+def cast_to_device(tensor: torch.Tensor, device: torch.device, dtype: torch.dtype, copy: bool = False) -> torch.Tensor:
+    """#### Cast a tensor to a device
+
+    #### Args:
+        - `tensor` (torch.Tensor): The tensor to cast
+        - `device` (torch.device): The device to cast the tensor to
+        - `dtype` (torch.dtype): The dtype precision to cast to
+        - `copy` (bool, optional): Whether to copy the tensor. Defaults to False.
+
+    #### Returns:
+        - `torch.Tensor`: The tensor cast to the device
+    """
+    device_supports_cast = False
+    if tensor.dtype == torch.float32 or tensor.dtype == torch.float16:
+        device_supports_cast = True
+    elif tensor.dtype == torch.bfloat16:
+        if hasattr(device, "type") and device.type.startswith("cuda"):
+            device_supports_cast = True
+        elif is_intel_xpu():
+            device_supports_cast = True
+
+    non_blocking = device_supports_non_blocking(device)
+
+    if device_supports_cast:
+        if copy:
+            if tensor.device == device:
+                return tensor.to(dtype, copy=copy, non_blocking=non_blocking)
+            return tensor.to(device, copy=copy, non_blocking=non_blocking).to(
+                dtype, non_blocking=non_blocking
+            )
+        else:
+            return tensor.to(device, non_blocking=non_blocking).to(
+                dtype, non_blocking=non_blocking
+            )
+    else:
+        return tensor.to(device, dtype, copy=copy, non_blocking=non_blocking)
+
+def pick_weight_dtype(dtype: torch.dtype, fallback_dtype: torch.dtype, device: torch.device) -> torch.dtype:
+    """#### Pick the weight dtype
+    
+    #### Args:
+        - `dtype`: The dtype
+        - `fallback_dtype`: The fallback dtype
+        - `device`: The device
+        
+    #### Returns:
+        - `torch.dtype`: The weight dtype
+    """
+    if dtype is None:
+        dtype = fallback_dtype
+    elif dtype_size(dtype) > dtype_size(fallback_dtype):
+        dtype = fallback_dtype
+
+    if not supports_cast(device, dtype):
+        dtype = fallback_dtype
+
+    return dtype
+
+def xformers_enabled() -> bool:
+    """#### Check if xformers is enabled
+    
+    #### Returns:
+        - `bool`: Whether xformers is enabled
+    """
+    global directml_enabled
+    global cpu_state
+    if cpu_state != CPUState.GPU:
+        return False
+    if is_intel_xpu():
+        return False
+    if directml_enabled:
+        return False
+    return XFORMERS_IS_AVAILABLE
+
+
+def xformers_enabled_vae() -> bool:
+    """#### Check if xformers is enabled for VAE
+    
+    #### Returns:
+        - `bool`: Whether xformers is enabled for VAE
+    """
+    enabled = xformers_enabled()
+    if not enabled:
+        return False
+
+    return XFORMERS_ENABLED_VAE
+
+
+def pytorch_attention_enabled() -> bool:
+    """#### Check if PyTorch attention is enabled
+    
+    #### Returns:
+        - `bool`: Whether PyTorch attention is enabled
+    """
+    global ENABLE_PYTORCH_ATTENTION
+    return ENABLE_PYTORCH_ATTENTION
+
+def pytorch_attention_flash_attention() -> bool:
+    """#### Check if PyTorch flash attention is enabled and supported.
+
+    #### Returns:
+        - `bool`: True if PyTorch flash attention is enabled and supported, False otherwise.
+    """
+    global ENABLE_PYTORCH_ATTENTION
+    if ENABLE_PYTORCH_ATTENTION:
+        if is_nvidia():  # pytorch flash attention only works on Nvidia
+            return True
+    return False
+
+
+def get_free_memory(dev: torch.device = None, torch_free_too: bool = False) -> Union[int, Tuple[int, int]]:
+    """#### Get the free memory available on the device.
+
+    #### Args:
+        - `dev` (torch.device, optional): The device to check memory for. Defaults to None.
+        - `torch_free_too` (bool, optional): Whether to return both total and torch free memory. Defaults to False.
+
+    #### Returns:
+        - `int` or `Tuple[int, int]`: The free memory available. If `torch_free_too` is True, returns a tuple of total and torch free memory.
+    """
+    global directml_enabled
+    if dev is None:
+        dev = get_torch_device()
+
+    if hasattr(dev, "type") and (dev.type == "cpu" or dev.type == "mps"):
+        mem_free_total = psutil.virtual_memory().available
+        mem_free_torch = mem_free_total
+    else:
+        if directml_enabled:
+            mem_free_total = 1024 * 1024 * 1024
+            mem_free_torch = mem_free_total
+        elif is_intel_xpu():
+            stats = torch.xpu.memory_stats(dev)
+            mem_active = stats["active_bytes.all.current"]
+            mem_reserved = stats["reserved_bytes.all.current"]
+            mem_free_torch = mem_reserved - mem_active
+            mem_free_xpu = (
+                torch.xpu.get_device_properties(dev).total_memory - mem_reserved
+            )
+            mem_free_total = mem_free_xpu + mem_free_torch
+        else:
+            stats = torch.cuda.memory_stats(dev)
+            mem_active = stats["active_bytes.all.current"]
+            mem_reserved = stats["reserved_bytes.all.current"]
+            mem_free_cuda, _ = torch.cuda.mem_get_info(dev)
+            mem_free_torch = mem_reserved - mem_active
+            mem_free_total = mem_free_cuda + mem_free_torch
+
+    if torch_free_too:
+        return (mem_free_total, mem_free_torch)
+    else:
+        return mem_free_total
+
+
+def cpu_mode() -> bool:
+    """#### Check if the current mode is CPU.
+
+    #### Returns:
+        - `bool`: True if the current mode is CPU, False otherwise.
+    """
+    global cpu_state
+    return cpu_state == CPUState.CPU
+
+
+def mps_mode() -> bool:
+    """#### Check if the current mode is MPS.
+
+    #### Returns:
+        - `bool`: True if the current mode is MPS, False otherwise.
+    """
+    global cpu_state
+    return cpu_state == CPUState.MPS
+
+
+def is_device_type(device: torch.device, type: str) -> bool:
+    """#### Check if the device is of a specific type.
+
+    #### Args:
+        - `device` (torch.device): The device to check.
+        - `type` (str): The type to check for.
+
+    #### Returns:
+        - `bool`: True if the device is of the specified type, False otherwise.
+    """
+    if hasattr(device, "type"):
+        if device.type == type:
+            return True
+    return False
+
+
+def is_device_cpu(device: torch.device) -> bool:
+    """#### Check if the device is a CPU.
+
+    #### Args:
+        - `device` (torch.device): The device to check.
+
+    #### Returns:
+        - `bool`: True if the device is a CPU, False otherwise.
+    """
+    return is_device_type(device, "cpu")
+
+
+def is_device_mps(device: torch.device) -> bool:
+    """#### Check if the device is an MPS.
+
+    #### Args:
+        - `device` (torch.device): The device to check.
+
+    #### Returns:
+        - `bool`: True if the device is an MPS, False otherwise.
+    """
+    return is_device_type(device, "mps")
+
+
+def is_device_cuda(device: torch.device) -> bool:
+    """#### Check if the device is a CUDA device.
+
+    #### Args:
+        - `device` (torch.device): The device to check.
+
+    #### Returns:
+        - `bool`: True if the device is a CUDA device, False otherwise.
+    """
+    return is_device_type(device, "cuda")
+
+
+def should_use_fp16(
+    device: torch.device = None, model_params: int = 0, prioritize_performance: bool = True, manual_cast: bool = False
+) -> bool:
+    """#### Determine if FP16 should be used.
+
+    #### Args:
+        - `device` (torch.device, optional): The device to check. Defaults to None.
+        - `model_params` (int, optional): The number of model parameters. Defaults to 0.
+        - `prioritize_performance` (bool, optional): Whether to prioritize performance. Defaults to True.
+        - `manual_cast` (bool, optional): Whether to manually cast. Defaults to False.
+
+    #### Returns:
+        - `bool`: True if FP16 should be used, False otherwise.
+    """
+    global directml_enabled
+
+    if device is not None:
+        if is_device_cpu(device):
+            return False
+
+    if FORCE_FP16:
+        return True
+
+    if device is not None:
+        if is_device_mps(device):
+            return True
+
+    if FORCE_FP32:
+        return False
+
+    if directml_enabled:
+        return False
+
+    if mps_mode():
+        return True
+
+    if cpu_mode():
+        return False
+
+    if is_intel_xpu():
+        return True
+
+    if torch.version.hip:
+        return True
+
+    props = torch.cuda.get_device_properties("cuda")
+    if props.major >= 8:
+        return True
+
+    if props.major < 6:
+        return False
+
+    fp16_works = False
+    nvidia_10_series = [
+        "1080",
+        "1070",
+        "titan x",
+        "p3000",
+        "p3200",
+        "p4000",
+        "p4200",
+        "p5000",
+        "p5200",
+        "p6000",
+        "1060",
+        "1050",
+        "p40",
+        "p100",
+        "p6",
+        "p4",
+    ]
+    for x in nvidia_10_series:
+        if x in props.name.lower():
+            fp16_works = True
+
+    if fp16_works or manual_cast:
+        free_model_memory = get_free_memory() * 0.9 - minimum_inference_memory()
+        if (not prioritize_performance) or model_params * 4 > free_model_memory:
+            return True
+
+    if props.major < 7:
+        return False
+
+    nvidia_16_series = [
+        "1660",
+        "1650",
+        "1630",
+        "T500",
+        "T550",
+        "T600",
+        "MX550",
+        "MX450",
+        "CMP 30HX",
+        "T2000",
+        "T1000",
+        "T1200",
+    ]
+    for x in nvidia_16_series:
+        if x in props.name:
+            return False
+
+    return True
+
+
+def should_use_bf16(
+    device: torch.device = None, model_params: int = 0, prioritize_performance: bool = True, manual_cast: bool = False
+) -> bool:
+    """#### Determine if BF16 should be used.
+
+    #### Args:
+        - `device` (torch.device, optional): The device to check. Defaults to None.
+        - `model_params` (int, optional): The number of model parameters. Defaults to 0.
+        - `prioritize_performance` (bool, optional): Whether to prioritize performance. Defaults to True.
+        - `manual_cast` (bool, optional): Whether to manually cast. Defaults to False.
+
+    #### Returns:
+        - `bool`: True if BF16 should be used, False otherwise.
+    """
+    if device is not None:
+        if is_device_cpu(device):
+            return False
+
+    if device is not None:
+        if is_device_mps(device):
+            return False
+
+    if FORCE_FP32:
+        return False
+
+    if directml_enabled:
+        return False
+
+    if cpu_mode() or mps_mode():
+        return False
+
+    if is_intel_xpu():
+        return True
+
+    if device is None:
+        device = torch.device("cuda")
+
+    props = torch.cuda.get_device_properties(device)
+    if props.major >= 8:
+        return True
+
+    bf16_works = torch.cuda.is_bf16_supported()
+
+    if bf16_works or manual_cast:
+        free_model_memory = get_free_memory() * 0.9 - minimum_inference_memory()
+        if (not prioritize_performance) or model_params * 4 > free_model_memory:
+            return True
+
+    return False
+
+
+def soft_empty_cache(force: bool = False) -> None:
+    """#### Softly empty the cache.
+
+    #### Args:
+        - `force` (bool, optional): Whether to force emptying the cache. Defaults to False.
+    """
+    global cpu_state
+    if cpu_state == CPUState.MPS:
+        torch.mps.empty_cache()
+    elif is_intel_xpu():
+        torch.xpu.empty_cache()
+    elif torch.cuda.is_available():
+        if (
+            force or is_nvidia()
+        ):  # This seems to make things worse on ROCm so I only do it for cuda
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+
+def unload_all_models() -> None:
+    """#### Unload all models."""
+    free_memory(1e30, get_torch_device())
+
+
+def resolve_lowvram_weight(weight: torch.Tensor, model: torch.nn.Module, key: str) -> torch.Tensor:
+    """#### Resolve low VRAM weight.
+
+    #### Args:
+        - `weight` (torch.Tensor): The weight tensor.
+        - `model` (torch.nn.Module): The model.
+        - `key` (str): The key.
+
+    #### Returns:
+        - `torch.Tensor`: The resolved weight tensor.
+    """
+    return weight

modules/FileManaging/Downloader.py

@@ -0,0 +1,116 @@
+import glob
+from huggingface_hub import hf_hub_download
+
+
+def CheckAndDownload():
+    """#### Check and download all the necessary safetensors and checkpoints models"""
+    if glob.glob("./_internal/checkpoints/*.safetensors") == []:
+
+        hf_hub_download(
+            repo_id="Meina/MeinaMix",
+            filename="Meina V10 - baked VAE.safetensors",
+            local_dir="./_internal/checkpoints/",
+        )
+        hf_hub_download(
+            repo_id="Lykon/DreamShaper",
+            filename="DreamShaper_8_pruned.safetensors",
+            local_dir="./_internal/checkpoints/",
+        )
+    if glob.glob("./_internal/yolos/*.pt") == []:
+
+        hf_hub_download(
+            repo_id="Bingsu/adetailer",
+            filename="hand_yolov9c.pt",
+            local_dir="./_internal/yolos/",
+        )
+        hf_hub_download(
+            repo_id="Bingsu/adetailer",
+            filename="face_yolov9c.pt",
+            local_dir="./_internal/yolos/",
+        )
+        hf_hub_download(
+            repo_id="Bingsu/adetailer",
+            filename="person_yolov8m-seg.pt",
+            local_dir="./_internal/yolos/",
+        )
+        hf_hub_download(
+            repo_id="segments-arnaud/sam_vit_b",
+            filename="sam_vit_b_01ec64.pth",
+            local_dir="./_internal/yolos/",
+        )
+    if glob.glob("./_internal/ESRGAN/*.pth") == []:
+
+        hf_hub_download(
+            repo_id="lllyasviel/Annotators",
+            filename="RealESRGAN_x4plus.pth",
+            local_dir="./_internal/ESRGAN/",
+        )
+    if glob.glob("./_internal/loras/*.safetensors") == []:
+
+        hf_hub_download(
+            repo_id="EvilEngine/add_detail",
+            filename="add_detail.safetensors",
+            local_dir="./_internal/loras/",
+        )
+    if glob.glob("./_internal/embeddings/*.pt") == []:
+
+        hf_hub_download(
+            repo_id="EvilEngine/badhandv4",
+            filename="badhandv4.pt",
+            local_dir="./_internal/embeddings/",
+        )
+        # hf_hub_download(
+        #     repo_id="segments-arnaud/sam_vit_b",
+        #     filename="EasyNegative.safetensors",
+        #     local_dir="./_internal/embeddings/",
+        # )
+    if glob.glob("./_internal/vae_approx/*.pth") == []:
+
+        hf_hub_download(
+            repo_id="madebyollin/taesd",
+            filename="taesd_decoder.safetensors",
+            local_dir="./_internal/vae_approx/",
+        )
+
+def CheckAndDownloadFlux():
+    """#### Check and download all the necessary safetensors and checkpoints models for FLUX"""
+    if glob.glob("./_internal/embeddings/*.pt") == []:
+        hf_hub_download(
+            repo_id="EvilEngine/badhandv4",
+            filename="badhandv4.pt",
+            local_dir="./_internal/embeddings",
+        )
+    if glob.glob("./_internal/unet/*.gguf") == []:
+
+        hf_hub_download(
+            repo_id="city96/FLUX.1-dev-gguf",
+            filename="flux1-dev-Q8_0.gguf",
+            local_dir="./_internal/unet",
+        )
+    if glob.glob("./_internal/clip/*.gguf") == []:
+
+        hf_hub_download(
+            repo_id="city96/t5-v1_1-xxl-encoder-gguf",
+            filename="t5-v1_1-xxl-encoder-Q8_0.gguf",
+            local_dir="./_internal/clip",
+        )
+        hf_hub_download(
+            repo_id="comfyanonymous/flux_text_encoders",
+            filename="clip_l.safetensors",
+            local_dir="./_internal/clip",
+        )
+    if glob.glob("./_internal/vae/*.safetensors") == []:
+
+        hf_hub_download(
+            repo_id="black-forest-labs/FLUX.1-schnell",
+            filename="ae.safetensors",
+            local_dir="./_internal/vae",
+        )
+
+    if glob.glob("./_internal/vae_approx/*.pth") == []:
+
+        hf_hub_download(
+            repo_id="madebyollin/taef1",
+            filename="diffusion_pytorch_model.safetensors",
+            local_dir="./_internal/vae_approx/",
+        )

modules/FileManaging/ImageSaver.py

@@ -0,0 +1,126 @@
+import os
+import numpy as np
+from PIL import Image
+
+output_directory = "./_internal/output"
+
+
+def get_output_directory() -> str:
+    """#### Get the output directory.
+
+    #### Returns:
+        - `str`: The output directory.
+    """
+    global output_directory
+    return output_directory
+
+
+def get_save_image_path(
+    filename_prefix: str, output_dir: str, image_width: int = 0, image_height: int = 0
+) -> tuple:
+    """#### Get the save image path.
+
+    #### Args:
+        - `filename_prefix` (str): The filename prefix.
+        - `output_dir` (str): The output directory.
+        - `image_width` (int, optional): The image width. Defaults to 0.
+        - `image_height` (int, optional): The image height. Defaults to 0.
+
+    #### Returns:
+        - `tuple`: The full output folder, filename, counter, subfolder, and filename prefix.
+    """
+
+    def map_filename(filename: str) -> tuple:
+        prefix_len = len(os.path.basename(filename_prefix))
+        prefix = filename[: prefix_len + 1]
+        try:
+            digits = int(filename[prefix_len + 1 :].split("_")[0])
+        except:
+            digits = 0
+        return (digits, prefix)
+
+    def compute_vars(input: str, image_width: int, image_height: int) -> str:
+        input = input.replace("%width%", str(image_width))
+        input = input.replace("%height%", str(image_height))
+        return input
+
+    filename_prefix = compute_vars(filename_prefix, image_width, image_height)
+
+    subfolder = os.path.dirname(os.path.normpath(filename_prefix))
+    filename = os.path.basename(os.path.normpath(filename_prefix))
+
+    full_output_folder = os.path.join(output_dir, subfolder)
+    try:
+        counter = (
+            max(
+                filter(
+                    lambda a: a[1][:-1] == filename and a[1][-1] == "_",
+                    map(map_filename, os.listdir(full_output_folder)),
+                )
+            )[0]
+            + 1
+        )
+    except ValueError:
+        counter = 1
+    except FileNotFoundError:
+        os.makedirs(full_output_folder, exist_ok=True)
+        counter = 1
+    return full_output_folder, filename, counter, subfolder, filename_prefix
+
+
+MAX_RESOLUTION = 16384
+
+
+class SaveImage:
+    """#### Class for saving images."""
+
+    def __init__(self):
+        """#### Initialize the SaveImage class."""
+        self.output_dir = get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+        self.compress_level = 4
+
+    def save_images(
+        self,
+        images: list,
+        filename_prefix: str = "LD",
+        prompt: str = None,
+        extra_pnginfo: dict = None,
+    ) -> dict:
+        """#### Save images to the output directory.
+
+        #### Args:
+            - `images` (list): The list of images.
+            - `filename_prefix` (str, optional): The filename prefix. Defaults to "LD".
+            - `prompt` (str, optional): The prompt. Defaults to None.
+            - `extra_pnginfo` (dict, optional): Additional PNG info. Defaults to None.
+
+        #### Returns:
+            - `dict`: The saved images information.
+        """
+        filename_prefix += self.prefix_append
+        full_output_folder, filename, counter, subfolder, filename_prefix = (
+            get_save_image_path(
+                filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]
+            )
+        )
+        results = list()
+        for batch_number, image in enumerate(images):
+            i = 255.0 * image.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            metadata = None
+
+            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
+            file = f"{filename_with_batch_num}_{counter:05}_.png"
+            img.save(
+                os.path.join(full_output_folder, file),
+                pnginfo=metadata,
+                compress_level=self.compress_level,
+            )
+            results.append(
+                {"filename": file, "subfolder": subfolder, "type": self.type}
+            )
+            counter += 1
+
+        return {"ui": {"images": results}}

modules/FileManaging/Loader.py

@@ -0,0 +1,138 @@
+import logging
+import torch
+from modules.Utilities import util
+from modules.AutoEncoders import VariationalAE
+from modules.Device import Device
+from modules.Model import ModelPatcher
+from modules.NeuralNetwork import unet
+from modules.clip import Clip
+
+
+def load_checkpoint_guess_config(
+    ckpt_path: str,
+    output_vae: bool = True,
+    output_clip: bool = True,
+    output_clipvision: bool = False,
+    embedding_directory: str = None,
+    output_model: bool = True,
+) -> tuple:
+    """#### Load a checkpoint and guess the configuration.
+
+    #### Args:
+        - `ckpt_path` (str): The path to the checkpoint file.
+        - `output_vae` (bool, optional): Whether to output the VAE. Defaults to True.
+        - `output_clip` (bool, optional): Whether to output the CLIP. Defaults to True.
+        - `output_clipvision` (bool, optional): Whether to output the CLIP vision. Defaults to False.
+        - `embedding_directory` (str, optional): The embedding directory. Defaults to None.
+        - `output_model` (bool, optional): Whether to output the model. Defaults to True.
+
+    #### Returns:
+        - `tuple`: The model patcher, CLIP, VAE, and CLIP vision.
+    """
+    sd = util.load_torch_file(ckpt_path)
+    sd.keys()
+    clip = None
+    clipvision = None
+    vae = None
+    model = None
+    model_patcher = None
+    clip_target = None
+
+    parameters = util.calculate_parameters(sd, "model.diffusion_model.")
+    load_device = Device.get_torch_device()
+
+    model_config = unet.model_config_from_unet(sd, "model.diffusion_model.")
+    unet_dtype = unet.unet_dtype1(
+        model_params=parameters,
+        supported_dtypes=model_config.supported_inference_dtypes,
+    )
+    manual_cast_dtype = Device.unet_manual_cast(
+        unet_dtype, load_device, model_config.supported_inference_dtypes
+    )
+    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+
+    if output_model:
+        inital_load_device = Device.unet_inital_load_device(parameters, unet_dtype)
+        Device.unet_offload_device()
+        model = model_config.get_model(
+            sd, "model.diffusion_model.", device=inital_load_device
+        )
+        model.load_model_weights(sd, "model.diffusion_model.")
+
+    if output_vae:
+        vae_sd = util.state_dict_prefix_replace(
+            sd, {k: "" for k in model_config.vae_key_prefix}, filter_keys=True
+        )
+        vae_sd = model_config.process_vae_state_dict(vae_sd)
+        vae = VariationalAE.VAE(sd=vae_sd)
+
+    if output_clip:
+        clip_target = model_config.clip_target()
+        if clip_target is not None:
+            clip_sd = model_config.process_clip_state_dict(sd)
+            if len(clip_sd) > 0:
+                clip = Clip.CLIP(clip_target, embedding_directory=embedding_directory)
+                m, u = clip.load_sd(clip_sd, full_model=True)
+                if len(m) > 0:
+                    m_filter = list(
+                        filter(
+                            lambda a: ".logit_scale" not in a
+                            and ".transformer.text_projection.weight" not in a,
+                            m,
+                        )
+                    )
+                    if len(m_filter) > 0:
+                        logging.warning("clip missing: {}".format(m))
+                    else:
+                        logging.debug("clip missing: {}".format(m))
+
+                if len(u) > 0:
+                    logging.debug("clip unexpected {}:".format(u))
+            else:
+                logging.warning(
+                    "no CLIP/text encoder weights in checkpoint, the text encoder model will not be loaded."
+                )
+
+    left_over = sd.keys()
+    if len(left_over) > 0:
+        logging.debug("left over keys: {}".format(left_over))
+
+    if output_model:
+        model_patcher = ModelPatcher.ModelPatcher(
+            model,
+            load_device=load_device,
+            offload_device=Device.unet_offload_device(),
+            current_device=inital_load_device,
+        )
+        if inital_load_device != torch.device("cpu"):
+            logging.info("loaded straight to GPU")
+            Device.load_model_gpu(model_patcher)
+
+    return (model_patcher, clip, vae, clipvision)
+
+
+class CheckpointLoaderSimple:
+    """#### Class for loading checkpoints."""
+
+    def load_checkpoint(
+        self, ckpt_name: str, output_vae: bool = True, output_clip: bool = True
+    ) -> tuple:
+        """#### Load a checkpoint.
+
+        #### Args:
+            - `ckpt_name` (str): The name of the checkpoint.
+            - `output_vae` (bool, optional): Whether to output the VAE. Defaults to True.
+            - `output_clip` (bool, optional): Whether to output the CLIP. Defaults to True.
+
+        #### Returns:
+            - `tuple`: The model patcher, CLIP, and VAE.
+        """
+        ckpt_path = f"{ckpt_name}"
+        out = load_checkpoint_guess_config(
+            ckpt_path,
+            output_vae=output_vae,
+            output_clip=output_clip,
+            embedding_directory="./_internal/embeddings/",
+        )
+        print("loading", ckpt_path)
+        return out[:3]

modules/Model/LoRas.py

@@ -0,0 +1,193 @@
+import torch
+from modules.Utilities import util
+from modules.NeuralNetwork import unet
+
+LORA_CLIP_MAP = {
+    "mlp.fc1": "mlp_fc1",
+    "mlp.fc2": "mlp_fc2",
+    "self_attn.k_proj": "self_attn_k_proj",
+    "self_attn.q_proj": "self_attn_q_proj",
+    "self_attn.v_proj": "self_attn_v_proj",
+    "self_attn.out_proj": "self_attn_out_proj",
+}
+
+
+def load_lora(lora: dict, to_load: dict) -> dict:
+    """#### Load a LoRA model.
+
+    #### Args:
+        - `lora` (dict): The LoRA model state dictionary.
+        - `to_load` (dict): The keys to load from the LoRA model.
+
+    #### Returns:
+        - `dict`: The loaded LoRA model.
+    """
+    patch_dict = {}
+    loaded_keys = set()
+    for x in to_load:
+        alpha_name = "{}.alpha".format(x)
+        alpha = None
+        if alpha_name in lora.keys():
+            alpha = lora[alpha_name].item()
+            loaded_keys.add(alpha_name)
+
+        "{}.dora_scale".format(x)
+        dora_scale = None
+
+        regular_lora = "{}.lora_up.weight".format(x)
+        "{}_lora.up.weight".format(x)
+        "{}.lora_linear_layer.up.weight".format(x)
+        A_name = None
+
+        if regular_lora in lora.keys():
+            A_name = regular_lora
+            B_name = "{}.lora_down.weight".format(x)
+            "{}.lora_mid.weight".format(x)
+
+        if A_name is not None:
+            mid = None
+            patch_dict[to_load[x]] = (
+                "lora",
+                (lora[A_name], lora[B_name], alpha, mid, dora_scale),
+            )
+            loaded_keys.add(A_name)
+            loaded_keys.add(B_name)
+    return patch_dict
+
+
+def model_lora_keys_clip(model: torch.nn.Module, key_map: dict = {}) -> dict:
+    """#### Get the keys for a LoRA model's CLIP component.
+
+    #### Args:
+        - `model` (torch.nn.Module): The LoRA model.
+        - `key_map` (dict, optional): The key map. Defaults to {}.
+
+    #### Returns:
+        - `dict`: The keys for the CLIP component.
+    """
+    sdk = model.state_dict().keys()
+
+    text_model_lora_key = "lora_te_text_model_encoder_layers_{}_{}"
+    for b in range(32):
+        for c in LORA_CLIP_MAP:
+            k = "clip_l.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)
+            if k in sdk:
+                lora_key = text_model_lora_key.format(b, LORA_CLIP_MAP[c])
+                key_map[lora_key] = k
+                lora_key = "lora_te1_text_model_encoder_layers_{}_{}".format(
+                    b, LORA_CLIP_MAP[c]
+                )  # SDXL base
+                key_map[lora_key] = k
+                lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(
+                    b, c
+                )  # diffusers lora
+                key_map[lora_key] = k
+    return key_map
+
+
+def model_lora_keys_unet(model: torch.nn.Module, key_map: dict = {}) -> dict:
+    """#### Get the keys for a LoRA model's UNet component.
+
+    #### Args:
+        - `model` (torch.nn.Module): The LoRA model.
+        - `key_map` (dict, optional): The key map. Defaults to {}.
+
+    #### Returns:
+        - `dict`: The keys for the UNet component.
+    """
+    sdk = model.state_dict().keys()
+
+    for k in sdk:
+        if k.startswith("diffusion_model.") and k.endswith(".weight"):
+            key_lora = k[len("diffusion_model.") : -len(".weight")].replace(".", "_")
+            key_map["lora_unet_{}".format(key_lora)] = k
+            key_map["lora_prior_unet_{}".format(key_lora)] = k  # cascade lora:
+
+    diffusers_keys = unet.unet_to_diffusers(model.model_config.unet_config)
+    for k in diffusers_keys:
+        if k.endswith(".weight"):
+            unet_key = "diffusion_model.{}".format(diffusers_keys[k])
+            key_lora = k[: -len(".weight")].replace(".", "_")
+            key_map["lora_unet_{}".format(key_lora)] = unet_key
+
+            diffusers_lora_prefix = ["", "unet."]
+            for p in diffusers_lora_prefix:
+                diffusers_lora_key = "{}{}".format(
+                    p, k[: -len(".weight")].replace(".to_", ".processor.to_")
+                )
+                if diffusers_lora_key.endswith(".to_out.0"):
+                    diffusers_lora_key = diffusers_lora_key[:-2]
+                key_map[diffusers_lora_key] = unet_key
+    return key_map
+
+
+def load_lora_for_models(
+    model: object, clip: object, lora: dict, strength_model: float, strength_clip: float
+) -> tuple:
+    """#### Load a LoRA model for the given models.
+
+    #### Args:
+        - `model` (object): The model.
+        - `clip` (object): The CLIP model.
+        - `lora` (dict): The LoRA model state dictionary.
+        - `strength_model` (float): The strength of the model.
+        - `strength_clip` (float): The strength of the CLIP model.
+
+    #### Returns:
+        - `tuple`: The new model patcher and CLIP model.
+    """
+    key_map = {}
+    if model is not None:
+        key_map = model_lora_keys_unet(model.model, key_map)
+    if clip is not None:
+        key_map = model_lora_keys_clip(clip.cond_stage_model, key_map)
+
+    loaded = load_lora(lora, key_map)
+    new_modelpatcher = model.clone()
+    k = new_modelpatcher.add_patches(loaded, strength_model)
+
+    new_clip = clip.clone()
+    k1 = new_clip.add_patches(loaded, strength_clip)
+    k = set(k)
+    k1 = set(k1)
+
+    return (new_modelpatcher, new_clip)
+
+
+class LoraLoader:
+    """#### Class for loading LoRA models."""
+
+    def __init__(self):
+        """#### Initialize the LoraLoader class."""
+        self.loaded_lora = None
+
+    def load_lora(
+        self,
+        model: object,
+        clip: object,
+        lora_name: str,
+        strength_model: float,
+        strength_clip: float,
+    ) -> tuple:
+        """#### Load a LoRA model.
+
+        #### Args:
+            - `model` (object): The model.
+            - `clip` (object): The CLIP model.
+            - `lora_name` (str): The name of the LoRA model.
+            - `strength_model` (float): The strength of the model.
+            - `strength_clip` (float): The strength of the CLIP model.
+
+        #### Returns:
+            - `tuple`: The new model patcher and CLIP model.
+        """
+        lora_path = util.get_full_path("loras", lora_name)
+        lora = None
+        if lora is None:
+            lora = util.load_torch_file(lora_path, safe_load=True)
+            self.loaded_lora = (lora_path, lora)
+
+        model_lora, clip_lora = load_lora_for_models(
+            model, clip, lora, strength_model, strength_clip
+        )
+        return (model_lora, clip_lora)

modules/Model/ModelBase.py

@@ -0,0 +1,363 @@
+import logging
+import math
+import torch
+
+from modules.Utilities import Latent
+from modules.Device import Device
+from modules.NeuralNetwork import unet
+from modules.cond import cast, cond
+from modules.sample import sampling
+
+
+class BaseModel(torch.nn.Module):
+    """#### Base class for models."""
+
+    def __init__(
+        self,
+        model_config: object,
+        model_type: sampling.ModelType = sampling.ModelType.EPS,
+        device: torch.device = None,
+        unet_model: object = unet.UNetModel1,
+        flux: bool = False,
+    ):
+        """#### Initialize the BaseModel class.
+
+        #### Args:
+            - `model_config` (object): The model configuration.
+            - `model_type` (sampling.ModelType, optional): The model type. Defaults to sampling.ModelType.EPS.
+            - `device` (torch.device, optional): The device to use. Defaults to None.
+            - `unet_model` (object, optional): The UNet model. Defaults to unet.UNetModel1.
+        """
+        super().__init__()
+
+        unet_config = model_config.unet_config
+        self.latent_format = model_config.latent_format
+        self.model_config = model_config
+        self.manual_cast_dtype = model_config.manual_cast_dtype
+        self.device = device
+        if flux:
+            if not unet_config.get("disable_unet_model_creation", False):
+                operations = model_config.custom_operations
+                self.diffusion_model = unet_model(
+                    **unet_config, device=device, operations=operations
+                )
+                logging.info(
+                    "model weight dtype {}, manual cast: {}".format(
+                        self.get_dtype(), self.manual_cast_dtype
+                    )
+                )
+        else:
+            if not unet_config.get("disable_unet_model_creation", False):
+                if self.manual_cast_dtype is not None:
+                    operations = cast.manual_cast
+                else:
+                    operations = cast.disable_weight_init
+                self.diffusion_model = unet_model(
+                    **unet_config, device=device, operations=operations
+                )
+        self.model_type = model_type
+        self.model_sampling = sampling.model_sampling(model_config, model_type, flux=flux)
+
+        self.adm_channels = unet_config.get("adm_in_channels", None)
+        if self.adm_channels is None:
+            self.adm_channels = 0
+
+        self.concat_keys = ()
+        logging.info("model_type {}".format(model_type.name))
+        logging.debug("adm {}".format(self.adm_channels))
+        self.memory_usage_factor = model_config.memory_usage_factor if flux else 2.0
+
+    def apply_model(
+        self,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        c_concat: torch.Tensor = None,
+        c_crossattn: torch.Tensor = None,
+        control: torch.Tensor = None,
+        transformer_options: dict = {},
+        **kwargs,
+    ) -> torch.Tensor:
+        """#### Apply the model to the input tensor.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `t` (torch.Tensor): The timestep tensor.
+            - `c_concat` (torch.Tensor, optional): The concatenated condition tensor. Defaults to None.
+            - `c_crossattn` (torch.Tensor, optional): The cross-attention condition tensor. Defaults to None.
+            - `control` (torch.Tensor, optional): The control tensor. Defaults to None.
+            - `transformer_options` (dict, optional): The transformer options. Defaults to {}.
+            - `**kwargs`: Additional keyword arguments.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        sigma = t
+        xc = self.model_sampling.calculate_input(sigma, x)
+        if c_concat is not None:
+            xc = torch.cat([xc] + [c_concat], dim=1)
+
+        context = c_crossattn
+        dtype = self.get_dtype()
+
+        if self.manual_cast_dtype is not None:
+            dtype = self.manual_cast_dtype
+
+        xc = xc.to(dtype)
+        t = self.model_sampling.timestep(t).float()
+        context = context.to(dtype)
+        extra_conds = {}
+        for o in kwargs:
+            extra = kwargs[o]
+            if hasattr(extra, "dtype"):
+                if extra.dtype != torch.int and extra.dtype != torch.long:
+                    extra = extra.to(dtype)
+            extra_conds[o] = extra
+
+        model_output = self.diffusion_model(
+            xc,
+            t,
+            context=context,
+            control=control,
+            transformer_options=transformer_options,
+            **extra_conds,
+        ).float()
+        return self.model_sampling.calculate_denoised(sigma, model_output, x)
+
+    def get_dtype(self) -> torch.dtype:
+        """#### Get the data type of the model.
+
+        #### Returns:
+            - `torch.dtype`: The data type.
+        """
+        return self.diffusion_model.dtype
+
+    def encode_adm(self, **kwargs) -> None:
+        """#### Encode the ADM.
+
+        #### Args:
+            - `**kwargs`: Additional keyword arguments.
+
+        #### Returns:
+            - `None`: The encoded ADM.
+        """
+        return None
+
+    def extra_conds(self, **kwargs) -> dict:
+        """#### Get the extra conditions.
+
+        #### Args:
+            - `**kwargs`: Additional keyword arguments.
+
+        #### Returns:
+            - `dict`: The extra conditions.
+        """
+        out = {}
+        adm = self.encode_adm(**kwargs)
+        if adm is not None:
+            out["y"] = cond.CONDRegular(adm)
+
+        cross_attn = kwargs.get("cross_attn", None)
+        if cross_attn is not None:
+            out["c_crossattn"] = cond.CONDCrossAttn(cross_attn)
+
+        cross_attn_cnet = kwargs.get("cross_attn_controlnet", None)
+        if cross_attn_cnet is not None:
+            out["crossattn_controlnet"] = cond.CONDCrossAttn(cross_attn_cnet)
+
+        return out
+
+    def load_model_weights(self, sd: dict, unet_prefix: str = "") -> "BaseModel":
+        """#### Load the model weights.
+
+        #### Args:
+            - `sd` (dict): The state dictionary.
+            - `unet_prefix` (str, optional): The UNet prefix. Defaults to "".
+
+        #### Returns:
+            - `BaseModel`: The model with loaded weights.
+        """
+        to_load = {}
+        keys = list(sd.keys())
+        for k in keys:
+            if k.startswith(unet_prefix):
+                to_load[k[len(unet_prefix) :]] = sd.pop(k)
+
+        to_load = self.model_config.process_unet_state_dict(to_load)
+        m, u = self.diffusion_model.load_state_dict(to_load, strict=False)
+        if len(m) > 0:
+            logging.warning("unet missing: {}".format(m))
+
+        if len(u) > 0:
+            logging.warning("unet unexpected: {}".format(u))
+        del to_load
+        return self
+
+    def process_latent_in(self, latent: torch.Tensor) -> torch.Tensor:
+        """#### Process the latent input.
+
+        #### Args:
+            - `latent` (torch.Tensor): The latent tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The processed latent tensor.
+        """
+        return self.latent_format.process_in(latent)
+
+    def process_latent_out(self, latent: torch.Tensor) -> torch.Tensor:
+        """#### Process the latent output.
+
+        #### Args:
+            - `latent` (torch.Tensor): The latent tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The processed latent tensor.
+        """
+        return self.latent_format.process_out(latent)
+
+    def memory_required(self, input_shape: tuple) -> float:
+        """#### Calculate the memory required for the model.
+
+        #### Args:
+            - `input_shape` (tuple): The input shape.
+
+        #### Returns:
+            - `float`: The memory required.
+        """
+        dtype = self.get_dtype()
+        if self.manual_cast_dtype is not None:
+            dtype = self.manual_cast_dtype
+        # TODO: this needs to be tweaked
+        area = input_shape[0] * math.prod(input_shape[2:])
+        return (area * Device.dtype_size(dtype) * 0.01 * self.memory_usage_factor) * (
+            1024 * 1024
+        )
+
+
+class BASE:
+    """#### Base class for model configurations."""
+
+    unet_config = {}
+    unet_extra_config = {
+        "num_heads": -1,
+        "num_head_channels": 64,
+    }
+
+    required_keys = {}
+
+    clip_prefix = []
+    clip_vision_prefix = None
+    noise_aug_config = None
+    sampling_settings = {}
+    latent_format = Latent.LatentFormat
+    vae_key_prefix = ["first_stage_model."]
+    text_encoder_key_prefix = ["cond_stage_model."]
+    supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+
+    memory_usage_factor = 2.0
+
+    manual_cast_dtype = None
+    custom_operations = None
+
+    @classmethod
+    def matches(cls, unet_config: dict, state_dict: dict = None) -> bool:
+        """#### Check if the UNet configuration matches.
+
+        #### Args:
+            - `unet_config` (dict): The UNet configuration.
+            - `state_dict` (dict, optional): The state dictionary. Defaults to None.
+
+        #### Returns:
+            - `bool`: Whether the configuration matches.
+        """
+        for k in cls.unet_config:
+            if k not in unet_config or cls.unet_config[k] != unet_config[k]:
+                return False
+        if state_dict is not None:
+            for k in cls.required_keys:
+                if k not in state_dict:
+                    return False
+        return True
+
+    def model_type(self, state_dict: dict, prefix: str = "") -> sampling.ModelType:
+        """#### Get the model type.
+
+        #### Args:
+            - `state_dict` (dict): The state dictionary.
+            - `prefix` (str, optional): The prefix. Defaults to "".
+
+        #### Returns:
+            - `sampling.ModelType`: The model type.
+        """
+        return sampling.ModelType.EPS
+
+    def inpaint_model(self) -> bool:
+        """#### Check if the model is an inpaint model.
+
+        #### Returns:
+            - `bool`: Whether the model is an inpaint model.
+        """
+        return self.unet_config["in_channels"] > 4
+
+    def __init__(self, unet_config: dict):
+        """#### Initialize the BASE class.
+
+        #### Args:
+            - `unet_config` (dict): The UNet configuration.
+        """
+        self.unet_config = unet_config.copy()
+        self.sampling_settings = self.sampling_settings.copy()
+        self.latent_format = self.latent_format()
+        for x in self.unet_extra_config:
+            self.unet_config[x] = self.unet_extra_config[x]
+
+    def get_model(
+        self, state_dict: dict, prefix: str = "", device: torch.device = None
+    ) -> BaseModel:
+        """#### Get the model.
+
+        #### Args:
+            - `state_dict` (dict): The state dictionary.
+            - `prefix` (str, optional): The prefix. Defaults to "".
+            - `device` (torch.device, optional): The device to use. Defaults to None.
+
+        #### Returns:
+            - `BaseModel`: The model.
+        """
+        out = BaseModel(
+            self, model_type=self.model_type(state_dict, prefix), device=device
+        )
+        return out
+
+    def process_unet_state_dict(self, state_dict: dict) -> dict:
+        """#### Process the UNet state dictionary.
+
+        #### Args:
+            - `state_dict` (dict): The state dictionary.
+
+        #### Returns:
+            - `dict`: The processed state dictionary.
+        """
+        return state_dict
+
+    def process_vae_state_dict(self, state_dict: dict) -> dict:
+        """#### Process the VAE state dictionary.
+
+        #### Args:
+            - `state_dict` (dict): The state dictionary.
+
+        #### Returns:
+            - `dict`: The processed state dictionary.
+        """
+        return state_dict
+
+    def set_inference_dtype(
+        self, dtype: torch.dtype, manual_cast_dtype: torch.dtype
+    ) -> None:
+        """#### Set the inference data type.
+
+        #### Args:
+            - `dtype` (torch.dtype): The data type.
+            - `manual_cast_dtype` (torch.dtype): The manual cast data type.
+        """
+        self.unet_config["dtype"] = dtype
+        self.manual_cast_dtype = manual_cast_dtype

modules/Model/ModelPatcher.py

@@ -0,0 +1,779 @@
+import copy
+import logging
+import uuid
+
+import torch
+
+from modules.NeuralNetwork import unet
+from modules.Utilities import util
+from modules.Device import Device
+
+def wipe_lowvram_weight(m):
+    if hasattr(m, "prev_comfy_cast_weights"):
+        m.comfy_cast_weights = m.prev_comfy_cast_weights
+        del m.prev_comfy_cast_weights
+    m.weight_function = None
+    m.bias_function = None
+
+class ModelPatcher:
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        load_device: torch.device,
+        offload_device: torch.device,
+        size: int = 0,
+        current_device: torch.device = None,
+        weight_inplace_update: bool = False,
+    ):
+        """#### Initialize the ModelPatcher class.
+
+        #### Args:
+            - `model` (torch.nn.Module): The model.
+            - `load_device` (torch.device): The device to load the model on.
+            - `offload_device` (torch.device): The device to offload the model to.
+            - `size` (int, optional): The size of the model. Defaults to 0.
+            - `current_device` (torch.device, optional): The current device. Defaults to None.
+            - `weight_inplace_update` (bool, optional): Whether to update weights in place. Defaults to False.
+        """
+        self.size = size
+        self.model = model
+        self.patches = {}
+        self.backup = {}
+        self.object_patches = {}
+        self.object_patches_backup = {}
+        self.model_options = {"transformer_options": {}}
+        self.model_size()
+        self.load_device = load_device
+        self.offload_device = offload_device
+        if current_device is None:
+            self.current_device = self.offload_device
+        else:
+            self.current_device = current_device
+
+        self.weight_inplace_update = weight_inplace_update
+        self.model_lowvram = False
+        self.lowvram_patch_counter = 0
+        self.patches_uuid = uuid.uuid4()
+
+        if not hasattr(self.model, "model_loaded_weight_memory"):
+            self.model.model_loaded_weight_memory = 0
+
+        if not hasattr(self.model, "model_lowvram"):
+            self.model.model_lowvram = False
+
+        if not hasattr(self.model, "lowvram_patch_counter"):
+            self.model.lowvram_patch_counter = 0
+
+    def loaded_size(self) -> int:
+        """#### Get the loaded size
+
+        #### Returns:
+            - `int`: The loaded size
+        """
+        return self.model.model_loaded_weight_memory
+
+    def model_size(self) -> int:
+        """#### Get the size of the model.
+
+        #### Returns:
+            - `int`: The size of the model.
+        """
+        if self.size > 0:
+            return self.size
+        model_sd = self.model.state_dict()
+        self.size = Device.module_size(self.model)
+        self.model_keys = set(model_sd.keys())
+        return self.size
+
+    def clone(self) -> "ModelPatcher":
+        """#### Clone the ModelPatcher object.
+
+        #### Returns:
+            - `ModelPatcher`: The cloned ModelPatcher object.
+        """
+        n = ModelPatcher(
+            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.patches_uuid = self.patches_uuid
+
+        n.object_patches = self.object_patches.copy()
+        n.model_options = copy.deepcopy(self.model_options)
+        n.model_keys = self.model_keys
+        n.backup = self.backup
+        n.object_patches_backup = self.object_patches_backup
+        return n
+
+    def is_clone(self, other: object) -> bool:
+        """#### Check if the object is a clone.
+
+        #### Args:
+            - `other` (object): The other object.
+
+        #### Returns:
+            - `bool`: Whether the object is a clone.
+        """
+        if hasattr(other, "model") and self.model is other.model:
+            return True
+        return False
+
+    def memory_required(self, input_shape: tuple) -> float:
+        """#### Calculate the memory required for the model.
+
+        #### Args:
+            - `input_shape` (tuple): The input shape.
+
+        #### Returns:
+            - `float`: The memory required.
+        """
+        return self.model.memory_required(input_shape=input_shape)
+
+    def set_model_unet_function_wrapper(self, unet_wrapper_function: callable) -> None:
+        """#### Set the UNet function wrapper for the model.
+
+        #### Args:
+            - `unet_wrapper_function` (callable): The UNet function wrapper.
+        """
+        self.model_options["model_function_wrapper"] = unet_wrapper_function
+
+    def set_model_denoise_mask_function(self, denoise_mask_function: callable) -> None:
+        """#### Set the denoise mask function for the model.
+
+        #### Args:
+            - `denoise_mask_function` (callable): The denoise mask function.
+        """
+        self.model_options["denoise_mask_function"] = denoise_mask_function
+
+    def get_model_object(self, name: str) -> object:
+        """#### Get an object from the model.
+
+        #### Args:
+            - `name` (str): The name of the object.
+
+        #### Returns:
+            - `object`: The object.
+        """
+        return util.get_attr(self.model, name)
+
+    def model_patches_to(self, device: torch.device) -> None:
+        """#### Move model patches to a device.
+
+        #### Args:
+            - `device` (torch.device): The device.
+        """
+        self.model_options["transformer_options"]
+        if "model_function_wrapper" in self.model_options:
+            wrap_func = self.model_options["model_function_wrapper"]
+            if hasattr(wrap_func, "to"):
+                self.model_options["model_function_wrapper"] = wrap_func.to(device)
+
+    def model_dtype(self) -> torch.dtype:
+        """#### Get the data type of the model.
+
+        #### Returns:
+            - `torch.dtype`: The data type.
+        """
+        if hasattr(self.model, "get_dtype"):
+            return self.model.get_dtype()
+
+    def add_patches(
+        self, patches: dict, strength_patch: float = 1.0, strength_model: float = 1.0
+    ) -> list:
+        """#### Add patches to the model.
+
+        #### Args:
+            - `patches` (dict): The patches to add.
+            - `strength_patch` (float, optional): The strength of the patches. Defaults to 1.0.
+            - `strength_model` (float, optional): The strength of the model. Defaults to 1.0.
+
+        #### Returns:
+            - `list`: The list of patched keys.
+        """
+        p = set()
+        for k in patches:
+            if k in self.model_keys:
+                p.add(k)
+                current_patches = self.patches.get(k, [])
+                current_patches.append((strength_patch, patches[k], strength_model))
+                self.patches[k] = current_patches
+
+        self.patches_uuid = uuid.uuid4()
+        return list(p)
+
+    def set_model_patch(self, patch: list, name: str):
+        """#### Set a patch for the model.
+
+        #### Args:
+            - `patch` (list): The patch.
+            - `name` (str): The name of the patch.
+        """
+        to = self.model_options["transformer_options"]
+        if "patches" not in to:
+            to["patches"] = {}
+        to["patches"][name] = to["patches"].get(name, []) + [patch]
+
+    def set_model_attn1_patch(self, patch: list):
+        """#### Set the attention 1 patch for the model.
+
+        #### Args:
+            - `patch` (list): The patch.
+        """
+        self.set_model_patch(patch, "attn1_patch")
+
+    def set_model_attn2_patch(self, patch: list):
+        """#### Set the attention 2 patch for the model.
+
+        #### Args:
+            - `patch` (list): The patch.
+        """
+        self.set_model_patch(patch, "attn2_patch")
+
+    def set_model_attn1_output_patch(self, patch: list):
+        """#### Set the attention 1 output patch for the model.
+
+        #### Args:
+            - `patch` (list): The patch.
+        """
+        self.set_model_patch(patch, "attn1_output_patch")
+
+    def set_model_attn2_output_patch(self, patch: list):
+        """#### Set the attention 2 output patch for the model.
+
+        #### Args:
+            - `patch` (list): The patch.
+        """
+        self.set_model_patch(patch, "attn2_output_patch")
+
+    def model_state_dict(self, filter_prefix: str = None) -> dict:
+        """#### Get the state dictionary of the model.
+
+        #### Args:
+            - `filter_prefix` (str, optional): The prefix to filter. Defaults to None.
+
+        #### Returns:
+            - `dict`: The state dictionary.
+        """
+        sd = self.model.state_dict()
+        list(sd.keys())
+        return sd
+
+    def patch_weight_to_device(self, key: str, device_to: torch.device = None) -> None:
+        """#### Patch the weight of a key to a device.
+
+        #### Args:
+            - `key` (str): The key.
+            - `device_to` (torch.device, optional): The device to patch to. Defaults to None.
+        """
+        if key not in self.patches:
+            return
+
+        weight = util.get_attr(self.model, key)
+
+        inplace_update = self.weight_inplace_update
+
+        if key not in self.backup:
+            self.backup[key] = weight.to(
+                device=self.offload_device, copy=inplace_update
+            )
+
+        if device_to is not None:
+            temp_weight = Device.cast_to_device(
+                weight, device_to, torch.float32, copy=True
+            )
+        else:
+            temp_weight = weight.to(torch.float32, copy=True)
+        out_weight = self.calculate_weight(self.patches[key], temp_weight, key).to(
+            weight.dtype
+        )
+        if inplace_update:
+            util.copy_to_param(self.model, key, out_weight)
+        else:
+            util.set_attr_param(self.model, key, out_weight)
+
+    def load(
+        self,
+        device_to: torch.device = None,
+        lowvram_model_memory: int = 0,
+        force_patch_weights: bool = False,
+        full_load: bool = False,
+    ):
+        """#### Load the model.
+
+        #### Args:
+            - `device_to` (torch.device, optional): The device to load to. Defaults to None.
+            - `lowvram_model_memory` (int, optional): The low VRAM model memory. Defaults to 0.
+            - `force_patch_weights` (bool, optional): Whether to force patch weights. Defaults to False.
+            - `full_load` (bool, optional): Whether to fully load the model. Defaults to False.
+        """
+        mem_counter = 0
+        patch_counter = 0
+        lowvram_counter = 0
+        loading = []
+        for n, m in self.model.named_modules():
+            if hasattr(m, "comfy_cast_weights") or hasattr(m, "weight"):
+                loading.append((Device.module_size(m), n, m))
+
+        load_completely = []
+        loading.sort(reverse=True)
+        for x in loading:
+            n = x[1]
+            m = x[2]
+            module_mem = x[0]
+
+            lowvram_weight = False
+
+            if not full_load and hasattr(m, "comfy_cast_weights"):
+                if mem_counter + module_mem >= lowvram_model_memory:
+                    lowvram_weight = True
+                    lowvram_counter += 1
+                    if hasattr(m, "prev_comfy_cast_weights"):  # Already lowvramed
+                        continue
+
+            weight_key = "{}.weight".format(n)
+            bias_key = "{}.bias".format(n)
+
+            if lowvram_weight:
+                if weight_key in self.patches:
+                    if force_patch_weights:
+                        self.patch_weight_to_device(weight_key)
+                if bias_key in self.patches:
+                    if force_patch_weights:
+                        self.patch_weight_to_device(bias_key)
+
+                m.prev_comfy_cast_weights = m.comfy_cast_weights
+                m.comfy_cast_weights = True
+            else:
+                if hasattr(m, "comfy_cast_weights"):
+                    if m.comfy_cast_weights:
+                        wipe_lowvram_weight(m)
+
+                if hasattr(m, "weight"):
+                    mem_counter += module_mem
+                    load_completely.append((module_mem, n, m))
+
+        load_completely.sort(reverse=True)
+        for x in load_completely:
+            n = x[1]
+            m = x[2]
+            weight_key = "{}.weight".format(n)
+            bias_key = "{}.bias".format(n)
+            if hasattr(m, "comfy_patched_weights"):
+                if m.comfy_patched_weights is True:
+                    continue
+
+            self.patch_weight_to_device(weight_key, device_to=device_to)
+            self.patch_weight_to_device(bias_key, device_to=device_to)
+            logging.debug("lowvram: loaded module regularly {} {}".format(n, m))
+            m.comfy_patched_weights = True
+
+        for x in load_completely:
+            x[2].to(device_to)
+
+        if lowvram_counter > 0:
+            logging.info(
+                "loaded partially {} {} {}".format(
+                    lowvram_model_memory / (1024 * 1024),
+                    mem_counter / (1024 * 1024),
+                    patch_counter,
+                )
+            )
+            self.model.model_lowvram = True
+        else:
+            logging.info(
+                "loaded completely {} {} {}".format(
+                    lowvram_model_memory / (1024 * 1024),
+                    mem_counter / (1024 * 1024),
+                    full_load,
+                )
+            )
+            self.model.model_lowvram = False
+            if full_load:
+                self.model.to(device_to)
+                mem_counter = self.model_size()
+
+
+        self.model.lowvram_patch_counter += patch_counter
+        self.model.device = device_to
+        self.model.model_loaded_weight_memory = mem_counter
+
+    def patch_model_flux(
+        self,
+        device_to: torch.device = None,
+        lowvram_model_memory: int =0,
+        load_weights: bool = True,
+        force_patch_weights: bool = False,
+    ):
+        """#### Patch the model.
+
+        #### Args:
+            - `device_to` (torch.device, optional): The device to patch to. Defaults to None.
+            - `lowvram_model_memory` (int, optional): The low VRAM model memory. Defaults to 0.
+            - `load_weights` (bool, optional): Whether to load weights. Defaults to True.
+            - `force_patch_weights` (bool, optional): Whether to force patch weights. Defaults to False.
+
+        #### Returns:
+            - `torch.nn.Module`: The patched model.
+        """
+        for k in self.object_patches:
+            old = util.set_attr(self.model, k, self.object_patches[k])
+            if k not in self.object_patches_backup:
+                self.object_patches_backup[k] = old
+
+        if lowvram_model_memory == 0:
+            full_load = True
+        else:
+            full_load = False
+
+        if load_weights:
+            self.load(
+                device_to,
+                lowvram_model_memory=lowvram_model_memory,
+                force_patch_weights=force_patch_weights,
+                full_load=full_load,
+            )
+        return self.model
+
+    def patch_model_lowvram_flux(
+        self,
+        device_to: torch.device = None,
+        lowvram_model_memory: int = 0,
+        force_patch_weights: bool = False,
+    ) -> torch.nn.Module:
+        """#### Patch the model for low VRAM.
+
+        #### Args:
+            - `device_to` (torch.device, optional): The device to patch to. Defaults to None.
+            - `lowvram_model_memory` (int, optional): The low VRAM model memory. Defaults to 0.
+            - `force_patch_weights` (bool, optional): Whether to force patch weights. Defaults to False.
+
+        #### Returns:
+            - `torch.nn.Module`: The patched model.
+        """
+        self.patch_model(device_to)
+
+        logging.info(
+            "loading in lowvram mode {}".format(lowvram_model_memory / (1024 * 1024))
+        )
+
+        class LowVramPatch:
+            def __init__(self, key: str, model_patcher: "ModelPatcher"):
+                self.key = key
+                self.model_patcher = model_patcher
+
+            def __call__(self, weight: torch.Tensor) -> torch.Tensor:
+                return self.model_patcher.calculate_weight(
+                    self.model_patcher.patches[self.key], weight, self.key
+                )
+
+        mem_counter = 0
+        patch_counter = 0
+        for n, m in self.model.named_modules():
+            lowvram_weight = False
+            if hasattr(m, "comfy_cast_weights"):
+                module_mem = Device.module_size(m)
+                if mem_counter + module_mem >= lowvram_model_memory:
+                    lowvram_weight = True
+
+            weight_key = "{}.weight".format(n)
+            bias_key = "{}.bias".format(n)
+
+            if lowvram_weight:
+                if weight_key in self.patches:
+                    if force_patch_weights:
+                        self.patch_weight_to_device(weight_key)
+                    else:
+                        m.weight_function = LowVramPatch(weight_key, self)
+                        patch_counter += 1
+                if bias_key in self.patches:
+                    if force_patch_weights:
+                        self.patch_weight_to_device(bias_key)
+                    else:
+                        m.bias_function = LowVramPatch(bias_key, self)
+                        patch_counter += 1
+
+                m.prev_comfy_cast_weights = m.comfy_cast_weights
+                m.comfy_cast_weights = True
+            else:
+                if hasattr(m, "weight"):
+                    self.patch_weight_to_device(weight_key, device_to)
+                    self.patch_weight_to_device(bias_key, device_to)
+                    m.to(device_to)
+                    mem_counter += Device.module_size(m)
+                    logging.debug("lowvram: loaded module regularly {}".format(m))
+
+        self.model_lowvram = True
+        self.lowvram_patch_counter = patch_counter
+        return self.model
+
+    def patch_model(
+        self, device_to: torch.device = None, patch_weights: bool = True
+    ) -> torch.nn.Module:
+        """#### Patch the model.
+
+        #### Args:
+            - `device_to` (torch.device, optional): The device to patch to. Defaults to None.
+            - `patch_weights` (bool, optional): Whether to patch weights. Defaults to True.
+
+        #### Returns:
+            - `torch.nn.Module`: The patched model.
+        """
+        for k in self.object_patches:
+            old = util.set_attr(self.model, k, self.object_patches[k])
+            if k not in self.object_patches_backup:
+                self.object_patches_backup[k] = old
+
+        if patch_weights:
+            model_sd = self.model_state_dict()
+            for key in self.patches:
+                if key not in model_sd:
+                    logging.warning(
+                        "could not patch. key doesn't exist in model: {}".format(key)
+                    )
+                    continue
+
+                self.patch_weight_to_device(key, device_to)
+
+            if device_to is not None:
+                self.model.to(device_to)
+                self.current_device = device_to
+
+        return self.model
+
+    def patch_model_lowvram(
+        self,
+        device_to: torch.device = None,
+        lowvram_model_memory: int = 0,
+        force_patch_weights: bool = False,
+    ) -> torch.nn.Module:
+        """#### Patch the model for low VRAM.
+
+        #### Args:
+            - `device_to` (torch.device, optional): The device to patch to. Defaults to None.
+            - `lowvram_model_memory` (int, optional): The low VRAM model memory. Defaults to 0.
+            - `force_patch_weights` (bool, optional): Whether to force patch weights. Defaults to False.
+
+        #### Returns:
+            - `torch.nn.Module`: The patched model.
+        """
+        self.patch_model(device_to, patch_weights=False)
+
+        logging.info(
+            "loading in lowvram mode {}".format(lowvram_model_memory / (1024 * 1024))
+        )
+
+        class LowVramPatch:
+            def __init__(self, key: str, model_patcher: "ModelPatcher"):
+                self.key = key
+                self.model_patcher = model_patcher
+
+            def __call__(self, weight: torch.Tensor) -> torch.Tensor:
+                return self.model_patcher.calculate_weight(
+                    self.model_patcher.patches[self.key], weight, self.key
+                )
+
+        mem_counter = 0
+        patch_counter = 0
+        for n, m in self.model.named_modules():
+            lowvram_weight = False
+            if hasattr(m, "comfy_cast_weights"):
+                module_mem = Device.module_size(m)
+                if mem_counter + module_mem >= lowvram_model_memory:
+                    lowvram_weight = True
+
+            weight_key = "{}.weight".format(n)
+            bias_key = "{}.bias".format(n)
+
+            if lowvram_weight:
+                if weight_key in self.patches:
+                    if force_patch_weights:
+                        self.patch_weight_to_device(weight_key)
+                    else:
+                        m.weight_function = LowVramPatch(weight_key, self)
+                        patch_counter += 1
+                if bias_key in self.patches:
+                    if force_patch_weights:
+                        self.patch_weight_to_device(bias_key)
+                    else:
+                        m.bias_function = LowVramPatch(bias_key, self)
+                        patch_counter += 1
+
+                m.prev_comfy_cast_weights = m.comfy_cast_weights
+                m.comfy_cast_weights = True
+            else:
+                if hasattr(m, "weight"):
+                    self.patch_weight_to_device(weight_key, device_to)
+                    self.patch_weight_to_device(bias_key, device_to)
+                    m.to(device_to)
+                    mem_counter += Device.module_size(m)
+                    logging.debug("lowvram: loaded module regularly {}".format(m))
+
+        self.model_lowvram = True
+        self.lowvram_patch_counter = patch_counter
+        return self.model
+
+    def calculate_weight(
+        self, patches: list, weight: torch.Tensor, key: str
+    ) -> torch.Tensor:
+        """#### Calculate the weight of a key.
+
+        #### Args:
+            - `patches` (list): The list of patches.
+            - `weight` (torch.Tensor): The weight tensor.
+            - `key` (str): The key.
+
+        #### Returns:
+            - `torch.Tensor`: The calculated weight.
+        """
+        for p in patches:
+            alpha = p[0]
+            v = p[1]
+            p[2]
+            v[0]
+            v = v[1]
+            mat1 = Device.cast_to_device(v[0], weight.device, torch.float32)
+            mat2 = Device.cast_to_device(v[1], weight.device, torch.float32)
+            v[4]
+            if v[2] is not None:
+                alpha *= v[2] / mat2.shape[0]
+            weight += (
+                (alpha * torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1)))
+                .reshape(weight.shape)
+                .type(weight.dtype)
+            )
+        return weight
+
+    def unpatch_model(
+        self, device_to: torch.device = None, unpatch_weights: bool = True
+    ) -> None:
+        """#### Unpatch the model.
+
+        #### Args:
+            - `device_to` (torch.device, optional): The device to unpatch to. Defaults to None.
+            - `unpatch_weights` (bool, optional): Whether to unpatch weights. Defaults to True.
+        """
+        if unpatch_weights:
+            keys = list(self.backup.keys())
+            for k in keys:
+                util.set_attr_param(self.model, k, self.backup[k])
+            self.backup.clear()
+            if device_to is not None:
+                self.model.to(device_to)
+                self.current_device = device_to
+
+        keys = list(self.object_patches_backup.keys())
+        self.object_patches_backup.clear()
+
+    def partially_load(self, device_to: torch.device, extra_memory: int = 0) -> int:
+        """#### Partially load the model.
+
+        #### Args:
+            - `device_to` (torch.device): The device to load to.
+            - `extra_memory` (int, optional): The extra memory. Defaults to 0.
+
+        #### Returns:
+            - `int`: The memory loaded.
+        """
+        self.unpatch_model(unpatch_weights=False)
+        self.patch_model(patch_weights=False)
+        full_load = False
+        if self.model.model_lowvram is False:
+            return 0
+        if self.model.model_loaded_weight_memory + extra_memory > self.model_size():
+            full_load = True
+        current_used = self.model.model_loaded_weight_memory
+        self.load(
+            device_to,
+            lowvram_model_memory=current_used + extra_memory,
+            full_load=full_load,
+        )
+        return self.model.model_loaded_weight_memory - current_used
+
+    def add_object_patch(self, name, obj):
+            self.object_patches[name] = obj
+
+def unet_prefix_from_state_dict(state_dict: dict) -> str:
+    """#### Get the UNet prefix from the state dictionary.
+
+    #### Args:
+        - `state_dict` (dict): The state dictionary.
+
+    #### Returns:
+        - `str`: The UNet prefix.
+    """
+    candidates = [
+        "model.diffusion_model.",  # ldm/sgm models
+        "model.model.",  # audio models
+    ]
+    counts = {k: 0 for k in candidates}
+    for k in state_dict:
+        for c in candidates:
+            if k.startswith(c):
+                counts[c] += 1
+                break
+
+    top = max(counts, key=counts.get)
+    if counts[top] > 5:
+        return top
+    else:
+        return "model."  # aura flow and others
+
+def load_diffusion_model_state_dict(
+    sd, model_options={}
+) -> ModelPatcher:
+    """#### Load the diffusion model state dictionary.
+
+    #### Args:
+        - `sd`: The state dictionary.
+        - `model_options` (dict, optional): The model options. Defaults to {}.
+
+    #### Returns:
+        - `ModelPatcher`: The model patcher.
+    """
+    # load unet in diffusers or regular format
+    dtype = model_options.get("dtype", None)
+
+    # Allow loading unets from checkpoint files
+    diffusion_model_prefix = unet_prefix_from_state_dict(sd)
+    temp_sd = util.state_dict_prefix_replace(
+        sd, {diffusion_model_prefix: ""}, filter_keys=True
+    )
+    if len(temp_sd) > 0:
+        sd = temp_sd
+
+    parameters = util.calculate_parameters(sd)
+    load_device = Device.get_torch_device()
+    model_config = unet.model_config_from_unet(sd, "")
+
+    if model_config is not None:
+        new_sd = sd
+
+    offload_device = Device.unet_offload_device()
+    if dtype is None:
+        unet_dtype2 = Device.unet_dtype(
+            model_params=parameters,
+            supported_dtypes=model_config.supported_inference_dtypes,
+        )
+    else:
+        unet_dtype2 = dtype
+
+    manual_cast_dtype = Device.unet_manual_cast(
+        unet_dtype2, load_device, model_config.supported_inference_dtypes
+    )
+    model_config.set_inference_dtype(unet_dtype2, manual_cast_dtype)
+    model_config.custom_operations = model_options.get(
+        "custom_operations", model_config.custom_operations
+    )
+    model = model_config.get_model(new_sd, "")
+    model = model.to(offload_device)
+    model.load_model_weights(new_sd, "")
+    left_over = sd.keys()
+    if len(left_over) > 0:
+        logging.info("left over keys in unet: {}".format(left_over))
+    return ModelPatcher(model, load_device=load_device, offload_device=offload_device)

modules/NeuralNetwork/transformer.py

@@ -0,0 +1,443 @@
+from einops import rearrange
+import torch
+from modules.Utilities import util
+import torch.nn as nn
+from modules.Attention import Attention
+from modules.Device import Device
+from modules.cond import Activation
+from modules.cond import cast
+from modules.sample import sampling_util
+
+if Device.xformers_enabled():
+    pass
+
+ops = cast.disable_weight_init
+
+_ATTN_PRECISION = "fp32"
+
+
+class FeedForward(nn.Module):
+    """#### FeedForward neural network module.
+
+    #### Args:
+        - `dim` (int): The input dimension.
+        - `dim_out` (int, optional): The output dimension. Defaults to None.
+        - `mult` (int, optional): The multiplier for the inner dimension. Defaults to 4.
+        - `glu` (bool, optional): Whether to use Gated Linear Units. Defaults to False.
+        - `dropout` (float, optional): The dropout rate. Defaults to 0.0.
+        - `dtype` (torch.dtype, optional): The data type. Defaults to None.
+        - `device` (torch.device, optional): The device. Defaults to None.
+        - `operations` (object, optional): The operations module. Defaults to `ops`.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int = None,
+        mult: int = 4,
+        glu: bool = False,
+        dropout: float = 0.0,
+        dtype: torch.dtype = None,
+        device: torch.device = None,
+        operations: object = ops,
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = util.default(dim_out, dim)
+        project_in = (
+            nn.Sequential(
+                operations.Linear(dim, inner_dim, dtype=dtype, device=device), nn.GELU()
+            )
+            if not glu
+            else Activation.GEGLU(dim, inner_dim)
+        )
+
+        self.net = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            operations.Linear(inner_dim, dim_out, dtype=dtype, device=device),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass of the FeedForward network.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        return self.net(x)
+
+
+class BasicTransformerBlock(nn.Module):
+    """#### Basic Transformer block.
+
+    #### Args:
+        - `dim` (int): The input dimension.
+        - `n_heads` (int): The number of attention heads.
+        - `d_head` (int): The dimension of each attention head.
+        - `dropout` (float, optional): The dropout rate. Defaults to 0.0.
+        - `context_dim` (int, optional): The context dimension. Defaults to None.
+        - `gated_ff` (bool, optional): Whether to use Gated FeedForward. Defaults to True.
+        - `checkpoint` (bool, optional): Whether to use checkpointing. Defaults to True.
+        - `ff_in` (bool, optional): Whether to use FeedForward input. Defaults to False.
+        - `inner_dim` (int, optional): The inner dimension. Defaults to None.
+        - `disable_self_attn` (bool, optional): Whether to disable self-attention. Defaults to False.
+        - `disable_temporal_crossattention` (bool, optional): Whether to disable temporal cross-attention. Defaults to False.
+        - `switch_temporal_ca_to_sa` (bool, optional): Whether to switch temporal cross-attention to self-attention. Defaults to False.
+        - `dtype` (torch.dtype, optional): The data type. Defaults to None.
+        - `device` (torch.device, optional): The device. Defaults to None.
+        - `operations` (object, optional): The operations module. Defaults to `ops`.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        d_head: int,
+        dropout: float = 0.0,
+        context_dim: int = None,
+        gated_ff: bool = True,
+        checkpoint: bool = True,
+        ff_in: bool = False,
+        inner_dim: int = None,
+        disable_self_attn: bool = False,
+        disable_temporal_crossattention: bool = False,
+        switch_temporal_ca_to_sa: bool = False,
+        dtype: torch.dtype = None,
+        device: torch.device = None,
+        operations: object = ops,
+    ):
+        super().__init__()
+
+        self.ff_in = ff_in or inner_dim is not None
+        if inner_dim is None:
+            inner_dim = dim
+
+        self.is_res = inner_dim == dim
+        self.disable_self_attn = disable_self_attn
+        self.attn1 = Attention.CrossAttention(
+            query_dim=inner_dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            context_dim=context_dim if self.disable_self_attn else None,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )  # is a self-attention if not self.disable_self_attn
+        self.ff = FeedForward(
+            inner_dim,
+            dim_out=dim,
+            dropout=dropout,
+            glu=gated_ff,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )
+
+        context_dim_attn2 = None
+        if not switch_temporal_ca_to_sa:
+            context_dim_attn2 = context_dim
+
+        self.attn2 = Attention.CrossAttention(
+            query_dim=inner_dim,
+            context_dim=context_dim_attn2,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )  # is self-attn if context is none
+        self.norm2 = operations.LayerNorm(inner_dim, dtype=dtype, device=device)
+
+        self.norm1 = operations.LayerNorm(inner_dim, dtype=dtype, device=device)
+        self.norm3 = operations.LayerNorm(inner_dim, dtype=dtype, device=device)
+        self.checkpoint = checkpoint
+        self.n_heads = n_heads
+        self.d_head = d_head
+        self.switch_temporal_ca_to_sa = switch_temporal_ca_to_sa
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor = None,
+        transformer_options: dict = {},
+    ) -> torch.Tensor:
+        """#### Forward pass of the Basic Transformer block.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `context` (torch.Tensor, optional): The context tensor. Defaults to None.
+            - `transformer_options` (dict, optional): Additional transformer options. Defaults to {}.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        return sampling_util.checkpoint(
+            self._forward,
+            (x, context, transformer_options),
+            self.parameters(),
+            self.checkpoint,
+        )
+
+    def _forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor = None,
+        transformer_options: dict = {},
+    ) -> torch.Tensor:
+        """#### Internal forward pass of the Basic Transformer block.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `context` (torch.Tensor, optional): The context tensor. Defaults to None.
+            - `transformer_options` (dict, optional): Additional transformer options. Defaults to {}.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        extra_options = {}
+        block = transformer_options.get("block", None)
+        block_index = transformer_options.get("block_index", 0)
+        transformer_patches_replace = {}
+
+        for k in transformer_options:
+            extra_options[k] = transformer_options[k]
+
+        extra_options["n_heads"] = self.n_heads
+        extra_options["dim_head"] = self.d_head
+
+        n = self.norm1(x)
+        context_attn1 = None
+        value_attn1 = None
+
+        transformer_block = (block[0], block[1], block_index)
+        attn1_replace_patch = transformer_patches_replace.get("attn1", {})
+        block_attn1 = transformer_block
+        if block_attn1 not in attn1_replace_patch:
+            block_attn1 = block
+
+        n = self.attn1(n, context=context_attn1, value=value_attn1)
+
+        x += n
+
+        if self.attn2 is not None:
+            n = self.norm2(x)
+            context_attn2 = context
+            value_attn2 = None
+
+            attn2_replace_patch = transformer_patches_replace.get("attn2", {})
+            block_attn2 = transformer_block
+            if block_attn2 not in attn2_replace_patch:
+                block_attn2 = block
+            n = self.attn2(n, context=context_attn2, value=value_attn2)
+
+        x += n
+        if self.is_res:
+            x_skip = x
+        x = self.ff(self.norm3(x))
+        if self.is_res:
+            x += x_skip
+
+        return x
+
+
+class SpatialTransformer(nn.Module):
+    """#### Spatial Transformer module.
+
+    #### Args:
+        - `in_channels` (int): The number of input channels.
+        - `n_heads` (int): The number of attention heads.
+        - `d_head` (int): The dimension of each attention head.
+        - `depth` (int, optional): The depth of the transformer. Defaults to 1.
+        - `dropout` (float, optional): The dropout rate. Defaults to 0.0.
+        - `context_dim` (int, optional): The context dimension. Defaults to None.
+        - `disable_self_attn` (bool, optional): Whether to disable self-attention. Defaults to False.
+        - `use_linear` (bool, optional): Whether to use linear projections. Defaults to False.
+        - `use_checkpoint` (bool, optional): Whether to use checkpointing. Defaults to True.
+        - `dtype` (torch.dtype, optional): The data type. Defaults to None.
+        - `device` (torch.device, optional): The device. Defaults to None.
+        - `operations` (object, optional): The operations module. Defaults to `ops`.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        n_heads: int,
+        d_head: int,
+        depth: int = 1,
+        dropout: float = 0.0,
+        context_dim: int = None,
+        disable_self_attn: bool = False,
+        use_linear: bool = False,
+        use_checkpoint: bool = True,
+        dtype: torch.dtype = None,
+        device: torch.device = None,
+        operations: object = ops,
+    ):
+        super().__init__()
+        if util.exists(context_dim) and not isinstance(context_dim, list):
+            context_dim = [context_dim] * depth
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = operations.GroupNorm(
+            num_groups=32,
+            num_channels=in_channels,
+            eps=1e-6,
+            affine=True,
+            dtype=dtype,
+            device=device,
+        )
+        if not use_linear:
+            self.proj_in = operations.Conv2d(
+                in_channels,
+                inner_dim,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                dtype=dtype,
+                device=device,
+            )
+        else:
+            self.proj_in = operations.Linear(
+                in_channels, inner_dim, dtype=dtype, device=device
+            )
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    n_heads,
+                    d_head,
+                    dropout=dropout,
+                    context_dim=context_dim[d],
+                    disable_self_attn=disable_self_attn,
+                    checkpoint=use_checkpoint,
+                    dtype=dtype,
+                    device=device,
+                    operations=operations,
+                )
+                for d in range(depth)
+            ]
+        )
+        if not use_linear:
+            self.proj_out = operations.Conv2d(
+                inner_dim,
+                in_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                dtype=dtype,
+                device=device,
+            )
+        else:
+            self.proj_out = operations.Linear(
+                in_channels, inner_dim, dtype=dtype, device=device
+            )
+        self.use_linear = use_linear
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor = None,
+        transformer_options: dict = {},
+    ) -> torch.Tensor:
+        """#### Forward pass of the Spatial Transformer.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `context` (torch.Tensor, optional): The context tensor. Defaults to None.
+            - `transformer_options` (dict, optional): Additional transformer options. Defaults to {}.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        # note: if no context is given, cross-attention defaults to self-attention
+        if not isinstance(context, list):
+            context = [context] * len(self.transformer_blocks)
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        if not self.use_linear:
+            x = self.proj_in(x)
+        x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+        if self.use_linear:
+            x = self.proj_in(x)
+        for i, block in enumerate(self.transformer_blocks):
+            transformer_options["block_index"] = i
+            x = block(x, context=context[i], transformer_options=transformer_options)
+        if self.use_linear:
+            x = self.proj_out(x)
+        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+        if not self.use_linear:
+            x = self.proj_out(x)
+        return x + x_in
+
+
+def count_blocks(state_dict_keys: list, prefix_string: str) -> int:
+    """#### Count the number of blocks in a state dictionary.
+
+    #### Args:
+        - `state_dict_keys` (list): The list of state dictionary keys.
+        - `prefix_string` (str): The prefix string to match.
+
+    #### Returns:
+        - `int`: The number of blocks.
+    """
+    count = 0
+    while True:
+        c = False
+        for k in state_dict_keys:
+            if k.startswith(prefix_string.format(count)):
+                c = True
+                break
+        if c is False:
+            break
+        count += 1
+    return count
+
+
+def calculate_transformer_depth(
+    prefix: str, state_dict_keys: list, state_dict: dict
+) -> tuple:
+    """#### Calculate the depth of a transformer.
+
+    #### Args:
+        - `prefix` (str): The prefix string.
+        - `state_dict_keys` (list): The list of state dictionary keys.
+        - `state_dict` (dict): The state dictionary.
+
+    #### Returns:
+        - `tuple`: The transformer depth, context dimension, use of linear in transformer, and time stack.
+    """
+    context_dim = None
+    use_linear_in_transformer = False
+
+    transformer_prefix = prefix + "1.transformer_blocks."
+    transformer_keys = sorted(
+        list(filter(lambda a: a.startswith(transformer_prefix), state_dict_keys))
+    )
+    if len(transformer_keys) > 0:
+        last_transformer_depth = count_blocks(
+            state_dict_keys, transformer_prefix + "{}"
+        )
+        context_dim = state_dict[
+            "{}0.attn2.to_k.weight".format(transformer_prefix)
+        ].shape[1]
+        use_linear_in_transformer = (
+            len(state_dict["{}1.proj_in.weight".format(prefix)].shape) == 2
+        )
+        time_stack = (
+            "{}1.time_stack.0.attn1.to_q.weight".format(prefix) in state_dict
+            or "{}1.time_mix_blocks.0.attn1.to_q.weight".format(prefix) in state_dict
+        )
+        return (
+            last_transformer_depth,
+            context_dim,
+            use_linear_in_transformer,
+            time_stack,
+        )
+    return None

modules/NeuralNetwork/unet.py

@@ -0,0 +1,1132 @@
+import logging
+import math
+from typing import Any, Dict, List, Optional
+import torch.nn as nn
+import torch as th
+import torch
+
+from modules.Utilities import util
+from modules.AutoEncoders import ResBlock
+from modules.NeuralNetwork import transformer
+from modules.cond import cast
+from modules.sample import sampling, sampling_util
+
+UNET_MAP_ATTENTIONS = {
+    "proj_in.weight",
+    "proj_in.bias",
+    "proj_out.weight",
+    "proj_out.bias",
+    "norm.weight",
+    "norm.bias",
+}
+
+TRANSFORMER_BLOCKS = {
+    "norm1.weight",
+    "norm1.bias",
+    "norm2.weight",
+    "norm2.bias",
+    "norm3.weight",
+    "norm3.bias",
+    "attn1.to_q.weight",
+    "attn1.to_k.weight",
+    "attn1.to_v.weight",
+    "attn1.to_out.0.weight",
+    "attn1.to_out.0.bias",
+    "attn2.to_q.weight",
+    "attn2.to_k.weight",
+    "attn2.to_v.weight",
+    "attn2.to_out.0.weight",
+    "attn2.to_out.0.bias",
+    "ff.net.0.proj.weight",
+    "ff.net.0.proj.bias",
+    "ff.net.2.weight",
+    "ff.net.2.bias",
+}
+
+UNET_MAP_RESNET = {
+    "in_layers.2.weight": "conv1.weight",
+    "in_layers.2.bias": "conv1.bias",
+    "emb_layers.1.weight": "time_emb_proj.weight",
+    "emb_layers.1.bias": "time_emb_proj.bias",
+    "out_layers.3.weight": "conv2.weight",
+    "out_layers.3.bias": "conv2.bias",
+    "skip_connection.weight": "conv_shortcut.weight",
+    "skip_connection.bias": "conv_shortcut.bias",
+    "in_layers.0.weight": "norm1.weight",
+    "in_layers.0.bias": "norm1.bias",
+    "out_layers.0.weight": "norm2.weight",
+    "out_layers.0.bias": "norm2.bias",
+}
+
+UNET_MAP_BASIC = {
+    ("label_emb.0.0.weight", "class_embedding.linear_1.weight"),
+    ("label_emb.0.0.bias", "class_embedding.linear_1.bias"),
+    ("label_emb.0.2.weight", "class_embedding.linear_2.weight"),
+    ("label_emb.0.2.bias", "class_embedding.linear_2.bias"),
+    ("label_emb.0.0.weight", "add_embedding.linear_1.weight"),
+    ("label_emb.0.0.bias", "add_embedding.linear_1.bias"),
+    ("label_emb.0.2.weight", "add_embedding.linear_2.weight"),
+    ("label_emb.0.2.bias", "add_embedding.linear_2.bias"),
+    ("input_blocks.0.0.weight", "conv_in.weight"),
+    ("input_blocks.0.0.bias", "conv_in.bias"),
+    ("out.0.weight", "conv_norm_out.weight"),
+    ("out.0.bias", "conv_norm_out.bias"),
+    ("out.2.weight", "conv_out.weight"),
+    ("out.2.bias", "conv_out.bias"),
+    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
+    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
+    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
+    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
+}
+
+# taken from https://github.com/TencentARC/T2I-Adapter
+
+
+def unet_to_diffusers(unet_config: dict) -> dict:
+    """#### Convert a UNet configuration to a diffusers configuration.
+
+    #### Args:
+        - `unet_config` (dict): The UNet configuration.
+
+    #### Returns:
+        - `dict`: The diffusers configuration.
+    """
+    if "num_res_blocks" not in unet_config:
+        return {}
+    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 UNET_MAP_RESNET:
+                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 UNET_MAP_ATTENTIONS:
+                    diffusers_unet_map[
+                        "down_blocks.{}.attentions.{}.{}".format(x, i, 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)
+            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 UNET_MAP_ATTENTIONS:
+        diffusers_unet_map["mid_block.attentions.{}.{}".format(i, 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 i, n in enumerate([0, 2]):
+        for b in UNET_MAP_RESNET:
+            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
+        length = num_res_blocks[x] + 1
+        for i in range(length):
+            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 == length - 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
+
+
+def apply_control1(h: th.Tensor, control: any, name: str) -> th.Tensor:
+    """#### Apply control to a tensor.
+
+    #### Args:
+        - `h` (torch.Tensor): The input tensor.
+        - `control` (any): The control to apply.
+        - `name` (str): The name of the control.
+
+    #### Returns:
+        - `torch.Tensor`: The controlled tensor.
+    """
+    return h
+
+
+oai_ops = cast.disable_weight_init
+
+
+class UNetModel1(nn.Module):
+    """#### UNet Model class."""
+
+    def __init__(
+        self,
+        image_size: int,
+        in_channels: int,
+        model_channels: int,
+        out_channels: int,
+        num_res_blocks: list,
+        dropout: float = 0,
+        channel_mult: tuple = (1, 2, 4, 8),
+        conv_resample: bool = True,
+        dims: int = 2,
+        num_classes: int = None,
+        use_checkpoint: bool = False,
+        dtype: th.dtype = th.float32,
+        num_heads: int = -1,
+        num_head_channels: int = -1,
+        num_heads_upsample: int = -1,
+        use_scale_shift_norm: bool = False,
+        resblock_updown: bool = False,
+        use_new_attention_order: bool = False,
+        use_spatial_transformer: bool = False,  # custom transformer support
+        transformer_depth: int = 1,  # custom transformer support
+        context_dim: int = None,  # custom transformer support
+        n_embed: int = None,  # custom support for prediction of discrete ids into codebook of first stage vq model
+        legacy: bool = True,
+        disable_self_attentions: list = None,
+        num_attention_blocks: list = None,
+        disable_middle_self_attn: bool = False,
+        use_linear_in_transformer: bool = False,
+        adm_in_channels: int = None,
+        transformer_depth_middle: int = None,
+        transformer_depth_output: list = None,
+        use_temporal_resblock: bool = False,
+        use_temporal_attention: bool = False,
+        time_context_dim: int = None,
+        extra_ff_mix_layer: bool = False,
+        use_spatial_context: bool = False,
+        merge_strategy: any = None,
+        merge_factor: float = 0.0,
+        video_kernel_size: int = None,
+        disable_temporal_crossattention: bool = False,
+        max_ddpm_temb_period: int = 10000,
+        device: th.device = None,
+        operations: any = oai_ops,
+    ):
+        """#### Initialize the UNetModel1 class.
+
+        #### Args:
+            - `image_size` (int): The size of the input image.
+            - `in_channels` (int): The number of input channels.
+            - `model_channels` (int): The number of model channels.
+            - `out_channels` (int): The number of output channels.
+            - `num_res_blocks` (list): The number of residual blocks.
+            - `dropout` (float, optional): The dropout rate. Defaults to 0.
+            - `channel_mult` (tuple, optional): The channel multiplier. Defaults to (1, 2, 4, 8).
+            - `conv_resample` (bool, optional): Whether to use convolutional resampling. Defaults to True.
+            - `dims` (int, optional): The number of dimensions. Defaults to 2.
+            - `num_classes` (int, optional): The number of classes. Defaults to None.
+            - `use_checkpoint` (bool, optional): Whether to use checkpointing. Defaults to False.
+            - `dtype` (torch.dtype, optional): The data type. Defaults to torch.float32.
+            - `num_heads` (int, optional): The number of heads. Defaults to -1.
+            - `num_head_channels` (int, optional): The number of head channels. Defaults to -1.
+            - `num_heads_upsample` (int, optional): The number of heads for upsampling. Defaults to -1.
+            - `use_scale_shift_norm` (bool, optional): Whether to use scale-shift normalization. Defaults to False.
+            - `resblock_updown` (bool, optional): Whether to use residual blocks for up/down sampling. Defaults to False.
+            - `use_new_attention_order` (bool, optional): Whether to use a new attention order. Defaults to False.
+            - `use_spatial_transformer` (bool, optional): Whether to use a spatial transformer. Defaults to False.
+            - `transformer_depth` (int, optional): The depth of the transformer. Defaults to 1.
+            - `context_dim` (int, optional): The context dimension. Defaults to None.
+            - `n_embed` (int, optional): The number of embeddings. Defaults to None.
+            - `legacy` (bool, optional): Whether to use legacy mode. Defaults to True.
+            - `disable_self_attentions` (list, optional): The list of self-attentions to disable. Defaults to None.
+            - `num_attention_blocks` (list, optional): The number of attention blocks. Defaults to None.
+            - `disable_middle_self_attn` (bool, optional): Whether to disable middle self-attention. Defaults to False.
+            - `use_linear_in_transformer` (bool, optional): Whether to use linear in transformer. Defaults to False.
+            - `adm_in_channels` (int, optional): The number of ADM input channels. Defaults to None.
+            - `transformer_depth_middle` (int, optional): The depth of the middle transformer. Defaults to None.
+            - `transformer_depth_output` (list, optional): The depth of the output transformer. Defaults to None.
+            - `use_temporal_resblock` (bool, optional): Whether to use temporal residual blocks. Defaults to False.
+            - `use_temporal_attention` (bool, optional): Whether to use temporal attention. Defaults to False.
+            - `time_context_dim` (int, optional): The time context dimension. Defaults to None.
+            - `extra_ff_mix_layer` (bool, optional): Whether to use an extra feed-forward mix layer. Defaults to False.
+            - `use_spatial_context` (bool, optional): Whether to use spatial context. Defaults to False.
+            - `merge_strategy` (any, optional): The merge strategy. Defaults to None.
+            - `merge_factor` (float, optional): The merge factor. Defaults to 0.0.
+            - `video_kernel_size` (int, optional): The video kernel size. Defaults to None.
+            - `disable_temporal_crossattention` (bool, optional): Whether to disable temporal cross-attention. Defaults to False.
+            - `max_ddpm_temb_period` (int, optional): The maximum DDPM temporal embedding period. Defaults to 10000.
+            - `device` (torch.device, optional): The device to use. Defaults to None.
+            - `operations` (any, optional): The operations to use. Defaults to oai_ops.
+        """
+        super().__init__()
+
+        if context_dim is not None:
+            self.context_dim = context_dim
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+        if num_head_channels == -1:
+            assert num_heads != -1, "Either num_heads or num_head_channels has to be set"
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+
+        transformer_depth = transformer_depth[:]
+        transformer_depth_output = transformer_depth_output[:]
+
+        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.use_temporal_resblocks = use_temporal_resblock
+        self.predict_codebook_ids = n_embed is not None
+
+        self.default_num_video_frames = 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
+            ),
+        )
+
+        self.input_blocks = nn.ModuleList(
+            [
+                sampling.TimestepEmbedSequential1(
+                    operations.conv_nd(
+                        dims,
+                        in_channels,
+                        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
+
+        def get_attention_layer(
+            ch: int,
+            num_heads: int,
+            dim_head: int,
+            depth: int = 1,
+            context_dim: int = None,
+            use_checkpoint: bool = False,
+            disable_self_attn: bool = False,
+        ) -> transformer.SpatialTransformer:
+            """#### Get an attention layer.
+
+            #### Args:
+                - `ch` (int): The number of channels.
+                - `num_heads` (int): The number of heads.
+                - `dim_head` (int): The dimension of each head.
+                - `depth` (int, optional): The depth of the transformer. Defaults to 1.
+                - `context_dim` (int, optional): The context dimension. Defaults to None.
+                - `use_checkpoint` (bool, optional): Whether to use checkpointing. Defaults to False.
+                - `disable_self_attn` (bool, optional): Whether to disable self-attention. Defaults to False.
+
+            #### Returns:
+                - `transformer.SpatialTransformer`: The attention layer.
+            """
+            return transformer.SpatialTransformer(
+                ch,
+                num_heads,
+                dim_head,
+                depth=depth,
+                context_dim=context_dim,
+                disable_self_attn=disable_self_attn,
+                use_linear=use_linear_in_transformer,
+                use_checkpoint=use_checkpoint,
+                dtype=self.dtype,
+                device=device,
+                operations=operations,
+            )
+
+        def get_resblock(
+            merge_factor: float,
+            merge_strategy: any,
+            video_kernel_size: int,
+            ch: int,
+            time_embed_dim: int,
+            dropout: float,
+            out_channels: int,
+            dims: int,
+            use_checkpoint: bool,
+            use_scale_shift_norm: bool,
+            down: bool = False,
+            up: bool = False,
+            dtype: th.dtype = None,
+            device: th.device = None,
+            operations: any = oai_ops,
+        ) -> ResBlock.ResBlock1:
+            """#### Get a residual block.
+
+            #### Args:
+                - `merge_factor` (float): The merge factor.
+                - `merge_strategy` (any): The merge strategy.
+                - `video_kernel_size` (int): The video kernel size.
+                - `ch` (int): The number of channels.
+                - `time_embed_dim` (int): The time embedding dimension.
+                - `dropout` (float): The dropout rate.
+                - `out_channels` (int): The number of output channels.
+                - `dims` (int): The number of dimensions.
+                - `use_checkpoint` (bool): Whether to use checkpointing.
+                - `use_scale_shift_norm` (bool): Whether to use scale-shift normalization.
+                - `down` (bool, optional): Whether to use downsampling. Defaults to False.
+                - `up` (bool, optional): Whether to use upsampling. Defaults to False.
+                - `dtype` (torch.dtype, optional): The data type. Defaults to None.
+                - `device` (torch.device, optional): The device. Defaults to None.
+                - `operations` (any, optional): The operations to use. Defaults to oai_ops.
+
+            #### Returns:
+                - `ResBlock.ResBlock1`: The residual block.
+            """
+            return ResBlock.ResBlock1(
+                channels=ch,
+                emb_channels=time_embed_dim,
+                dropout=dropout,
+                out_channels=out_channels,
+                use_checkpoint=use_checkpoint,
+                dims=dims,
+                use_scale_shift_norm=use_scale_shift_norm,
+                down=down,
+                up=up,
+                dtype=dtype,
+                device=device,
+                operations=operations,
+            )
+
+        self.double_blocks = nn.ModuleList()
+        for level, mult in enumerate(channel_mult):
+            for nr in range(self.num_res_blocks[level]):
+                layers = [
+                    get_resblock(
+                        merge_factor=merge_factor,
+                        merge_strategy=merge_strategy,
+                        video_kernel_size=video_kernel_size,
+                        ch=ch,
+                        time_embed_dim=time_embed_dim,
+                        dropout=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,
+                    )
+                ]
+                ch = mult * model_channels
+                num_transformers = transformer_depth.pop(0)
+                if num_transformers > 0:
+                    dim_head = ch // num_heads
+                    disabled_sa = False
+
+                    if (
+                        not util.exists(num_attention_blocks)
+                        or nr < num_attention_blocks[level]
+                    ):
+                        layers.append(
+                            get_attention_layer(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                depth=num_transformers,
+                                context_dim=context_dim,
+                                disable_self_attn=disabled_sa,
+                                use_checkpoint=use_checkpoint,
+                            )
+                        )
+                self.input_blocks.append(sampling.TimestepEmbedSequential1(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    sampling.TimestepEmbedSequential1(
+                        get_resblock(
+                            merge_factor=merge_factor,
+                            merge_strategy=merge_strategy,
+                            video_kernel_size=video_kernel_size,
+                            ch=ch,
+                            time_embed_dim=time_embed_dim,
+                            dropout=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,
+                        )
+                        if resblock_updown
+                        else ResBlock.Downsample1(
+                            ch,
+                            conv_resample,
+                            dims=dims,
+                            out_channels=out_ch,
+                            dtype=self.dtype,
+                            device=device,
+                            operations=operations,
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+
+        dim_head = ch // num_heads
+        mid_block = [
+            get_resblock(
+                merge_factor=merge_factor,
+                merge_strategy=merge_strategy,
+                video_kernel_size=video_kernel_size,
+                ch=ch,
+                time_embed_dim=time_embed_dim,
+                dropout=dropout,
+                out_channels=None,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+                dtype=self.dtype,
+                device=device,
+                operations=operations,
+            )
+        ]
+
+        self.middle_block = None
+        if transformer_depth_middle >= -1:
+            if transformer_depth_middle >= 0:
+                mid_block += [
+                    get_attention_layer(  # 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_checkpoint=use_checkpoint,
+                    ),
+                    get_resblock(
+                        merge_factor=merge_factor,
+                        merge_strategy=merge_strategy,
+                        video_kernel_size=video_kernel_size,
+                        ch=ch,
+                        time_embed_dim=time_embed_dim,
+                        dropout=dropout,
+                        out_channels=None,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                        dtype=self.dtype,
+                        device=device,
+                        operations=operations,
+                    ),
+                ]
+            self.middle_block = sampling.TimestepEmbedSequential1(*mid_block)
+        self._feature_size += ch
+
+        self.output_blocks = nn.ModuleList([])
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(self.num_res_blocks[level] + 1):
+                ich = input_block_chans.pop()
+                layers = [
+                    get_resblock(
+                        merge_factor=merge_factor,
+                        merge_strategy=merge_strategy,
+                        video_kernel_size=video_kernel_size,
+                        ch=ch + ich,
+                        time_embed_dim=time_embed_dim,
+                        dropout=dropout,
+                        out_channels=model_channels * mult,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                        dtype=self.dtype,
+                        device=device,
+                        operations=operations,
+                    )
+                ]
+                ch = model_channels * mult
+                num_transformers = transformer_depth_output.pop()
+                if num_transformers > 0:
+                    dim_head = ch // num_heads
+                    disabled_sa = False
+
+                    if (
+                        not util.exists(num_attention_blocks)
+                        or i < num_attention_blocks[level]
+                    ):
+                        layers.append(
+                            get_attention_layer(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                depth=num_transformers,
+                                context_dim=context_dim,
+                                disable_self_attn=disabled_sa,
+                                use_checkpoint=use_checkpoint,
+                            )
+                        )
+                if level and i == self.num_res_blocks[level]:
+                    out_ch = ch
+                    layers.append(
+                        get_resblock(
+                            merge_factor=merge_factor,
+                            merge_strategy=merge_strategy,
+                            video_kernel_size=video_kernel_size,
+                            ch=ch,
+                            time_embed_dim=time_embed_dim,
+                            dropout=dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            up=True,
+                            dtype=self.dtype,
+                            device=device,
+                            operations=operations,
+                        )
+                        if resblock_updown
+                        else ResBlock.Upsample1(
+                            ch,
+                            conv_resample,
+                            dims=dims,
+                            out_channels=out_ch,
+                            dtype=self.dtype,
+                            device=device,
+                            operations=operations,
+                        )
+                    )
+                    ds //= 2
+                self.output_blocks.append(sampling.TimestepEmbedSequential1(*layers))
+                self._feature_size += ch
+
+        self.out = nn.Sequential(
+            operations.GroupNorm(32, ch, dtype=self.dtype, device=device),
+            nn.SiLU(),
+            util.zero_module(
+                operations.conv_nd(
+                    dims,
+                    model_channels,
+                    out_channels,
+                    3,
+                    padding=1,
+                    dtype=self.dtype,
+                    device=device,
+                )
+            ),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        timesteps: Optional[torch.Tensor] = None,
+        context: Optional[torch.Tensor] = None,
+        y: Optional[torch.Tensor] = None,
+        control: Optional[torch.Tensor] = None,
+        transformer_options: Dict[str, Any] = {},
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        """#### Forward pass of the UNet model.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+            - `timesteps` (Optional[torch.Tensor], optional): The timesteps tensor. Defaults to None.
+            - `context` (Optional[torch.Tensor], optional): The context tensor. Defaults to None.
+            - `y` (Optional[torch.Tensor], optional): The class labels tensor. Defaults to None.
+            - `control` (Optional[torch.Tensor], optional): The control tensor. Defaults to None.
+            - `transformer_options` (Dict[str, Any], optional): Options for the transformer. Defaults to {}.
+            - `**kwargs` (Any): Additional keyword arguments.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        transformer_options["original_shape"] = list(x.shape)
+        transformer_options["transformer_index"] = 0
+        transformer_patches = transformer_options.get("patches", {})
+
+        num_video_frames = kwargs.get("num_video_frames", self.default_num_video_frames)
+        image_only_indicator = kwargs.get("image_only_indicator", None)
+        time_context = kwargs.get("time_context", None)
+
+        assert (y is not None) == (
+            self.num_classes is not None
+        ), "must specify y if and only if the model is class-conditional"
+        hs = []
+        t_emb = sampling_util.timestep_embedding(
+            timesteps, self.model_channels
+        ).to(x.dtype)
+        emb = self.time_embed(t_emb)
+        h = x
+        for id, module in enumerate(self.input_blocks):
+            transformer_options["block"] = ("input", id)
+            h = ResBlock.forward_timestep_embed1(
+                module,
+                h,
+                emb,
+                context,
+                transformer_options,
+                time_context=time_context,
+                num_video_frames=num_video_frames,
+                image_only_indicator=image_only_indicator,
+            )
+            h = apply_control1(h, control, "input")
+            hs.append(h)
+
+        transformer_options["block"] = ("middle", 0)
+        if self.middle_block is not None:
+            h = ResBlock.forward_timestep_embed1(
+                self.middle_block,
+                h,
+                emb,
+                context,
+                transformer_options,
+                time_context=time_context,
+                num_video_frames=num_video_frames,
+                image_only_indicator=image_only_indicator,
+            )
+        h = apply_control1(h, control, "middle")
+
+        for id, module in enumerate(self.output_blocks):
+            transformer_options["block"] = ("output", id)
+            hsp = hs.pop()
+            hsp = apply_control1(hsp, control, "output")
+
+            h = torch.cat([h, hsp], dim=1)
+            del hsp
+            if len(hs) > 0:
+                output_shape = hs[-1].shape
+            else:
+                output_shape = None
+            h = ResBlock.forward_timestep_embed1(
+                module,
+                h,
+                emb,
+                context,
+                transformer_options,
+                output_shape,
+                time_context=time_context,
+                num_video_frames=num_video_frames,
+                image_only_indicator=image_only_indicator,
+            )
+        h = h.type(x.dtype)
+        return self.out(h)
+
+
+def detect_unet_config(state_dict: Dict[str, torch.Tensor], key_prefix: str) -> Dict[str, Any]:
+    """#### Detect the UNet configuration from a state dictionary.
+
+    #### Args:
+        - `state_dict` (Dict[str, torch.Tensor]): The state dictionary.
+        - `key_prefix` (str): The key prefix.
+
+    #### Returns:
+        - `Dict[str, Any]`: The detected UNet configuration.
+    """
+    state_dict_keys = list(state_dict.keys())
+
+    if (
+        "{}joint_blocks.0.context_block.attn.qkv.weight".format(key_prefix)
+        in state_dict_keys
+    ):  # mmdit model
+        unet_config = {}
+        unet_config["in_channels"] = state_dict[
+            "{}x_embedder.proj.weight".format(key_prefix)
+        ].shape[1]
+        patch_size = state_dict["{}x_embedder.proj.weight".format(key_prefix)].shape[2]
+        unet_config["patch_size"] = patch_size
+        final_layer = "{}final_layer.linear.weight".format(key_prefix)
+        if final_layer in state_dict:
+            unet_config["out_channels"] = state_dict[final_layer].shape[0] // (
+                patch_size * patch_size
+            )
+
+        unet_config["depth"] = (
+            state_dict["{}x_embedder.proj.weight".format(key_prefix)].shape[0] // 64
+        )
+        unet_config["input_size"] = None
+        y_key = "{}y_embedder.mlp.0.weight".format(key_prefix)
+        if y_key in state_dict_keys:
+            unet_config["adm_in_channels"] = state_dict[y_key].shape[1]
+
+        context_key = "{}context_embedder.weight".format(key_prefix)
+        if context_key in state_dict_keys:
+            in_features = state_dict[context_key].shape[1]
+            out_features = state_dict[context_key].shape[0]
+            unet_config["context_embedder_config"] = {
+                "target": "torch.nn.Linear",
+                "params": {"in_features": in_features, "out_features": out_features},
+            }
+        num_patches_key = "{}pos_embed".format(key_prefix)
+        if num_patches_key in state_dict_keys:
+            num_patches = state_dict[num_patches_key].shape[1]
+            unet_config["num_patches"] = num_patches
+            unet_config["pos_embed_max_size"] = round(math.sqrt(num_patches))
+
+        rms_qk = "{}joint_blocks.0.context_block.attn.ln_q.weight".format(key_prefix)
+        if rms_qk in state_dict_keys:
+            unet_config["qk_norm"] = "rms"
+
+        unet_config["pos_embed_scaling_factor"] = None  # unused for inference
+        context_processor = "{}context_processor.layers.0.attn.qkv.weight".format(
+            key_prefix
+        )
+        if context_processor in state_dict_keys:
+            unet_config["context_processor_layers"] = transformer.count_blocks(
+                state_dict_keys,
+                "{}context_processor.layers.".format(key_prefix) + "{}.",
+            )
+        return unet_config
+
+    if "{}clf.1.weight".format(key_prefix) in state_dict_keys:  # stable cascade
+        unet_config = {}
+        text_mapper_name = "{}clip_txt_mapper.weight".format(key_prefix)
+        if text_mapper_name in state_dict_keys:
+            unet_config["stable_cascade_stage"] = "c"
+            w = state_dict[text_mapper_name]
+            if w.shape[0] == 1536:  # stage c lite
+                unet_config["c_cond"] = 1536
+                unet_config["c_hidden"] = [1536, 1536]
+                unet_config["nhead"] = [24, 24]
+                unet_config["blocks"] = [[4, 12], [12, 4]]
+            elif w.shape[0] == 2048:  # stage c full
+                unet_config["c_cond"] = 2048
+        elif "{}clip_mapper.weight".format(key_prefix) in state_dict_keys:
+            unet_config["stable_cascade_stage"] = "b"
+            w = state_dict["{}down_blocks.1.0.channelwise.0.weight".format(key_prefix)]
+            if w.shape[-1] == 640:
+                unet_config["c_hidden"] = [320, 640, 1280, 1280]
+                unet_config["nhead"] = [-1, -1, 20, 20]
+                unet_config["blocks"] = [[2, 6, 28, 6], [6, 28, 6, 2]]
+                unet_config["block_repeat"] = [[1, 1, 1, 1], [3, 3, 2, 2]]
+            elif w.shape[-1] == 576:  # stage b lite
+                unet_config["c_hidden"] = [320, 576, 1152, 1152]
+                unet_config["nhead"] = [-1, 9, 18, 18]
+                unet_config["blocks"] = [[2, 4, 14, 4], [4, 14, 4, 2]]
+                unet_config["block_repeat"] = [[1, 1, 1, 1], [2, 2, 2, 2]]
+        return unet_config
+
+    if (
+        "{}transformer.rotary_pos_emb.inv_freq".format(key_prefix) in state_dict_keys
+    ):  # stable audio dit
+        unet_config = {}
+        unet_config["audio_model"] = "dit1.0"
+        return unet_config
+
+    if (
+        "{}double_layers.0.attn.w1q.weight".format(key_prefix) in state_dict_keys
+    ):  # aura flow dit
+        unet_config = {}
+        unet_config["max_seq"] = state_dict[
+            "{}positional_encoding".format(key_prefix)
+        ].shape[1]
+        unet_config["cond_seq_dim"] = state_dict[
+            "{}cond_seq_linear.weight".format(key_prefix)
+        ].shape[1]
+        double_layers = transformer.count_blocks(
+            state_dict_keys, "{}double_layers.".format(key_prefix) + "{}."
+        )
+        single_layers = transformer.count_blocks(
+            state_dict_keys, "{}single_layers.".format(key_prefix) + "{}."
+        )
+        unet_config["n_double_layers"] = double_layers
+        unet_config["n_layers"] = double_layers + single_layers
+        return unet_config
+
+    if "{}mlp_t5.0.weight".format(key_prefix) in state_dict_keys:  # Hunyuan DiT
+        unet_config = {}
+        unet_config["image_model"] = "hydit"
+        unet_config["depth"] = transformer.count_blocks(
+            state_dict_keys, "{}blocks.".format(key_prefix) + "{}."
+        )
+        unet_config["hidden_size"] = state_dict[
+            "{}x_embedder.proj.weight".format(key_prefix)
+        ].shape[0]
+        if unet_config["hidden_size"] == 1408 and unet_config["depth"] == 40:  # DiT-g/2
+            unet_config["mlp_ratio"] = 4.3637
+        if state_dict["{}extra_embedder.0.weight".format(key_prefix)].shape[1] == 3968:
+            unet_config["size_cond"] = True
+            unet_config["use_style_cond"] = True
+            unet_config["image_model"] = "hydit1"
+        return unet_config
+
+    if (
+        "{}double_blocks.0.img_attn.norm.key_norm.scale".format(key_prefix)
+        in state_dict_keys
+    ):  # Flux
+        dit_config = {}
+        dit_config["image_model"] = "flux"
+        dit_config["in_channels"] = 16
+        dit_config["vec_in_dim"] = 768
+        dit_config["context_in_dim"] = 4096
+        dit_config["hidden_size"] = 3072
+        dit_config["mlp_ratio"] = 4.0
+        dit_config["num_heads"] = 24
+        dit_config["depth"] = transformer.count_blocks(
+            state_dict_keys, "{}double_blocks.".format(key_prefix) + "{}."
+        )
+        dit_config["depth_single_blocks"] = transformer.count_blocks(
+            state_dict_keys, "{}single_blocks.".format(key_prefix) + "{}."
+        )
+        dit_config["axes_dim"] = [16, 56, 56]
+        dit_config["theta"] = 10000
+        dit_config["qkv_bias"] = True
+        dit_config["guidance_embed"] = (
+            "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys
+        )
+        return dit_config
+
+    if "{}input_blocks.0.0.weight".format(key_prefix) not in state_dict_keys:
+        return None
+
+    unet_config = {
+        "use_checkpoint": False,
+        "image_size": 32,
+        "use_spatial_transformer": True,
+        "legacy": False,
+    }
+
+    y_input = "{}label_emb.0.0.weight".format(key_prefix)
+    if y_input in state_dict_keys:
+        unet_config["num_classes"] = "sequential"
+        unet_config["adm_in_channels"] = state_dict[y_input].shape[1]
+    else:
+        unet_config["adm_in_channels"] = None
+
+    model_channels = state_dict["{}input_blocks.0.0.weight".format(key_prefix)].shape[0]
+    in_channels = state_dict["{}input_blocks.0.0.weight".format(key_prefix)].shape[1]
+
+    out_key = "{}out.2.weight".format(key_prefix)
+    if out_key in state_dict:
+        out_channels = state_dict[out_key].shape[0]
+    else:
+        out_channels = 4
+
+    num_res_blocks = []
+    channel_mult = []
+    transformer_depth = []
+    transformer_depth_output = []
+    context_dim = None
+    use_linear_in_transformer = False
+
+    video_model = False
+    video_model_cross = False
+
+    current_res = 1
+    count = 0
+
+    last_res_blocks = 0
+    last_channel_mult = 0
+
+    input_block_count = transformer.count_blocks(
+        state_dict_keys, "{}input_blocks".format(key_prefix) + ".{}."
+    )
+    for count in range(input_block_count):
+        prefix = "{}input_blocks.{}.".format(key_prefix, count)
+        prefix_output = "{}output_blocks.{}.".format(
+            key_prefix, input_block_count - count - 1
+        )
+
+        block_keys = sorted(
+            list(filter(lambda a: a.startswith(prefix), state_dict_keys))
+        )
+        if len(block_keys) == 0:
+            break
+
+        block_keys_output = sorted(
+            list(filter(lambda a: a.startswith(prefix_output), state_dict_keys))
+        )
+
+        if "{}0.op.weight".format(prefix) in block_keys:  # new layer
+            num_res_blocks.append(last_res_blocks)
+            channel_mult.append(last_channel_mult)
+
+            current_res *= 2
+            last_res_blocks = 0
+            last_channel_mult = 0
+            out = transformer.calculate_transformer_depth(
+                prefix_output, state_dict_keys, state_dict
+            )
+            if out is not None:
+                transformer_depth_output.append(out[0])
+            else:
+                transformer_depth_output.append(0)
+        else:
+            res_block_prefix = "{}0.in_layers.0.weight".format(prefix)
+            if res_block_prefix in block_keys:
+                last_res_blocks += 1
+                last_channel_mult = (
+                    state_dict["{}0.out_layers.3.weight".format(prefix)].shape[0]
+                    // model_channels
+                )
+
+                out = transformer.calculate_transformer_depth(prefix, state_dict_keys, state_dict)
+                if out is not None:
+                    transformer_depth.append(out[0])
+                    if context_dim is None:
+                        context_dim = out[1]
+                        use_linear_in_transformer = out[2]
+                        out[3]
+                else:
+                    transformer_depth.append(0)
+
+            res_block_prefix = "{}0.in_layers.0.weight".format(prefix_output)
+            if res_block_prefix in block_keys_output:
+                out = transformer.calculate_transformer_depth(
+                    prefix_output, state_dict_keys, state_dict
+                )
+                if out is not None:
+                    transformer_depth_output.append(out[0])
+                else:
+                    transformer_depth_output.append(0)
+
+    num_res_blocks.append(last_res_blocks)
+    channel_mult.append(last_channel_mult)
+    if "{}middle_block.1.proj_in.weight".format(key_prefix) in state_dict_keys:
+        transformer_depth_middle = transformer.count_blocks(
+            state_dict_keys,
+            "{}middle_block.1.transformer_blocks.".format(key_prefix) + "{}",
+        )
+    elif "{}middle_block.0.in_layers.0.weight".format(key_prefix) in state_dict_keys:
+        transformer_depth_middle = -1
+    else:
+        transformer_depth_middle = -2
+
+    unet_config["in_channels"] = in_channels
+    unet_config["out_channels"] = out_channels
+    unet_config["model_channels"] = model_channels
+    unet_config["num_res_blocks"] = num_res_blocks
+    unet_config["transformer_depth"] = transformer_depth
+    unet_config["transformer_depth_output"] = transformer_depth_output
+    unet_config["channel_mult"] = channel_mult
+    unet_config["transformer_depth_middle"] = transformer_depth_middle
+    unet_config["use_linear_in_transformer"] = use_linear_in_transformer
+    unet_config["context_dim"] = context_dim
+
+    if video_model:
+        unet_config["extra_ff_mix_layer"] = True
+        unet_config["use_spatial_context"] = True
+        unet_config["merge_strategy"] = "learned_with_images"
+        unet_config["merge_factor"] = 0.0
+        unet_config["video_kernel_size"] = [3, 1, 1]
+        unet_config["use_temporal_resblock"] = True
+        unet_config["use_temporal_attention"] = True
+        unet_config["disable_temporal_crossattention"] = not video_model_cross
+    else:
+        unet_config["use_temporal_resblock"] = False
+        unet_config["use_temporal_attention"] = False
+
+    return unet_config
+
+
+def model_config_from_unet_config(unet_config: Dict[str, Any], state_dict: Optional[Dict[str, torch.Tensor]] = None) -> Any:
+    """#### Get the model configuration from a UNet configuration.
+
+    #### Args:
+        - `unet_config` (Dict[str, Any]): The UNet configuration.
+        - `state_dict` (Optional[Dict[str, torch.Tensor]], optional): The state dictionary. Defaults to None.
+
+    #### Returns:
+        - `Any`: The model configuration.
+    """
+    from modules.SD15 import SD15
+
+    for model_config in SD15.models:
+        if model_config.matches(unet_config, state_dict):
+            return model_config(unet_config)
+
+    logging.error("no match {}".format(unet_config))
+    return None
+
+
+def model_config_from_unet(state_dict: Dict[str, torch.Tensor], unet_key_prefix: str, use_base_if_no_match: bool = False) -> Any:
+    """#### Get the model configuration from a UNet state dictionary.
+
+    #### Args:
+        - `state_dict` (Dict[str, torch.Tensor]): The state dictionary.
+        - `unet_key_prefix` (str): The UNet key prefix.
+        - `use_base_if_no_match` (bool, optional): Whether to use the base configuration if no match is found. Defaults to False.
+
+    #### Returns:
+        - `Any`: The model configuration.
+    """
+    unet_config = detect_unet_config(state_dict, unet_key_prefix)
+    if unet_config is None:
+        return None
+    model_config = model_config_from_unet_config(unet_config, state_dict)
+    return model_config
+
+
+def unet_dtype1(
+    device: Optional[torch.device] = None,
+    model_params: int = 0,
+    supported_dtypes: List[torch.dtype] = [torch.float16, torch.bfloat16, torch.float32],
+) -> torch.dtype:
+    """#### Get the dtype for the UNet model.
+
+    #### Args:
+        - `device` (Optional[torch.device], optional): The device. Defaults to None.
+        - `model_params` (int, optional): The model parameters. Defaults to 0.
+        - `supported_dtypes` (List[torch.dtype], optional): The supported dtypes. Defaults to [torch.float16, torch.bfloat16, torch.float32].
+
+    #### Returns:
+        - `torch.dtype`: The dtype for the UNet model.
+    """
+    return torch.float16

modules/Quantize/Quantizer.py

@@ -0,0 +1,1012 @@
+import copy
+import logging
+import gguf
+import torch
+
+from modules.Device import Device
+from modules.Model import ModelPatcher
+from modules.Utilities import util
+from modules.clip import Clip
+from modules.cond import cast
+
+# Constants for torch-compatible quantization types
+TORCH_COMPATIBLE_QTYPES = {
+    None,
+    gguf.GGMLQuantizationType.F32,
+    gguf.GGMLQuantizationType.F16,
+}
+
+
+def is_torch_compatible(tensor: torch.Tensor) -> bool:
+    """#### Check if a tensor is compatible with PyTorch operations.
+
+    #### Args:
+        - `tensor` (torch.Tensor): The tensor to check.
+
+    #### Returns:
+        - `bool`: Whether the tensor is torch-compatible.
+    """
+    return (
+        tensor is None
+        or getattr(tensor, "tensor_type", None) in TORCH_COMPATIBLE_QTYPES
+    )
+
+
+def is_quantized(tensor: torch.Tensor) -> bool:
+    """#### Check if a tensor is quantized.
+
+    #### Args:
+        - `tensor` (torch.Tensor): The tensor to check.
+
+    #### Returns:
+        - `bool`: Whether the tensor is quantized.
+    """
+    return not is_torch_compatible(tensor)
+
+
+def dequantize(
+    data: torch.Tensor,
+    qtype: gguf.GGMLQuantizationType,
+    oshape: tuple,
+    dtype: torch.dtype = None,
+) -> torch.Tensor:
+    """#### Dequantize tensor back to usable shape/dtype.
+
+    #### Args:
+        - `data` (torch.Tensor): The quantized data.
+        - `qtype` (gguf.GGMLQuantizationType): The quantization type.
+        - `oshape` (tuple): The output shape.
+        - `dtype` (torch.dtype, optional): The output dtype. Defaults to None.
+
+    #### Returns:
+        - `torch.Tensor`: The dequantized tensor.
+    """
+    # Get block size and type size for quantization format
+    block_size, type_size = gguf.GGML_QUANT_SIZES[qtype]
+    dequantize_blocks = dequantize_functions[qtype]
+
+    # Reshape data into blocks
+    rows = data.reshape((-1, data.shape[-1])).view(torch.uint8)
+    n_blocks = rows.numel() // type_size
+    blocks = rows.reshape((n_blocks, type_size))
+
+    # Dequantize blocks and reshape to target shape
+    blocks = dequantize_blocks(blocks, block_size, type_size, dtype)
+    return blocks.reshape(oshape)
+
+
+def split_block_dims(blocks: torch.Tensor, *args) -> list:
+    """#### Split blocks into dimensions.
+
+    #### Args:
+        - `blocks` (torch.Tensor): The blocks to split.
+        - `*args`: The dimensions to split into.
+
+    #### Returns:
+        - `list`: The split blocks.
+    """
+    n_max = blocks.shape[1]
+    dims = list(args) + [n_max - sum(args)]
+    return torch.split(blocks, dims, dim=1)
+
+
+# Legacy Quantization Functions
+def dequantize_blocks_Q8_0(
+    blocks: torch.Tensor, block_size: int, type_size: int, dtype: torch.dtype = None
+) -> torch.Tensor:
+    """#### Dequantize Q8_0 quantized blocks.
+
+    #### Args:
+        - `blocks` (torch.Tensor): The quantized blocks.
+        - `block_size` (int): The block size.
+        - `type_size` (int): The type size.
+        - `dtype` (torch.dtype, optional): The output dtype. Defaults to None.
+
+    #### Returns:
+        - `torch.Tensor`: The dequantized blocks.
+    """
+    # Split blocks into scale and quantized values
+    d, x = split_block_dims(blocks, 2)
+    d = d.view(torch.float16).to(dtype)
+    x = x.view(torch.int8)
+    return d * x
+
+
+# K Quants #
+QK_K = 256
+K_SCALE_SIZE = 12
+
+# Mapping of quantization types to dequantization functions
+dequantize_functions = {
+    gguf.GGMLQuantizationType.Q8_0: dequantize_blocks_Q8_0,
+}
+
+
+def dequantize_tensor(
+    tensor: torch.Tensor, dtype: torch.dtype = None, dequant_dtype: torch.dtype = None
+) -> torch.Tensor:
+    """#### Dequantize a potentially quantized tensor.
+
+    #### Args:
+        - `tensor` (torch.Tensor): The tensor to dequantize.
+        - `dtype` (torch.dtype, optional): Target dtype. Defaults to None.
+        - `dequant_dtype` (torch.dtype, optional): Intermediate dequantization dtype. Defaults to None.
+
+    #### Returns:
+        - `torch.Tensor`: The dequantized tensor.
+    """
+    qtype = getattr(tensor, "tensor_type", None)
+    oshape = getattr(tensor, "tensor_shape", tensor.shape)
+
+    if qtype in TORCH_COMPATIBLE_QTYPES:
+        return tensor.to(dtype)
+    elif qtype in dequantize_functions:
+        dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype
+        return dequantize(tensor.data, qtype, oshape, dtype=dequant_dtype).to(dtype)
+
+
+class GGMLLayer(torch.nn.Module):
+    """#### Base class for GGML quantized layers.
+
+    Handles dynamic dequantization of weights during forward pass.
+    """
+
+    comfy_cast_weights: bool = True
+    dequant_dtype: torch.dtype = None
+    patch_dtype: torch.dtype = None
+    torch_compatible_tensor_types: set = {
+        None,
+        gguf.GGMLQuantizationType.F32,
+        gguf.GGMLQuantizationType.F16,
+    }
+
+    def is_ggml_quantized(
+        self, *, weight: torch.Tensor = None, bias: torch.Tensor = None
+    ) -> bool:
+        """#### Check if layer weights are GGML quantized.
+
+        #### Args:
+            - `weight` (torch.Tensor, optional): Weight tensor to check. Defaults to self.weight.
+            - `bias` (torch.Tensor, optional): Bias tensor to check. Defaults to self.bias.
+
+        #### Returns:
+            - `bool`: Whether weights are quantized.
+        """
+        if weight is None:
+            weight = self.weight
+        if bias is None:
+            bias = self.bias
+        return is_quantized(weight) or is_quantized(bias)
+
+    def _load_from_state_dict(
+        self, state_dict: dict, prefix: str, *args, **kwargs
+    ) -> None:
+        """#### Load quantized weights from state dict.
+
+        #### Args:
+            - `state_dict` (dict): State dictionary.
+            - `prefix` (str): Key prefix.
+            - `*args`: Additional arguments.
+            - `**kwargs`: Additional keyword arguments.
+        """
+        weight = state_dict.get(f"{prefix}weight")
+        bias = state_dict.get(f"{prefix}bias")
+        # Use modified loader for quantized or linear layers
+        if self.is_ggml_quantized(weight=weight, bias=bias) or isinstance(
+            self, torch.nn.Linear
+        ):
+            return self.ggml_load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+    def ggml_load_from_state_dict(
+        self,
+        state_dict: dict,
+        prefix: str,
+        local_metadata: dict,
+        strict: bool,
+        missing_keys: list,
+        unexpected_keys: list,
+        error_msgs: list,
+    ) -> None:
+        """#### Load GGML quantized weights from state dict.
+
+        #### Args:
+            - `state_dict` (dict): State dictionary.
+            - `prefix` (str): Key prefix.
+            - `local_metadata` (dict): Local metadata.
+            - `strict` (bool): Strict loading mode.
+            - `missing_keys` (list): Keys missing from state dict.
+            - `unexpected_keys` (list): Unexpected keys found.
+            - `error_msgs` (list): Error messages.
+        """
+        prefix_len = len(prefix)
+        for k, v in state_dict.items():
+            if k[prefix_len:] == "weight":
+                self.weight = torch.nn.Parameter(v, requires_grad=False)
+            elif k[prefix_len:] == "bias" and v is not None:
+                self.bias = torch.nn.Parameter(v, requires_grad=False)
+            else:
+                missing_keys.append(k)
+
+    def _save_to_state_dict(self, *args, **kwargs) -> None:
+        """#### Save layer state to state dict.
+
+        #### Args:
+            - `*args`: Additional arguments.
+            - `**kwargs`: Additional keyword arguments.
+        """
+        if self.is_ggml_quantized():
+            return self.ggml_save_to_state_dict(*args, **kwargs)
+        return super()._save_to_state_dict(*args, **kwargs)
+
+    def ggml_save_to_state_dict(
+        self, destination: dict, prefix: str, keep_vars: bool
+    ) -> None:
+        """#### Save GGML layer state to state dict.
+
+        #### Args:
+            - `destination` (dict): Destination dictionary.
+            - `prefix` (str): Key prefix.
+            - `keep_vars` (bool): Whether to keep variables.
+        """
+        # Create fake tensors for VRAM estimation
+        weight = torch.zeros_like(self.weight, device=torch.device("meta"))
+        destination[prefix + "weight"] = weight
+        if self.bias is not None:
+            bias = torch.zeros_like(self.bias, device=torch.device("meta"))
+            destination[prefix + "bias"] = bias
+        return
+
+    def get_weight(self, tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+        """#### Get dequantized weight tensor.
+
+        #### Args:
+            - `tensor` (torch.Tensor): Input tensor.
+            - `dtype` (torch.dtype): Target dtype.
+
+        #### Returns:
+            - `torch.Tensor`: Dequantized tensor.
+        """
+        if tensor is None:
+            return
+
+        # Consolidate and load patches to GPU asynchronously
+        patch_list = []
+        device = tensor.device
+        for function, patches, key in getattr(tensor, "patches", []):
+            patch_list += move_patch_to_device(patches, device)
+
+        # Dequantize tensor while patches load
+        weight = dequantize_tensor(tensor, dtype, self.dequant_dtype)
+
+        # Apply patches
+        if patch_list:
+            if self.patch_dtype is None:
+                weight = function(patch_list, weight, key)
+            else:
+                # For testing, may degrade image quality
+                patch_dtype = (
+                    dtype if self.patch_dtype == "target" else self.patch_dtype
+                )
+                weight = function(patch_list, weight, key, patch_dtype)
+        return weight
+
+    def cast_bias_weight(
+        self,
+        input: torch.Tensor = None,
+        dtype: torch.dtype = None,
+        device: torch.device = None,
+        bias_dtype: torch.dtype = None,
+    ) -> tuple:
+        """#### Cast layer weights and bias to target dtype/device.
+
+        #### Args:
+            - `input` (torch.Tensor, optional): Input tensor for type/device inference.
+            - `dtype` (torch.dtype, optional): Target dtype.
+            - `device` (torch.device, optional): Target device.
+            - `bias_dtype` (torch.dtype, optional): Target bias dtype.
+
+        #### Returns:
+            - `tuple`: (cast_weight, cast_bias)
+        """
+        if input is not None:
+            if dtype is None:
+                dtype = getattr(input, "dtype", torch.float32)
+            if bias_dtype is None:
+                bias_dtype = dtype
+            if device is None:
+                device = input.device
+
+        bias = None
+        non_blocking = Device.device_supports_non_blocking(device)
+        if self.bias is not None:
+            bias = self.get_weight(self.bias.to(device), dtype)
+            bias = cast.cast_to(
+                bias, bias_dtype, device, non_blocking=non_blocking, copy=False
+            )
+
+        weight = self.get_weight(self.weight.to(device), dtype)
+        weight = cast.cast_to(
+            weight, dtype, device, non_blocking=non_blocking, copy=False
+        )
+        return weight, bias
+
+    def forward_comfy_cast_weights(
+        self, input: torch.Tensor, *args, **kwargs
+    ) -> torch.Tensor:
+        """#### Forward pass with weight casting.
+
+        #### Args:
+            - `input` (torch.Tensor): Input tensor.
+            - `*args`: Additional arguments.
+            - `**kwargs`: Additional keyword arguments.
+
+        #### Returns:
+            - `torch.Tensor`: Output tensor.
+        """
+        if self.is_ggml_quantized():
+            return self.forward_ggml_cast_weights(input, *args, **kwargs)
+        return super().forward_comfy_cast_weights(input, *args, **kwargs)
+
+
+class GGMLOps(cast.manual_cast):
+    """
+    Dequantize weights on the fly before doing the compute
+    """
+
+    class Linear(GGMLLayer, cast.manual_cast.Linear):
+        def __init__(
+            self, in_features, out_features, bias=True, device=None, dtype=None
+        ):
+            """
+            Initialize the Linear layer.
+
+            Args:
+                in_features (int): Number of input features.
+                out_features (int): Number of output features.
+                bias (bool, optional): If set to False, the layer will not learn an additive bias. Defaults to True.
+                device (torch.device, optional): The device to store the layer's parameters. Defaults to None.
+                dtype (torch.dtype, optional): The data type of the layer's parameters. Defaults to None.
+            """
+            torch.nn.Module.__init__(self)
+            # TODO: better workaround for reserved memory spike on windows
+            # Issue is with `torch.empty` still reserving the full memory for the layer
+            # Windows doesn't over-commit memory so without this 24GB+ of pagefile is used
+            self.in_features = in_features
+            self.out_features = out_features
+            self.weight = None
+            self.bias = None
+
+        def forward_ggml_cast_weights(self, input: torch.Tensor) -> torch.Tensor:
+            """
+            Forward pass with GGML cast weights.
+
+            Args:
+                input (torch.Tensor): The input tensor.
+
+            Returns:
+                torch.Tensor: The output tensor.
+            """
+            weight, bias = self.cast_bias_weight(input)
+            return torch.nn.functional.linear(input, weight, bias)
+
+    class Embedding(GGMLLayer, cast.manual_cast.Embedding):
+        def forward_ggml_cast_weights(
+            self, input: torch.Tensor, out_dtype: torch.dtype = None
+        ) -> torch.Tensor:
+            """
+            Forward pass with GGML cast weights for embedding.
+
+            Args:
+                input (torch.Tensor): The input tensor.
+                out_dtype (torch.dtype, optional): The output data type. Defaults to None.
+
+            Returns:
+                torch.Tensor: The output tensor.
+            """
+            output_dtype = out_dtype
+            if (
+                self.weight.dtype == torch.float16
+                or self.weight.dtype == torch.bfloat16
+            ):
+                out_dtype = None
+            weight, _bias = self.cast_bias_weight(
+                self, device=input.device, dtype=out_dtype
+            )
+            return torch.nn.functional.embedding(
+                input,
+                weight,
+                self.padding_idx,
+                self.max_norm,
+                self.norm_type,
+                self.scale_grad_by_freq,
+                self.sparse,
+            ).to(dtype=output_dtype)
+
+
+def gguf_sd_loader_get_orig_shape(
+    reader: gguf.GGUFReader, tensor_name: str
+) -> torch.Size:
+    """#### Get the original shape of a tensor from a GGUF reader.
+
+    #### Args:
+        - `reader` (gguf.GGUFReader): The GGUF reader.
+        - `tensor_name` (str): The name of the tensor.
+
+    #### Returns:
+        - `torch.Size`: The original shape of the tensor.
+    """
+    field_key = f"comfy.gguf.orig_shape.{tensor_name}"
+    field = reader.get_field(field_key)
+    if field is None:
+        return None
+    # Has original shape metadata, so we try to decode it.
+    if (
+        len(field.types) != 2
+        or field.types[0] != gguf.GGUFValueType.ARRAY
+        or field.types[1] != gguf.GGUFValueType.INT32
+    ):
+        raise TypeError(
+            f"Bad original shape metadata for {field_key}: Expected ARRAY of INT32, got {field.types}"
+        )
+    return torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data))
+
+
+class GGMLTensor(torch.Tensor):
+    """
+    Main tensor-like class for storing quantized weights
+    """
+
+    def __init__(self, *args, tensor_type, tensor_shape, patches=[], **kwargs):
+        """
+        Initialize the GGMLTensor.
+
+        Args:
+            *args: Variable length argument list.
+            tensor_type: The type of the tensor.
+            tensor_shape: The shape of the tensor.
+            patches (list, optional): List of patches. Defaults to [].
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super().__init__()
+        self.tensor_type = tensor_type
+        self.tensor_shape = tensor_shape
+        self.patches = patches
+
+    def __new__(cls, *args, tensor_type, tensor_shape, patches=[], **kwargs):
+        """
+        Create a new instance of GGMLTensor.
+
+        Args:
+            *args: Variable length argument list.
+            tensor_type: The type of the tensor.
+            tensor_shape: The shape of the tensor.
+            patches (list, optional): List of patches. Defaults to [].
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            GGMLTensor: A new instance of GGMLTensor.
+        """
+        return super().__new__(cls, *args, **kwargs)
+
+    def to(self, *args, **kwargs):
+        """
+        Convert the tensor to a specified device and/or dtype.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            GGMLTensor: The converted tensor.
+        """
+        new = super().to(*args, **kwargs)
+        new.tensor_type = getattr(self, "tensor_type", None)
+        new.tensor_shape = getattr(self, "tensor_shape", new.data.shape)
+        new.patches = getattr(self, "patches", []).copy()
+        return new
+
+    def clone(self, *args, **kwargs):
+        """
+        Clone the tensor.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            GGMLTensor: The cloned tensor.
+        """
+        return self
+
+    def detach(self, *args, **kwargs):
+        """
+        Detach the tensor from the computation graph.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            GGMLTensor: The detached tensor.
+        """
+        return self
+
+    def copy_(self, *args, **kwargs):
+        """
+        Copy the values from another tensor into this tensor.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            GGMLTensor: The tensor with copied values.
+        """
+        try:
+            return super().copy_(*args, **kwargs)
+        except Exception as e:
+            print(f"ignoring 'copy_' on tensor: {e}")
+
+    def __deepcopy__(self, *args, **kwargs):
+        """
+        Create a deep copy of the tensor.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            GGMLTensor: The deep copied tensor.
+        """
+        new = super().__deepcopy__(*args, **kwargs)
+        new.tensor_type = getattr(self, "tensor_type", None)
+        new.tensor_shape = getattr(self, "tensor_shape", new.data.shape)
+        new.patches = getattr(self, "patches", []).copy()
+        return new
+
+    @property
+    def shape(self):
+        """
+        Get the shape of the tensor.
+
+        Returns:
+            torch.Size: The shape of the tensor.
+        """
+        if not hasattr(self, "tensor_shape"):
+            self.tensor_shape = self.size()
+        return self.tensor_shape
+
+
+def gguf_sd_loader(path: str, handle_prefix: str = "model.diffusion_model."):
+    """#### Load a GGUF file into a state dict.
+
+    #### Args:
+        - `path` (str): The path to the GGUF file.
+        - `handle_prefix` (str, optional): The prefix to handle. Defaults to "model.diffusion_model.".
+
+    #### Returns:
+        - `dict`: The loaded state dict.
+    """
+    reader = gguf.GGUFReader(path)
+
+    # filter and strip prefix
+    has_prefix = False
+    if handle_prefix is not None:
+        prefix_len = len(handle_prefix)
+        tensor_names = set(tensor.name for tensor in reader.tensors)
+        has_prefix = any(s.startswith(handle_prefix) for s in tensor_names)
+
+    tensors = []
+    for tensor in reader.tensors:
+        sd_key = tensor_name = tensor.name
+        if has_prefix:
+            if not tensor_name.startswith(handle_prefix):
+                continue
+            sd_key = tensor_name[prefix_len:]
+        tensors.append((sd_key, tensor))
+
+    # detect and verify architecture
+    compat = None
+    arch_str = None
+    arch_field = reader.get_field("general.architecture")
+    if arch_field is not None:
+        if (
+            len(arch_field.types) != 1
+            or arch_field.types[0] != gguf.GGUFValueType.STRING
+        ):
+            raise TypeError(
+                f"Bad type for GGUF general.architecture key: expected string, got {arch_field.types!r}"
+            )
+        arch_str = str(arch_field.parts[arch_field.data[-1]], encoding="utf-8")
+        if arch_str not in {"flux", "sd1", "sdxl", "t5", "t5encoder"}:
+            raise ValueError(
+                f"Unexpected architecture type in GGUF file, expected one of flux, sd1, sdxl, t5encoder but got {arch_str!r}"
+            )
+
+    # main loading loop
+    state_dict = {}
+    qtype_dict = {}
+    for sd_key, tensor in tensors:
+        tensor_name = tensor.name
+        tensor_type_str = str(tensor.tensor_type)
+        torch_tensor = torch.from_numpy(tensor.data)  # mmap
+
+        shape = gguf_sd_loader_get_orig_shape(reader, tensor_name)
+        if shape is None:
+            shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
+            # Workaround for stable-diffusion.cpp SDXL detection.
+            if compat == "sd.cpp" and arch_str == "sdxl":
+                if any(
+                    [
+                        tensor_name.endswith(x)
+                        for x in (".proj_in.weight", ".proj_out.weight")
+                    ]
+                ):
+                    while len(shape) > 2 and shape[-1] == 1:
+                        shape = shape[:-1]
+
+        # add to state dict
+        if tensor.tensor_type in {
+            gguf.GGMLQuantizationType.F32,
+            gguf.GGMLQuantizationType.F16,
+        }:
+            torch_tensor = torch_tensor.view(*shape)
+        state_dict[sd_key] = GGMLTensor(
+            torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape
+        )
+        qtype_dict[tensor_type_str] = qtype_dict.get(tensor_type_str, 0) + 1
+
+    # sanity check debug print
+    print("\nggml_sd_loader:")
+    for k, v in qtype_dict.items():
+        print(f" {k:30}{v:3}")
+
+    return state_dict
+
+
+class GGUFModelPatcher(ModelPatcher.ModelPatcher):
+    patch_on_device = False
+
+    def unpatch_model(self, device_to=None, unpatch_weights=True):
+        """
+        Unpatch the model.
+
+        Args:
+            device_to (torch.device, optional): The device to move the model to. Defaults to None.
+            unpatch_weights (bool, optional): Whether to unpatch the weights. Defaults to True.
+
+        Returns:
+            GGUFModelPatcher: The unpatched model.
+        """
+        if unpatch_weights:
+            for p in self.model.parameters():
+                if is_torch_compatible(p):
+                    continue
+                patches = getattr(p, "patches", [])
+                if len(patches) > 0:
+                    p.patches = []
+        self.object_patches = {}
+        # TODO: Find another way to not unload after patches
+        return super().unpatch_model(
+            device_to=device_to, unpatch_weights=unpatch_weights
+        )
+
+    mmap_released = False
+
+    def load(self, *args, force_patch_weights=False, **kwargs):
+        """
+        Load the model.
+
+        Args:
+            *args: Variable length argument list.
+            force_patch_weights (bool, optional): Whether to force patch weights. Defaults to False.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super().load(*args, force_patch_weights=True, **kwargs)
+
+        # make sure nothing stays linked to mmap after first load
+        if not self.mmap_released:
+            linked = []
+            if kwargs.get("lowvram_model_memory", 0) > 0:
+                for n, m in self.model.named_modules():
+                    if hasattr(m, "weight"):
+                        device = getattr(m.weight, "device", None)
+                        if device == self.offload_device:
+                            linked.append((n, m))
+                            continue
+                    if hasattr(m, "bias"):
+                        device = getattr(m.bias, "device", None)
+                        if device == self.offload_device:
+                            linked.append((n, m))
+                            continue
+            if linked:
+                print(f"Attempting to release mmap ({len(linked)})")
+                for n, m in linked:
+                    # TODO: possible to OOM, find better way to detach
+                    m.to(self.load_device).to(self.offload_device)
+            self.mmap_released = True
+
+    def add_object_patch(self, name, obj):
+            self.object_patches[name] = obj
+
+    def clone(self, *args, **kwargs):
+        """
+        Clone the model patcher.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            GGUFModelPatcher: The cloned model patcher.
+        """
+        n = GGUFModelPatcher(
+            self.model,
+            self.load_device,
+            self.offload_device,
+            self.size,
+            weight_inplace_update=self.weight_inplace_update,
+        )
+        n.patches = {}
+        for k in self.patches:
+            n.patches[k] = self.patches[k][:]
+        n.patches_uuid = self.patches_uuid
+
+        n.object_patches = self.object_patches.copy()
+        n.model_options = copy.deepcopy(self.model_options)
+        n.backup = self.backup
+        n.object_patches_backup = self.object_patches_backup
+        n.patch_on_device = getattr(self, "patch_on_device", False)
+        return n
+
+
+class UnetLoaderGGUF:
+    def load_unet(
+        self,
+        unet_name: str,
+        dequant_dtype: str = None,
+        patch_dtype: str = None,
+        patch_on_device: bool = None,
+    ) -> tuple:
+        """
+        Load the UNet model.
+
+        Args:
+            unet_name (str): The name of the UNet model.
+            dequant_dtype (str, optional): The dequantization data type. Defaults to None.
+            patch_dtype (str, optional): The patch data type. Defaults to None.
+            patch_on_device (bool, optional): Whether to patch on device. Defaults to None.
+
+        Returns:
+            tuple: The loaded model.
+        """
+        ops = GGMLOps()
+
+        if dequant_dtype in ("default", None):
+            ops.Linear.dequant_dtype = None
+        elif dequant_dtype in ["target"]:
+            ops.Linear.dequant_dtype = dequant_dtype
+        else:
+            ops.Linear.dequant_dtype = getattr(torch, dequant_dtype)
+
+        if patch_dtype in ("default", None):
+            ops.Linear.patch_dtype = None
+        elif patch_dtype in ["target"]:
+            ops.Linear.patch_dtype = patch_dtype
+        else:
+            ops.Linear.patch_dtype = getattr(torch, patch_dtype)
+
+        unet_path = "./_internal/unet/" + unet_name
+        sd = gguf_sd_loader(unet_path)
+        model = ModelPatcher.load_diffusion_model_state_dict(
+            sd, model_options={"custom_operations": ops}
+        )
+        if model is None:
+            logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path))
+            raise RuntimeError(
+                "ERROR: Could not detect model type of: {}".format(unet_path)
+            )
+        model = GGUFModelPatcher.clone(model)
+        model.patch_on_device = patch_on_device
+        return (model,)
+
+
+clip_sd_map = {
+    "enc.": "encoder.",
+    ".blk.": ".block.",
+    "token_embd": "shared",
+    "output_norm": "final_layer_norm",
+    "attn_q": "layer.0.SelfAttention.q",
+    "attn_k": "layer.0.SelfAttention.k",
+    "attn_v": "layer.0.SelfAttention.v",
+    "attn_o": "layer.0.SelfAttention.o",
+    "attn_norm": "layer.0.layer_norm",
+    "attn_rel_b": "layer.0.SelfAttention.relative_attention_bias",
+    "ffn_up": "layer.1.DenseReluDense.wi_1",
+    "ffn_down": "layer.1.DenseReluDense.wo",
+    "ffn_gate": "layer.1.DenseReluDense.wi_0",
+    "ffn_norm": "layer.1.layer_norm",
+}
+
+clip_name_dict = {
+    "stable_diffusion": Clip.CLIPType.STABLE_DIFFUSION,
+    "sdxl": Clip.CLIPType.STABLE_DIFFUSION,
+    "sd3": Clip.CLIPType.SD3,
+    "flux": Clip.CLIPType.FLUX,
+}
+
+
+def gguf_clip_loader(path: str) -> dict:
+    """#### Load a CLIP model from a GGUF file.
+
+    #### Args:
+        - `path` (str): The path to the GGUF file.
+
+    #### Returns:
+        - `dict`: The loaded CLIP model.
+    """
+    raw_sd = gguf_sd_loader(path)
+    assert "enc.blk.23.ffn_up.weight" in raw_sd, "Invalid Text Encoder!"
+    sd = {}
+    for k, v in raw_sd.items():
+        for s, d in clip_sd_map.items():
+            k = k.replace(s, d)
+        sd[k] = v
+    return sd
+
+
+class CLIPLoaderGGUF:
+    def load_data(self, ckpt_paths: list) -> list:
+        """
+        Load data from checkpoint paths.
+
+        Args:
+            ckpt_paths (list): List of checkpoint paths.
+
+        Returns:
+            list: List of loaded data.
+        """
+        clip_data = []
+        for p in ckpt_paths:
+            if p.endswith(".gguf"):
+                clip_data.append(gguf_clip_loader(p))
+            else:
+                sd = util.load_torch_file(p, safe_load=True)
+                clip_data.append(
+                    {
+                        k: GGMLTensor(
+                            v,
+                            tensor_type=gguf.GGMLQuantizationType.F16,
+                            tensor_shape=v.shape,
+                        )
+                        for k, v in sd.items()
+                    }
+                )
+        return clip_data
+
+    def load_patcher(self, clip_paths: list, clip_type: str, clip_data: list) -> Clip:
+        """
+        Load the model patcher.
+
+        Args:
+            clip_paths (list): List of clip paths.
+            clip_type (str): The type of the clip.
+            clip_data (list): List of clip data.
+
+        Returns:
+            Clip: The loaded clip.
+        """
+        clip = Clip.load_text_encoder_state_dicts(
+            clip_type=clip_type,
+            state_dicts=clip_data,
+            model_options={
+                "custom_operations": GGMLOps,
+                "initial_device": Device.text_encoder_offload_device(),
+            },
+            embedding_directory="models/embeddings",
+        )
+        clip.patcher = GGUFModelPatcher.clone(clip.patcher)
+
+        # for some reason this is just missing in some SAI checkpoints
+        if getattr(clip.cond_stage_model, "clip_l", None) is not None:
+            if (
+                getattr(
+                    clip.cond_stage_model.clip_l.transformer.text_projection.weight,
+                    "tensor_shape",
+                    None,
+                )
+                is None
+            ):
+                clip.cond_stage_model.clip_l.transformer.text_projection = (
+                    cast.manual_cast.Linear(768, 768)
+                )
+        if getattr(clip.cond_stage_model, "clip_g", None) is not None:
+            if (
+                getattr(
+                    clip.cond_stage_model.clip_g.transformer.text_projection.weight,
+                    "tensor_shape",
+                    None,
+                )
+                is None
+            ):
+                clip.cond_stage_model.clip_g.transformer.text_projection = (
+                    cast.manual_cast.Linear(1280, 1280)
+                )
+
+        return clip
+
+
+class DualCLIPLoaderGGUF(CLIPLoaderGGUF):
+    def load_clip(self, clip_name1: str, clip_name2: str, type: str) -> tuple:
+        """
+        Load dual clips.
+
+        Args:
+            clip_name1 (str): The name of the first clip.
+            clip_name2 (str): The name of the second clip.
+            type (str): The type of the clip.
+
+        Returns:
+            tuple: The loaded clips.
+        """
+        clip_path1 = "./_internal/clip/" + clip_name1
+        clip_path2 = "./_internal/clip/" + clip_name2
+        clip_paths = (clip_path1, clip_path2)
+        clip_type = clip_name_dict.get(type, Clip.CLIPType.STABLE_DIFFUSION)
+        return (self.load_patcher(clip_paths, clip_type, self.load_data(clip_paths)),)
+
+
+class CLIPTextEncodeFlux:
+    def encode(
+        self,
+        clip: Clip,
+        clip_l: str,
+        t5xxl: str,
+        guidance: str,
+        flux_enabled: bool = False,
+    ) -> tuple:
+        """
+        Encode text using CLIP and T5XXL.
+
+        Args:
+            clip (Clip): The clip object.
+            clip_l (str): The clip text.
+            t5xxl (str): The T5XXL text.
+            guidance (str): The guidance text.
+            flux_enabled (bool, optional): Whether flux is enabled. Defaults to False.
+
+        Returns:
+            tuple: The encoded text.
+        """
+        tokens = clip.tokenize(clip_l)
+        tokens["t5xxl"] = clip.tokenize(t5xxl)["t5xxl"]
+
+        output = clip.encode_from_tokens(
+            tokens, return_pooled=True, return_dict=True, flux_enabled=flux_enabled
+        )
+        cond = output.pop("cond")
+        output["guidance"] = guidance
+        return ([[cond, output]],)
+
+
+class ConditioningZeroOut:
+    def zero_out(self, conditioning: list) -> list:
+        """
+        Zero out the conditioning.
+
+        Args:
+            conditioning (list): The conditioning list.
+
+        Returns:
+            list: The zeroed out conditioning.
+        """
+        c = []
+        for t in conditioning:
+            d = t[1].copy()
+            pooled_output = d.get("pooled_output", None)
+            if pooled_output is not None:
+                d["pooled_output"] = torch.zeros_like(pooled_output)
+            n = [torch.zeros_like(t[0]), d]
+            c.append(n)
+        return (c,)

modules/SD15/SD15.py

@@ -0,0 +1,81 @@
+import torch
+from modules.BlackForest import Flux
+from modules.Utilities import util
+from modules.Model import ModelBase
+from modules.SD15 import SDClip, SDToken
+from modules.Utilities import Latent
+from modules.clip import Clip
+
+
+class sm_SD15(ModelBase.BASE):
+    """#### Class representing the SD15 model.
+
+    #### Args:
+        - `ModelBase.BASE` (ModelBase.BASE): The base model class.
+    """
+
+    unet_config: dict = {
+        "context_dim": 768,
+        "model_channels": 320,
+        "use_linear_in_transformer": False,
+        "adm_in_channels": None,
+        "use_temporal_attention": False,
+    }
+
+    unet_extra_config: dict = {
+        "num_heads": 8,
+        "num_head_channels": -1,
+    }
+
+    latent_format: Latent.SD15 = Latent.SD15
+
+    def process_clip_state_dict(self, state_dict: dict) -> dict:
+        """#### Process the state dictionary for the CLIP model.
+
+        #### Args:
+            - `state_dict` (dict): The state dictionary.
+
+        #### Returns:
+            - `dict`: The processed state dictionary.
+        """
+        k = list(state_dict.keys())
+        for x in k:
+            if x.startswith("cond_stage_model.transformer.") and not x.startswith(
+                "cond_stage_model.transformer.text_model."
+            ):
+                y = x.replace(
+                    "cond_stage_model.transformer.",
+                    "cond_stage_model.transformer.text_model.",
+                )
+                state_dict[y] = state_dict.pop(x)
+
+        if (
+            "cond_stage_model.transformer.text_model.embeddings.position_ids"
+            in state_dict
+        ):
+            ids = state_dict[
+                "cond_stage_model.transformer.text_model.embeddings.position_ids"
+            ]
+            if ids.dtype == torch.float32:
+                state_dict[
+                    "cond_stage_model.transformer.text_model.embeddings.position_ids"
+                ] = ids.round()
+
+        replace_prefix = {}
+        replace_prefix["cond_stage_model."] = "clip_l."
+        state_dict = util.state_dict_prefix_replace(
+            state_dict, replace_prefix, filter_keys=True
+        )
+        return state_dict
+
+    def clip_target(self) -> Clip.ClipTarget:
+        """#### Get the target CLIP model.
+
+        #### Returns:
+            - `Clip.ClipTarget`: The target CLIP model.
+        """
+        return Clip.ClipTarget(SDToken.SD1Tokenizer, SDClip.SD1ClipModel)
+    
+models = [
+    sm_SD15, Flux.Flux
+]

modules/SD15/SDClip.py

@@ -0,0 +1,403 @@
+import json
+import logging
+import numbers
+import torch
+from modules.Device import Device
+from modules.cond import cast
+from modules.clip.CLIPTextModel import CLIPTextModel
+
+
+
+def gen_empty_tokens(special_tokens: dict, length: int) -> list:
+    """#### Generate a list of empty tokens.
+
+    #### Args:
+        - `special_tokens` (dict): The special tokens.
+        - `length` (int): The length of the token list.
+
+    #### Returns:
+        - `list`: The list of empty tokens.
+    """
+    start_token = special_tokens.get("start", None)
+    end_token = special_tokens.get("end", None)
+    pad_token = special_tokens.get("pad")
+    output = []
+    if start_token is not None:
+        output.append(start_token)
+    if end_token is not None:
+        output.append(end_token)
+    output += [pad_token] * (length - len(output))
+    return output
+
+
+class ClipTokenWeightEncoder:
+    """#### Class representing a CLIP token weight encoder."""
+
+    def encode_token_weights(self, token_weight_pairs: list) -> tuple:
+        """#### Encode token weights.
+
+        #### Args:
+            - `token_weight_pairs` (list): The token weight pairs.
+
+        #### Returns:
+            - `tuple`: The encoded tokens and the pooled output.
+        """
+        to_encode = list()
+        max_token_len = 0
+        has_weights = False
+        for x in token_weight_pairs:
+            tokens = list(map(lambda a: a[0], x))
+            max_token_len = max(len(tokens), max_token_len)
+            has_weights = has_weights or not all(map(lambda a: a[1] == 1.0, x))
+            to_encode.append(tokens)
+
+        sections = len(to_encode)
+        if has_weights or sections == 0:
+            to_encode.append(gen_empty_tokens(self.special_tokens, max_token_len))
+
+        o = self.encode(to_encode)
+        out, pooled = o[:2]
+
+        if pooled is not None:
+            first_pooled = pooled[0:1].to(Device.intermediate_device())
+        else:
+            first_pooled = pooled
+
+        output = []
+        for k in range(0, sections):
+            z = out[k : k + 1]
+            if has_weights:
+                z_empty = out[-1]
+                for i in range(len(z)):
+                    for j in range(len(z[i])):
+                        weight = token_weight_pairs[k][j][1]
+                        if weight != 1.0:
+                            z[i][j] = (z[i][j] - z_empty[j]) * weight + z_empty[j]
+            output.append(z)
+
+        if len(output) == 0:
+            r = (out[-1:].to(Device.intermediate_device()), first_pooled)
+        else:
+            r = (torch.cat(output, dim=-2).to(Device.intermediate_device()), first_pooled)
+
+        if len(o) > 2:
+            extra = {}
+            for k in o[2]:
+                v = o[2][k]
+                if k == "attention_mask":
+                    v = (
+                        v[:sections]
+                        .flatten()
+                        .unsqueeze(dim=0)
+                        .to(Device.intermediate_device())
+                    )
+                extra[k] = v
+
+            r = r + (extra,)
+        return r
+
+class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
+    """#### Uses the CLIP transformer encoder for text (from huggingface)."""
+
+    LAYERS = ["last", "pooled", "hidden"]
+
+    def __init__(
+        self,
+        version: str = "openai/clip-vit-large-patch14",
+        device: str = "cpu",
+        max_length: int = 77,
+        freeze: bool = True,
+        layer: str = "last",
+        layer_idx: int = None,
+        textmodel_json_config: str = None,
+        dtype: torch.dtype = None,
+        model_class: type = CLIPTextModel,
+        special_tokens: dict = {"start": 49406, "end": 49407, "pad": 49407},
+        layer_norm_hidden_state: bool = True,
+        enable_attention_masks: bool = False,
+        zero_out_masked:bool = False,
+        return_projected_pooled: bool = True,
+        return_attention_masks: bool = False,
+        model_options={},
+    ):
+        """#### Initialize the SDClipModel.
+
+        #### Args:
+            - `version` (str, optional): The version of the model. Defaults to "openai/clip-vit-large-patch14".
+            - `device` (str, optional): The device to use. Defaults to "cpu".
+            - `max_length` (int, optional): The maximum length of the input. Defaults to 77.
+            - `freeze` (bool, optional): Whether to freeze the model parameters. Defaults to True.
+            - `layer` (str, optional): The layer to use. Defaults to "last".
+            - `layer_idx` (int, optional): The index of the layer. Defaults to None.
+            - `textmodel_json_config` (str, optional): The path to the JSON config file. Defaults to None.
+            - `dtype` (torch.dtype, optional): The data type. Defaults to None.
+            - `model_class` (type, optional): The model class. Defaults to CLIPTextModel.
+            - `special_tokens` (dict, optional): The special tokens. Defaults to {"start": 49406, "end": 49407, "pad": 49407}.
+            - `layer_norm_hidden_state` (bool, optional): Whether to normalize the hidden state. Defaults to True.
+            - `enable_attention_masks` (bool, optional): Whether to enable attention masks. Defaults to False.
+            - `zero_out_masked` (bool, optional): Whether to zero out masked tokens. Defaults to False.
+            - `return_projected_pooled` (bool, optional): Whether to return the projected pooled output. Defaults to True.
+            - `return_attention_masks` (bool, optional): Whether to return the attention masks. Defaults to False.
+            - `model_options` (dict, optional): Additional model options. Defaults to {}.
+        """
+        super().__init__()
+        assert layer in self.LAYERS
+
+        if textmodel_json_config is None:
+            textmodel_json_config = "./_internal/clip/sd1_clip_config.json"
+
+        with open(textmodel_json_config) as f:
+            config = json.load(f)
+
+        operations = model_options.get("custom_operations", None)
+        if operations is None:
+            operations = cast.manual_cast
+
+        self.operations = operations
+        self.transformer = model_class(config, dtype, device, self.operations)
+        self.num_layers = self.transformer.num_layers
+
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+        self.layer = layer
+        self.layer_idx = None
+        self.special_tokens = special_tokens
+
+        self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
+        self.enable_attention_masks = enable_attention_masks
+        self.zero_out_masked = zero_out_masked
+
+        self.layer_norm_hidden_state = layer_norm_hidden_state
+        self.return_projected_pooled = return_projected_pooled
+        self.return_attention_masks = return_attention_masks
+
+        if layer == "hidden":
+            assert layer_idx is not None
+            assert abs(layer_idx) < self.num_layers
+            self.set_clip_options({"layer": layer_idx})
+        self.options_default = (
+            self.layer,
+            self.layer_idx,
+            self.return_projected_pooled,
+        )
+
+    def freeze(self) -> None:
+        """#### Freeze the model parameters."""
+        self.transformer = self.transformer.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def set_clip_options(self, options: dict) -> None:
+        """#### Set the CLIP options.
+
+        #### Args:
+            - `options` (dict): The options to set.
+        """
+        layer_idx = options.get("layer", self.layer_idx)
+        self.return_projected_pooled = options.get(
+            "projected_pooled", self.return_projected_pooled
+        )
+        if layer_idx is None or abs(layer_idx) > self.num_layers:
+            self.layer = "last"
+        else:
+            self.layer = "hidden"
+            self.layer_idx = layer_idx
+
+    def reset_clip_options(self) -> None:
+        """#### Reset the CLIP options to default."""
+        self.layer = self.options_default[0]
+        self.layer_idx = self.options_default[1]
+        self.return_projected_pooled = self.options_default[2]
+
+    def set_up_textual_embeddings(self, tokens: list, current_embeds: torch.nn.Embedding) -> list:
+        """#### Set up the textual embeddings.
+
+        #### Args:
+            - `tokens` (list): The input tokens.
+            - `current_embeds` (torch.nn.Embedding): The current embeddings.
+
+        #### Returns:
+            - `list`: The processed tokens.
+        """
+        out_tokens = []
+        next_new_token = token_dict_size = current_embeds.weight.shape[0]
+        embedding_weights = []
+
+        for x in tokens:
+            tokens_temp = []
+            for y in x:
+                if isinstance(y, numbers.Integral):
+                    tokens_temp += [int(y)]
+                else:
+                    if y.shape[0] == current_embeds.weight.shape[1]:
+                        embedding_weights += [y]
+                        tokens_temp += [next_new_token]
+                        next_new_token += 1
+                    else:
+                        logging.warning(
+                            "WARNING: shape mismatch when trying to apply embedding, embedding will be ignored {} != {}".format(
+                                y.shape[0], current_embeds.weight.shape[1]
+                            )
+                        )
+            while len(tokens_temp) < len(x):
+                tokens_temp += [self.special_tokens["pad"]]
+            out_tokens += [tokens_temp]
+
+        n = token_dict_size
+        if len(embedding_weights) > 0:
+            new_embedding = self.operations.Embedding(
+                next_new_token + 1,
+                current_embeds.weight.shape[1],
+                device=current_embeds.weight.device,
+                dtype=current_embeds.weight.dtype,
+            )
+            new_embedding.weight[:token_dict_size] = current_embeds.weight
+            for x in embedding_weights:
+                new_embedding.weight[n] = x
+                n += 1
+            self.transformer.set_input_embeddings(new_embedding)
+
+        processed_tokens = []
+        for x in out_tokens:
+            processed_tokens += [
+                list(map(lambda a: n if a == -1 else a, x))
+            ]  # The EOS token should always be the largest one
+
+        return processed_tokens
+
+    def forward(self, tokens: list) -> tuple:
+        """#### Forward pass of the model.
+
+        #### Args:
+            - `tokens` (list): The input tokens.
+
+        #### Returns:
+            - `tuple`: The output and the pooled output.
+        """
+        backup_embeds = self.transformer.get_input_embeddings()
+        device = backup_embeds.weight.device
+        tokens = self.set_up_textual_embeddings(tokens, backup_embeds)
+        tokens = torch.LongTensor(tokens).to(device)
+
+        attention_mask = None
+        if (
+            self.enable_attention_masks
+            or self.zero_out_masked
+            or self.return_attention_masks
+        ):
+            attention_mask = torch.zeros_like(tokens)
+            end_token = self.special_tokens.get("end", -1)
+            for x in range(attention_mask.shape[0]):
+                for y in range(attention_mask.shape[1]):
+                    attention_mask[x, y] = 1
+                    if tokens[x, y] == end_token:
+                        break
+
+        attention_mask_model = None
+        if self.enable_attention_masks:
+            attention_mask_model = attention_mask
+
+        outputs = self.transformer(
+            tokens,
+            attention_mask_model,
+            intermediate_output=self.layer_idx,
+            final_layer_norm_intermediate=self.layer_norm_hidden_state,
+            dtype=torch.float32,
+        )
+        self.transformer.set_input_embeddings(backup_embeds)
+
+        if self.layer == "last":
+            z = outputs[0].float()
+        else:
+            z = outputs[1].float()
+
+        if self.zero_out_masked:
+            z *= attention_mask.unsqueeze(-1).float()
+
+        pooled_output = None
+        if len(outputs) >= 3:
+            if (
+                not self.return_projected_pooled
+                and len(outputs) >= 4
+                and outputs[3] is not None
+            ):
+                pooled_output = outputs[3].float()
+            elif outputs[2] is not None:
+                pooled_output = outputs[2].float()
+
+        extra = {}
+        if self.return_attention_masks:
+            extra["attention_mask"] = attention_mask
+
+        if len(extra) > 0:
+            return z, pooled_output, extra
+
+        return z, pooled_output
+
+    def encode(self, tokens: list) -> tuple:
+        """#### Encode the input tokens.
+
+        #### Args:
+            - `tokens` (list): The input tokens.
+
+        #### Returns:
+            - `tuple`: The encoded tokens and the pooled output.
+        """
+        return self(tokens)
+
+    def load_sd(self, sd: dict) -> None:
+        """#### Load the state dictionary.
+
+        #### Args:
+            - `sd` (dict): The state dictionary.
+        """
+        return self.transformer.load_state_dict(sd, strict=False)
+
+
+class SD1ClipModel(torch.nn.Module):
+    """#### Class representing the SD1ClipModel."""
+
+    def __init__(
+        self, device: str = "cpu", dtype: torch.dtype = None, clip_name: str = "l", clip_model: type = SDClipModel, **kwargs
+    ):
+        """#### Initialize the SD1ClipModel.
+
+        #### Args:
+            - `device` (str, optional): The device to use. Defaults to "cpu".
+            - `dtype` (torch.dtype, optional): The data type. Defaults to None.
+            - `clip_name` (str, optional): The name of the CLIP model. Defaults to "l".
+            - `clip_model` (type, optional): The CLIP model class. Defaults to SDClipModel.
+            - `**kwargs`: Additional keyword arguments.
+        """
+        super().__init__()
+        self.clip_name = clip_name
+        self.clip = "clip_{}".format(self.clip_name)
+        self.lowvram_patch_counter = 0
+        self.model_loaded_weight_memory = 0
+        setattr(self, self.clip, clip_model(device=device, dtype=dtype, **kwargs))
+
+    def set_clip_options(self, options: dict) -> None:
+        """#### Set the CLIP options.
+
+        #### Args:
+            - `options` (dict): The options to set.
+        """
+        getattr(self, self.clip).set_clip_options(options)
+
+    def reset_clip_options(self) -> None:
+        """#### Reset the CLIP options to default."""
+        getattr(self, self.clip).reset_clip_options()
+
+    def encode_token_weights(self, token_weight_pairs: dict) -> tuple:
+        """#### Encode token weights.
+
+        #### Args:
+            - `token_weight_pairs` (dict): The token weight pairs.
+
+        #### Returns:
+            - `tuple`: The encoded tokens and the pooled output.
+        """
+        token_weight_pairs = token_weight_pairs[self.clip_name]
+        out, pooled = getattr(self, self.clip).encode_token_weights(token_weight_pairs)
+        return out, pooled

modules/SD15/SDToken.py

@@ -0,0 +1,450 @@
+import logging
+import os
+import traceback
+import torch
+from transformers import CLIPTokenizerFast
+
+def model_options_long_clip(sd, tokenizer_data, model_options):
+    w = sd.get("clip_l.text_model.embeddings.position_embedding.weight", None)
+    if w is None:
+        w = sd.get("text_model.embeddings.position_embedding.weight", None)
+    return tokenizer_data, model_options
+
+def parse_parentheses(string: str) -> list:
+    """#### Parse a string with nested parentheses.
+
+    #### Args:
+        - `string` (str): The input string.
+
+    #### Returns:
+        - `list`: The parsed list of strings.
+    """
+    result = []
+    current_item = ""
+    nesting_level = 0
+    for char in string:
+        if char == "(":
+            if nesting_level == 0:
+                if current_item:
+                    result.append(current_item)
+                    current_item = "("
+                else:
+                    current_item = "("
+            else:
+                current_item += char
+            nesting_level += 1
+        elif char == ")":
+            nesting_level -= 1
+            if nesting_level == 0:
+                result.append(current_item + ")")
+                current_item = ""
+            else:
+                current_item += char
+        else:
+            current_item += char
+    if current_item:
+        result.append(current_item)
+    return result
+
+
+def token_weights(string: str, current_weight: float) -> list:
+    """#### Parse a string into tokens with weights.
+
+    #### Args:
+        - `string` (str): The input string.
+        - `current_weight` (float): The current weight.
+
+    #### Returns:
+        - `list`: The list of token-weight pairs.
+    """
+    a = parse_parentheses(string)
+    out = []
+    for x in a:
+        weight = current_weight
+        if len(x) >= 2 and x[-1] == ")" and x[0] == "(":
+            x = x[1:-1]
+            xx = x.rfind(":")
+            weight *= 1.1
+            if xx > 0:
+                try:
+                    weight = float(x[xx + 1 :])
+                    x = x[:xx]
+                except:
+                    pass
+            out += token_weights(x, weight)
+        else:
+            out += [(x, current_weight)]
+    return out
+
+
+def escape_important(text: str) -> str:
+    """#### Escape important characters in a string.
+
+    #### Args:
+        - `text` (str): The input text.
+
+    #### Returns:
+        - `str`: The escaped text.
+    """
+    text = text.replace("\\)", "\0\1")
+    text = text.replace("\\(", "\0\2")
+    return text
+
+
+def unescape_important(text: str) -> str:
+    """#### Unescape important characters in a string.
+
+    #### Args:
+        - `text` (str): The input text.
+
+    #### Returns:
+        - `str`: The unescaped text.
+    """
+    text = text.replace("\0\1", ")")
+    text = text.replace("\0\2", "(")
+    return text
+
+
+def expand_directory_list(directories: list) -> list:
+    """#### Expand a list of directories to include all subdirectories.
+
+    #### Args:
+        - `directories` (list): The list of directories.
+
+    #### Returns:
+        - `list`: The expanded list of directories.
+    """
+    dirs = set()
+    for x in directories:
+        dirs.add(x)
+        for root, subdir, file in os.walk(x, followlinks=True):
+            dirs.add(root)
+    return list(dirs)
+
+
+def load_embed(embedding_name: str, embedding_directory: list, embedding_size: int, embed_key: str = None) -> torch.Tensor:
+    """#### Load an embedding from a directory.
+
+    #### Args:
+        - `embedding_name` (str): The name of the embedding.
+        - `embedding_directory` (list): The list of directories to search.
+        - `embedding_size` (int): The size of the embedding.
+        - `embed_key` (str, optional): The key for the embedding. Defaults to None.
+
+    #### Returns:
+        - `torch.Tensor`: The loaded embedding.
+    """
+    if isinstance(embedding_directory, str):
+        embedding_directory = [embedding_directory]
+
+    embedding_directory = expand_directory_list(embedding_directory)
+
+    valid_file = None
+    for embed_dir in embedding_directory:
+        embed_path = os.path.abspath(os.path.join(embed_dir, embedding_name))
+        embed_dir = os.path.abspath(embed_dir)
+        try:
+            if os.path.commonpath((embed_dir, embed_path)) != embed_dir:
+                continue
+        except:
+            continue
+        if not os.path.isfile(embed_path):
+            extensions = [".safetensors", ".pt", ".bin"]
+            for x in extensions:
+                t = embed_path + x
+                if os.path.isfile(t):
+                    valid_file = t
+                    break
+        else:
+            valid_file = embed_path
+        if valid_file is not None:
+            break
+
+    if valid_file is None:
+        return None
+
+    embed_path = valid_file
+
+    embed_out = None
+
+    try:
+        if embed_path.lower().endswith(".safetensors"):
+            import safetensors.torch
+
+            embed = safetensors.torch.load_file(embed_path, device="cpu")
+        else:
+            if "weights_only" in torch.load.__code__.co_varnames:
+                embed = torch.load(embed_path, weights_only=True, map_location="cpu")
+            else:
+                embed = torch.load(embed_path, map_location="cpu")
+    except Exception:
+        logging.warning(
+            "{}\n\nerror loading embedding, skipping loading: {}".format(
+                traceback.format_exc(), embedding_name
+            )
+        )
+        return None
+
+    if embed_out is None:
+        if "string_to_param" in embed:
+            values = embed["string_to_param"].values()
+            embed_out = next(iter(values))
+        elif isinstance(embed, list):
+            out_list = []
+            for x in range(len(embed)):
+                for k in embed[x]:
+                    t = embed[x][k]
+                    if t.shape[-1] != embedding_size:
+                        continue
+                    out_list.append(t.reshape(-1, t.shape[-1]))
+            embed_out = torch.cat(out_list, dim=0)
+        elif embed_key is not None and embed_key in embed:
+            embed_out = embed[embed_key]
+        else:
+            values = embed.values()
+            embed_out = next(iter(values))
+    return embed_out
+
+
+class SDTokenizer:
+    """#### Class representing a Stable Diffusion tokenizer."""
+
+    def __init__(
+        self,
+        tokenizer_path: str = None,
+        max_length: int = 77,
+        pad_with_end: bool = True,
+        embedding_directory: str = None,
+        embedding_size: int = 768,
+        embedding_key: str = "clip_l",
+        tokenizer_class: type = CLIPTokenizerFast,
+        has_start_token: bool = True,
+        pad_to_max_length: bool = True,
+        min_length: int = None,
+    ):
+        """#### Initialize the SDTokenizer.
+
+        #### Args:
+            - `tokenizer_path` (str, optional): The path to the tokenizer. Defaults to None.
+            - `max_length` (int, optional): The maximum length of the input. Defaults to 77.
+            - `pad_with_end` (bool, optional): Whether to pad with the end token. Defaults to True.
+            - `embedding_directory` (str, optional): The directory for embeddings. Defaults to None.
+            - `embedding_size` (int, optional): The size of the embeddings. Defaults to 768.
+            - `embedding_key` (str, optional): The key for the embeddings. Defaults to "clip_l".
+            - `tokenizer_class` (type, optional): The tokenizer class. Defaults to CLIPTokenizer.
+            - `has_start_token` (bool, optional): Whether the tokenizer has a start token. Defaults to True.
+            - `pad_to_max_length` (bool, optional): Whether to pad to the maximum length. Defaults to True.
+            - `min_length` (int, optional): The minimum length of the input. Defaults to None.
+        """
+        if tokenizer_path is None:
+            tokenizer_path = "_internal/sd1_tokenizer/"
+        self.tokenizer = tokenizer_class.from_pretrained(tokenizer_path)
+        self.max_length = max_length
+        self.min_length = min_length
+
+        empty = self.tokenizer("")["input_ids"]
+        if has_start_token:
+            self.tokens_start = 1
+            self.start_token = empty[0]
+            self.end_token = empty[1]
+        else:
+            self.tokens_start = 0
+            self.start_token = None
+            self.end_token = empty[0]
+        self.pad_with_end = pad_with_end
+        self.pad_to_max_length = pad_to_max_length
+
+        vocab = self.tokenizer.get_vocab()
+        self.inv_vocab = {v: k for k, v in vocab.items()}
+        self.embedding_directory = embedding_directory
+        self.max_word_length = 8
+        self.embedding_identifier = "embedding:"
+        self.embedding_size = embedding_size
+        self.embedding_key = embedding_key
+
+    def _try_get_embedding(self, embedding_name: str) -> tuple:
+        """#### Try to get an embedding.
+
+        #### Args:
+            - `embedding_name` (str): The name of the embedding.
+
+        #### Returns:
+            - `tuple`: The embedding and any leftover text.
+        """
+        embed = load_embed(
+            embedding_name,
+            self.embedding_directory,
+            self.embedding_size,
+            self.embedding_key,
+        )
+        if embed is None:
+            stripped = embedding_name.strip(",")
+            if len(stripped) < len(embedding_name):
+                embed = load_embed(
+                    stripped,
+                    self.embedding_directory,
+                    self.embedding_size,
+                    self.embedding_key,
+                )
+                return (embed, embedding_name[len(stripped) :])
+        return (embed, "")
+
+    def tokenize_with_weights(self, text: str, return_word_ids: bool = False) -> list:
+        """#### Tokenize text with weights.
+
+        #### Args:
+            - `text` (str): The input text.
+            - `return_word_ids` (bool, optional): Whether to return word IDs. Defaults to False.
+
+        #### Returns:
+            - `list`: The tokenized text with weights.
+        """
+        if self.pad_with_end:
+            pad_token = self.end_token
+        else:
+            pad_token = 0
+
+        text = escape_important(text)
+        parsed_weights = token_weights(text, 1.0)
+
+        # tokenize words
+        tokens = []
+        for weighted_segment, weight in parsed_weights:
+            to_tokenize = (
+                unescape_important(weighted_segment).replace("\n", " ").split(" ")
+            )
+            to_tokenize = [x for x in to_tokenize if x != ""]
+            for word in to_tokenize:
+                # if we find an embedding, deal with the embedding
+                if (
+                    word.startswith(self.embedding_identifier)
+                    and self.embedding_directory is not None
+                ):
+                    embedding_name = word[len(self.embedding_identifier) :].strip("\n")
+                    embed, leftover = self._try_get_embedding(embedding_name)
+                    if embed is None:
+                        logging.warning(
+                            f"warning, embedding:{embedding_name} does not exist, ignoring"
+                        )
+                    else:
+                        if len(embed.shape) == 1:
+                            tokens.append([(embed, weight)])
+                        else:
+                            tokens.append(
+                                [(embed[x], weight) for x in range(embed.shape[0])]
+                            )
+                        print("loading ", embedding_name)
+                    # if we accidentally have leftover text, continue parsing using leftover, else move on to next word
+                    if leftover != "":
+                        word = leftover
+                    else:
+                        continue
+                # parse word
+                tokens.append(
+                    [
+                        (t, weight)
+                        for t in self.tokenizer(word)["input_ids"][
+                            self.tokens_start : -1
+                        ]
+                    ]
+                )
+
+        # reshape token array to CLIP input size
+        batched_tokens = []
+        batch = []
+        if self.start_token is not None:
+            batch.append((self.start_token, 1.0, 0))
+        batched_tokens.append(batch)
+        for i, t_group in enumerate(tokens):
+            # determine if we're going to try and keep the tokens in a single batch
+            is_large = len(t_group) >= self.max_word_length
+
+            while len(t_group) > 0:
+                if len(t_group) + len(batch) > self.max_length - 1:
+                    remaining_length = self.max_length - len(batch) - 1
+                    # break word in two and add end token
+                    if is_large:
+                        batch.extend(
+                            [(t, w, i + 1) for t, w in t_group[:remaining_length]]
+                        )
+                        batch.append((self.end_token, 1.0, 0))
+                        t_group = t_group[remaining_length:]
+                    # add end token and pad
+                    else:
+                        batch.append((self.end_token, 1.0, 0))
+                        if self.pad_to_max_length:
+                            batch.extend([(pad_token, 1.0, 0)] * (remaining_length))
+                    # start new batch
+                    batch = []
+                    if self.start_token is not None:
+                        batch.append((self.start_token, 1.0, 0))
+                    batched_tokens.append(batch)
+                else:
+                    batch.extend([(t, w, i + 1) for t, w in t_group])
+                    t_group = []
+
+        # fill last batch
+        batch.append((self.end_token, 1.0, 0))
+        if self.pad_to_max_length:
+            batch.extend([(pad_token, 1.0, 0)] * (self.max_length - len(batch)))
+        if self.min_length is not None and len(batch) < self.min_length:
+            batch.extend([(pad_token, 1.0, 0)] * (self.min_length - len(batch)))
+
+        if not return_word_ids:
+            batched_tokens = [[(t, w) for t, w, _ in x] for x in batched_tokens]
+
+        return batched_tokens
+
+    def untokenize(self, token_weight_pair: list) -> list:
+        """#### Untokenize a list of token-weight pairs.
+
+        #### Args:
+            - `token_weight_pair` (list): The list of token-weight pairs.
+
+        #### Returns:
+            - `list`: The untokenized list.
+        """
+        return list(map(lambda a: (a, self.inv_vocab[a[0]]), token_weight_pair))
+
+
+class SD1Tokenizer:
+    """#### Class representing the SD1Tokenizer."""
+
+    def __init__(self, embedding_directory: str = None, clip_name: str = "l", tokenizer: type = SDTokenizer):
+        """#### Initialize the SD1Tokenizer.
+
+        #### Args:
+            - `embedding_directory` (str, optional): The directory for embeddings. Defaults to None.
+            - `clip_name` (str, optional): The name of the CLIP model. Defaults to "l".
+            - `tokenizer` (type, optional): The tokenizer class. Defaults to SDTokenizer.
+        """
+        self.clip_name = clip_name
+        self.clip = "clip_{}".format(self.clip_name)
+        setattr(self, self.clip, tokenizer(embedding_directory=embedding_directory))
+
+    def tokenize_with_weights(self, text: str, return_word_ids: bool = False) -> dict:
+        """#### Tokenize text with weights.
+
+        #### Args:
+            - `text` (str): The input text.
+            - `return_word_ids` (bool, optional): Whether to return word IDs. Defaults to False.
+
+        #### Returns:
+            - `dict`: The tokenized text with weights.
+        """
+        out = {}
+        out[self.clip_name] = getattr(self, self.clip).tokenize_with_weights(
+            text, return_word_ids
+        )
+        return out
+
+    def untokenize(self, token_weight_pair: list) -> list:
+        """#### Untokenize a list of token-weight pairs.
+
+        #### Args:
+            - `token_weight_pair` (list): The list of token-weight pairs.
+
+        #### Returns:
+            - `list`: The untokenized list.
+        """
+        return getattr(self, self.clip).untokenize(token_weight_pair)

modules/StableFast/StableFast.py

@@ -0,0 +1,274 @@
+import contextlib
+import functools
+import logging
+from dataclasses import dataclass
+
+import torch
+
+try:
+    from sfast.compilers.diffusion_pipeline_compiler import CompilationConfig
+    from sfast.compilers.diffusion_pipeline_compiler import (
+        _enable_xformers,
+        _modify_model,
+    )
+    from sfast.cuda.graphs import make_dynamic_graphed_callable
+    from sfast.jit import utils as jit_utils
+    from sfast.jit.trace_helper import trace_with_kwargs
+except:
+    pass
+
+
+def hash_arg(arg):
+    # micro optimization: bool obj is an instance of int
+    if isinstance(arg, (str, int, float, bytes)):
+        return arg
+    if isinstance(arg, (tuple, list)):
+        return tuple(map(hash_arg, arg))
+    if isinstance(arg, dict):
+        return tuple(
+            sorted(
+                ((hash_arg(k), hash_arg(v)) for k, v in arg.items()), key=lambda x: x[0]
+            )
+        )
+    return type(arg)
+
+
+class ModuleFactory:
+    def get_converted_kwargs(self):
+        return self.converted_kwargs
+
+
+import torch as th
+import torch.nn as nn
+import copy
+
+
+class BaseModelApplyModelModule(torch.nn.Module):
+    def __init__(self, func, module):
+        super().__init__()
+        self.func = func
+        self.module = module
+
+    def forward(
+        self,
+        input_x,
+        timestep,
+        c_concat=None,
+        c_crossattn=None,
+        y=None,
+        control=None,
+        transformer_options={},
+    ):
+        kwargs = {"y": y}
+
+        new_transformer_options = {}
+
+        return self.func(
+            input_x,
+            timestep,
+            c_concat=c_concat,
+            c_crossattn=c_crossattn,
+            control=control,
+            transformer_options=new_transformer_options,
+            **kwargs,
+        )
+
+
+class BaseModelApplyModelModuleFactory(ModuleFactory):
+    kwargs_name = (
+        "input_x",
+        "timestep",
+        "c_concat",
+        "c_crossattn",
+        "y",
+        "control",
+    )
+
+    def __init__(self, callable, kwargs) -> None:
+        self.callable = callable
+        self.unet_config = callable.__self__.model_config.unet_config
+        self.kwargs = kwargs
+        self.patch_module = {}
+        self.patch_module_parameter = {}
+        self.converted_kwargs = self.gen_converted_kwargs()
+
+    def gen_converted_kwargs(self):
+        converted_kwargs = {}
+        for arg_name, arg in self.kwargs.items():
+            if arg_name in self.kwargs_name:
+                converted_kwargs[arg_name] = arg
+
+        transformer_options = self.kwargs.get("transformer_options", {})
+        patches = transformer_options.get("patches", {})
+
+        patch_module = {}
+        patch_module_parameter = {}
+
+        new_transformer_options = {}
+        new_transformer_options["patches"] = patch_module_parameter
+
+        self.patch_module = patch_module
+        self.patch_module_parameter = patch_module_parameter
+        return converted_kwargs
+
+    def gen_cache_key(self):
+        key_kwargs = {}
+        for k, v in self.converted_kwargs.items():
+            key_kwargs[k] = v
+
+        patch_module_cache_key = {}
+        return (
+            self.callable.__class__.__qualname__,
+            hash_arg(self.unet_config),
+            hash_arg(key_kwargs),
+            hash_arg(patch_module_cache_key),
+        )
+
+    @contextlib.contextmanager
+    def converted_module_context(self):
+        module = BaseModelApplyModelModule(self.callable, self.callable.__self__)
+        yield (module, self.converted_kwargs)
+
+
+logger = logging.getLogger()
+
+
+@dataclass
+class TracedModuleCacheItem:
+    module: object
+    patch_id: int
+    device: str
+
+
+class LazyTraceModule:
+    traced_modules = {}
+
+    def __init__(self, config=None, patch_id=None, **kwargs_) -> None:
+        self.config = config
+        self.patch_id = patch_id
+        self.kwargs_ = kwargs_
+        self.modify_model = functools.partial(
+            _modify_model,
+            enable_cnn_optimization=config.enable_cnn_optimization,
+            prefer_lowp_gemm=config.prefer_lowp_gemm,
+            enable_triton=config.enable_triton,
+            enable_triton_reshape=config.enable_triton,
+            memory_format=config.memory_format,
+        )
+        self.cuda_graph_modules = {}
+
+    def ts_compiler(
+        self,
+        m,
+    ):
+        with torch.jit.optimized_execution(True):
+            if self.config.enable_jit_freeze:
+                # raw freeze causes Tensor reference leak
+                # because the constant Tensors in the GraphFunction of
+                # the compilation unit are never freed.
+                m.eval()
+                m = jit_utils.better_freeze(m)
+            self.modify_model(m)
+
+        if self.config.enable_cuda_graph:
+            m = make_dynamic_graphed_callable(m)
+        return m
+
+    def __call__(self, model_function, /, **kwargs):
+        module_factory = BaseModelApplyModelModuleFactory(model_function, kwargs)
+        kwargs = module_factory.get_converted_kwargs()
+        key = module_factory.gen_cache_key()
+
+        traced_module = self.cuda_graph_modules.get(key)
+        if traced_module is None:
+            with module_factory.converted_module_context() as (m_model, m_kwargs):
+                logger.info(
+                    f'Tracing {getattr(m_model, "__name__", m_model.__class__.__name__)}'
+                )
+                traced_m, call_helper = trace_with_kwargs(
+                    m_model, None, m_kwargs, **self.kwargs_
+                )
+
+            traced_m = self.ts_compiler(traced_m)
+            traced_module = call_helper(traced_m)
+            self.cuda_graph_modules[key] = traced_module
+
+        return traced_module(**kwargs)
+
+
+def build_lazy_trace_module(config, device, patch_id):
+    config.enable_cuda_graph = config.enable_cuda_graph and device.type == "cuda"
+
+    if config.enable_xformers:
+        _enable_xformers(None)
+
+    return LazyTraceModule(
+        config=config,
+        patch_id=patch_id,
+        check_trace=True,
+        strict=True,
+    )
+
+
+def gen_stable_fast_config():
+    config = CompilationConfig.Default()
+    try:
+        import xformers
+
+        config.enable_xformers = True
+    except ImportError:
+        print("xformers not installed, skip")
+
+    # CUDA Graph is suggested for small batch sizes.
+    # After capturing, the model only accepts one fixed image size.
+    # If you want the model to be dynamic, don't enable it.
+    config.enable_cuda_graph = False
+    # config.enable_jit_freeze = False
+    return config
+
+
+class StableFastPatch:
+    def __init__(self, model, config):
+        self.model = model
+        self.config = config
+        self.stable_fast_model = None
+
+    def __call__(self, model_function, params):
+        input_x = params.get("input")
+        timestep_ = params.get("timestep")
+        c = params.get("c")
+
+        if self.stable_fast_model is None:
+            self.stable_fast_model = build_lazy_trace_module(
+                self.config,
+                input_x.device,
+                id(self),
+            )
+
+        return self.stable_fast_model(
+            model_function, input_x=input_x, timestep=timestep_, **c
+        )
+
+    def to(self, device):
+        if type(device) == torch.device:
+            if self.config.enable_cuda_graph or self.config.enable_jit_freeze:
+                if device.type == "cpu":
+                    del self.stable_fast_model
+                    self.stable_fast_model = None
+                    print(
+                        "\33[93mWarning: Your graphics card doesn't have enough video memory to keep the model. If you experience a noticeable delay every time you start sampling, please consider disable enable_cuda_graph.\33[0m"
+                    )
+        return self
+
+
+class ApplyStableFastUnet:
+    def apply_stable_fast(self, model, enable_cuda_graph):
+        config = gen_stable_fast_config()
+
+        if config.memory_format is not None:
+            model.model.to(memory_format=config.memory_format)
+
+        patch = StableFastPatch(model, config)
+        model_stable_fast = model.clone()
+        model_stable_fast.set_model_unet_function_wrapper(patch)
+        return (model_stable_fast,)

modules/UltimateSDUpscale/RDRB.py

@@ -0,0 +1,471 @@
+from collections import OrderedDict
+import functools
+import math
+import re
+from typing import Union, Dict
+import torch
+import torch.nn as nn
+from modules.UltimateSDUpscale import USDU_util
+
+
+class RRDB(nn.Module):
+    """#### Residual in Residual Dense Block (RRDB) class.
+
+    #### Args:
+        - `nf` (int): Number of filters.
+        - `kernel_size` (int, optional): Kernel size. Defaults to 3.
+        - `gc` (int, optional): Growth channel. Defaults to 32.
+        - `stride` (int, optional): Stride. Defaults to 1.
+        - `bias` (bool, optional): Whether to use bias. Defaults to True.
+        - `pad_type` (str, optional): Padding type. Defaults to "zero".
+        - `norm_type` (str, optional): Normalization type. Defaults to None.
+        - `act_type` (str, optional): Activation type. Defaults to "leakyrelu".
+        - `mode` (USDU_util.ConvMode, optional): Convolution mode. Defaults to "CNA".
+        - `_convtype` (str, optional): Convolution type. Defaults to "Conv2D".
+        - `_spectral_norm` (bool, optional): Whether to use spectral normalization. Defaults to False.
+        - `plus` (bool, optional): Whether to use the plus variant. Defaults to False.
+        - `c2x2` (bool, optional): Whether to use 2x2 convolution. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        nf: int,
+        kernel_size: int = 3,
+        gc: int = 32,
+        stride: int = 1,
+        bias: bool = True,
+        pad_type: str = "zero",
+        norm_type: str = None,
+        act_type: str = "leakyrelu",
+        mode: USDU_util.ConvMode = "CNA",
+        _convtype: str = "Conv2D",
+        _spectral_norm: bool = False,
+        plus: bool = False,
+        c2x2: bool = False,
+    ) -> None:
+        super(RRDB, self).__init__()
+        self.RDB1 = ResidualDenseBlock_5C(
+            nf,
+            kernel_size,
+            gc,
+            stride,
+            bias,
+            pad_type,
+            norm_type,
+            act_type,
+            mode,
+            plus=plus,
+            c2x2=c2x2,
+        )
+        self.RDB2 = ResidualDenseBlock_5C(
+            nf,
+            kernel_size,
+            gc,
+            stride,
+            bias,
+            pad_type,
+            norm_type,
+            act_type,
+            mode,
+            plus=plus,
+            c2x2=c2x2,
+        )
+        self.RDB3 = ResidualDenseBlock_5C(
+            nf,
+            kernel_size,
+            gc,
+            stride,
+            bias,
+            pad_type,
+            norm_type,
+            act_type,
+            mode,
+            plus=plus,
+            c2x2=c2x2,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass of the RRDB.
+
+        #### Args:
+            - `x` (torch.Tensor): Input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: Output tensor.
+        """
+        out = self.RDB1(x)
+        out = self.RDB2(out)
+        out = self.RDB3(out)
+        return out * 0.2 + x
+
+
+class ResidualDenseBlock_5C(nn.Module):
+    """#### Residual Dense Block with 5 Convolutions (ResidualDenseBlock_5C) class.
+
+    #### Args:
+        - `nf` (int, optional): Number of filters. Defaults to 64.
+        - `kernel_size` (int, optional): Kernel size. Defaults to 3.
+        - `gc` (int, optional): Growth channel. Defaults to 32.
+        - `stride` (int, optional): Stride. Defaults to 1.
+        - `bias` (bool, optional): Whether to use bias. Defaults to True.
+        - `pad_type` (str, optional): Padding type. Defaults to "zero".
+        - `norm_type` (str, optional): Normalization type. Defaults to None.
+        - `act_type` (str, optional): Activation type. Defaults to "leakyrelu".
+        - `mode` (USDU_util.ConvMode, optional): Convolution mode. Defaults to "CNA".
+        - `plus` (bool, optional): Whether to use the plus variant. Defaults to False.
+        - `c2x2` (bool, optional): Whether to use 2x2 convolution. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        nf: int = 64,
+        kernel_size: int = 3,
+        gc: int = 32,
+        stride: int = 1,
+        bias: bool = True,
+        pad_type: str = "zero",
+        norm_type: str = None,
+        act_type: str = "leakyrelu",
+        mode: USDU_util.ConvMode = "CNA",
+        plus: bool = False,
+        c2x2: bool = False,
+    ) -> None:
+        super(ResidualDenseBlock_5C, self).__init__()
+
+        self.conv1x1 = None
+
+        self.conv1 = USDU_util.conv_block(
+            nf,
+            gc,
+            kernel_size,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=act_type,
+            mode=mode,
+            c2x2=c2x2,
+        )
+        self.conv2 = USDU_util.conv_block(
+            nf + gc,
+            gc,
+            kernel_size,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=act_type,
+            mode=mode,
+            c2x2=c2x2,
+        )
+        self.conv3 = USDU_util.conv_block(
+            nf + 2 * gc,
+            gc,
+            kernel_size,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=act_type,
+            mode=mode,
+            c2x2=c2x2,
+        )
+        self.conv4 = USDU_util.conv_block(
+            nf + 3 * gc,
+            gc,
+            kernel_size,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=act_type,
+            mode=mode,
+            c2x2=c2x2,
+        )
+        last_act = None
+        self.conv5 = USDU_util.conv_block(
+            nf + 4 * gc,
+            nf,
+            3,
+            stride,
+            bias=bias,
+            pad_type=pad_type,
+            norm_type=norm_type,
+            act_type=last_act,
+            mode=mode,
+            c2x2=c2x2,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass of the ResidualDenseBlock_5C.
+
+        #### Args:
+            - `x` (torch.Tensor): Input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: Output tensor.
+        """
+        x1 = self.conv1(x)
+        x2 = self.conv2(torch.cat((x, x1), 1))
+        x3 = self.conv3(torch.cat((x, x1, x2), 1))
+        x4 = self.conv4(torch.cat((x, x1, x2, x3), 1))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5 * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    """#### Residual in Residual Dense Block Network (RRDBNet) class.
+
+    #### Args:
+        - `state_dict` (dict): State dictionary.
+        - `norm` (str, optional): Normalization type. Defaults to None.
+        - `act` (str, optional): Activation type. Defaults to "leakyrelu".
+        - `upsampler` (str, optional): Upsampler type. Defaults to "upconv".
+        - `mode` (USDU_util.ConvMode, optional): Convolution mode. Defaults to "CNA".
+    """
+
+    def __init__(
+        self,
+        state_dict: Dict[str, torch.Tensor],
+        norm: str = None,
+        act: str = "leakyrelu",
+        upsampler: str = "upconv",
+        mode: USDU_util.ConvMode = "CNA",
+    ) -> None:
+        super(RRDBNet, self).__init__()
+        self.model_arch = "ESRGAN"
+        self.sub_type = "SR"
+
+        self.state = state_dict
+        self.norm = norm
+        self.act = act
+        self.upsampler = upsampler
+        self.mode = mode
+
+        self.state_map = {
+            # currently supports old, new, and newer RRDBNet arch _internal
+            # ESRGAN, BSRGAN/RealSR, Real-ESRGAN
+            "model.0.weight": ("conv_first.weight",),
+            "model.0.bias": ("conv_first.bias",),
+            "model.1.sub./NB/.weight": ("trunk_conv.weight", "conv_body.weight"),
+            "model.1.sub./NB/.bias": ("trunk_conv.bias", "conv_body.bias"),
+            r"model.1.sub.\1.RDB\2.conv\3.0.\4": (
+                r"RRDB_trunk\.(\d+)\.RDB(\d)\.conv(\d+)\.(weight|bias)",
+                r"body\.(\d+)\.rdb(\d)\.conv(\d+)\.(weight|bias)",
+            ),
+        }
+        self.num_blocks = self.get_num_blocks()
+        self.plus = any("conv1x1" in k for k in self.state.keys())
+
+        self.state = self.new_to_old_arch(self.state)
+
+        self.key_arr = list(self.state.keys())
+
+        self.in_nc: int = self.state[self.key_arr[0]].shape[1]
+        self.out_nc: int = self.state[self.key_arr[-1]].shape[0]
+
+        self.scale: int = self.get_scale()
+        self.num_filters: int = self.state[self.key_arr[0]].shape[0]
+
+        c2x2 = False
+
+        self.supports_fp16 = True
+        self.supports_bfp16 = True
+        self.min_size_restriction = None
+
+        self.shuffle_factor = None
+
+        upsample_block = {
+            "upconv": USDU_util.upconv_block,
+        }.get(self.upsampler)
+        upsample_blocks = [
+            upsample_block(
+                in_nc=self.num_filters,
+                out_nc=self.num_filters,
+                act_type=self.act,
+                c2x2=c2x2,
+            )
+            for _ in range(int(math.log(self.scale, 2)))
+        ]
+
+        self.model = USDU_util.sequential(
+            # fea conv
+            USDU_util.conv_block(
+                in_nc=self.in_nc,
+                out_nc=self.num_filters,
+                kernel_size=3,
+                norm_type=None,
+                act_type=None,
+                c2x2=c2x2,
+            ),
+            USDU_util.ShortcutBlock(
+                USDU_util.sequential(
+                    # rrdb blocks
+                    *[
+                        RRDB(
+                            nf=self.num_filters,
+                            kernel_size=3,
+                            gc=32,
+                            stride=1,
+                            bias=True,
+                            pad_type="zero",
+                            norm_type=self.norm,
+                            act_type=self.act,
+                            mode="CNA",
+                            plus=self.plus,
+                            c2x2=c2x2,
+                        )
+                        for _ in range(self.num_blocks)
+                    ],
+                    # lr conv
+                    USDU_util.conv_block(
+                        in_nc=self.num_filters,
+                        out_nc=self.num_filters,
+                        kernel_size=3,
+                        norm_type=self.norm,
+                        act_type=None,
+                        mode=self.mode,
+                        c2x2=c2x2,
+                    ),
+                )
+            ),
+            *upsample_blocks,
+            # hr_conv0
+            USDU_util.conv_block(
+                in_nc=self.num_filters,
+                out_nc=self.num_filters,
+                kernel_size=3,
+                norm_type=None,
+                act_type=self.act,
+                c2x2=c2x2,
+            ),
+            # hr_conv1
+            USDU_util.conv_block(
+                in_nc=self.num_filters,
+                out_nc=self.out_nc,
+                kernel_size=3,
+                norm_type=None,
+                act_type=None,
+                c2x2=c2x2,
+            ),
+        )
+
+        self.load_state_dict(self.state, strict=False)
+
+    def new_to_old_arch(self, state: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """#### Convert new architecture state dictionary to old architecture.
+
+        #### Args:
+            - `state` (dict): State dictionary.
+
+        #### Returns:
+            - `dict`: Converted state dictionary.
+        """
+        # add nb to state keys
+        for kind in ("weight", "bias"):
+            self.state_map[f"model.1.sub.{self.num_blocks}.{kind}"] = self.state_map[
+                f"model.1.sub./NB/.{kind}"
+            ]
+            del self.state_map[f"model.1.sub./NB/.{kind}"]
+
+        old_state = OrderedDict()
+        for old_key, new_keys in self.state_map.items():
+            for new_key in new_keys:
+                if r"\1" in old_key:
+                    for k, v in state.items():
+                        sub = re.sub(new_key, old_key, k)
+                        if sub != k:
+                            old_state[sub] = v
+                else:
+                    if new_key in state:
+                        old_state[old_key] = state[new_key]
+
+        # upconv layers
+        max_upconv = 0
+        for key in state.keys():
+            match = re.match(r"(upconv|conv_up)(\d)\.(weight|bias)", key)
+            if match is not None:
+                _, key_num, key_type = match.groups()
+                old_state[f"model.{int(key_num) * 3}.{key_type}"] = state[key]
+                max_upconv = max(max_upconv, int(key_num) * 3)
+
+        # final layers
+        for key in state.keys():
+            if key in ("HRconv.weight", "conv_hr.weight"):
+                old_state[f"model.{max_upconv + 2}.weight"] = state[key]
+            elif key in ("HRconv.bias", "conv_hr.bias"):
+                old_state[f"model.{max_upconv + 2}.bias"] = state[key]
+            elif key in ("conv_last.weight",):
+                old_state[f"model.{max_upconv + 4}.weight"] = state[key]
+            elif key in ("conv_last.bias",):
+                old_state[f"model.{max_upconv + 4}.bias"] = state[key]
+
+        # Sort by first numeric value of each layer
+        def compare(item1: str, item2: str) -> int:
+            parts1 = item1.split(".")
+            parts2 = item2.split(".")
+            int1 = int(parts1[1])
+            int2 = int(parts2[1])
+            return int1 - int2
+
+        sorted_keys = sorted(old_state.keys(), key=functools.cmp_to_key(compare))
+
+        # Rebuild the output dict in the right order
+        out_dict = OrderedDict((k, old_state[k]) for k in sorted_keys)
+
+        return out_dict
+
+    def get_scale(self, min_part: int = 6) -> int:
+        """#### Get the scale factor.
+
+        #### Args:
+            - `min_part` (int, optional): Minimum part. Defaults to 6.
+
+        #### Returns:
+            - `int`: Scale factor.
+        """
+        n = 0
+        for part in list(self.state):
+            parts = part.split(".")[1:]
+            if len(parts) == 2:
+                part_num = int(parts[0])
+                if part_num > min_part and parts[1] == "weight":
+                    n += 1
+        return 2**n
+
+    def get_num_blocks(self) -> int:
+        """#### Get the number of blocks.
+
+        #### Returns:
+            - `int`: Number of blocks.
+        """
+        nbs = []
+        state_keys = self.state_map[r"model.1.sub.\1.RDB\2.conv\3.0.\4"] + (
+            r"model\.\d+\.sub\.(\d+)\.RDB(\d+)\.conv(\d+)\.0\.(weight|bias)",
+        )
+        for state_key in state_keys:
+            for k in self.state:
+                m = re.search(state_key, k)
+                if m:
+                    nbs.append(int(m.group(1)))
+            if nbs:
+                break
+        return max(*nbs) + 1
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass of the RRDBNet.
+
+        #### Args:
+            - `x` (torch.Tensor): Input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: Output tensor.
+        """
+        return self.model(x)
+
+
+PyTorchSRModels = (RRDBNet,)
+PyTorchSRModel = Union[RRDBNet,]
+
+PyTorchModels = (*PyTorchSRModels,)
+PyTorchModel = Union[PyTorchSRModel]

modules/UltimateSDUpscale/USDU_upscaler.py

@@ -0,0 +1,182 @@
+import logging as logger
+import torch
+from PIL import Image
+
+from modules.Device import Device
+from modules.UltimateSDUpscale import RDRB
+from modules.UltimateSDUpscale import image_util
+from modules.Utilities import util
+
+
+def load_state_dict(state_dict: dict) -> RDRB.PyTorchModel:
+    """#### Load a state dictionary into a PyTorch model.
+
+    #### Args:
+        - `state_dict` (dict): The state dictionary.
+
+    #### Returns:
+        - `RDRB.PyTorchModel`: The loaded PyTorch model.
+    """
+    logger.debug("Loading state dict into pytorch model arch")
+    state_dict_keys = list(state_dict.keys())
+    if "params_ema" in state_dict_keys:
+        state_dict = state_dict["params_ema"]
+    model = RDRB.RRDBNet(state_dict)
+    return model
+
+
+class UpscaleModelLoader:
+    """#### Class for loading upscale models."""
+
+    def load_model(self, model_name: str) -> tuple:
+        """#### Load an upscale model.
+
+        #### Args:
+            - `model_name` (str): The name of the model.
+
+        #### Returns:
+            - `tuple`: The loaded model.
+        """
+        model_path = f"_internal/ESRGAN/{model_name}"
+        sd = util.load_torch_file(model_path, safe_load=True)
+        if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
+            sd = util.state_dict_prefix_replace(sd, {"module.": ""})
+        out = load_state_dict(sd).eval()
+        return (out,)
+
+
+class ImageUpscaleWithModel:
+    """#### Class for upscaling images with a model."""
+
+    def upscale(self, upscale_model: torch.nn.Module, image: torch.Tensor) -> tuple:
+        """#### Upscale an image using a model.
+
+        #### Args:
+            - `upscale_model` (torch.nn.Module): The upscale model.
+            - `image` (torch.Tensor): The input image tensor.
+
+        #### Returns:
+            - `tuple`: The upscaled image tensor.
+        """
+        device = torch.device(torch.cuda.current_device())
+        upscale_model.to(device)
+        in_img = image.movedim(-1, -3).to(device)
+        Device.get_free_memory(device)
+
+        tile = 512
+        overlap = 32
+
+        oom = True
+        while oom:
+            steps = in_img.shape[0] * image_util.get_tiled_scale_steps(
+                in_img.shape[3],
+                in_img.shape[2],
+                tile_x=tile,
+                tile_y=tile,
+                overlap=overlap,
+            )
+            pbar = util.ProgressBar(steps)
+            s = image_util.tiled_scale(
+                in_img,
+                lambda a: upscale_model(a),
+                tile_x=tile,
+                tile_y=tile,
+                overlap=overlap,
+                upscale_amount=upscale_model.scale,
+                pbar=pbar,
+            )
+            oom = False
+
+        upscale_model.cpu()
+        s = torch.clamp(s.movedim(-3, -1), min=0, max=1.0)
+        return (s,)
+
+
+def torch_gc() -> None:
+    """#### Perform garbage collection for PyTorch."""
+    pass
+
+
+class Script:
+    """#### Class representing a script."""
+    pass
+
+
+class Options:
+    """#### Class representing options."""
+
+    img2img_background_color: str = "#ffffff"  # Set to white for now
+
+
+class State:
+    """#### Class representing the state."""
+
+    interrupted: bool = False
+
+    def begin(self) -> None:
+        """#### Begin the state."""
+        pass
+
+    def end(self) -> None:
+        """#### End the state."""
+        pass
+
+
+opts = Options()
+state = State()
+
+# Will only ever hold 1 upscaler
+sd_upscalers = [None]
+actual_upscaler = None
+
+# Batch of images to upscale
+batch = None
+
+
+if not hasattr(Image, "Resampling"):  # For older versions of Pillow
+    Image.Resampling = Image
+
+
+class Upscaler:
+    """#### Class for upscaling images."""
+
+    def _upscale(self, img: Image.Image, scale: float) -> Image.Image:
+        """#### Upscale an image.
+
+        #### Args:
+            - `img` (Image.Image): The input image.
+            - `scale` (float): The scale factor.
+
+        #### Returns:
+            - `Image.Image`: The upscaled image.
+        """
+        global actual_upscaler
+        tensor = image_util.pil_to_tensor(img)
+        image_upscale_node = ImageUpscaleWithModel()
+        (upscaled,) = image_upscale_node.upscale(actual_upscaler, tensor)
+        return image_util.tensor_to_pil(upscaled)
+
+    def upscale(self, img: Image.Image, scale: float, selected_model: str = None) -> Image.Image:
+        """#### Upscale an image with a selected model.
+
+        #### Args:
+            - `img` (Image.Image): The input image.
+            - `scale` (float): The scale factor.
+            - `selected_model` (str, optional): The selected model. Defaults to None.
+
+        #### Returns:
+            - `Image.Image`: The upscaled image.
+        """
+        global batch
+        batch = [self._upscale(img, scale) for img in batch]
+        return batch[0]
+
+
+class UpscalerData:
+    """#### Class for storing upscaler data."""
+
+    name: str = ""
+    data_path: str = ""
+
+    def __init__(self):
+        self.scaler = Upscaler()

modules/UltimateSDUpscale/USDU_util.py

@@ -0,0 +1,173 @@
+from typing import Literal
+import torch
+import torch.nn as nn
+
+ConvMode = Literal["CNA", "NAC", "CNAC"]
+
+def act(act_type: str, inplace: bool = True, neg_slope: float = 0.2, n_prelu: int = 1) -> nn.Module:
+    """#### Get the activation layer.
+
+    #### Args:
+        - `act_type` (str): The type of activation.
+        - `inplace` (bool, optional): Whether to perform the operation in-place. Defaults to True.
+        - `neg_slope` (float, optional): The negative slope for LeakyReLU. Defaults to 0.2.
+        - `n_prelu` (int, optional): The number of PReLU parameters. Defaults to 1.
+
+    #### Returns:
+        - `nn.Module`: The activation layer.
+    """
+    act_type = act_type.lower()
+    layer = nn.LeakyReLU(neg_slope, inplace)
+    return layer
+
+def get_valid_padding(kernel_size: int, dilation: int) -> int:
+    """#### Get the valid padding for a convolutional layer.
+
+    #### Args:
+        - `kernel_size` (int): The size of the kernel.
+        - `dilation` (int): The dilation rate.
+
+    #### Returns:
+        - `int`: The valid padding.
+    """
+    kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1)
+    padding = (kernel_size - 1) // 2
+    return padding
+
+def sequential(*args: nn.Module) -> nn.Sequential:
+    """#### Create a sequential container.
+
+    #### Args:
+        - `*args` (nn.Module): The modules to include in the sequential container.
+
+    #### Returns:
+        - `nn.Sequential`: The sequential container.
+    """
+    modules = []
+    for module in args:
+        if isinstance(module, nn.Sequential):
+            for submodule in module.children():
+                modules.append(submodule)
+        elif isinstance(module, nn.Module):
+            modules.append(module)
+    return nn.Sequential(*modules)
+
+def conv_block(
+    in_nc: int,
+    out_nc: int,
+    kernel_size: int,
+    stride: int = 1,
+    dilation: int = 1,
+    groups: int = 1,
+    bias: bool = True,
+    pad_type: str = "zero",
+    norm_type: str | None = None,
+    act_type: str | None = "relu",
+    mode: ConvMode = "CNA",
+    c2x2: bool = False,
+) -> nn.Sequential:
+    """#### Create a convolutional block.
+
+    #### Args:
+        - `in_nc` (int): The number of input channels.
+        - `out_nc` (int): The number of output channels.
+        - `kernel_size` (int): The size of the kernel.
+        - `stride` (int, optional): The stride of the convolution. Defaults to 1.
+        - `dilation` (int, optional): The dilation rate. Defaults to 1.
+        - `groups` (int, optional): The number of groups. Defaults to 1.
+        - `bias` (bool, optional): Whether to include a bias term. Defaults to True.
+        - `pad_type` (str, optional): The type of padding. Defaults to "zero".
+        - `norm_type` (str | None, optional): The type of normalization. Defaults to None.
+        - `act_type` (str | None, optional): The type of activation. Defaults to "relu".
+        - `mode` (ConvMode, optional): The mode of the convolution. Defaults to "CNA".
+        - `c2x2` (bool, optional): Whether to use 2x2 convolutions. Defaults to False.
+
+    #### Returns:
+        - `nn.Sequential`: The convolutional block.
+    """
+    assert mode in ("CNA", "NAC", "CNAC"), "Wrong conv mode [{:s}]".format(mode)
+    padding = get_valid_padding(kernel_size, dilation)
+    padding = padding if pad_type == "zero" else 0
+
+    c = nn.Conv2d(
+        in_nc,
+        out_nc,
+        kernel_size=kernel_size,
+        stride=stride,
+        padding=padding,
+        dilation=dilation,
+        bias=bias,
+        groups=groups,
+    )
+    a = act(act_type) if act_type else None
+    if mode in ("CNA", "CNAC"):
+        return sequential(None, c, None, a)
+
+def upconv_block(
+    in_nc: int,
+    out_nc: int,
+    upscale_factor: int = 2,
+    kernel_size: int = 3,
+    stride: int = 1,
+    bias: bool = True,
+    pad_type: str = "zero",
+    norm_type: str | None = None,
+    act_type: str = "relu",
+    mode: str = "nearest",
+    c2x2: bool = False,
+) -> nn.Sequential:
+    """#### Create an upsampling convolutional block.
+
+    #### Args:
+        - `in_nc` (int): The number of input channels.
+        - `out_nc` (int): The number of output channels.
+        - `upscale_factor` (int, optional): The upscale factor. Defaults to 2.
+        - `kernel_size` (int, optional): The size of the kernel. Defaults to 3.
+        - `stride` (int, optional): The stride of the convolution. Defaults to 1.
+        - `bias` (bool, optional): Whether to include a bias term. Defaults to True.
+        - `pad_type` (str, optional): The type of padding. Defaults to "zero".
+        - `norm_type` (str | None, optional): The type of normalization. Defaults to None.
+        - `act_type` (str, optional): The type of activation. Defaults to "relu".
+        - `mode` (str, optional): The mode of upsampling. Defaults to "nearest".
+        - `c2x2` (bool, optional): Whether to use 2x2 convolutions. Defaults to False.
+
+    #### Returns:
+        - `nn.Sequential`: The upsampling convolutional block.
+    """
+    upsample = nn.Upsample(scale_factor=upscale_factor, mode=mode)
+    conv = conv_block(
+        in_nc,
+        out_nc,
+        kernel_size,
+        stride,
+        bias=bias,
+        pad_type=pad_type,
+        norm_type=norm_type,
+        act_type=act_type,
+        c2x2=c2x2,
+    )
+    return sequential(upsample, conv)
+
+class ShortcutBlock(nn.Module):
+    """#### Elementwise sum the output of a submodule to its input."""
+
+    def __init__(self, submodule: nn.Module):
+        """#### Initialize the ShortcutBlock.
+
+        #### Args:
+            - `submodule` (nn.Module): The submodule to apply.
+        """
+        super(ShortcutBlock, self).__init__()
+        self.sub = submodule
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """#### Forward pass.
+
+        #### Args:
+            - `x` (torch.Tensor): The input tensor.
+
+        #### Returns:
+            - `torch.Tensor`: The output tensor.
+        """
+        output = x + self.sub(x)
+        return output

modules/UltimateSDUpscale/UltimateSDUpscale.py

@@ -0,0 +1,1019 @@
+from modules.AutoEncoders import VariationalAE
+from modules.sample import sampling
+from modules.UltimateSDUpscale import USDU_upscaler, image_util
+import torch
+from PIL import ImageFilter, ImageDraw, Image
+from enum import Enum
+import math
+
+# taken from https://github.com/ssitu/ComfyUI_UltimateSDUpscale
+
+state = USDU_upscaler.state
+
+class UnsupportedModel(Exception):
+    """#### Exception raised for unsupported models."""
+    pass
+
+
+class StableDiffusionProcessing:
+    """#### Class representing the processing of Stable Diffusion images."""
+
+    def __init__(
+        self,
+        init_img: Image.Image,
+        model: torch.nn.Module,
+        positive: str,
+        negative: str,
+        vae: VariationalAE.VAE,
+        seed: int,
+        steps: int,
+        cfg: float,
+        sampler_name: str,
+        scheduler: str,
+        denoise: float,
+        upscale_by: float,
+        uniform_tile_mode: bool,
+    ):
+        """
+        #### Initialize the StableDiffusionProcessing class.
+
+        #### Args:
+            - `init_img` (Image.Image): The initial image.
+            - `model` (torch.nn.Module): The model.
+            - `positive` (str): The positive prompt.
+            - `negative` (str): The negative prompt.
+            - `vae` (VariationalAE.VAE): The variational autoencoder.
+            - `seed` (int): The seed.
+            - `steps` (int): The number of steps.
+            - `cfg` (float): The CFG scale.
+            - `sampler_name` (str): The sampler name.
+            - `scheduler` (str): The scheduler.
+            - `denoise` (float): The denoise strength.
+            - `upscale_by` (float): The upscale factor.
+            - `uniform_tile_mode` (bool): Whether to use uniform tile mode.
+        """
+        # Variables used by the USDU script
+        self.init_images = [init_img]
+        self.image_mask = None
+        self.mask_blur = 0
+        self.inpaint_full_res_padding = 0
+        self.width = init_img.width
+        self.height = init_img.height
+
+        self.model = model
+        self.positive = positive
+        self.negative = negative
+        self.vae = vae
+        self.seed = seed
+        self.steps = steps
+        self.cfg = cfg
+        self.sampler_name = sampler_name
+        self.scheduler = scheduler
+        self.denoise = denoise
+
+        # Variables used only by this script
+        self.init_size = init_img.width, init_img.height
+        self.upscale_by = upscale_by
+        self.uniform_tile_mode = uniform_tile_mode
+
+        # Other required A1111 variables for the USDU script that is currently unused in this script
+        self.extra_generation_params = {}
+
+
+class Processed:
+    """#### Class representing the processed images."""
+
+    def __init__(
+        self, p: StableDiffusionProcessing, images: list, seed: int, info: str
+    ):
+        """
+        #### Initialize the Processed class.
+
+        #### Args:
+            - `p` (StableDiffusionProcessing): The processing object.
+            - `images` (list): The list of images.
+            - `seed` (int): The seed.
+            - `info` (str): The information string.
+        """
+        self.images = images
+        self.seed = seed
+        self.info = info
+
+    def infotext(self, p: StableDiffusionProcessing, index: int) -> str:
+        """
+        #### Get the information text.
+
+        #### Args:
+            - `p` (StableDiffusionProcessing): The processing object.
+            - `index` (int): The index.
+
+        #### Returns:
+            - `str`: The information text.
+        """
+        return None
+
+
+def fix_seed(p: StableDiffusionProcessing) -> None:
+    """
+    #### Fix the seed for reproducibility.
+
+    #### Args:
+        - `p` (StableDiffusionProcessing): The processing object.
+    """
+    pass
+
+
+def process_images(p: StableDiffusionProcessing, pipeline: bool = False) -> Processed:
+    """
+    #### Process the images.
+
+    #### Args:
+        - `p` (StableDiffusionProcessing): The processing object.
+
+    #### Returns:
+        - `Processed`: The processed images.
+    """
+    # Where the main image generation happens in A1111
+
+    # Setup
+    image_mask = p.image_mask.convert("L")
+    init_image = p.init_images[0]
+
+    # Locate the white region of the mask outlining the tile and add padding
+    crop_region = image_util.get_crop_region(image_mask, p.inpaint_full_res_padding)
+
+    x1, y1, x2, y2 = crop_region
+    crop_width = x2 - x1
+    crop_height = y2 - y1
+    crop_ratio = crop_width / crop_height
+    p_ratio = p.width / p.height
+    if crop_ratio > p_ratio:
+        target_width = crop_width
+        target_height = round(crop_width / p_ratio)
+    else:
+        target_width = round(crop_height * p_ratio)
+        target_height = crop_height
+    crop_region, _ = image_util.expand_crop(
+        crop_region,
+        image_mask.width,
+        image_mask.height,
+        target_width,
+        target_height,
+    )
+    tile_size = p.width, p.height
+
+    # Blur the mask
+    if p.mask_blur > 0:
+        image_mask = image_mask.filter(ImageFilter.GaussianBlur(p.mask_blur))
+
+    # Crop the images to get the tiles that will be used for generation
+    tiles = [img.crop(crop_region) for img in USDU_upscaler.batch]
+
+    # Assume the same size for all images in the batch
+    initial_tile_size = tiles[0].size
+
+    # Resize if necessary
+    for i, tile in enumerate(tiles):
+        if tile.size != tile_size:
+            tiles[i] = tile.resize(tile_size, Image.Resampling.LANCZOS)
+
+    # Crop conditioning
+    positive_cropped = image_util.crop_cond(
+        p.positive, crop_region, p.init_size, init_image.size, tile_size
+    )
+    negative_cropped = image_util.crop_cond(
+        p.negative, crop_region, p.init_size, init_image.size, tile_size
+    )
+
+    # Encode the image
+    vae_encoder = VariationalAE.VAEEncode()
+    batched_tiles = torch.cat([image_util.pil_to_tensor(tile) for tile in tiles], dim=0)
+    (latent,) = vae_encoder.encode(p.vae, batched_tiles)
+
+    # Generate samples
+    (samples,) = sampling.common_ksampler(
+        p.model,
+        p.seed,
+        p.steps,
+        p.cfg,
+        p.sampler_name,
+        p.scheduler,
+        positive_cropped,
+        negative_cropped,
+        latent,
+        denoise=p.denoise,
+        pipeline=pipeline
+    )
+
+    # Decode the sample
+    vae_decoder = VariationalAE.VAEDecode()
+    (decoded,) = vae_decoder.decode(p.vae, samples)
+
+    # Convert the sample to a PIL image
+    tiles_sampled = [image_util.tensor_to_pil(decoded, i) for i in range(len(decoded))]
+
+    for i, tile_sampled in enumerate(tiles_sampled):
+        init_image = USDU_upscaler.batch[i]
+
+        # Resize back to the original size
+        if tile_sampled.size != initial_tile_size:
+            tile_sampled = tile_sampled.resize(
+                initial_tile_size, Image.Resampling.LANCZOS
+            )
+
+        # Put the tile into position
+        image_tile_only = Image.new("RGBA", init_image.size)
+        image_tile_only.paste(tile_sampled, crop_region[:2])
+
+        # Add the mask as an alpha channel
+        # Must make a copy due to the possibility of an edge becoming black
+        temp = image_tile_only.copy()
+        image_mask = image_mask.resize(temp.size)
+        temp.putalpha(image_mask)
+        temp.putalpha(image_mask)
+        image_tile_only.paste(temp, image_tile_only)
+
+        # Add back the tile to the initial image according to the mask in the alpha channel
+        result = init_image.convert("RGBA")
+        result.alpha_composite(image_tile_only)
+
+        # Convert back to RGB
+        result = result.convert("RGB")
+        USDU_upscaler.batch[i] = result
+
+    processed = Processed(p, [USDU_upscaler.batch[0]], p.seed, None)
+    return processed
+
+
+class USDUMode(Enum):
+    """#### Enum representing the modes for Ultimate SD Upscale."""
+    LINEAR = 0
+    CHESS = 1
+    NONE = 2
+
+
+class USDUSFMode(Enum):
+    """#### Enum representing the seam fix modes for Ultimate SD Upscale."""
+    NONE = 0
+    BAND_PASS = 1
+    HALF_TILE = 2
+    HALF_TILE_PLUS_INTERSECTIONS = 3
+
+
+class USDUpscaler:
+    """#### Class representing the Ultimate SD Upscaler."""
+
+    def __init__(
+        self,
+        p: StableDiffusionProcessing,
+        image: Image.Image,
+        upscaler_index: int,
+        save_redraw: bool,
+        save_seams_fix: bool,
+        tile_width: int,
+        tile_height: int,
+    ) -> None:
+        """
+        #### Initialize the USDUpscaler class.
+
+        #### Args:
+            - `p` (StableDiffusionProcessing): The processing object.
+            - `image` (Image.Image): The image.
+            - `upscaler_index` (int): The upscaler index.
+            - `save_redraw` (bool): Whether to save the redraw.
+            - `save_seams_fix` (bool): Whether to save the seams fix.
+            - `tile_width` (int): The tile width.
+            - `tile_height` (int): The tile height.
+        """
+        self.p: StableDiffusionProcessing = p
+        self.image: Image = image
+        self.scale_factor = math.ceil(
+            max(p.width, p.height) / max(image.width, image.height)
+        )
+        self.upscaler = USDU_upscaler.sd_upscalers[upscaler_index]
+        self.redraw = USDURedraw()
+        self.redraw.save = save_redraw
+        self.redraw.tile_width = tile_width if tile_width > 0 else tile_height
+        self.redraw.tile_height = tile_height if tile_height > 0 else tile_width
+        self.seams_fix = USDUSeamsFix()
+        self.seams_fix.save = save_seams_fix
+        self.seams_fix.tile_width = tile_width if tile_width > 0 else tile_height
+        self.seams_fix.tile_height = tile_height if tile_height > 0 else tile_width
+        self.initial_info = None
+        self.rows = math.ceil(self.p.height / self.redraw.tile_height)
+        self.cols = math.ceil(self.p.width / self.redraw.tile_width)
+
+    def get_factor(self, num: int) -> int:
+        """
+        #### Get the factor for a given number.
+
+        #### Args:
+            - `num` (int): The number.
+
+        #### Returns:
+            - `int`: The factor.
+        """
+        if num == 1:
+            return 2
+        if num % 4 == 0:
+            return 4
+        if num % 3 == 0:
+            return 3
+        if num % 2 == 0:
+            return 2
+        return 0
+
+    def get_factors(self) -> None:
+        """
+        #### Get the list of scale factors.
+        """
+        scales = []
+        current_scale = 1
+        current_scale_factor = self.get_factor(self.scale_factor)
+        while current_scale < self.scale_factor:
+            current_scale_factor = self.get_factor(self.scale_factor // current_scale)
+            scales.append(current_scale_factor)
+            current_scale = current_scale * current_scale_factor
+        self.scales = enumerate(scales)
+
+    def upscale(self) -> None:
+        """
+        #### Upscale the image.
+        """
+        # Log info
+        print(f"Canva size: {self.p.width}x{self.p.height}")
+        print(f"Image size: {self.image.width}x{self.image.height}")
+        print(f"Scale factor: {self.scale_factor}")
+        # Get list with scale factors
+        self.get_factors()
+        # Upscaling image over all factors
+        for index, value in self.scales:
+            print(f"Upscaling iteration {index + 1} with scale factor {value}")
+            self.image = self.upscaler.scaler.upscale(
+                self.image, value, self.upscaler.data_path
+            )
+        # Resize image to set values
+        self.image = self.image.resize(
+            (self.p.width, self.p.height), resample=Image.LANCZOS
+        )
+
+    def setup_redraw(self, redraw_mode: int, padding: int, mask_blur: int) -> None:
+        """
+        #### Set up the redraw.
+
+        #### Args:
+            - `redraw_mode` (int): The redraw mode.
+            - `padding` (int): The padding.
+            - `mask_blur` (int): The mask blur.
+        """
+        self.redraw.mode = USDUMode(redraw_mode)
+        self.redraw.enabled = self.redraw.mode != USDUMode.NONE
+        self.redraw.padding = padding
+        self.p.mask_blur = mask_blur
+
+    def setup_seams_fix(
+        self, padding: int, denoise: float, mask_blur: int, width: int, mode: int
+    ) -> None:
+        """
+        #### Set up the seams fix.
+
+        #### Args:
+            - `padding` (int): The padding.
+            - `denoise` (float): The denoise strength.
+            - `mask_blur` (int): The mask blur.
+            - `width` (int): The width.
+            - `mode` (int): The mode.
+        """
+        self.seams_fix.padding = padding
+        self.seams_fix.denoise = denoise
+        self.seams_fix.mask_blur = mask_blur
+        self.seams_fix.width = width
+        self.seams_fix.mode = USDUSFMode(mode)
+        self.seams_fix.enabled = self.seams_fix.mode != USDUSFMode.NONE
+
+    def calc_jobs_count(self) -> None:
+        """
+        #### Calculate the number of jobs.
+        """
+        redraw_job_count = (self.rows * self.cols) if self.redraw.enabled else 0
+        seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols
+        global state
+        state.job_count = redraw_job_count + seams_job_count
+
+    def print_info(self) -> None:
+        """
+        #### Print the information.
+        """
+        print(f"Tile size: {self.redraw.tile_width}x{self.redraw.tile_height}")
+        print(f"Tiles amount: {self.rows * self.cols}")
+        print(f"Grid: {self.rows}x{self.cols}")
+        print(f"Redraw enabled: {self.redraw.enabled}")
+        print(f"Seams fix mode: {self.seams_fix.mode.name}")
+
+    def add_extra_info(self) -> None:
+        """
+        #### Add extra information.
+        """
+        self.p.extra_generation_params["Ultimate SD upscale upscaler"] = (
+            self.upscaler.name
+        )
+        self.p.extra_generation_params["Ultimate SD upscale tile_width"] = (
+            self.redraw.tile_width
+        )
+        self.p.extra_generation_params["Ultimate SD upscale tile_height"] = (
+            self.redraw.tile_height
+        )
+        self.p.extra_generation_params["Ultimate SD upscale mask_blur"] = (
+            self.p.mask_blur
+        )
+        self.p.extra_generation_params["Ultimate SD upscale padding"] = (
+            self.redraw.padding
+        )
+
+    def process(self, pipeline) -> None:
+        """
+        #### Process the image.
+        """
+        USDU_upscaler.state.begin()
+        self.calc_jobs_count()
+        self.result_images = []
+        if self.redraw.enabled:
+            self.image = self.redraw.start(self.p, self.image, self.rows, self.cols, pipeline)
+            self.initial_info = self.redraw.initial_info
+        self.result_images.append(self.image)
+
+        if self.seams_fix.enabled:
+            self.image = self.seams_fix.start(self.p, self.image, self.rows, self.cols, pipeline)
+            self.initial_info = self.seams_fix.initial_info
+            self.result_images.append(self.image)
+        USDU_upscaler.state.end()
+
+
+class USDURedraw:
+    """#### Class representing the redraw functionality for Ultimate SD Upscale."""
+
+    def init_draw(self, p: StableDiffusionProcessing, width: int, height: int) -> tuple:
+        """
+        #### Initialize the draw.
+
+        #### Args:
+            - `p` (StableDiffusionProcessing): The processing object.
+            - `width` (int): The width.
+            - `height` (int): The height.
+
+        #### Returns:
+            - `tuple`: The mask and draw objects.
+        """
+        p.inpaint_full_res = True
+        p.inpaint_full_res_padding = self.padding
+        p.width = math.ceil((self.tile_width + self.padding) / 64) * 64
+        p.height = math.ceil((self.tile_height + self.padding) / 64) * 64
+        mask = Image.new("L", (width, height), "black")
+        draw = ImageDraw.Draw(mask)
+        return mask, draw
+
+    def calc_rectangle(self, xi: int, yi: int) -> tuple:
+        """
+        #### Calculate the rectangle coordinates.
+
+        #### Args:
+            - `xi` (int): The x index.
+            - `yi` (int): The y index.
+
+        #### Returns:
+            - `tuple`: The rectangle coordinates.
+        """
+        x1 = xi * self.tile_width
+        y1 = yi * self.tile_height
+        x2 = xi * self.tile_width + self.tile_width
+        y2 = yi * self.tile_height + self.tile_height
+
+        return x1, y1, x2, y2
+
+    def linear_process(
+        self, p: StableDiffusionProcessing, image: Image.Image, rows: int, cols: int, pipeline: bool = False
+    ) -> Image.Image:
+        """
+        #### Perform linear processing.
+
+        #### Args:
+            - `p` (StableDiffusionProcessing): The processing object.
+            - `image` (Image.Image): The image.
+            - `rows` (int): The number of rows.
+            - `cols` (int): The number of columns.
+
+        #### Returns:
+            - `Image.Image`: The processed image.
+        """
+        global state
+        mask, draw = self.init_draw(p, image.width, image.height)
+        for yi in range(rows):
+            for xi in range(cols):
+                if state.interrupted:
+                    break
+                draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
+                p.init_images = [image]
+                p.image_mask = mask
+                processed = process_images(p, pipeline)
+                draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
+                if len(processed.images) > 0:
+                    image = processed.images[0]
+
+        p.width = image.width
+        p.height = image.height
+        self.initial_info = processed.infotext(p, 0)
+
+        return image
+
+    def start(self, p: StableDiffusionProcessing, image: Image.Image, rows: int, cols: int, pipeline: bool = False) -> Image.Image:
+        """#### Start the redraw.
+
+        #### Args:
+            - `p` (StableDiffusionProcessing): The processing object.
+            - `image` (Image.Image): The image.
+            - `rows` (int): The number of rows.
+            - `cols` (int): The number of columns.
+            
+        #### Returns:
+            - `Image.Image`: The processed image.
+        """
+        self.initial_info = None
+        return self.linear_process(p, image, rows, cols, pipeline=pipeline)
+
+
+class USDUSeamsFix:
+    """#### Class representing the seams fix functionality for Ultimate SD Upscale."""
+
+    def init_draw(self, p: StableDiffusionProcessing) -> None:
+        """#### Initialize the draw.
+
+        #### Args:
+            - `p` (StableDiffusionProcessing): The processing object.
+        """
+        self.initial_info = None
+        p.width = math.ceil((self.tile_width + self.padding) / 64) * 64
+        p.height = math.ceil((self.tile_height + self.padding) / 64) * 64
+
+    def half_tile_process(
+        self, p: StableDiffusionProcessing, image: Image.Image, rows: int, cols: int, pipeline: bool = False
+    ) -> Image.Image:
+        """#### Perform half-tile processing.
+
+        #### Args:
+            - `p` (StableDiffusionProcessing): The processing object.
+            - `image` (Image.Image): The image.
+            - `rows` (int): The number of rows.
+            - `cols` (int): The number of columns.
+
+        #### Returns:
+            - `Image.Image`: The processed image.
+        """
+        global state
+        self.init_draw(p)
+        processed = None
+
+        gradient = Image.linear_gradient("L")
+        row_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
+        row_gradient.paste(
+            gradient.resize(
+                (self.tile_width, self.tile_height // 2), resample=Image.BICUBIC
+            ),
+            (0, 0),
+        )
+        row_gradient.paste(
+            gradient.rotate(180).resize(
+                (self.tile_width, self.tile_height // 2), resample=Image.BICUBIC
+            ),
+            (0, self.tile_height // 2),
+        )
+        col_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
+        col_gradient.paste(
+            gradient.rotate(90).resize(
+                (self.tile_width // 2, self.tile_height), resample=Image.BICUBIC
+            ),
+            (0, 0),
+        )
+        col_gradient.paste(
+            gradient.rotate(270).resize(
+                (self.tile_width // 2, self.tile_height), resample=Image.BICUBIC
+            ),
+            (self.tile_width // 2, 0),
+        )
+
+        p.denoising_strength = self.denoise
+        p.mask_blur = self.mask_blur
+
+        for yi in range(rows - 1):
+            for xi in range(cols):
+                p.width = self.tile_width
+                p.height = self.tile_height
+                p.inpaint_full_res = True
+                p.inpaint_full_res_padding = self.padding
+                mask = Image.new("L", (image.width, image.height), "black")
+                mask.paste(
+                    row_gradient,
+                    (
+                        xi * self.tile_width,
+                        yi * self.tile_height + self.tile_height // 2,
+                    ),
+                )
+
+                p.init_images = [image]
+                p.image_mask = mask
+                processed = process_images(p, pipeline)
+                if len(processed.images) > 0:
+                    image = processed.images[0]
+
+        for yi in range(rows):
+            for xi in range(cols - 1):
+                p.width = self.tile_width
+                p.height = self.tile_height
+                p.inpaint_full_res = True
+                p.inpaint_full_res_padding = self.padding
+                mask = Image.new("L", (image.width, image.height), "black")
+                mask.paste(
+                    col_gradient,
+                    (
+                        xi * self.tile_width + self.tile_width // 2,
+                        yi * self.tile_height,
+                    ),
+                )
+
+                p.init_images = [image]
+                p.image_mask = mask
+                processed = process_images(p, pipeline)
+                if len(processed.images) > 0:
+                    image = processed.images[0]
+
+        p.width = image.width
+        p.height = image.height
+        if processed is not None:
+            self.initial_info = processed.infotext(p, 0)
+
+        return image
+
+    def start(
+        self, p: StableDiffusionProcessing, image: Image.Image, rows: int, cols: int, pipeline: bool = False
+    ) -> Image.Image:
+        """#### Start the seams fix process.
+
+        #### Args:
+            - `p` (StableDiffusionProcessing): The processing object.
+            - `image` (Image.Image): The image.
+            - `rows` (int): The number of rows.
+            - `cols` (int): The number of columns.
+
+        #### Returns:
+            - `Image.Image`: The processed image.
+        """
+        return self.half_tile_process(p, image, rows, cols, pipeline=pipeline)
+
+
+class Script(USDU_upscaler.Script):
+    """#### Class representing the script for Ultimate SD Upscale."""
+
+    def run(
+        self,
+        p: StableDiffusionProcessing,
+        _: None,
+        tile_width: int,
+        tile_height: int,
+        mask_blur: int,
+        padding: int,
+        seams_fix_width: int,
+        seams_fix_denoise: float,
+        seams_fix_padding: int,
+        upscaler_index: int,
+        save_upscaled_image: bool,
+        redraw_mode: int,
+        save_seams_fix_image: bool,
+        seams_fix_mask_blur: int,
+        seams_fix_type: int,
+        target_size_type: int,
+        custom_width: int,
+        custom_height: int,
+        custom_scale: float,
+        pipeline: bool = False,
+    ) -> Processed:
+        """#### Run the script.
+
+        #### Args:
+            - `p` (StableDiffusionProcessing): The processing object.
+            - `_` (None): Unused parameter.
+            - `tile_width` (int): The tile width.
+            - `tile_height` (int): The tile height.
+            - `mask_blur` (int): The mask blur.
+            - `padding` (int): The padding.
+            - `seams_fix_width` (int): The seams fix width.
+            - `seams_fix_denoise` (float): The seams fix denoise strength.
+            - `seams_fix_padding` (int): The seams fix padding.
+            - `upscaler_index` (int): The upscaler index.
+            - `save_upscaled_image` (bool): Whether to save the upscaled image.
+            - `redraw_mode` (int): The redraw mode.
+            - `save_seams_fix_image` (bool): Whether to save the seams fix image.
+            - `seams_fix_mask_blur` (int): The seams fix mask blur.
+            - `seams_fix_type` (int): The seams fix type.
+            - `target_size_type` (int): The target size type.
+            - `custom_width` (int): The custom width.
+            - `custom_height` (int): The custom height.
+            - `custom_scale` (float): The custom scale.
+
+        #### Returns:
+            - `Processed`: The processed images.
+        """
+        # Init
+        fix_seed(p)
+        USDU_upscaler.torch_gc()
+
+        p.do_not_save_grid = True
+        p.do_not_save_samples = True
+        p.inpaint_full_res = False
+
+        p.inpainting_fill = 1
+        p.n_iter = 1
+        p.batch_size = 1
+
+        seed = p.seed
+
+        # Init image
+        init_img = p.init_images[0]
+        init_img = image_util.flatten(
+            init_img, USDU_upscaler.opts.img2img_background_color
+        )
+
+        p.width = math.ceil((init_img.width * custom_scale) / 64) * 64
+        p.height = math.ceil((init_img.height * custom_scale) / 64) * 64
+
+        # Upscaling
+        upscaler = USDUpscaler(
+            p,
+            init_img,
+            upscaler_index,
+            save_upscaled_image,
+            save_seams_fix_image,
+            tile_width,
+            tile_height,
+        )
+        upscaler.upscale()
+
+        # Drawing
+        upscaler.setup_redraw(redraw_mode, padding, mask_blur)
+        upscaler.setup_seams_fix(
+            seams_fix_padding,
+            seams_fix_denoise,
+            seams_fix_mask_blur,
+            seams_fix_width,
+            seams_fix_type,
+        )
+        upscaler.print_info()
+        upscaler.add_extra_info()
+        upscaler.process(pipeline=pipeline)
+        result_images = upscaler.result_images
+
+        return Processed(
+            p,
+            result_images,
+            seed,
+            upscaler.initial_info if upscaler.initial_info is not None else "",
+        )
+
+
+# Upscaler
+old_init = USDUpscaler.__init__
+
+
+def new_init(
+    self: USDUpscaler,
+    p: StableDiffusionProcessing,
+    image: Image.Image,
+    upscaler_index: int,
+    save_redraw: bool,
+    save_seams_fix: bool,
+    tile_width: int,
+    tile_height: int,
+) -> None:
+    """#### Initialize the USDUpscaler class with new settings.
+
+    #### Args:
+        - `self` (USDUpscaler): The USDUpscaler instance.
+        - `p` (StableDiffusionProcessing): The processing object.
+        - `image` (Image.Image): The image.
+        - `upscaler_index` (int): The upscaler index.
+        - `save_redraw` (bool): Whether to save the redraw.
+        - `save_seams_fix` (bool): Whether to save the seams fix.
+        - `tile_width` (int): The tile width.
+        - `tile_height` (int): The tile height.
+    """
+    p.width = math.ceil((image.width * p.upscale_by) / 8) * 8
+    p.height = math.ceil((image.height * p.upscale_by) / 8) * 8
+    old_init(
+        self,
+        p,
+        image,
+        upscaler_index,
+        save_redraw,
+        save_seams_fix,
+        tile_width,
+        tile_height,
+    )
+
+
+USDUpscaler.__init__ = new_init
+
+# Redraw
+old_setup_redraw = USDURedraw.init_draw
+
+
+def new_setup_redraw(
+    self: USDURedraw, p: StableDiffusionProcessing, width: int, height: int
+) -> tuple:
+    """#### Set up the redraw with new settings.
+
+    #### Args:
+        - `self` (USDURedraw): The USDURedraw instance.
+        - `p` (StableDiffusionProcessing): The processing object.
+        - `width` (int): The width.
+        - `height` (int): The height.
+
+    #### Returns:
+        - `tuple`: The mask and draw objects.
+    """
+    mask, draw = old_setup_redraw(self, p, width, height)
+    p.width = math.ceil((self.tile_width + self.padding) / 8) * 8
+    p.height = math.ceil((self.tile_height + self.padding) / 8) * 8
+    return mask, draw
+
+
+USDURedraw.init_draw = new_setup_redraw
+
+# Seams fix
+old_setup_seams_fix = USDUSeamsFix.init_draw
+
+
+def new_setup_seams_fix(self: USDUSeamsFix, p: StableDiffusionProcessing) -> None:
+    """#### Set up the seams fix with new settings.
+
+    #### Args:
+        - `self` (USDUSeamsFix): The USDUSeamsFix instance.
+        - `p` (StableDiffusionProcessing): The processing object.
+    """
+    old_setup_seams_fix(self, p)
+    p.width = math.ceil((self.tile_width + self.padding) / 8) * 8
+    p.height = math.ceil((self.tile_height + self.padding) / 8) * 8
+
+
+USDUSeamsFix.init_draw = new_setup_seams_fix
+
+# Make the script upscale on a batch of images instead of one image
+old_upscale = USDUpscaler.upscale
+
+
+def new_upscale(self: USDUpscaler) -> None:
+    """#### Upscale a batch of images.
+
+    #### Args:
+        - `self` (USDUpscaler): The USDUpscaler instance.
+    """
+    old_upscale(self)
+    USDU_upscaler.batch = [self.image] + [
+        img.resize((self.p.width, self.p.height), resample=Image.LANCZOS)
+        for img in USDU_upscaler.batch[1:]
+    ]
+
+
+USDUpscaler.upscale = new_upscale
+MAX_RESOLUTION = 8192
+# The modes available for Ultimate SD Upscale
+MODES = {
+    "Linear": USDUMode.LINEAR,
+    "Chess": USDUMode.CHESS,
+    "None": USDUMode.NONE,
+}
+# The seam fix modes
+SEAM_FIX_MODES = {
+    "None": USDUSFMode.NONE,
+    "Band Pass": USDUSFMode.BAND_PASS,
+    "Half Tile": USDUSFMode.HALF_TILE,
+    "Half Tile + Intersections": USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS,
+}
+
+
+class UltimateSDUpscale:
+    """#### Class representing the Ultimate SD Upscale functionality."""
+
+    def upscale(
+        self,
+        image: torch.Tensor,
+        model: torch.nn.Module,
+        positive: str,
+        negative: str,
+        vae: VariationalAE.VAE,
+        upscale_by: float,
+        seed: int,
+        steps: int,
+        cfg: float,
+        sampler_name: str,
+        scheduler: str,
+        denoise: float,
+        upscale_model: any,
+        mode_type: str,
+        tile_width: int,
+        tile_height: int,
+        mask_blur: int,
+        tile_padding: int,
+        seam_fix_mode: str,
+        seam_fix_denoise: float,
+        seam_fix_mask_blur: int,
+        seam_fix_width: int,
+        seam_fix_padding: int,
+        force_uniform_tiles: bool,
+        pipeline: bool = False,
+    ) -> tuple:
+        """#### Upscale the image.
+
+        #### Args:
+            - `image` (torch.Tensor): The image tensor.
+            - `model` (torch.nn.Module): The model.
+            - `positive` (str): The positive prompt.
+            - `negative` (str): The negative prompt.
+            - `vae` (VariationalAE.VAE): The variational autoencoder.
+            - `upscale_by` (float): The upscale factor.
+            - `seed` (int): The seed.
+            - `steps` (int): The number of steps.
+            - `cfg` (float): The CFG scale.
+            - `sampler_name` (str): The sampler name.
+            - `scheduler` (str): The scheduler.
+            - `denoise` (float): The denoise strength.
+            - `upscale_model` (any): The upscale model.
+            - `mode_type` (str): The mode type.
+            - `tile_width` (int): The tile width.
+            - `tile_height` (int): The tile height.
+            - `mask_blur` (int): The mask blur.
+            - `tile_padding` (int): The tile padding.
+            - `seam_fix_mode` (str): The seam fix mode.
+            - `seam_fix_denoise` (float): The seam fix denoise strength.
+            - `seam_fix_mask_blur` (int): The seam fix mask blur.
+            - `seam_fix_width` (int): The seam fix width.
+            - `seam_fix_padding` (int): The seam fix padding.
+            - `force_uniform_tiles` (bool): Whether to force uniform tiles.
+
+        #### Returns:
+            - `tuple`: The resulting tensor.
+        """
+        # Set up A1111 patches
+
+        # Upscaler
+        # An object that the script works with
+        USDU_upscaler.sd_upscalers[0] = USDU_upscaler.UpscalerData()
+        # Where the actual upscaler is stored, will be used when the script upscales using the Upscaler in UpscalerData
+        USDU_upscaler.actual_upscaler = upscale_model
+
+        # Set the batch of images
+        USDU_upscaler.batch = [image_util.tensor_to_pil(image, i) for i in range(len(image))]
+
+        # Processing
+        sdprocessing = StableDiffusionProcessing(
+            image_util.tensor_to_pil(image),
+            model,
+            positive,
+            negative,
+            vae,
+            seed,
+            steps,
+            cfg,
+            sampler_name,
+            scheduler,
+            denoise,
+            upscale_by,
+            force_uniform_tiles,
+        )
+
+        # Running the script
+        script = Script()
+        script.run(
+            p=sdprocessing,
+            _=None,
+            tile_width=tile_width,
+            tile_height=tile_height,
+            mask_blur=mask_blur,
+            padding=tile_padding,
+            seams_fix_width=seam_fix_width,
+            seams_fix_denoise=seam_fix_denoise,
+            seams_fix_padding=seam_fix_padding,
+            upscaler_index=0,
+            save_upscaled_image=False,
+            redraw_mode=MODES[mode_type],
+            save_seams_fix_image=False,
+            seams_fix_mask_blur=seam_fix_mask_blur,
+            seams_fix_type=SEAM_FIX_MODES[seam_fix_mode],
+            target_size_type=2,
+            custom_width=None,
+            custom_height=None,
+            custom_scale=upscale_by,
+            pipeline=pipeline,
+        )
+
+        # Return the resulting images
+        images = [image_util.pil_to_tensor(img) for img in USDU_upscaler.batch]
+        tensor = torch.cat(images, dim=0)
+        return (tensor,)

modules/UltimateSDUpscale/image_util.py

@@ -0,0 +1,265 @@
+import math
+import numpy as np
+import torch
+from PIL import Image
+
+
+def get_tiled_scale_steps(width: int, height: int, tile_x: int, tile_y: int, overlap: int) -> int:
+    """#### Calculate the number of steps required for tiled scaling.
+
+    #### Args:
+        - `width` (int): The width of the image.
+        - `height` (int): The height of the image.
+        - `tile_x` (int): The width of each tile.
+        - `tile_y` (int): The height of each tile.
+        - `overlap` (int): The overlap between tiles.
+
+    #### Returns:
+        - `int`: The number of steps required for tiled scaling.
+    """
+    return math.ceil((height / (tile_y - overlap))) * math.ceil(
+        (width / (tile_x - overlap))
+    )
+
+
+@torch.inference_mode()
+def tiled_scale(
+    samples: torch.Tensor,
+    function: callable,
+    tile_x: int = 64,
+    tile_y: int = 64,
+    overlap: int = 8,
+    upscale_amount: float = 4,
+    out_channels: int = 3,
+    pbar: any = None,
+) -> torch.Tensor:
+    """#### Perform tiled scaling on a batch of samples.
+
+    #### Args:
+        - `samples` (torch.Tensor): The input samples.
+        - `function` (callable): The function to apply to each tile.
+        - `tile_x` (int, optional): The width of each tile. Defaults to 64.
+        - `tile_y` (int, optional): The height of each tile. Defaults to 64.
+        - `overlap` (int, optional): The overlap between tiles. Defaults to 8.
+        - `upscale_amount` (float, optional): The upscale amount. Defaults to 4.
+        - `out_channels` (int, optional): The number of output channels. Defaults to 3.
+        - `pbar` (any, optional): The progress bar. Defaults to None.
+
+    #### Returns:
+        - `torch.Tensor`: The scaled output tensor.
+    """
+    output = torch.empty(
+        (
+            samples.shape[0],
+            out_channels,
+            round(samples.shape[2] * upscale_amount),
+            round(samples.shape[3] * upscale_amount),
+        ),
+        device="cpu",
+    )
+    for b in range(samples.shape[0]):
+        s = samples[b : b + 1]
+        out = torch.zeros(
+            (
+                s.shape[0],
+                out_channels,
+                round(s.shape[2] * upscale_amount),
+                round(s.shape[3] * upscale_amount),
+            ),
+            device="cpu",
+        )
+        out_div = torch.zeros(
+            (
+                s.shape[0],
+                out_channels,
+                round(s.shape[2] * upscale_amount),
+                round(s.shape[3] * upscale_amount),
+            ),
+            device="cpu",
+        )
+        for y in range(0, s.shape[2], tile_y - overlap):
+            for x in range(0, s.shape[3], tile_x - overlap):
+                s_in = s[:, :, y : y + tile_y, x : x + tile_x]
+
+                ps = function(s_in).cpu()
+                mask = torch.ones_like(ps)
+                feather = round(overlap * upscale_amount)
+                for t in range(feather):
+                    mask[:, :, t : 1 + t, :] *= (1.0 / feather) * (t + 1)
+                    mask[:, :, mask.shape[2] - 1 - t : mask.shape[2] - t, :] *= (
+                        1.0 / feather
+                    ) * (t + 1)
+                    mask[:, :, :, t : 1 + t] *= (1.0 / feather) * (t + 1)
+                    mask[:, :, :, mask.shape[3] - 1 - t : mask.shape[3] - t] *= (
+                        1.0 / feather
+                    ) * (t + 1)
+                out[
+                    :,
+                    :,
+                    round(y * upscale_amount) : round((y + tile_y) * upscale_amount),
+                    round(x * upscale_amount) : round((x + tile_x) * upscale_amount),
+                ] += ps * mask
+                out_div[
+                    :,
+                    :,
+                    round(y * upscale_amount) : round((y + tile_y) * upscale_amount),
+                    round(x * upscale_amount) : round((x + tile_x) * upscale_amount),
+                ] += mask
+
+        output[b : b + 1] = out / out_div
+    return output
+
+
+def flatten(img: Image.Image, bgcolor: str) -> Image.Image:
+    """#### Replace transparency with a background color.
+
+    #### Args:
+        - `img` (Image.Image): The input image.
+        - `bgcolor` (str): The background color.
+
+    #### Returns:
+        - `Image.Image`: The image with transparency replaced by the background color.
+    """
+    if img.mode in ("RGB"):
+        return img
+    return Image.alpha_composite(Image.new("RGBA", img.size, bgcolor), img).convert(
+        "RGB"
+    )
+
+
+BLUR_KERNEL_SIZE = 15
+
+
+def tensor_to_pil(img_tensor: torch.Tensor, batch_index: int = 0) -> Image.Image:
+    """#### Convert a tensor to a PIL image.
+
+    #### Args:
+        - `img_tensor` (torch.Tensor): The input tensor.
+        - `batch_index` (int, optional): The batch index. Defaults to 0.
+
+    #### Returns:
+        - `Image.Image`: The converted PIL image.
+    """
+    img_tensor = img_tensor[batch_index].unsqueeze(0)
+    i = 255.0 * img_tensor.cpu().numpy()
+    img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8).squeeze())
+    return img
+
+
+def pil_to_tensor(image: Image.Image) -> torch.Tensor:
+    """#### Convert a PIL image to a tensor.
+
+    #### Args:
+        - `image` (Image.Image): The input PIL image.
+
+    #### Returns:
+        - `torch.Tensor`: The converted tensor.
+    """
+    image = np.array(image).astype(np.float32) / 255.0
+    image = torch.from_numpy(image).unsqueeze(0)
+    return image
+
+
+def get_crop_region(mask: Image.Image, pad: int = 0) -> tuple:
+    """#### Get the coordinates of the white rectangular mask region.
+
+    #### Args:
+        - `mask` (Image.Image): The input mask image in 'L' mode.
+        - `pad` (int, optional): The padding to apply. Defaults to 0.
+
+    #### Returns:
+        - `tuple`: The coordinates of the crop region.
+    """
+    coordinates = mask.getbbox()
+    if coordinates is not None:
+        x1, y1, x2, y2 = coordinates
+    else:
+        x1, y1, x2, y2 = mask.width, mask.height, 0, 0
+    # Apply padding
+    x1 = max(x1 - pad, 0)
+    y1 = max(y1 - pad, 0)
+    x2 = min(x2 + pad, mask.width)
+    y2 = min(y2 + pad, mask.height)
+    return fix_crop_region((x1, y1, x2, y2), (mask.width, mask.height))
+
+
+def fix_crop_region(region: tuple, image_size: tuple) -> tuple:
+    """#### Remove the extra pixel added by the get_crop_region function.
+
+    #### Args:
+        - `region` (tuple): The crop region coordinates.
+        - `image_size` (tuple): The size of the image.
+
+    #### Returns:
+        - `tuple`: The fixed crop region coordinates.
+    """
+    image_width, image_height = image_size
+    x1, y1, x2, y2 = region
+    if x2 < image_width:
+        x2 -= 1
+    if y2 < image_height:
+        y2 -= 1
+    return x1, y1, x2, y2
+
+
+def expand_crop(region: tuple, width: int, height: int, target_width: int, target_height: int) -> tuple:
+    """#### Expand a crop region to a specified target size.
+
+    #### Args:
+        - `region` (tuple): The crop region coordinates.
+        - `width` (int): The width of the image.
+        - `height` (int): The height of the image.
+        - `target_width` (int): The desired width of the crop region.
+        - `target_height` (int): The desired height of the crop region.
+
+    #### Returns:
+        - `tuple`: The expanded crop region coordinates and the target size.
+    """
+    x1, y1, x2, y2 = region
+    actual_width = x2 - x1
+    actual_height = y2 - y1
+
+    # Try to expand region to the right of half the difference
+    width_diff = target_width - actual_width
+    x2 = min(x2 + width_diff // 2, width)
+    # Expand region to the left of the difference including the pixels that could not be expanded to the right
+    width_diff = target_width - (x2 - x1)
+    x1 = max(x1 - width_diff, 0)
+    # Try the right again
+    width_diff = target_width - (x2 - x1)
+    x2 = min(x2 + width_diff, width)
+
+    # Try to expand region to the bottom of half the difference
+    height_diff = target_height - actual_height
+    y2 = min(y2 + height_diff // 2, height)
+    # Expand region to the top of the difference including the pixels that could not be expanded to the bottom
+    height_diff = target_height - (y2 - y1)
+    y1 = max(y1 - height_diff, 0)
+    # Try the bottom again
+    height_diff = target_height - (y2 - y1)
+    y2 = min(y2 + height_diff, height)
+
+    return (x1, y1, x2, y2), (target_width, target_height)
+
+
+def crop_cond(cond: list, region: tuple, init_size: tuple, canvas_size: tuple, tile_size: tuple, w_pad: int = 0, h_pad: int = 0) -> list:
+    """#### Crop conditioning data to match a specific region.
+
+    #### Args:
+        - `cond` (list): The conditioning data.
+        - `region` (tuple): The crop region coordinates.
+        - `init_size` (tuple): The initial size of the image.
+        - `canvas_size` (tuple): The size of the canvas.
+        - `tile_size` (tuple): The size of the tile.
+        - `w_pad` (int, optional): The width padding. Defaults to 0.
+        - `h_pad` (int, optional): The height padding. Defaults to 0.
+
+    #### Returns:
+        - `list`: The cropped conditioning data.
+    """
+    cropped = []
+    for emb, x in cond:
+        cond_dict = x.copy()
+        n = [emb, cond_dict]
+        cropped.append(n)
+    return cropped